Added a multibody collision movie script

examples/Fragmentation/.gitignore

@@ -1,3 +1,4 @@
 *
 !.gitignore
 !Fragmentation_Movie.py
+!Multibody_Movie.py

examples/Fragmentation/Multibody_Movie.py

@@ -0,0 +1,288 @@
+#!/usr/bin/env python3
+"""
+ Copyright 2022 - David Minton, Carlisle Wishard, Jennifer Pouplin, Jake Elliott, & Dana Singh
+ This file is part of Swiftest.
+ Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License 
+ as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
+ Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty 
+ of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
+ You should have received a copy of the GNU General Public License along with Swiftest. 
+ If not, see: https://www.gnu.org/licenses. 
+"""
+
+"""
+Generates a movie of a multi-body fragmentation event from set of Swiftest output files.
+
+Inputs
+_______
+param.in : ASCII text file
+    Swiftest parameter input file.
+out.nc   : NetCDF file
+    Swiftest output file.
+
+Returns
+-------
+fragmentation.mp4 : mp4 movie file
+    Movie of a fragmentation event.
+"""
+
+import swiftest
+import numpy as np
+import xarray as xr
+import matplotlib.pyplot as plt
+import matplotlib.animation as animation
+from scipy.spatial.transform import Rotation as R
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Define the names and initial conditions of the various fragmentation simulation types
+# ----------------------------------------------------------------------------------------------------------------------
+available_movie_styles = ["supercatastrophic_multi"]
+movie_title_list = ["Multi-body Supercatastrophic"]
+movie_titles = dict(zip(available_movie_styles, movie_title_list))
+num_movie_frames = 1000
+
+# These initial conditions were generated by trial and error
+names = ["Body1","Body2","Body3","Body4"]
+
+rdistance = 1e-4/np.sqrt(2.0)
+vrel = 2.0/np.sqrt(2.0)
+
+rcb = np.array([1.0, 0.0,  0.0])
+vcb = np.array([0.0, 6.28, 0.0])
+
+r1 = rcb + np.array([-rdistance,-rdistance,0.0])
+r2 = rcb + np.array([ rdistance,-rdistance,0.0])
+r3 = rcb + np.array([ rdistance, rdistance,0.0])
+r4 = rcb + np.array([-rdistance, rdistance,0.0])
+
+v1 = vcb + np.array([ vrel, vrel, 0.0])
+v2 = vcb + np.array([-vrel, vrel, 0.0])
+v3 = vcb + np.array([-vrel,-vrel, 0.0])
+v4 = vcb + np.array([ vrel,-vrel, 0.0])
+
+rot1 = np.array([0.0,0.0,1e5])
+rot2 = np.array([0.0,0.0,-2e5])
+rot3 = np.array([0.0,0.0,3e5])
+rot4 = np.array([0.0,0.0,0.0])
+
+m1 = 2e-7
+m2 = 2e-7
+m3 = 2e-7
+m4 = 2e-7
+
+pos_vectors = {"supercatastrophic_multi"   : [r1, r2, r3, r4],
+               }
+
+vel_vectors = {"supercatastrophic_multi":   [v1, v2, v3, v4]
+               }
+
+rot_vectors = { "supercatastrophic_multi":   [rot1, rot2, rot3, rot4]
+               }
+
+body_Gmass = {"supercatastrophic_multi":   [m1, m2, m3, m4]
+               }
+
+tstop = {"supercatastrophic_multi"   : 2.0e-3,
+         }
+
+density = 3000 * swiftest.AU2M**3 / swiftest.MSun
+GU = swiftest.GMSun * swiftest.YR2S**2 / swiftest.AU2M**3
+body_radius = body_Gmass.copy()
+for k,v in body_Gmass.items():
+    body_radius[k] = [((Gmass/GU)/(4./3.*np.pi*density))**(1./3.) for Gmass in v]
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Define the animation class that will generate the movies of the fragmentation outcomes
+# ----------------------------------------------------------------------------------------------------------------------
+def encounter_combiner(sim):
+    """
+    Combines simulation data with encounter data to produce a dataset that contains the position,
+    mass, radius, etc. of both. It will interpolate over empty time values to fill in gaps.
+    """
+
+    # Only keep a minimal subset of necessary data from the simulation and encounter datasets
+    keep_vars = ['name','rh','vh','Gmass','radius','rot']
+    data = sim.data[keep_vars]
+    enc = sim.encounters[keep_vars].load()
+
+    # Remove any encounter data at the same time steps that appear in the data to prevent duplicates
+    t_not_duplicate = ~enc['time'].isin(data['time'])
+    enc = enc.where(t_not_duplicate,drop=True)
+    tgood=enc.time.where(~np.isnan(enc.time),drop=True)
+    enc = enc.sel(time=tgood)
+
+    # The following will combine the two datasets along the time dimension, sort the time dimension, and then fill in any time gaps with interpolation
+    ds = xr.combine_nested([data,enc],concat_dim='time').sortby("time").interpolate_na(dim="time")
+
+    # Interpolate in time to make a smooth, constant time step dataset 
+    # Add a bit of padding to the time, otherwise there are some issues with the interpolation in the last few frames.
+    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',
+                      'Supercatastrophic': 'xkcd:shocking pink',
+                      'Hit and run fragmention': 'xkcd:blue with a hint of purple',
+                      'Central body': 'xkcd:almost black'}
+
+        # Set up the figure and axes...
+        self.figsize = (4,4)
+        self.fig, self.ax = self.setup_plot()
+        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}")
+
+    def setup_plot(self):
+        fig = plt.figure(figsize=self.figsize, dpi=300)
+        plt.tight_layout(pad=0)
+
+        # Calculate the distance along the y-axis between the colliding bodies at the start of the simulation.
+        # This will be used to scale the axis limits on the movie.
+        rhy1 = self.ds['rh'].sel(name="Body1",space='y').isel(time=0).values[()]
+        rhy2 = self.ds['rh'].sel(name="Body3",space='y').isel(time=0).values[()]
+
+        scale_frame =   abs(rhy1) + abs(rhy2)
+
+        ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8])
+        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([])
+        ax.set_yticks([])
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        ax.set_title(self.title)
+        fig.add_axes(ax)
+
+        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
+        t, Gmass, rh, radius, rotangle = next(self.data_stream(frame))
+        x_ref, y_ref = center(Gmass, rh[:,0], rh[:,1]) 
+        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)
+            radius = ds['radius'].values
+            Gmass = ds['Gmass'].values
+            rh = ds['rh'].values
+            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.")
+    print("1. Multi-body supercatastrophic")
+    print("2. All of the above")
+    user_selection = int(input("? "))
+
+    if user_selection > 0 and user_selection < 2:
+        movie_styles = [available_movie_styles[user_selection-1]]
+    else:
+        print("Generating all movie styles")
+        movie_styles = available_movie_styles.copy()
+
+    for style in movie_styles:
+        print(f"Generating {movie_titles[style]}")
+        movie_filename = f"{style}.mp4"
+        # Pull in the Swiftest output data from the parameter file and store it as a Xarray dataset.
+        sim = swiftest.Simulation(simdir=style, rotation=True, init_cond_format = "XV", compute_conservation_values=True)
+        sim.add_solar_system_body("Sun")
+        sim.add_body(name=names, Gmass=body_Gmass[style], radius=body_radius[style], rh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style])
+
+        # Set fragmentation parameters
+        minimum_fragment_gmass = 0.05 * body_Gmass[style][1] 
+        gmtiny = 0.10 * body_Gmass[style][1] 
+        sim.set_parameter(collision_model="fraggle", encounter_save="both", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False)
+        sim.run(dt=5e-4, tstop=tstop[style], istep_out=1, dump_cadence=0)
+
+        print("Generating animation")
+        anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1)