From 4d4417cf1a81bd27ed2c42f90e0a70abc98fb96e Mon Sep 17 00:00:00 2001 From: David A Minton Date: Wed, 7 Dec 2022 16:46:13 -0500 Subject: [PATCH] Added in code that can merge together the encounter and simulation data --- examples/Fragmentation/Fragmentation_Movie.py | 39 ++++++++++++++----- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/examples/Fragmentation/Fragmentation_Movie.py b/examples/Fragmentation/Fragmentation_Movie.py index 4e2ada041..329a6f343 100644 --- a/examples/Fragmentation/Fragmentation_Movie.py +++ b/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 @@ -75,21 +76,37 @@ 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, style, nskip=1): - if style == "disruption_headon" or style == "hitandrun": - tgood = sim.enc.where(sim.enc['Gmass'] > 0.8 * body_Gmass[style][0]).time - else: - tgood = sim.enc.where(sim.enc['Gmass'] > 0.01 * body_Gmass[style][1]).time - self.ds = sim.enc[['rh','vh','Gmass','radius']].sel(time=tgood).load().interpolate_na(dim="time") + self.ds = encounter_combiner(sim) nframes = int(self.ds['time'].size) self.sim = sim @@ -101,6 +118,7 @@ def __init__(self, sim, animfile, title, style, nskip=1): '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. @@ -110,7 +128,7 @@ def __init__(self, sim, animfile, title, style, 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. @@ -120,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([]) @@ -185,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=4.0e-3, istep_out=1, dump_cadence=0) + sim.run(dt=1e-5, tstop=4.0e-3, istep_out=10, dump_cadence=0) + print("Generating animation") anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1) \ No newline at end of file