From 8d01494c6d3d48a0773c933c984a9dfac30d69f5 Mon Sep 17 00:00:00 2001 From: David A Minton Date: Thu, 19 Jan 2023 01:40:32 -0500 Subject: [PATCH] Updated movies with rotation indicated and did some more updates to the rotation component of Fraggle to spread out rotations among bodies --- .gitignore | 1 + examples/Fragmentation/Fragmentation_Movie.py | 137 ++++++++++++------ src/fraggle/fraggle_generate.f90 | 69 +++++---- 3 files changed, 130 insertions(+), 77 deletions(-) diff --git a/.gitignore b/.gitignore index 2233adc6b..523f9aafb 100644 --- a/.gitignore +++ b/.gitignore @@ -30,3 +30,4 @@ dump* bin/ build/* +disruption_headon/swiftest_driver.sh diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index f0c25d196..c77c69da6 100755 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/examples/Fragmentation/Fragmentation_Movie.py @@ -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], @@ -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', @@ -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}") @@ -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([]) @@ -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.") diff --git a/src/fraggle/fraggle_generate.f90 b/src/fraggle/fraggle_generate.f90 index 43e2467b9..ada685da7 100644 --- a/src/fraggle/fraggle_generate.f90 +++ b/src/fraggle/fraggle_generate.f90 @@ -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") @@ -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