Merge remote-tracking branch 'origin/debug' into debug

examples/.gitignore

@@ -1,6 +1,6 @@
 *
 !.gitignore
-!Basic_Simulation/*
-!Fragmentation/*
-!helio_gr_test/*
-!whm_gr_test/*
+!Basic_Simulation
+!Fragmentation
+!helio_gr_test
+!whm_gr_test

examples/Fragmentation/Fragmentation_Movie.py

@@ -28,6 +28,7 @@
 
 import swiftest
 import numpy as np
+import xarray as xr
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 from pathlib import Path
@@ -40,8 +41,8 @@
 movie_titles = dict(zip(available_movie_styles, movie_title_list))
 
 # These initial conditions were generated by trial and error
-pos_vectors = {"disruption_headon"         : [np.array([1.0, -2.807993e-05, 0.0]),
-                                              np.array([1.0,  2.807993e-05 ,0.0])],
+pos_vectors = {"disruption_headon"         : [np.array([1.0, -5.0e-05, 0.0]),
+                                              np.array([1.0,  5.0e-05 ,0.0])],
                "supercatastrophic_off_axis": [np.array([1.0, -4.2e-05,      0.0]),
                                               np.array([1.0,  4.2e-05,      0.0])],
                "hitandrun"                 : [np.array([1.0, -2.0e-05,      0.0]),
@@ -75,27 +76,49 @@
 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
 # ----------------------------------------------------------------------------------------------------------------------
-figsize = (4,4)
+
+
+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 = ['rh','Gmass','radius']
+    data = sim.data[keep_vars]
+    enc = sim.enc[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)
+
+    # 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")
+
+    return ds
+
 class AnimatedScatter(object):
     """An animated scatter plot using matplotlib.animations.FuncAnimation."""
 
-    def __init__(self, sim, animfile, title, nskip=1):
-        tgood = sim.enc.where(sim.enc['Gmass'] > 9e-8).time
-        self.ds = sim.enc.sel(time=tgood)
+    def __init__(self, sim, animfile, title, style, nskip=1):
+
+        self.ds = encounter_combiner(sim)
+
         nframes = int(self.ds['time'].size)
         self.sim = sim
         self.title = title
         self.body_color_list = {'Initial conditions': 'xkcd:windows blue',
                       'Disruption': 'xkcd:baby poop',
                       'Supercatastrophic': 'xkcd:shocking pink',
-                      'Hit and run fragment': 'xkcd:blue with a hint of purple',
+                      '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()
 
         # Then setup FuncAnimation.
@@ -105,7 +128,7 @@ def __init__(self, sim, animfile, title, nskip=1):
         print(f"Finished writing {animfile}")
 
     def setup_plot(self):
-        fig = plt.figure(figsize=figsize, dpi=300)
+        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.
@@ -115,7 +138,7 @@ def setup_plot(self):
 
         scale_frame =   abs(rhy1) + abs(rhy2)
         ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8])
-        self.ax_pt_size = figsize[0] * 0.8 *  72 / (2 * scale_frame)
+        self.ax_pt_size = self.figsize[0] * 0.8 *  72 / (2 * scale_frame)
         ax.set_xlim(-scale_frame, scale_frame)
         ax.set_ylim(-scale_frame, scale_frame)
         ax.set_xticks([])
@@ -134,13 +157,12 @@ def center(Gmass, x, y):
             x = x[~np.isnan(x)]
             y = y[~np.isnan(y)]
             Gmass = Gmass[~np.isnan(Gmass)]
-            x = x[Gmass]
             x_com = np.sum(Gmass * x) / np.sum(Gmass)
-            y_com = #np.sum(Gmass * y) / np.sum(Gmass)
+            y_com = np.sum(Gmass * y) / np.sum(Gmass)
             return x_com, y_com
 
         Gmass, rh, point_rad = next(self.data_stream(frame))
-        #x_com, y_com = center(Gmass, rh[:,0], rh[:,1])
+        x_com, y_com = center(Gmass, rh[:,0], rh[:,1])
         self.scatter_artist.set_offsets(np.c_[rh[:,0] - x_com, rh[:,1] - y_com])
         self.scatter_artist.set_sizes(point_rad**2)
         return self.scatter_artist,
@@ -181,6 +203,7 @@ def data_stream(self, frame=0):
         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, fragmentation_save="TRAJECTORY", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False)
-        sim.run(dt=1e-4, tstop=2.0e-3, istep_out=1, dump_cadence=0)
+        sim.run(dt=1e-4, tstop=1.0e-3, istep_out=1, dump_cadence=0)
 
-        anim = AnimatedScatter(sim,movie_filename,movie_titles[style],nskip=1)
+        print("Generating animation")
+        anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1)

python/swiftest/swiftest/io.py

@@ -817,11 +817,12 @@ def process_netcdf_input(ds, param):
     ds : xarray dataset
     """
 
+    ds = ds.where(~np.isnan(ds.id) ,drop=True)
     if param['OUT_TYPE'] == "NETCDF_DOUBLE":
         ds = fix_types(ds,ftype=np.float64)
     elif param['OUT_TYPE'] == "NETCDF_FLOAT":
         ds = fix_types(ds,ftype=np.float32)
-    ds = ds.where(ds.id >=0 ,drop=True)
+
     # Check if the name variable contains unique values. If so, make name the dimension instead of id
     if "id" in ds.dims:
         if len(np.unique(ds['name'])) == len(ds['name']):

python/swiftest/swiftest/simulation_class.py

@@ -17,6 +17,7 @@
 from swiftest import __file__ as _pyfile
 import json
 import os
+from glob import glob
 from pathlib import Path
 import datetime
 import xarray as xr
@@ -323,12 +324,14 @@ def __init__(self,read_param: bool = False, read_old_output_file: bool = False,
         self.simdir = Path(simdir)
         if self.simdir.exists():
             if not self.simdir.is_dir():
-                msg = f"Cannot create the {self.simdir} directory: File exists."
+                msg = f"Cannot create the {self.simdir.resolve()} directory: File exists."
                 msg += "\nDelete the file or change the location of param_file"
-                warnings.warn(msg,stacklevel=2)
+                raise NotADirectoryError(msg)
         else:
-            self.simdir.mkdir(parents=True, exist_ok=False)
-
+            if read_old_output_file or read_param:
+                raise NotADirectoryError(f"Cannot find directory {self.simdir.resolve()} ")
+            else:
+                self.simdir.mkdir(parents=True, exist_ok=False)
 
         # Set the location of the parameter input file, choosing the default if it isn't specified.
         param_file = kwargs.pop("param_file",Path.cwd() / self.simdir / "param.in")
@@ -376,7 +379,7 @@ def __init__(self,read_param: bool = False, read_old_output_file: bool = False,
             if os.path.exists(binpath):
                 self.read_output_file()
             else:
-                warnings.warn(f"BIN_OUT file {binpath} not found.",stacklevel=2)
+                raise FileNotFoundError(f"BIN_OUT file {binpath} not found.")
         return
 
     def _run_swiftest_driver(self):
@@ -2665,23 +2668,23 @@ def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", t
 
         Parameters
         ----------
-           param_file : string
-                File name of the input parameter file
-            newcodename : string
-                Name of the desired format (Swift/Swifter/Swiftest)
-            plname : string
-                File name of the massive body input file
-            tpname : string
-                File name of the test particle input file
-            cbname : string
-                File name of the central body input file
-            conversion_questions : dictronary
-                Dictionary of additional parameters required to convert between formats
+        param_file : string
+        File name of the input parameter file
+        newcodename : string
+        Name of the desired format (Swift/Swifter/Swiftest)
+        plname : string
+        File name of the massive body input file
+        tpname : string
+        File name of the test particle input file
+        cbname : string
+        File name of the central body input file
+        conversion_questions : dictronary
+        Dictionary of additional parameters required to convert between formats
 
         Returns
         -------
-            oldparam : xarray dataset
-                The old parameter configuration.
+        oldparam : xarray dataset
+        The old parameter configuration.
         """
         oldparam = self.param
         if self.codename == newcodename:
@@ -2718,12 +2721,12 @@ def read_output_file(self,read_init_cond : bool = True):
 
         Parameters
         ----------
-            read_init_cond : bool
-                Read in an initial conditions file along with the output file. Default is True
+        read_init_cond : bool
+        Read in an initial conditions file along with the output file. Default is True
 
         Returns
         -------
-            self.data : xarray dataset
+        self.data : xarray dataset
         """
 
         # Make a temporary copy of the parameter dictionary so we can supply the absolute path of the binary file
@@ -2763,15 +2766,19 @@ def read_output_file(self,read_init_cond : bool = True):
     def read_encounter(self):
         if self.verbose:
             print("Reading encounter history file as .enc")
-        enc_files = self.simdir.glob("**/encounter_*.nc")
+        enc_files = glob(f"{self.simdir}{os.path.sep}encounter_*.nc")
+        enc_files.sort()
 
         # 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.enc = xr.open_mfdataset(enc_files,parallel=True,preprocess=partial_func,mask_and_scale=True)
+        self.enc = xr.open_mfdataset(enc_files,parallel=True,combine="nested",concat_dim="time",join="left",preprocess=partial_func,mask_and_scale=True)
         self.enc = io.process_netcdf_input(self.enc, self.param)
+        # Remove any overlapping time values
+        tgood,tid = np.unique(self.enc.time,return_index=True)
+        self.enc = self.enc.isel(time=tid)
 
         return
 

src/CMakeLists.txt

@@ -30,6 +30,7 @@ SET(FAST_MATH_FILES
             ${SRC}/encounter/encounter_check.f90
             ${SRC}/encounter/encounter_setup.f90
             ${SRC}/encounter/encounter_util.f90
+            ${SRC}/encounter/encounter_io.f90
             ${SRC}/fraggle/fraggle_generate.f90
             ${SRC}/fraggle/fraggle_io.f90
             ${SRC}/fraggle/fraggle_placeholder.f90
@@ -74,6 +75,7 @@ SET(FAST_MATH_FILES
             ${SRC}/util/util_dealloc.f90
             ${SRC}/util/util_exit.f90
             ${SRC}/util/util_fill.f90
+            ${SRC}/util/util_final.f90
             ${SRC}/util/util_flatten.f90
             ${SRC}/util/util_get_energy_momentum.f90
             ${SRC}/util/util_index_array.f90