More updates for preparing to output encounter data

examples/Fragmentation/Fragmentation_Movie.py

@@ -36,7 +36,7 @@
 # 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_title_list = ["Head-on Disruption", "Off-axis Supercatastrophic", "Hit and Run"]
 movie_titles = dict(zip(available_movie_styles, movie_title_list))
 
 # These initial conditions were generated by trial and error
@@ -135,8 +135,8 @@ def center(Gmass, x, y):
             y_com = np.sum(Gmass * y) / np.sum(Gmass)
             return x_com, y_com
 
-        Gmass, x, y, point_rad = next(self.data_stream(frame))
-        x_com, y_com = center(Gmass, x, y)
+        Gmass, rh, point_rad = next(self.data_stream(frame))
+        x_com, y_com = center(Gmass, rh[:,1], rh[:,2])
         self.scatter_artist.set_offsets(np.c_[x - x_com, y - y_com])
         self.scatter_artist.set_sizes(point_rad)
         return self.scatter_artist,
@@ -147,10 +147,9 @@ def data_stream(self, frame=0):
             ds = ds.where(ds['name'] != "Sun", drop=True)
             radius = ds['radius'].values
             Gmass = ds['Gmass'].values
-            x = ds['xhx'].values
-            y = ds['xhy'].values
+            rh = ds['rh'].values
             point_rad = 2 * radius * self.ax_pt_size
-            yield Gmass, x, y, point_rad
+            yield Gmass, rh, point_rad
 
 if __name__ == "__main__":
 
@@ -188,7 +187,7 @@ def data_stream(self, frame=0):
        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], rh=pos_vectors[style], vh=vel_vectors[style], rot=rot_vectors[style])
+           sim.add_body(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.2 * body_Gmass[style][1] # Make the minimum fragment mass a fraction of the smallest body

src/modules/symba_classes.f90

@@ -177,6 +177,12 @@ module symba_classes
       procedure :: resolve_collision      => symba_collision_resolve_plplenc              !! Process the pl-pl collision list, then modifiy the massive bodies based on the outcome of the c
    end type symba_plplenc
 
+   type, extends(helio_nbody_system) :: symba_system_snapshot
+   contains
+      procedure :: snapshot => symba_util_take_system_snapshot
+      final :: symba_util_final_snapshot
+   end type 
+
    !********************************************************************************************************************************
    !  symba_nbody_system class definitions and method interfaces
    !********************************************************************************************************************************
@@ -188,6 +194,7 @@ module symba_classes
       integer(I4B)                                    :: irec               !! System recursion level
       type(encounter_storage(nframes=:)), allocatable :: encounter_history  !! Stores encounter history for later retrieval and saving to file
       integer(I4B)                                    :: ienc_frame = 0     !! Encounter history frame number
+      type(symba_system_snapshot)                     :: snapshot
    contains
       procedure :: write_discard    => symba_io_write_discard             !! Write out information about discarded and merged planets and test particles in SyMBA
       procedure :: initialize       => symba_setup_initialize_system      !! Performs SyMBA-specific initilization steps
@@ -201,11 +208,6 @@ module symba_classes
       final     :: symba_util_final_system                                !! Finalizes the SyMBA nbody system object - deallocates all allocatables
    end type symba_nbody_system
 
-   type, extends(helio_nbody_system) :: symba_system_snapshot
-   contains
-      procedure :: snapshot => symba_util_take_system_snapshot
-      final :: symba_util_final_snapshot
-   end type 
 
    interface
 

src/setup/setup.f90

@@ -68,6 +68,8 @@ module subroutine setup_construct_system(system, param)
             allocate(symba_pltpenc :: system%pltpenc_list)
             allocate(symba_plplenc :: system%plplenc_list)
             allocate(symba_plplenc :: system%plplcollision_list)
+            allocate(symba_pl :: system%snapshot%pl)
+            allocate(symba_tp :: system%snapshot%tp)
          end select
       case (RINGMOONS)
          write(*,*) 'RINGMOONS-SyMBA integrator not yet enabled'

src/symba/symba_step.f90

@@ -218,6 +218,8 @@ recursive module subroutine symba_step_recur_system(self, param, t, ireci)
                   call pl%drift(system, param, dtl)
                   call tp%drift(system, param, dtl)
 
+                  !call system%
+
                   if (lencounter) call system%recursive_step(param, t+dth,irecp)
                   system%irec = ireci
 

src/symba/symba_util.f90

@@ -1300,10 +1300,15 @@ module subroutine symba_util_take_system_snapshot(self, system, param, t)
       !! Takes a minimal snapshot of the state of the system during an encounter so that the trajectories
       !! Can be played back through the encounter
       implicit none
+      ! Internals
       class(symba_system_snapshot),    intent(inout) :: self   !! SyMBA nbody system snapshot object
       class(symba_nbody_system),       intent(in)    :: system !! SyMBA nbody system object
       class(symba_parameters),         intent(in)    :: param  !! Current run configuration parameters 
       real(DP),                        intent(in)    :: t      !! current time
+      ! Arguments
+      logical, dimension(:), allocatable :: lmask
+
+      !if (system%pl)
 
       return
    end subroutine symba_util_take_system_snapshot