From d6572812d208361dabd4a51d8325ad5d4cf294ab Mon Sep 17 00:00:00 2001 From: David Minton Date: Tue, 10 Aug 2021 10:06:43 -0400 Subject: [PATCH] Removed need for sla library and also changed Mass to GMass for consistency --- .../symba_energy_momentum/collision_movie.py | 313 ++++++++++++++++++ python/swiftest/swiftest/io.py | 38 +-- python/swiftest/swiftest/simulation_class.py | 9 + python/swiftest/swiftest/tool.py | 16 +- 4 files changed, 350 insertions(+), 26 deletions(-) create mode 100644 examples/symba_energy_momentum/collision_movie.py diff --git a/examples/symba_energy_momentum/collision_movie.py b/examples/symba_energy_momentum/collision_movie.py new file mode 100644 index 000000000..c3c1cf112 --- /dev/null +++ b/examples/symba_energy_momentum/collision_movie.py @@ -0,0 +1,313 @@ +import swiftest +import numpy as np +import matplotlib.pyplot as plt +from matplotlib import animation +import matplotlib.collections as clt +from scipy.spatial.transform import Rotation as R + +xmin = -20.0 +xmax = 20.0 +ymin = -20.0 +ymax = 20.0 + +#cases = ['supercat_head', 'supercat_off', 'disruption_head', 'disruption_off'] +cases = ['disruption_off'] + +def scale_sim(ds, param): + + dsscale = ds + + dsscale['Mass'] = ds['Mass'] / param['GU'] + Mtot = dsscale['Mass'].sum(skipna=True, dim="id").isel(time=0) + rscale = sum(ds['Radius'].sel(id=[2, 3], time=0)).item() + ds['Radius'] /= rscale + + dsscale['radmarker'] = dsscale['Radius'].fillna(0) + + dsscale['px'] /= rscale + dsscale['py'] /= rscale + dsscale['pz'] /= rscale + + mpx = dsscale['Mass'] * dsscale['px'] + mpy = dsscale['Mass'] * dsscale['py'] + mpz = dsscale['Mass'] * dsscale['pz'] + xbsys = mpx.sum(skipna=True, dim="id") / Mtot + ybsys = mpy.sum(skipna=True, dim="id") / Mtot + zbsys = mpz.sum(skipna=True, dim="id") / Mtot + + mvx = dsscale['Mass'] * dsscale['vx'] + mvy = dsscale['Mass'] * dsscale['vy'] + mvz = dsscale['Mass'] * dsscale['vz'] + vxbsys = mvx.sum(skipna=True, dim="id") / Mtot + vybsys = mvy.sum(skipna=True, dim="id") / Mtot + vzbsys = mvz.sum(skipna=True, dim="id") / Mtot + + dsscale['pxb'] = dsscale['px'] - xbsys + dsscale['pyb'] = dsscale['py'] - ybsys + dsscale['pzb'] = dsscale['pz'] - zbsys + + dsscale['vxb'] = dsscale['vx'] - vxbsys + dsscale['vyb'] = dsscale['vy'] - vybsys + dsscale['vzb'] = dsscale['vz'] - vzbsys + + return dsscale + +class UpdatablePatchCollection(clt.PatchCollection): + def __init__(self, patches, *args, **kwargs): + self.patches = patches + clt.PatchCollection.__init__(self, patches, *args, **kwargs) + + def get_paths(self): + self.set_paths(self.patches) + return self._paths + +class AnimatedScatter(object): + """An animated scatter plot using matplotlib.animations.FuncAnimation.""" + + def __init__(self, ds, param): + + frame = 0 + nframes = ds['time'].size + self.ds = scale_sim(ds, param) + self.param = param + self.rot_angle = {} + + self.clist = {'Initial conditions' : 'xkcd:windows blue', + 'Disruption' : 'xkcd:baby poop', + 'Supercatastrophic' : 'xkcd:shocking pink', + 'Hit and run fragment' : 'xkcd:blue with a hint of purple', + 'Central body' : 'xkcd:almost black'} + + self.stream = self.data_stream(frame) + # Setup the figure and axes... + self.fig, self.ax = plt.subplots(figsize=(8,8)) + # Then setup FuncAnimation. + self.ani = animation.FuncAnimation(self.fig, self.update, interval=1, frames=nframes, + init_func=self.setup_plot, blit=False) + self.ani.save(animfile, fps=60, dpi=300, + extra_args=['-vcodec', 'mpeg4']) + + def plot_pl_circles(self, pl, radmarker): + patches = [] + for i in range(pl.shape[0]): + s = plt.Circle((pl[i, 0], pl[i, 1]), radmarker[i]) + patches.append(s) + return patches + + def vec_props(self, c): + arrowprops = { + 'arrowstyle': '<|-', + 'mutation_scale': 20, + 'connectionstyle': 'arc3', + } + + arrow_args = { + 'xycoords': 'data', + 'textcoords': 'data', + 'arrowprops': arrowprops, + 'annotation_clip': True, + 'zorder': 100, + 'animated' : True + } + aarg = arrow_args.copy() + aprop = arrowprops.copy() + aprop['color'] = c + aarg['arrowprops'] = aprop + aarg['color'] = c + return aarg + + def plot_pl_vectors(self, pl, cval, r): + varrowend, varrowtip = self.velocity_vectors(pl, r) + arrows = [] + for i in range(pl.shape[0]): + aarg = self.vec_props(cval[i]) + a = self.ax.annotate("",xy=varrowend[i],xytext=varrowtip[i], **aarg) + arrows.append(a) + return arrows + + def plot_pl_spins(self, pl, id, cval, len): + sarrowend, sarrowtip = self.spin_arrows(pl, id, len) + arrows = [] + for i in range(pl.shape[0]): + aarg = self.vec_props(cval[i]) + aarg['arrowprops']['mutation_scale'] = 5 + aarg['arrowprops']['arrowstyle'] = "simple" + a = self.ax.annotate("",xy=sarrowend[i],xytext=sarrowtip[i], **aarg) + arrows.append(a) + return arrows + + def origin_to_color(self, origin): + cval = [] + for o in origin: + c = self.clist[o] + cval.append(c) + + return cval + + def velocity_vectors(self, pl, r): + px = pl[:, 0] + py = pl[:, 1] + vx = pl[:, 2] + vy = pl[:, 3] + vmag = np.sqrt(vx ** 2 + vy ** 2) + ux = np.zeros_like(vx) + uy = np.zeros_like(vx) + goodv = vmag > 0.0 + ux[goodv] = vx[goodv] / vmag[goodv] + uy[goodv] = vy[goodv] / vmag[goodv] + varrowend = [] + varrowtip = [] + for i in range(pl.shape[0]): + vend = (px[i], py[i]) + vtip = (px[i] + vx[i] * self.v_length, py[i] + vy[i] * self.v_length) + varrowend.append(vend) + varrowtip.append(vtip) + return varrowend, varrowtip + + def spin_arrows(self, pl, id, len): + px = pl[:, 0] + py = pl[:, 1] + sarrowend = [] + sarrowtip = [] + for i in range(pl.shape[0]): + endrel = np.array([0.0, len[i], 0.0]) + tiprel = np.array([0.0, -len[i], 0.0]) + r = R.from_rotvec(self.rot_angle[id[i]]) + endrel = r.apply(endrel) + tiprel = r.apply(tiprel) + send = (px[i] + endrel[0], py[i] + endrel[1]) + stip = (px[i] + tiprel[0], py[i] + tiprel[1]) + sarrowend.append(send) + sarrowtip.append(stip) + return sarrowend, sarrowtip + + def setup_plot(self): + # First frame + """Initial drawing of the scatter plot.""" + t, name, Mass, Radius, npl, pl, radmarker, origin = next(self.data_stream(0)) + + cval = self.origin_to_color(origin) + # set up the figure + self.ax = plt.axes(xlim=(xmin, xmax), ylim=(ymin, ymax)) + plt.axis('off') + plt.tight_layout(pad=0) + self.ax.set_aspect(1) + self.ax.get_xaxis().set_visible(False) + self.ax.get_yaxis().set_visible(False) + + # Scale markers to the size of the system + self.v_length = 0.50 # Length of arrow as fraction of velocity + + self.ax.margins(x=1, y=1) + self.ax.set_xlabel('x distance / ($R_1 + R_2$)', fontsize='16', labelpad=1) + self.ax.set_ylabel('y distance / ($R_1 + R_2$)', fontsize='16', labelpad=1) + + self.title = self.ax.text(0.50, 0.90, "", bbox={'facecolor': 'w', 'pad': 5}, transform=self.ax.transAxes, + ha="center", zorder=1000) + + self.title.set_text(titletext) + self.patches = self.plot_pl_circles(pl, radmarker) + + self.collection = UpdatablePatchCollection(self.patches, color=cval, alpha=0.5, zorder=50) + self.ax.add_collection(self.collection) + #self.varrows = self.plot_pl_vectors(pl, cval, radmarker) + self.sarrows = self.plot_pl_spins(pl, name, cval, radmarker) + + return self.collection, self.sarrows + + def update(self,frame): + """Update the scatter plot.""" + t, name, Mass, Radius, npl, pl, radmarker, origin = next(self.data_stream(frame)) + cval = self.origin_to_color(origin) + #varrowend, varrowtip = self.velocity_vectors(pl, radmarker) + sarrowend, sarrowtip = self.spin_arrows(pl, name, radmarker) + for i, p in enumerate(self.patches): + p.set_center((pl[i, 0], pl[i,1])) + p.set_radius(radmarker[i]) + p.set_color(cval[i]) + #self.varrows[i].set_position(varrowtip[i]) + #self.varrows[i].xy = varrowend[i] + self.sarrows[i].set_position(sarrowtip[i]) + self.sarrows[i].xy = sarrowend[i] + + self.collection.set_paths(self.patches) + return self.collection, self.sarrows + + def data_stream(self, frame=0): + while True: + d = self.ds.isel(time=frame) + Radius = d['radmarker'].values + Mass = d['Mass'].values + x = d['pxb'].values + y = d['pyb'].values + vx = d['vxb'].values + vy = d['vyb'].values + name = d['id'].values + npl = d['npl'].values + id = d['id'].values + rotx = d['rot_x'].values + roty = d['rot_y'].values + rotz = d['rot_z'].values + + radmarker = d['radmarker'].values + origin = d['origin_type'].values + + t = self.ds.coords['time'].values[frame] + self.mask = np.logical_not(np.isnan(x)) + + x = np.nan_to_num(x, copy=False) + y = np.nan_to_num(y, copy=False) + vx = np.nan_to_num(vx, copy=False) + vy = np.nan_to_num(vy, copy=False) + radmarker = np.nan_to_num(radmarker, copy=False) + Mass = np.nan_to_num(Mass, copy=False) + Radius = np.nan_to_num(Radius, copy=False) + rotx = np.nan_to_num(rotx, copy=False) + roty = np.nan_to_num(roty, copy=False) + rotz = np.nan_to_num(rotz, copy=False) + rotvec = np.array([rotx, roty, rotz]) + self.rotvec = dict(zip(id, zip(*rotvec))) + + if frame == 0: + tmp = np.zeros_like(rotvec) + self.rot_angle = dict(zip(id, zip(*tmp))) + else: + t0 = self.ds.coords['time'].values[frame-1] + dt = t - t0 + idxactive = np.arange(id.size)[self.mask] + for i in id[idxactive]: + self.rot_angle[i] = self.rot_angle[i] + dt * np.array(self.rotvec[i]) + frame += 1 + yield t, name, Mass, Radius, npl, np.c_[x, y, vx, vy], radmarker, origin + +for case in cases: + if case == 'supercat_off': + animfile = 'movies/supercat_off_axis.mp4' + titletext = "Supercatastrophic - Off Axis" + paramfile = 'param.supercatastrophic_off_axis.in' + elif case == 'supercat_head': + animfile = 'movies/supercat_headon.mp4' + titletext = "Supercatastrophic - Head on" + paramfile = 'param.supercatastrophic_headon.in' + elif case == 'disruption_off': + animfile = 'movies/disruption_off_axis.mp4' + titletext = "Disruption - Off Axis" + paramfile = 'param.disruption_off_axis.in' + elif case == 'disruption_head': + animfile = 'movies/disruption_headon.mp4' + titletext = "Disruption- Head on" + paramfile = 'param.disruption_headon.in' + elif case == 'merger': + animfile = 'movies/merger.mp4' + titletext = "Merger" + paramfile = 'param.merger.in' + else: + print(f'{case} is an unknown case') + exit(-1) + sim = swiftest.Simulation(param_file=paramfile) + sim.bin2xr() + ds = sim.ds + print('Making animation') + anim = AnimatedScatter(ds,sim.param) + print('Animation finished') + plt.close(fig='all') diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 81840ba05..37f3370fd 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -318,7 +318,7 @@ def swifter_stream(f, param): plid : int array IDs of massive bodies pvec : float array - (npl,N) - vector of N quantities or each particle (6 of XV/EL + Mass, Radius, etc) + (npl,N) - vector of N quantities or each particle (6 of XV/EL + GMass, Radius, etc) plab : string list Labels for the pvec data ntp : int @@ -376,7 +376,7 @@ def swifter_stream(f, param): tlab.append('omega') tlab.append('capm') plab = tlab.copy() - plab.append('Mass') + plab.append('GMass') plab.append('Radius') pvec = np.vstack([pvec, Mpl, Rpl]) @@ -401,11 +401,11 @@ def make_swiftest_labels(param): tlab.append('omega') tlab.append('capm') plab = tlab.copy() - plab.append('Mass') + plab.append('GMass') plab.append('Radius') if param['RHILL_PRESENT'] == 'YES': plab.append('Rhill') - clab = ['Mass', 'Radius', 'J_2', 'J_4'] + clab = ['GMass', 'Radius', 'J_2', 'J_4'] if param['ROTATION'] == 'YES': clab.append('Ip_x') clab.append('Ip_y') @@ -443,13 +443,13 @@ def swiftest_stream(f, param): cbid : int array ID of central body (always returns 0) cvec : float array - (npl,1) - vector of quantities for the massive body (Mass, Radius, J2, J4, etc) + (npl,1) - vector of quantities for the massive body (GMass, Radius, J2, J4, etc) npl : int Number of massive bodies plid : int array IDs of massive bodies pvec : float array - (npl,N) - vector of N quantities or each particle (6 of XV/EL + Mass, Radius, etc) + (npl,N) - vector of N quantities or each particle (6 of XV/EL + GMass, Radius, etc) plab : string list Labels for the pvec data ntp : int @@ -656,10 +656,10 @@ def swiftest_xr2infile(ds, param, framenum=-1): frame = ds.isel(time=framenum) cb = frame.where(frame.id == 0, drop=True) pl = frame.where(frame.id > 0, drop=True) - pl = pl.where(np.invert(np.isnan(pl['Mass'])), drop=True).drop_vars(['J_2', 'J_4']) - tp = frame.where(np.isnan(frame['Mass']), drop=True).drop_vars(['Mass', 'Radius', 'J_2', 'J_4']) + pl = pl.where(np.invert(np.isnan(pl['GMass'])), drop=True).drop_vars(['J_2', 'J_4']) + tp = frame.where(np.isnan(frame['GMass']), drop=True).drop_vars(['GMass', 'Radius', 'J_2', 'J_4']) - GMSun = np.double(cb['Mass']) + GMSun = np.double(cb['GMass']) RSun = np.double(cb['Radius']) J2 = np.double(cb['J_2']) J4 = np.double(cb['J_4']) @@ -680,9 +680,9 @@ def swiftest_xr2infile(ds, param, framenum=-1): for i in pl.id: pli = pl.sel(id=i) if param['RHILL_PRESENT'] == 'YES': - print(i.values, pli['Mass'].values, pli['Rhill'].values, file=plfile) + print(i.values, pli['GMass'].values, pli['Rhill'].values, file=plfile) else: - print(i.values, pli['Mass'].values, file=plfile) + print(i.values, pli['GMass'].values, file=plfile) print(pli['Radius'].values, file=plfile) print(pli['px'].values, pli['py'].values, pli['pz'].values, file=plfile) print(pli['vx'].values, pli['vy'].values, pli['vz'].values, file=plfile) @@ -716,7 +716,7 @@ def swiftest_xr2infile(ds, param, framenum=-1): vx = pl['vx'].values vy = pl['vy'].values vz = pl['vz'].values - mass = pl['Mass'].values + Gmass = pl['GMass'].values radius = pl['Radius'].values plfile.write_record(npl) @@ -727,7 +727,7 @@ def swiftest_xr2infile(ds, param, framenum=-1): plfile.write_record(vx) plfile.write_record(vy) plfile.write_record(vz) - plfile.write_record(mass) + plfile.write_record(Gmass) if param['RHILL_PRESENT'] == 'YES': rhill = pl['Rhill'].values plfile.write_record(rhill) @@ -774,10 +774,10 @@ def swifter_xr2infile(ds, param, framenum=-1): frame = ds.isel(time=framenum) cb = frame.where(frame.id == 0, drop=True) pl = frame.where(frame.id > 0, drop=True) - pl = pl.where(np.invert(np.isnan(pl['Mass'])), drop=True).drop_vars(['J_2', 'J_4']) - tp = frame.where(np.isnan(frame['Mass']), drop=True).drop_vars(['Mass', 'Radius', 'J_2', 'J_4']) + pl = pl.where(np.invert(np.isnan(pl['GMass'])), drop=True).drop_vars(['J_2', 'J_4']) + tp = frame.where(np.isnan(frame['GMass']), drop=True).drop_vars(['GMass', 'Radius', 'J_2', 'J_4']) - GMSun = np.double(cb['Mass']) + GMSun = np.double(cb['GMass']) RSun = np.double(cb['Radius']) param['J2'] = np.double(cb['J_2']) param['J4'] = np.double(cb['J_4']) @@ -786,15 +786,15 @@ def swifter_xr2infile(ds, param, framenum=-1): # Swiftest Central body file plfile = open(param['PL_IN'], 'w') print(pl.id.count().values + 1, file=plfile) - print(cb.id.values[0], cb['Mass'].values[0], file=plfile) + print(cb.id.values[0], cb['GMass'].values[0], file=plfile) print('0.0 0.0 0.0', file=plfile) print('0.0 0.0 0.0', file=plfile) for i in pl.id: pli = pl.sel(id=i) if param['RHILL_PRESENT'] == "YES": - print(i.values, pli['Mass'].values, pli['Rhill'].values, file=plfile) + print(i.values, pli['GMass'].values, pli['Rhill'].values, file=plfile) else: - print(i.values, pli['Mass'].values, file=plfile) + print(i.values, pli['GMass'].values, file=plfile) if param['CHK_CLOSE'] == "YES": print(pli['Radius'].values, file=plfile) print(pli['px'].values, pli['py'].values, pli['pz'].values, file=plfile) diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 05b6896b1..bea59ae5f 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -52,6 +52,7 @@ def __init__(self, codename="Swiftest", param_file=""): self.read_param(param_file, codename) return + def add(self, plname, date=date.today().isoformat(), idval=None): """ Adds a solar system body to an existing simulation DataSet. @@ -69,6 +70,7 @@ def add(self, plname, date=date.today().isoformat(), idval=None): self.ds = init_cond.solar_system_horizons(plname, idval, self.param, date, self.ds) return + def read_param(self, param_file, codename="Swiftest"): if codename == "Swiftest": self.param = io.read_swiftest_param(param_file, self.param) @@ -84,6 +86,7 @@ def read_param(self, param_file, codename="Swiftest"): self.codename = "Unknown" return + def write_param(self, param_file, param=None): if param is None: param = self.param @@ -97,6 +100,7 @@ def write_param(self, param_file, param=None): print('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".') return + def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", tpname="tp.swiftest.in", cbname="cb.swiftest.in", conversion_questions={}): """ Converts simulation input files from one code type to another (Swift, Swifter, or Swiftest). Returns the old parameter configuration. @@ -130,6 +134,7 @@ def convert(self, param_file, newcodename="Swiftest", plname="pl.swiftest.in", t print(f"Conversion from {self.codename} to {newcodename} is not supported.") return oldparam + def bin2xr(self): if self.codename == "Swiftest": self.ds = io.swiftest2xr(self.param) @@ -143,6 +148,7 @@ def bin2xr(self): print('Cannot process unknown code type. Call the read_param method with a valid code name. Valid options are "Swiftest", "Swifter", or "Swift".') return + def follow(self, codestyle="Swifter"): if self.ds is None: self.bin2xr() @@ -163,10 +169,13 @@ def follow(self, codestyle="Swifter"): ifol = None nskp = None fol = tool.follow_swift(self.ds, ifol=ifol, nskp=nskp) + else: + fol = None print('follow.out written') return fol + def save(self, param_file, framenum=-1, codename="Swiftest"): if codename == "Swiftest": io.swiftest_xr2infile(self.ds, self.param, framenum) diff --git a/python/swiftest/swiftest/tool.py b/python/swiftest/swiftest/tool.py index 741dc0f1b..a96610bc2 100644 --- a/python/swiftest/swiftest/tool.py +++ b/python/swiftest/swiftest/tool.py @@ -2,15 +2,17 @@ import numpy as np import os import glob -from pyslalib import slalib import xarray as xr """ Functions that recreate the Swift/Swifter tool programs """ -def sla_dranrm(angle): - func = np.vectorize(slalib.sla_dranrm) - return xr.apply_ufunc(func, angle) +def wrap_angle(angle): + while np.any(angle >= 2 * np.pi): + angle[angle >= 2 * np.pi] -= 2 * np.pi + while np.any(angle < 0.0): + angle[angle < 0.0] += 2 * np.pi + return angle def follow_swift(ds, ifol=None, nskp=None): """ @@ -36,11 +38,11 @@ def follow_swift(ds, ifol=None, nskp=None): ifol = int(intxt) print(f"Following particle {ifol}") if ifol < 0: # Negative numbers are planets - fol = ds.where(np.invert(np.isnan(ds['Mass'])), drop=True) + fol = ds.where(np.invert(np.isnan(ds['GMass'])), drop=True) fol = fol.where(np.invert(np.isnan(fol['a'])), drop=True) # Remove times where this body doesn't exist (but this also gets rid of the central body) fol = fol.isel(id = -ifol - 2) # Take 1 off for 0-indexed arrays in Python, and take 1 more off because the central body is gone elif ifol > 0: # Positive numbers are test particles - fol = ds.where(np.isnan(ds['Mass']), drop=True).drop_vars(['Mass', 'Radius']) + fol = ds.where(np.isnan(ds['GMass']), drop=True).drop_vars(['GMass', 'Radius']) fol = fol.where(np.invert(np.isnan(fol['a'])), drop=True) fol = fol.isel(id = ifol - 1) # Take 1 off for 0-indexed arrays in Python @@ -51,7 +53,7 @@ def follow_swift(ds, ifol=None, nskp=None): dr = 180.0 / np.pi fol['obar'] = fol['capom'] + fol['omega'] fol['obar'] = fol['obar'].fillna(0) - fol['obar'] = sla_dranrm(fol['obar']) + fol['obar'] = wrap_angle(fol['obar']) fol['obar'] = fol['obar'] * dr fol['inc'] = fol['inc'] * dr fol['capom'] = fol['capom'] * dr