Skip to content
This repository was archived by the owner on Aug 28, 2024. It is now read-only.

Commit

Permalink
Merge branch 'debug'
Browse files Browse the repository at this point in the history
  • Loading branch information
@daminton
daminton committed Jan 19, 2023
2 parents 42234ad + 8d01494 commit e254f38
Show file tree
Hide file tree
Showing 3 changed files with 130 additions and 77 deletions.
1 change: 1 addition & 0 deletions .gitignore
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,3 +30,4 @@ dump*

bin/
build/*
disruption_headon/swiftest_driver.sh
137 changes: 95 additions & 42 deletions examples/Fragmentation/Fragmentation_Movie.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -71,24 +71,24 @@
np.array([-1.45, -6.28, 0.0])],
"hitandrun_pure" : [np.array([ 0.00, 6.28, 0.0]),
np.array([-1.52, -6.28, 0.0])],
"merge" : [np.array([ 0.05, 6.28, 0.0]),
"merge" : [np.array([ 0.04, 6.28, 0.0]),
np.array([ 0.05, 6.18, 0.0])]
}

rot_vectors = {"disruption_headon" : [np.array([1.0e3, 0.0, 1.0e2]),
np.array([-1.0e2, -1.0e3, 0.0])],
"disruption_off_axis": [np.array([0.0, 0.0, 2.0e3]),
rot_vectors = {"disruption_headon" : [np.array([0.0, 0.0, 1.0e5]),
np.array([0.0, 0.0, -5e5])],
"disruption_off_axis": [np.array([0.0, 0.0, 2.0e5]),
np.array([0.0, 0.0, -1.0e5])],
"supercatastrophic_headon": [np.array([0.0, 0.0, 1.0e5]),
np.array([0.0, 0.0, -5.0e5])],
"supercatastrophic_off_axis": [np.array([0.0, 0.0, 1.0e5]),
np.array([0.0, 0.0, -1.0e4])],
"supercatastrophic_headon": [np.array([3.0e2, 3.0e2, 1.0e3]),
np.array([1.0e3, -3.0e2, -5.0e3])],
"supercatastrophic_off_axis": [np.array([3.0e2, 3.0e2, 1.0e3]),
np.array([1.0e3, 3.0e3, -1.0e4])],
"hitandrun_disrupt" : [np.array([0.0, 0.0, 6.0e3]),
np.array([0.0, 0.0, 1.0e4])],
"hitandrun_disrupt" : [np.array([0.0, 0.0, -6.0e5]),
np.array([0.0, 0.0, 1.0e5])],
"hitandrun_pure" : [np.array([0.0, 0.0, 6.0e3]),
np.array([0.0, 0.0, 1.0e4])],
"merge" : [np.array([0.0, 0.0, -6.0e3]),
np.array([0.0, 0.0, 1.0e4])]
"merge" : [np.array([0.0, 0.0, 0.0]),
np.array([0.0, 0.0, 0.0])]
}

body_Gmass = {"disruption_headon" : [1e-7, 1e-9],
Expand Down Expand Up @@ -160,15 +160,24 @@ def encounter_combiner(sim):
smooth_time = np.linspace(start=tgood.isel(time=0), stop=ds.time[-1], num=int(1.2*num_movie_frames))
ds = ds.interp(time=smooth_time)
ds['rotangle'] = xr.zeros_like(ds['rot'])
ds['rot'] = ds['rot'].fillna(0.0)

return ds

def center(Gmass, x, y):
x = x[~np.isnan(x)]
y = y[~np.isnan(y)]
Gmass = Gmass[~np.isnan(Gmass)]
x_com = np.sum(Gmass * x) / np.sum(Gmass)
y_com = np.sum(Gmass * y) / np.sum(Gmass)
return x_com, y_com
class AnimatedScatter(object):
"""An animated scatter plot using matplotlib.animations.FuncAnimation."""

def __init__(self, sim, animfile, title, style, nskip=1):

self.ds = encounter_combiner(sim)
self.npl = len(self.ds['name']) - 1
self.title = title
self.body_color_list = {'Initial conditions': 'xkcd:windows blue',
'Disruption': 'xkcd:baby poop',
Expand All @@ -179,9 +188,7 @@ def __init__(self, sim, animfile, title, style, nskip=1):
# Set up the figure and axes...
self.figsize = (4,4)
self.fig, self.ax = self.setup_plot()

# Then setup FuncAnimation.
self.ani = animation.FuncAnimation(self.fig, self.update_plot, interval=1, frames=range(0,num_movie_frames,nskip), blit=True)
self.ani = animation.FuncAnimation(self.fig, self.update_plot, init_func=self.init_func, interval=1, frames=range(0,num_movie_frames,nskip), blit=True)
self.ani.save(animfile, fps=60, dpi=300, extra_args=['-vcodec', 'libx264'])
print(f"Finished writing {animfile}")

Expand All @@ -199,7 +206,7 @@ def setup_plot(self):
scale_frame *= 2

ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8])
self.ax_pt_size = self.figsize[0] * 72 / scale_frame
self.ax_pt_size = self.figsize[0] * 72 / scale_frame * 0.7
ax.set_xlim(-scale_frame, scale_frame)
ax.set_ylim(-scale_frame, scale_frame)
ax.set_xticks([])
Expand All @@ -209,46 +216,92 @@ def setup_plot(self):
ax.set_title(self.title)
fig.add_axes(ax)

self.scatter_artist = ax.scatter([], [], animated=True, c='k', edgecolors='face')
return fig, ax

def init_func(self):
self.artists = []
aarg = self.vec_props('xkcd:beige')
for i in range(self.npl):
self.artists.append(self.ax.annotate("",xy=(0,0),**aarg))

self.artists.append(self.ax.scatter([],[],c='k', animated=True, zorder=10))
return self.artists

def update_plot(self, frame):

# Define a function to calculate a reference frame for the animation
# This will be based on the initial velocity of the Target body
def center(Gmass, x, y):
x = x[~np.isnan(x)]
y = y[~np.isnan(y)]
Gmass = Gmass[~np.isnan(Gmass)]
x_com = np.sum(Gmass * x) / np.sum(Gmass)
y_com = np.sum(Gmass * y) / np.sum(Gmass)
return x_com, y_com

t, Gmass, rh, point_rad = next(self.data_stream(frame))

t, Gmass, rh, radius, rotangle = next(self.data_stream(frame))
x_ref, y_ref = center(Gmass, rh[:,0], rh[:,1])
self.scatter_artist.set_offsets(np.c_[rh[:,0] - x_ref, rh[:,1] - y_ref])
self.scatter_artist.set_sizes(point_rad**2)
return self.scatter_artist,
rh = np.c_[rh[:,0] - x_ref, rh[:,1] - y_ref]
self.artists[-1].set_offsets(rh)
point_rad = radius * self.ax_pt_size
self.artists[-1].set_sizes(point_rad**2)

sarrowend, sarrowtip = self.spin_arrows(rh, rotangle, 1.1*radius)
for i, s in enumerate(self.artists[:-1]):
self.artists[i].set_position(sarrowtip[i])
self.artists[i].xy = sarrowend[i]

return self.artists

def data_stream(self, frame=0):
while True:
t = self.ds.isel(time=frame)['time'].values[()]

if frame > 0:
dsold = self.ds.isel(time=frame-1)
told = dsold['time'].values[()]
dt = t - told
self.ds['rotangle'][dict(time=frame)] = dsold['rotangle'] + dt * dsold['rot']

ds = self.ds.isel(time=frame)
ds = ds.where(ds['name'] != "Sun", drop=True)
t = ds['time'].values[()]
radius = ds['radius'].values
Gmass = ds['Gmass'].values
rh = ds['rh'].values
point_rad = radius * self.ax_pt_size
if frame > 0:
dsold = self.ds.isel(time=frame-1)
told = dsold['time'].values[()]
dt = t - told
ds['rotangle'] = dsold['rotangle'] + dt * dsold['rot']
ds['rotangle'] = np.mod(ds['rotangle'], 360.0)
self.ds[dict(time=frame)]['rotangle'] = ds['rotangle']

yield t, Gmass, rh, point_rad

rotangle = ds['rotangle'].values

yield t, Gmass, rh, radius, rotangle

def spin_arrows(self, rh, rotangle, rotlen):
px = rh[:, 0]
py = rh[:, 1]
sarrowend = []
sarrowtip = []
for i in range(rh.shape[0]):
endrel = np.array([0.0, -rotlen[i], 0.0])
tiprel = np.array([0.0, rotlen[i], 0.0])
r = R.from_rotvec(rotangle[i,:], degrees=True)
endrel = r.apply(endrel)
tiprel = r.apply(tiprel)
send = (px[i] + endrel[0], py[i] + endrel[1])
stip = (px[i] + tiprel[0], py[i] + tiprel[1])
sarrowend.append(send)
sarrowtip.append(stip)
return sarrowend, sarrowtip

def vec_props(self, c):
arrowprops = {
'arrowstyle': '-',
'linewidth' : 1,
}

arrow_args = {
'xycoords': 'data',
'textcoords': 'data',
'arrowprops': arrowprops,
'annotation_clip': True,
'zorder': 100,
'animated' : True
}
aarg = arrow_args.copy()
aprop = arrowprops.copy()
aprop['color'] = c
aarg['arrowprops'] = aprop
aarg['color'] = c
return aarg

if __name__ == "__main__":

print("Select a fragmentation movie to generate.")
Expand Down
69 changes: 34 additions & 35 deletions src/fraggle/fraggle_generate.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -562,6 +562,14 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu
call fragments%set_coordinate_system()
ke_min = 0.5_DP * fragments%mtot * vesc**2

! Try to put as much of the residual angular momentum into the spin of the fragments before the target body
call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))
do i = istart, fragments%nbody
fragments%L_spin(:,i) = fragments%L_spin(:,i) - L_residual(:) * fragments%mass(i) / fragments%mtot
fragments%rot(:,i) = fragments%L_spin(:,i) / (fragments%mass(i) * fragments%radius(i)**2 * fragments%Ip(:,i))
end do

do loop = 1, MAXLOOP
nsteps = loop * try
call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
Expand All @@ -570,48 +578,39 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu
ke_avail = ke_avail + 0.5_DP * fragments%mass(i) * max(fragments%vmag(i) - vesc,0.0_DP)**2
end do

! Check for any residual angular momentum, and if there is any, put it into spin/shear depending on the type of collision (hit and runs prioritize velocity shear over spin)
L_residual(:) = collider_local%L_total(:,2) - collider_local%L_total(:,1)
if (lhitandrun .or. (ke_avail < epsilon(1.0_DP))) then
! Start by putting residual angular momentum into velocity shear
mfrag = sum(fragments%mass(istart:fragments%nbody))
L_residual_unit(:) = .unit. L_residual(:)
fragments%r_unit(:,:) = .unit.fragments%rc(:,:)
do i = 1, fragments%nbody
r_lever(:) = (L_residual_unit(:) .cross. fragments%r_unit(:,i))
rmag = .mag.r_lever(:)
if (rmag > epsilon(1.0_DP)) then
vdir(:) = -.unit. r_lever(:)
vmag = .mag.L_residual(:) / (fragments%mtot * .mag.r_lever(:) * fragments%rmag(i))
Li(:) = fragments%mass(i) * fragments%rc(:,i) .cross. (vmag * vdir(:))
Lrat(:) = L_residual(:) / Li(:)
fragments%vc(:,i) = fragments%vc(:,i) + vmag * vdir(:)
end if
end do
fragments%vmag(:) = .mag.fragments%vc(:,:)

! Update the coordinate system now that the velocities have changed
call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc)
call fragments%set_coordinate_system()

! Try to put as much of the residual angular momentum into the spin of the fragments before the target body
call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))
do i = istart, fragments%nbody
fragments%L_spin(:,i) = fragments%L_spin(:,i) - L_residual(:) * fragments%mass(i) / fragments%mtot
fragments%rot(:,i) = fragments%L_spin(:,i) / (fragments%mass(i) * fragments%radius(i)**2 * fragments%Ip(:,i))
end do
call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))
end if
! Check for any residual angular momentum, and put it into spin and shear
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))

! Start by putting residual angular momentum into velocity shear
mfrag = sum(fragments%mass(istart:fragments%nbody))
L_residual_unit(:) = .unit. L_residual(:)
fragments%r_unit(:,:) = .unit.fragments%rc(:,:)
do i = 1, fragments%nbody
r_lever(:) = (L_residual_unit(:) .cross. fragments%r_unit(:,i))
rmag = .mag.r_lever(:)
if (rmag > epsilon(1.0_DP)) then
vdir(:) = -.unit. r_lever(:)
vmag = .mag.L_residual(:) / (fragments%mtot * .mag.r_lever(:) * fragments%rmag(i))
Li(:) = fragments%mass(i) * fragments%rc(:,i) .cross. (vmag * vdir(:))
Lrat(:) = L_residual(:) / Li(:)
fragments%vc(:,i) = fragments%vc(:,i) + vmag * vdir(:)
end if
end do
fragments%vmag(:) = .mag.fragments%vc(:,:)

! Update the coordinate system now that the velocities have changed
call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc)
call fragments%set_coordinate_system()
call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))

! Put any remaining residual angular momentum into the spin of the target body
fragments%L_spin(:,1) = fragments%L_spin(:,1) - L_residual(:)
fragments%rot(:,1) = fragments%L_spin(:,1) / (fragments%mass(1) * fragments%radius(1)**2 * fragments%Ip(:,1))
fragments%rotmag(:) = .mag.fragments%rot(:,:)

call collider_local%get_energy_and_momentum(nbody_system, param, phase="after")
L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1))

dE = collider_local%te(2) - collider_local%te(1)
E_residual = dE + impactors%Qloss

Expand Down

0 comments on commit e254f38

Please sign in to comment.