Merge branch 'RevampedRestart' into coarray_test_particles

examples/.gitignore

@@ -3,4 +3,5 @@
 !Basic_Simulation
 !Fragmentation
 !helio_gr_test
-!whm_gr_test
+!whm_gr_test
+!Chambers2013

examples/Chambers2013/.gitignore

@@ -0,0 +1,4 @@
+*
+!.gitignore
+!init_cond.py
+!aescattermovie.py

examples/Chambers2013/aescattermovie.py

@@ -0,0 +1,121 @@
+#!/usr/bin/env python3
+import swiftest 
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib import animation
+import matplotlib.colors as mcolors
+
+titletext = "Chambers (2013)"
+radscale = 2000
+xmin = 0.0
+xmax = 2.20
+ymin = 0.0 
+ymax = 1.0
+framejump = 1
+
+class AnimatedScatter():
+    """An animated scatter plot using matplotlib.animations.FuncAnimation."""
+    def __init__(self, ds, param):
+
+        frame = 0
+        nframes = int(ds['time'].size / framejump)
+        self.ds = ds
+        self.param = param
+        self.Rcb = self.ds['radius'].sel(name="Sun").isel(time=0).values[()]
+
+        self.clist = {'Initial conditions' : 'xkcd:faded blue',
+                      'Disruption' : 'xkcd:marigold',
+                      'Supercatastrophic' : 'xkcd:shocking pink',
+                      'Hit and run fragmentation' : 'xkcd:baby poop green'}
+
+        # Setup the figure and axes...
+        fig = plt.figure(figsize=(8,4.5), dpi=300)
+        plt.tight_layout(pad=0)
+        # set up the figure
+        self.ax = plt.Axes(fig, [0.1, 0.15, 0.8, 0.75])
+        self.ax.set_xlim(xmin, xmax)
+        self.ax.set_ylim(ymin, ymax)
+        fig.add_axes(self.ax)
+        self.ani = animation.FuncAnimation(fig, self.update, interval=1, frames=nframes, init_func=self.setup_plot, blit=True)
+        self.ani.save('aescatter.mp4', fps=60, dpi=300, extra_args=['-vcodec', 'libx264'])
+        print('Finished writing aescattter.mp4')
+
+    def scatters(self, pl, radmarker, origin):
+        scat = []
+        for key, value in self.clist.items():
+            idx = origin == key
+            s = self.ax.scatter(pl[idx, 0], pl[idx, 1], marker='o', s=radmarker[idx], c=value, alpha=0.75, label=key)
+            scat.append(s)
+        return scat
+
+    def setup_plot(self):
+        # First frame
+        """Initial drawing of the scatter plot."""
+        t, name, Gmass, radius, npl, pl, radmarker, origin = next(self.data_stream(0))
+
+        # set up the figure
+        self.ax.margins(x=10, y=1)
+        self.ax.set_xlabel("Semimajor Axis (AU)", fontsize='16', labelpad=1)
+        self.ax.set_ylabel("Eccentricity", fontsize='16', labelpad=1)
+
+        self.title = self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes,
+                        ha="center")
+
+        self.title.set_text(f"{titletext} -  Time = ${t*1e-6:6.2f}$ My with ${npl:4.0f}$ particles")
+        slist = self.scatters(pl, radmarker, origin)
+        self.s0 = slist[0]
+        self.s1 = slist[1]
+        self.s2 = slist[2]
+        self.s3 = slist[3]
+        leg = plt.legend(loc="upper right", scatterpoints=1, fontsize=10)
+        for i,l in enumerate(leg.legendHandles):
+           leg.legendHandles[i]._sizes = [20]
+        return self.s0, self.s1, self.s2, self.s3, self.title
+
+    def data_stream(self, frame=0):
+        while True:
+            d = self.ds.isel(time = frame)
+            name_good = d.name.where(d['status'] != 1, drop=True)
+            name_good = name_good.where(name_good != "Sun", drop=True)
+            d = d.sel(name=name_good) 
+            d['radmarker'] = (d['radius'] / self.Rcb) * radscale
+            radius = d['radmarker'].values
+
+            radius = d['radmarker'].values
+            Gmass = d['Gmass'].values
+            a = d['a'].values 
+            e = d['e'].values
+            name = d['name'].values
+            npl = d['npl'].values
+            radmarker = d['radmarker']
+            origin = d['origin_type']
+
+            t = self.ds.coords['time'].values[frame]
+
+            yield t, name, Gmass, radius, npl, np.c_[a, e], radmarker, origin
+
+    def update(self,frame):
+        """Update the scatter plot."""
+        t, name, Gmass, radius, npl, pl, radmarker, origin = next(self.data_stream(framejump * frame))
+
+        self.title.set_text(f"{titletext} - Time = ${t*1e-6:6.3f}$ My with ${npl:4.0f}$ particles")
+
+        # We need to return the updated artist for FuncAnimation to draw..
+        # Note that it expects a sequence of artists, thus the trailing comma.
+        s = [self.s0, self.s1, self.s2, self.s3]
+        for i, (key, value) in enumerate(self.clist.items()):
+            idx = origin == key
+            s[i].set_sizes(radmarker[idx])
+            s[i].set_offsets(pl[idx,:])
+            s[i].set_facecolor(value)
+
+        self.s0 = s[0]
+        self.s1 = s[1]
+        self.s2 = s[2]
+        self.s3 = s[3]
+        return self.s0, self.s1, self.s2, self.s3, self.title,
+
+sim = swiftest.Simulation(simdir="fragglesim",read_old_output=True)
+print('Making animation')
+anim = AnimatedScatter(sim.data,sim.param)
+print('Animation finished')

examples/Chambers2013/init_cond.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+import swiftest
+import numpy as np
+from numpy.random import default_rng
+
+# Initialize simulation object
+sim = swiftest.Simulation(simdir="fragglesim")
+
+sim.set_parameter(compute_conservation_values=True,  rotation=True, init_cond_format="EL",collision_model="fraggle",encounter_save="none")
+
+# Add bodies described in Chambers (2013) Sec. 2.1, with the uniform spatial distribution and two bodies sizes (big and small)
+Nb = 14
+Ns = 140
+Mb = 2.8e-7 * 14 / Nb
+Ms = 2.8e-8 * 140 / Ns
+dens = 3000.0 / (sim.param['MU2KG'] / sim.param['DU2M']**3)
+Rb = (3 * Mb / (4 * np.pi * dens) )**(1.0 / 3.0)
+Rs = (3 * Ms / (4 * np.pi * dens) )**(1.0 / 3.0)
+mtiny = 1e-2 * Ms
+mininum_fragment_mass = 1e-4 * Ms
+rng = default_rng(seed=880334)
+
+# Define the initial orbital elements of the big and small bodies
+avalb = rng.uniform(0.3, 2.0, Nb)
+avals = rng.uniform(0.3, 2.0, Ns)
+evalb = rng.uniform(0.0, 0.01, Nb)
+evals = rng.uniform(0.0, 0.01, Ns)
+incvalb = rng.uniform(0.0, 0.005 * 180 / np.pi, Nb)
+incvals = rng.uniform(0.0, 0.005 * 180 / np.pi, Ns)
+capomvalb = rng.uniform(0.0, 360.0, Nb)
+capomvals = rng.uniform(0.0, 360.0, Ns)
+omegavalb = rng.uniform(0.0, 360.0, Nb)
+omegavals = rng.uniform(0.0, 360.0, Ns)
+capmvalb = rng.uniform(0.0, 360.0, Nb)
+capmvals = rng.uniform(0.0, 360.0, Ns)
+Ipvalb = np.full((Nb,3), 0.4)
+Ipvals = np.full((Ns,3), 0.4)
+rotvalb = np.zeros_like(Ipvalb)
+rotvals = np.zeros_like(Ipvals)
+GMvalb = np.full(Nb, Mb * sim.GU)
+GMvals = np.full(Ns, Ms * sim.GU)
+Rvalb = np.full(Nb, Rb)
+Rvals = np.full(Ns, Rs)
+
+# Give the bodies unique names
+nameb = [f"Big{i:03}" for i in range(Nb)]
+names = [f"Small{i:03}" for i in range(Ns)]
+
+# Add the modern planets and the Sun using the JPL Horizons Database.
+sim.add_solar_system_body(["Sun","Jupiter","Saturn","Uranus","Neptune"])
+sim.add_body(name=nameb, a=avalb, e=evalb, inc=incvalb, capom=capomvalb, omega=omegavalb, capm=capmvalb, Gmass=GMvalb, radius=Rvalb, rot=rotvalb, Ip=Ipvalb)
+sim.add_body(name=names, a=avals, e=evals, inc=incvals, capom=capomvals, omega=omegavals, capm=capmvals, Gmass=GMvals, radius=Rvals, rot=rotvals, Ip=Ipvals)
+sim.set_parameter(mtiny=mtiny, minimum_fragment_mass=mininum_fragment_mass)
+
+sim.run(tstop=3e8, dt=6.0875/365.25, istep_out=60000, dump_cadence=1)
+

examples/Fragmentation/swiftest_fragmentation.py

@@ -29,8 +29,8 @@
 Each subdirectory contains:
 data.nc             : A NetCDF file containing the simulation output.
 init_cond.nc        : A NetCDF file containing the initial conditions for the simulation.
-collision_000001.nc : A NetCDF file containing the data for the collision.
-encounter_000001.nc : A NetCDF file containing the data for the close encounter.
+collision_000001.nc : A NetCDF file containing the data for the collisions.
+encounter_000001.nc : A NetCDF file containing the data for the close encounters.
 dump_bin1.nc        : A NetCDF file containing the necessary inputs to restart a simulation from t!=0.
 dump_bin2.nc        : A NetCDF file containing the necessary inputs to restart a simulation from t!=0.
 dump_param1.in      : An ASCII file containing the necessary parameters to restart a simulation.

python/swiftest/swiftest/io.py

@@ -63,7 +63,7 @@
 # handles strings differently than Python's Xarray.
 string_varnames = ["name", "particle_type", "status", "origin_type", "stage", "regime"]
 char_varnames = ["space"]
-int_varnames = ["id", "ntp", "npl", "nplm", "discard_body_id", "collision_id", "loopnum"]
+int_varnames = ["id", "ntp", "npl", "nplm", "discard_body_id", "collision_id", "status"]
 
 def bool2yesno(boolval):
     """
@@ -804,8 +804,6 @@ def process_netcdf_input(ds, param):
     Performs several tasks to convert raw NetCDF files output by the Fortran program into a form that
     is used by the Python side. These include:
     - Ensuring all types are correct
-    - Removing any bad id values (empty id slots)
-    - Swapping the id and name dimension if the names are unique
 
     Parameters
     ----------

python/swiftest/swiftest/simulation_class.py

@@ -1039,7 +1039,7 @@ def set_feature(self,
             in initial conditions.
         encounter_save : {"NONE","TRAJECTORY","CLOSEST","BOTH"}, default "NONE"
             Indicate if and how encounter data should be saved. If set to "TRAJECTORY", the position and velocity vectors
-            of all bodies undergoing close encounters are saved at each intermediate step to the encounter files.
+            of all bodies undergoing close encounter are saved at each intermediate step to the encounter files.
             If set to "CLOSEST", the position  and velocities at the point of closest approach between pairs of bodies are 
             computed and stored to the encounter files. If set to "BOTH", then this stores the values that would be computed
             in "TRAJECTORY" and "CLOSEST". If set to "NONE" no trajectory information is saved.
@@ -2738,8 +2738,8 @@ def read_output_file(self,read_init_cond : bool = True):
                 else:
                     self.init_cond = self.data.isel(time=0)
 
-            self.read_encounters()
-            self.read_collisions()
+            self.read_encounter()
+            self.read_collision()
             if self.verbose:
                 print("Finished reading Swiftest dataset files.")
 
@@ -2752,23 +2752,14 @@ def read_output_file(self,read_init_cond : bool = True):
             warnings.warn('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".',stacklevel=2)
         return
 
-    def read_encounters(self):
-        enc_files = glob(f"{self.simdir}{os.path.sep}encounter_*.nc")
-        if len(enc_files) == 0:
-            return
-
-        if self.verbose:
-            print("Reading encounter history file as .encounters")
-
-        enc_files.sort()
+    def read_encounter(self):
+        enc_file = self.simdir / "encounters.nc"
+        if not os.path.exists(enc_file):
+           return
 
-        # This is needed in order to pass the param argument down to the io.process_netcdf_input function
-        def _preprocess(ds, param):
-            return io.process_netcdf_input(ds,param)
-        partial_func = partial(_preprocess, param=self.param)
-
-        self.encounters = xr.open_mfdataset(enc_files,parallel=True,combine="nested",concat_dim="time",join="left",preprocess=partial_func,mask_and_scale=True)
+        self.encounters = xr.open_dataset(enc_file)
         self.encounters = io.process_netcdf_input(self.encounters, self.param)
+
         # Remove any overlapping time values
         tgood,tid = np.unique(self.encounters.time,return_index=True)
         self.encounters = self.encounters.isel(time=tid)
@@ -2778,21 +2769,16 @@ def _preprocess(ds, param):
 
         return
 
-    def read_collisions(self):
-        col_files = glob(f"{self.simdir}{os.path.sep}collision_*.nc")
-        if len(col_files) == 0:
-            return
+    def read_collision(self):
+
+        col_file = self.simdir / "collisions.nc"
+        if not os.path.exists(col_file):
+           return
 
-        col_files.sort()
         if self.verbose:
-                print("Reading collision history file as .collisions")
-
-        # This is needed in order to pass the param argument down to the io.process_netcdf_input function
-        def _preprocess(ds, param):
-            return io.process_netcdf_input(ds,param)
-        partial_func = partial(_preprocess, param=self.param)
+                print("Reading collisions history file as .collisions")
 
-        self.collisions = xr.open_mfdataset(col_files,parallel=True,  combine="nested", concat_dim="collision", preprocess=partial_func,mask_and_scale=True)
+        self.collisions = xr.open_dataset(col_file)
         self.collisions = io.process_netcdf_input(self.collisions, self.param)
 
         return
@@ -2964,17 +2950,13 @@ def clean(self):
         old_files = [self.simdir / self.param['BIN_OUT'],
                      self.simdir / "fraggle.log",
                      self.simdir / "swiftest.log",
+                     self.simdir / "collisions.log",
+                     self.simdir / "collisions.nc",
+                     self.simdir / "encounters.nc",
+                     self.simdir / "param.restart.in",
                      ]
-        glob_files = [self.simdir.glob("**/dump_param?.in")] \
-                     + [self.simdir.glob("**/dump_bin?.nc")] \
-                     + [self.simdir.glob("**/encounter_*.nc")] \
-                     + [self.simdir.glob("**/collision_*.nc")]
 
         for f in old_files:
             if f.exists():
                 os.remove(f)
-        for g in glob_files:
-            for f in g:
-                if f.exists():
-                    os.remove(f)
         return

src/base/base_module.f90

@@ -20,7 +20,7 @@ module base
    !> User defined parameters that are read in from the parameters input file. 
    !>    Each paramter is initialized to a default values. 
    type, abstract :: base_parameters
-      character(len=:), allocatable           :: integrator                             !! Symbolic name of the nbody integrator  used
+      character(len=:), allocatable           :: integrator                             !! Name of the nbody integrator used
       character(len=:), allocatable           :: param_file_name                        !! The name of the parameter file
       integer(I4B)                            :: maxid                = -1              !! The current maximum particle id number 
       integer(I4B)                            :: maxid_collision      = 0               !! The current maximum collision id number
@@ -29,7 +29,6 @@ module base
       real(DP)                                :: tstop                = -1.0_DP         !! Integration stop time
       real(DP)                                :: dt                   = -1.0_DP         !! Time step
       integer(I8B)                            :: iloop                = 0_I8B           !! Main loop counter
-      integer(I4B)                            :: ioutput              = 1               !! Output counter
       character(STRMAX)                       :: incbfile             = CB_INFILE       !! Name of input file for the central body
       character(STRMAX)                       :: inplfile             = PL_INFILE       !! Name of input file for massive bodies
       character(STRMAX)                       :: intpfile             = TP_INFILE       !! Name of input file for test particles
@@ -85,15 +84,15 @@ module base
       logical :: ltides         = .false. !! Include tidal dissipation 
 
       ! Initial values to pass to the energy report subroutine (usually only used in the case of a restart, otherwise these will be updated with initial conditions values)
-      real(DP)                  :: Eorbit_orig  = 0.0_DP  !! Initial orbital energy
+      real(DP)                  :: E_orbit_orig  = 0.0_DP  !! Initial orbital energy
       real(DP)                  :: GMtot_orig   = 0.0_DP  !! Initial system mass
-      real(DP), dimension(NDIM) :: Ltot_orig    = 0.0_DP  !! Initial total angular momentum vector
-      real(DP), dimension(NDIM) :: Lorbit_orig  = 0.0_DP  !! Initial orbital angular momentum
-      real(DP), dimension(NDIM) :: Lspin_orig   = 0.0_DP  !! Initial spin angular momentum vector
-      real(DP), dimension(NDIM) :: Lescape      = 0.0_DP  !! Angular momentum of bodies that escaped the system (used for bookeeping)
+      real(DP), dimension(NDIM) :: L_total_orig    = 0.0_DP  !! Initial total angular momentum vector
+      real(DP), dimension(NDIM) :: L_orbit_orig  = 0.0_DP  !! Initial orbital angular momentum
+      real(DP), dimension(NDIM) :: L_spin_orig   = 0.0_DP  !! Initial spin angular momentum vector
+      real(DP), dimension(NDIM) :: L_escape      = 0.0_DP  !! Angular momentum of bodies that escaped the system (used for bookeeping)
       real(DP)                  :: GMescape     = 0.0_DP  !! Mass of bodies that escaped the system (used for bookeeping)
-      real(DP)                  :: Ecollisions  = 0.0_DP  !! Energy lost from system due to collisions
-      real(DP)                  :: Euntracked   = 0.0_DP  !! Energy gained from system due to escaped bodies
+      real(DP)                  :: E_collisions  = 0.0_DP  !! Energy lost from system due to collisions
+      real(DP)                  :: E_untracked   = 0.0_DP  !! Energy gained from system due to escaped bodies
       logical                   :: lfirstenergy = .true.  !! This is the first time computing energe
       logical                   :: lfirstkick   = .true.  !! Initiate the first kick in a symplectic step
       logical                   :: lrestart     = .false. !! Indicates whether or not this is a restarted run

src/collision/collision_check.f90

@@ -244,7 +244,7 @@ module subroutine collision_check_pltp(self, nbody_system, param, t, dt, irec, l
                      write(message, *) "Particle " // trim(adjustl(tp%info(j)%name)) // " ("  // trim(adjustl(idstrj)) // ")" &
                            //  " collided with massive body " // trim(adjustl(pl%info(i)%name)) // " (" // trim(adjustl(idstri)) // ")" &
                            //  " at t = " // trim(adjustl(timestr))
-                     !call swiftest_io_log_one_message(COLLISION_LOG_OUT, message)
+                     call swiftest_io_log_one_message(COLLISION_LOG_OUT, message)
                   end if
                end do