Restructured movie code a bit and added some more clarifying comments…

…. Fixed bug that was causing bodies not to be drawn if there was a fragmentation event

examples/Fragmentation/Fragmentation_Movie.py

@@ -32,13 +32,9 @@
 import matplotlib.animation as animation
 from pathlib import Path
 
-print("Select a fragmentation movie to generate.")
-print("1. Head-on disruption")
-print("2. Off-axis supercatastrophic")
-print("3. Hit and run")
-print("4. All of the above")
-user_selection = int(input("? "))
-
+# ----------------------------------------------------------------------------------------------------------------------
+# Define the names and initial conditions of the various fragmentation simulation types
+# ----------------------------------------------------------------------------------------------------------------------
 available_movie_styles = ["disruption_headon", "supercatastrophic_off_axis", "hitandrun"]
 movie_title_list = ["Head-on Disrutption", "Off-axis Supercatastrophic", "Hit and Run"]
 movie_titles = dict(zip(available_movie_styles, movie_title_list))
@@ -79,12 +75,10 @@
 for k,v in body_Gmass.items():
     body_radius[k] = [((Gmass/GU)/(4./3.*np.pi*density))**(1./3.) for Gmass in v]
 
-if user_selection > 0 and user_selection < 4:
-   movie_styles = [available_movie_styles[user_selection-1]]
-else:
-   print("Generating all movie styles")
-   movie_styles = available_movie_styles.copy()
 
+# ----------------------------------------------------------------------------------------------------------------------
+# Define the animation class that will generate the movies of the fragmentation outcomes
+# ----------------------------------------------------------------------------------------------------------------------
 figsize = (4,4)
 class AnimatedScatter(object):
     """An animated scatter plot using matplotlib.animations.FuncAnimation."""
@@ -133,17 +127,16 @@ def setup_plot(self):
 
     def update_plot(self, frame):
         # Define a function to calculate the center of mass of the system.
-        def center(frame):
-            ds = self.sim.data.isel(time=frame).where(self.sim.data['name'] != "Sun", drop=True)
-            Gmass = ds['Gmass'].values
-            xhx = ds['xhx'].values
-            xhy = ds['xhy'].values
-            x_com = np.sum(Gmass * xhx) / np.sum(Gmass)
-            y_com = np.sum(Gmass * xhy) / np.sum(Gmass)
+        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
 
-        x, y, point_rad = next(self.data_stream(frame))
-        x_com, y_com = center(frame)
+        Gmass, x, y, point_rad = next(self.data_stream(frame))
+        x_com, y_com = center(Gmass, x, y)
         self.scatter_artist.set_offsets(np.c_[x - x_com, y - y_com])
         self.scatter_artist.set_sizes(point_rad)
         return self.scatter_artist,
@@ -153,27 +146,56 @@ def data_stream(self, frame=0):
             ds = self.sim.data.isel(time=frame)
             ds = ds.where(ds['name'] != "Sun", drop=True)
             radius = ds['radius'].values
+            Gmass = ds['Gmass'].values
             x = ds['xhx'].values
             y = ds['xhy'].values
             point_rad = 2 * radius * self.ax_pt_size
-            yield x, y, point_rad
-
-
-
-for style in movie_styles:
-    param_file = Path(style) / "param.in"
-
-    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(param_file=param_file, rotation=True, init_cond_format = "XV", compute_conservation_values=True)
-    sim.add_solar_system_body("Sun")
-    sim.add_body(Gmass=body_Gmass[style], radius=body_radius[style], xh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style])
-
-    # Set fragmentation parameters
-    minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body
-    gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades
-    sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass)
-    sim.run(dt=1e-8, tstop=2.e-5)
-
-    anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=10)
+            yield Gmass, x, y, point_rad
+
+if __name__ == "__main__":
+
+   print("Select a fragmentation movie to generate.")
+   print("1. Head-on disruption")
+   print("2. Off-axis supercatastrophic")
+   print("3. Hit and run")
+   print("4. All of the above")
+   user_selection = int(input("? "))
+
+   if user_selection > 0 and user_selection < 4:
+      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:
+       param_file = Path(style) / "param.in"
+       bin_file = Path(style) / "bin.nc"
+       if bin_file.exists():
+           user_selection = input(f"An older simulation of {movie_titles[style]} exists. Do you want to re-run it? [y/N] ")
+           if user_selection == "":
+               run_new = False
+           else:
+               try:
+                   run_new = swiftest.io.str2bool(user_selection)
+               except:
+                   run_new = False
+       else:
+           run_new = True
+
+       movie_filename = f"{style}.mp4"
+
+       # Pull in the Swiftest output data from the parameter file and store it as a Xarray dataset.
+       if run_new:
+           sim = swiftest.Simulation(param_file=param_file, rotation=True, init_cond_format = "XV", compute_conservation_values=True)
+           sim.add_solar_system_body("Sun")
+           sim.add_body(Gmass=body_Gmass[style], radius=body_radius[style], xh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style])
+
+           # Set fragmentation parameters
+           minimum_fragment_gmass = 0.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body
+           gmtiny = 0.99 * body_Gmass[style][1] # Make GMTINY just smaller than the smallest original body. This will prevent runaway collisional cascades
+           sim.set_parameter(fragmentation = True, gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False)
+           sim.run(dt=1e-8, tstop=2.e-5)
+       else:
+           sim = swiftest.Simulation(param_file=param_file, read_old_output_file=True)
+
+       anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=10)