diff --git a/examples/Chambers2013/scattermovie.py b/examples/Chambers2013/scattermovie.py index 2ef59c867..09b29072a 100755 --- a/examples/Chambers2013/scattermovie.py +++ b/examples/Chambers2013/scattermovie.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt from matplotlib import animation import matplotlib.colors as mcolors +from collections import namedtuple plt.switch_backend('agg') titletext = "Chambers (2013)" @@ -22,6 +23,8 @@ framejump = 1 animation_file = f"Chambers2013-{plot_style}.mp4" +origin_types = ["Initial conditions", "Merger", "Disruption", "Supercatastrophic", "Hit and run fragmentation"] + class AnimatedScatter(object): """An animated scatter plot using matplotlib.animations.FuncAnimation.""" def __init__(self, ds, param): @@ -31,91 +34,85 @@ def __init__(self, ds, param): self.ds = ds self.param = param self.Rcb = self.ds['radius'].sel(name="Sun").isel(time=0).values[()] + colors = ["k", "xkcd:faded blue", "xkcd:marigold", "xkcd:shocking pink", "xkcd:baby poop green"] + self.clist = dict(zip(origin_types,colors)) - self.clist = {'Initial conditions' : 'k', - 'Merger' : '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(xlim[plot_style]) - self.ax.set_ylim(ylim[plot_style]) fig.add_axes(self.ax) + + self.make_artists() + + self.ani = animation.FuncAnimation(fig, func=self.update, interval=1, frames=nframes, init_func=self.init_plot, blit=True) + self.ani.save(animation_file, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) + print(f'Finished writing {animation_file}') + + def make_artists(self): + scatter_names = [f"s{i}" for i,k in enumerate(origin_types)] + self.scatter_artist_names = dict(zip(origin_types,scatter_names)) + + animated_elements = [self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes,ha="center", animated=True)] + element_names = ["title"] + for key, value in self.clist.items(): + animated_elements.append(self.ax.scatter([], [], marker='o', s=[], c=value, alpha=0.75, label=key, animated=True)) + element_names.append(self.scatter_artist_names[key]) - # First frame - """Initial drawing of the scatter plot.""" - t, npl, pl, radmarker, origin = next(self.data_stream(0)) + Artists = namedtuple("Artists",tuple(element_names)) + self.artists = Artists(*animated_elements) + return + def init_plot(self): + self.ax.set_xlim(xlim[plot_style]) + self.ax.set_ylim(ylim[plot_style]) + # set up the figure self.ax.margins(x=10, y=1) self.ax.set_xlabel(xlabel[plot_style], fontsize='16', labelpad=1) - self.ax.set_ylabel(ylabel[plot_style], fontsize='16', labelpad=1) - - title = self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes, - ha="center", animated=True) - - title.set_text(f"{titletext} - Time = ${t*1e-6:6.2f}$ My with ${npl:4.0f}$ particles") - self.slist = self.scatters(pl, radmarker, origin) - self.slist.append(title) + self.ax.set_ylabel(ylabel[plot_style], fontsize='16', labelpad=1) leg = plt.legend(loc="upper left", scatterpoints=1, fontsize=10) for i,l in enumerate(leg.legendHandles): - leg.legendHandles[i]._sizes = [20] + leg.legendHandles[i]._sizes = [20] - self.ani = animation.FuncAnimation(fig, self.update, interval=1, frames=nframes, init_func=self.setup_plot, blit=True) - self.ani.save(animation_file, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) - print(f'Finished writing {animation_file}') + return self.artists + + def get_data(self, frame=0): + d = self.ds.isel(time = frame) + n=len(d['name']) + d = d.isel(name=range(1,n)) + d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale + + t = d['time'].values + npl = d['npl'].values + radmarker = d['radmarker'].values + origin = d['origin_type'].values + + if plot_style == "aescatter": + pl = np.c_[d['a'].values,d['e'].values] + elif plot_style == "aiscatter": + pl = np.c_[d['a'].values,d['inc'].values] - 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, animated=True) - scat.append(s) - return scat - - def setup_plot(self): - return self.slist - - def data_stream(self, frame=0): - while True: - d = self.ds.isel(time = frame) - n=len(d['name']) - d = d.isel(name=range(1,n)) - d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale - - t = d['time'].values - npl = d['npl'].values - radmarker = d['radmarker'].values - origin = d['origin_type'].values - - if plot_style == "aescatter": - pl = np.c_[d['a'].values,d['e'].values] - elif plot_style == "aiscatter": - pl = np.c_[d['a'].values,d['inc'].values] - - yield t, npl, pl, radmarker, origin + return t, npl, pl, radmarker, origin def update(self,frame): """Update the scatter plot.""" - t, npl, pl, radmarker, origin = next(self.data_stream(framejump * frame)) + t, npl, pl, radmarker, origin = self.get_data(framejump * frame) - self.slist[-1].set_text(f"{titletext} - Time = ${t*1e-6:6.3f}$ My with ${npl:4.0f}$ particles") + self.artists.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. - for i, (key, value) in enumerate(self.clist.items()): + for key,name in self.scatter_artist_names.items(): idx = origin == key - self.slist[i].set_sizes(radmarker[idx]) - self.slist[i].set_offsets(pl[idx,:]) - self.slist[i].set_facecolor(value) + if any(idx) and any(~np.isnan(radmarker[idx])): + scatter = self.artists._asdict()[name] + scatter.set_sizes(radmarker[idx]) + scatter.set_offsets(pl[idx,:]) + scatter.set_facecolor(self.clist[key]) - return self.slist + return self.artists sim = swiftest.Simulation(read_data=True) print('Making animation') diff --git a/examples/Fragmentation/.gitignore b/examples/Fragmentation/.gitignore index ff09bd225..2105da1cb 100644 --- a/examples/Fragmentation/.gitignore +++ b/examples/Fragmentation/.gitignore @@ -1,3 +1,4 @@ * !.gitignore !Fragmentation_Movie.py +!Multibody_Movie.py diff --git a/examples/Fragmentation/Multibody_Movie.py b/examples/Fragmentation/Multibody_Movie.py new file mode 100755 index 000000000..6751573b3 --- /dev/null +++ b/examples/Fragmentation/Multibody_Movie.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 +""" + Copyright 2022 - David Minton, Carlisle Wishard, Jennifer Pouplin, Jake Elliott, & Dana Singh + This file is part of Swiftest. + Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License + as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. + Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty + of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. + You should have received a copy of the GNU General Public License along with Swiftest. + If not, see: https://www.gnu.org/licenses. +""" + +""" +Generates a movie of a multi-body fragmentation event from set of Swiftest output files. + +Inputs +_______ +param.in : ASCII text file + Swiftest parameter input file. +out.nc : NetCDF file + Swiftest output file. + +Returns +------- +fragmentation.mp4 : mp4 movie file + Movie of a fragmentation event. +""" + +import swiftest +import numpy as np +import xarray as xr +import matplotlib.pyplot as plt +import matplotlib.animation as animation +from scipy.spatial.transform import Rotation as R + +# ---------------------------------------------------------------------------------------------------------------------- +# Define the names and initial conditions of the various fragmentation simulation types +# ---------------------------------------------------------------------------------------------------------------------- +available_movie_styles = ["supercatastrophic_multi"] +movie_title_list = ["Multi-body Supercatastrophic"] +movie_titles = dict(zip(available_movie_styles, movie_title_list)) +num_movie_frames = 1000 + +# These initial conditions were generated by trial and error +names = ["Body1","Body2","Body3","Body4"] + +rdistance = 1e-4/np.sqrt(2.0) +vrel = 2.0/np.sqrt(2.0) + +rcb = np.array([1.0, 0.0, 0.0]) +vcb = np.array([0.0, 6.28, 0.0]) + +r1 = rcb + np.array([-rdistance,-rdistance,0.0]) +r2 = rcb + np.array([ rdistance,-rdistance,0.0]) +r3 = rcb + np.array([ rdistance, rdistance,0.0]) +r4 = rcb + np.array([-rdistance, rdistance,0.0]) + +v1 = vcb + np.array([ vrel, vrel, 0.0]) +v2 = vcb + np.array([-vrel, vrel, 0.0]) +v3 = vcb + np.array([-vrel,-vrel, 0.0]) +v4 = vcb + np.array([ vrel,-vrel, 0.0]) + +rot1 = np.array([0.0,0.0,1e5]) +rot2 = np.array([0.0,0.0,-2e5]) +rot3 = np.array([0.0,0.0,3e5]) +rot4 = np.array([0.0,0.0,0.0]) + +m1 = 2e-7 +m2 = 2e-7 +m3 = 2e-7 +m4 = 2e-7 + +pos_vectors = {"supercatastrophic_multi" : [r1, r2, r3, r4], + } + +vel_vectors = {"supercatastrophic_multi": [v1, v2, v3, v4] + } + +rot_vectors = { "supercatastrophic_multi": [rot1, rot2, rot3, rot4] + } + +body_Gmass = {"supercatastrophic_multi": [m1, m2, m3, m4] + } + +tstop = {"supercatastrophic_multi" : 2.0e-3, + } + +density = 3000 * swiftest.AU2M**3 / swiftest.MSun +GU = swiftest.GMSun * swiftest.YR2S**2 / swiftest.AU2M**3 +body_radius = body_Gmass.copy() +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 +# ---------------------------------------------------------------------------------------------------------------------- +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 = ['name','rh','vh','Gmass','radius','rot'] + data = sim.data[keep_vars] + enc = sim.encounters[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) + tgood=enc.time.where(~np.isnan(enc.time),drop=True) + enc = enc.sel(time=tgood) + + # 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") + + # Interpolate in time to make a smooth, constant time step dataset + # Add a bit of padding to the time, otherwise there are some issues with the interpolation in the last few frames. + smooth_time = np.linspace(start=tgood.isel(time=0), stop=ds.time[-1], num=int(1.2*num_movie_frames)) + ds = ds.interp(time=smooth_time) + ds['rotangle'] = xr.zeros_like(ds['rot']) + ds['rot'] = ds['rot'].fillna(0.0) + + return ds + +def center(Gmass, x, y): + x = x[~np.isnan(x)] + y = y[~np.isnan(y)] + Gmass = Gmass[~np.isnan(Gmass)] + x_com = np.sum(Gmass * x) / np.sum(Gmass) + y_com = np.sum(Gmass * y) / np.sum(Gmass) + return x_com, y_com +class AnimatedScatter(object): + """An animated scatter plot using matplotlib.animations.FuncAnimation.""" + + def __init__(self, sim, animfile, title, style, nskip=1): + + self.ds = encounter_combiner(sim) + self.npl = len(self.ds['name']) - 1 + self.title = title + self.body_color_list = {'Initial conditions': 'xkcd:windows blue', + 'Disruption': 'xkcd:baby poop', + 'Supercatastrophic': 'xkcd:shocking pink', + '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() + self.ani = animation.FuncAnimation(self.fig, self.update_plot, init_func=self.init_func, interval=1, frames=range(0,num_movie_frames,nskip), blit=True) + self.ani.save(animfile, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) + print(f"Finished writing {animfile}") + + def setup_plot(self): + 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. + # This will be used to scale the axis limits on the movie. + rhy1 = self.ds['rh'].sel(name="Body1",space='y').isel(time=0).values[()] + rhy2 = self.ds['rh'].sel(name="Body3",space='y').isel(time=0).values[()] + + scale_frame = abs(rhy1) + abs(rhy2) + + ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8]) + self.ax_pt_size = self.figsize[0] * 72 / scale_frame * 0.7 + ax.set_xlim(-scale_frame, scale_frame) + ax.set_ylim(-scale_frame, scale_frame) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_xlabel("x") + ax.set_ylabel("y") + ax.set_title(self.title) + fig.add_axes(ax) + + return fig, ax + + def init_func(self): + self.artists = [] + aarg = self.vec_props('xkcd:beige') + for i in range(self.npl): + self.artists.append(self.ax.annotate("",xy=(0,0),**aarg)) + + self.artists.append(self.ax.scatter([],[],c='k', animated=True, zorder=10)) + return self.artists + + def update_plot(self, frame): + # Define a function to calculate a reference frame for the animation + t, Gmass, rh, radius, rotangle = next(self.data_stream(frame)) + x_ref, y_ref = center(Gmass, rh[:,0], rh[:,1]) + rh = np.c_[rh[:,0] - x_ref, rh[:,1] - y_ref] + self.artists[-1].set_offsets(rh) + point_rad = radius * self.ax_pt_size + self.artists[-1].set_sizes(point_rad**2) + + sarrowend, sarrowtip = self.spin_arrows(rh, rotangle, 1.1*radius) + for i, s in enumerate(self.artists[:-1]): + self.artists[i].set_position(sarrowtip[i]) + self.artists[i].xy = sarrowend[i] + + return self.artists + + def data_stream(self, frame=0): + while True: + t = self.ds.isel(time=frame)['time'].values[()] + + if frame > 0: + dsold = self.ds.isel(time=frame-1) + told = dsold['time'].values[()] + dt = t - told + self.ds['rotangle'][dict(time=frame)] = dsold['rotangle'] + dt * dsold['rot'] + + ds = self.ds.isel(time=frame) + ds = ds.where(ds['name'] != "Sun", drop=True) + radius = ds['radius'].values + Gmass = ds['Gmass'].values + rh = ds['rh'].values + rotangle = ds['rotangle'].values + + yield t, Gmass, rh, radius, rotangle + + def spin_arrows(self, rh, rotangle, rotlen): + px = rh[:, 0] + py = rh[:, 1] + sarrowend = [] + sarrowtip = [] + for i in range(rh.shape[0]): + endrel = np.array([0.0, -rotlen[i], 0.0]) + tiprel = np.array([0.0, rotlen[i], 0.0]) + r = R.from_rotvec(rotangle[i,:], degrees=True) + 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 vec_props(self, c): + arrowprops = { + 'arrowstyle': '-', + 'linewidth' : 1, + } + + 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 + +if __name__ == "__main__": + + print("Select a fragmentation movie to generate.") + print("1. Multi-body supercatastrophic") + print("2. All of the above") + user_selection = int(input("? ")) + + if user_selection > 0 and user_selection < 2: + movie_styles = [available_movie_styles[user_selection-1]] + else: + print("Generating all movie styles") + movie_styles = available_movie_styles.copy() + + for style in movie_styles: + print(f"Generating {movie_titles[style]}") + movie_filename = f"{style}.mp4" + # Pull in the Swiftest output data from the parameter file and store it as a Xarray dataset. + sim = swiftest.Simulation(simdir=style, rotation=True, init_cond_format = "XV", compute_conservation_values=True) + sim.add_solar_system_body("Sun") + sim.add_body(name=names, 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.05 * body_Gmass[style][1] + gmtiny = 0.10 * body_Gmass[style][1] + sim.set_parameter(collision_model="fraggle", encounter_save="both", gmtiny=gmtiny, minimum_fragment_gmass=minimum_fragment_gmass, verbose=False) + sim.run(dt=5e-4, tstop=tstop[style], istep_out=1, dump_cadence=0) + + print("Generating animation") + anim = AnimatedScatter(sim,movie_filename,movie_titles[style],style,nskip=1) diff --git a/python/swiftest/swiftest/io.py b/python/swiftest/swiftest/io.py index 0f9941066..98e6c154a 100644 --- a/python/swiftest/swiftest/io.py +++ b/python/swiftest/swiftest/io.py @@ -34,8 +34,6 @@ "DUMP_CADENCE", "ENCOUNTER_SAVE") - - # This list defines features that are booleans, so must be converted to/from string when writing/reading from file bool_param = ["RESTART", "CHK_CLOSE", @@ -83,7 +81,8 @@ def bool2yesno(boolval): return "YES" else: return "NO" - + + def bool2tf(boolval): """ Converts a boolean into a string of either "T" or "F". @@ -103,6 +102,7 @@ def bool2tf(boolval): else: return "F" + def str2bool(input_str): """ Converts a string into an equivalent boolean. @@ -801,6 +801,7 @@ def swifter2xr(param, verbose=True): if verbose: print(f"Successfully converted {ds.sizes['time']} output frames.") return ds + def process_netcdf_input(ds, param): """ Performs several tasks to convert raw NetCDF files output by the Fortran program into a form that @@ -823,6 +824,7 @@ def process_netcdf_input(ds, param): return ds + def swiftest2xr(param, verbose=True): """ Converts a Swiftest binary data file into an xarray DataSet. @@ -849,6 +851,7 @@ def swiftest2xr(param, verbose=True): return ds + def xstrip_nonstr(a): """ Cleans up the string values in the DataSet to remove extra white space @@ -864,6 +867,7 @@ def xstrip_nonstr(a): func = lambda x: np.char.strip(x) return xr.apply_ufunc(func, a.str.decode(encoding='utf-8'),dask='parallelized') + def xstrip_str(a): """ Cleans up the string values in the DataSet to remove extra white space @@ -898,6 +902,7 @@ def string_converter(da): return da + def char_converter(da): """` Converts a string to a unicode string @@ -916,6 +921,7 @@ def char_converter(da): da = xstrip_nonstr(da) return da + def clean_string_values(ds): """ Cleans up the string values in the DataSet that have artifacts as a result of coming from NetCDF Fortran @@ -963,6 +969,7 @@ def unclean_string_values(ds): ds[c] = n.str.ljust(1).str.encode('utf-8') return ds + def fix_types(ds,itype=np.int64,ftype=np.float64): ds = clean_string_values(ds) @@ -1054,6 +1061,7 @@ def swiftest_particle_2xr(param): return infoxr + def select_active_from_frame(ds, param, framenum=-1): """ Selects a particular frame from a DataSet and returns only the active particles in that frame @@ -1096,6 +1104,7 @@ def select_active_from_frame(ds, param, framenum=-1): return frame + def swiftest_xr2infile(ds, param, in_type="NETCDF_DOUBLE", infile_name=None,framenum=-1,verbose=True): """ Writes a set of Swiftest input files from a single frame of a Swiftest xarray dataset @@ -1281,6 +1290,7 @@ def swifter_xr2infile(ds, param, framenum=-1): return + def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): """ Converts from a Swift run to a Swifter run @@ -1501,6 +1511,7 @@ def swift2swifter(swift_param, plname="", tpname="", conversion_questions={}): return swifter_param + def swifter2swiftest(swifter_param, plname="", tpname="", cbname="", conversion_questions={}): """ Converts from a Swifter run to a Swiftest run @@ -1753,6 +1764,7 @@ def swifter2swiftest(swifter_param, plname="", tpname="", cbname="", conversion_ swiftest_param['! VERSION'] = "Swiftest parameter file converted from Swifter" return swiftest_param + def swift2swiftest(swift_param, plname="", tpname="", cbname="", conversion_questions={}): """ Converts from a Swift run to a Swiftest run @@ -1798,6 +1810,7 @@ def swift2swiftest(swift_param, plname="", tpname="", cbname="", conversion_ques swiftest_param['! VERSION'] = "Swiftest parameter file converted from Swift" return swiftest_param + def swiftest2swifter_param(swiftest_param, J2=0.0, J4=0.0): """ Converts from a Swiftest run to a Swifter run diff --git a/python/swiftest/swiftest/simulation_class.py b/python/swiftest/swiftest/simulation_class.py index 82fd6e70a..6e80909c4 100644 --- a/python/swiftest/swiftest/simulation_class.py +++ b/python/swiftest/swiftest/simulation_class.py @@ -2999,5 +2999,3 @@ def clean(self): os.remove(f) return - - diff --git a/python/swiftest/swiftest/tool.py b/python/swiftest/swiftest/tool.py index 1b53b2c7a..7fe2c8884 100644 --- a/python/swiftest/swiftest/tool.py +++ b/python/swiftest/swiftest/tool.py @@ -9,10 +9,7 @@ If not, see: https://www.gnu.org/licenses. """ -import swiftest import numpy as np -import os -import glob import xarray as xr """ Functions that recreate the Swift/Swifter tool programs diff --git a/python/swiftest/swiftest/visualize.py b/python/swiftest/swiftest/visualize.py new file mode 100644 index 000000000..bf095fc2f --- /dev/null +++ b/python/swiftest/swiftest/visualize.py @@ -0,0 +1,42 @@ +""" + Copyright 2023 - David Minton, + This file is part of Swiftest. + Swiftest is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License + as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. + Swiftest is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty + of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. + You should have received a copy of the GNU General Public License along with Swiftest. + If not, see: https://www.gnu.org/licenses. +""" +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +def _square_plot(): + figsize = (4,4) + fig = plt.figure(figsize=figsize, dpi=300) + plt.tight_layout(pad=0) + + ax = plt.Axes(fig, [0.1, 0.1, 0.8, 0.8]) + ax.set_xticks([]) + ax.set_yticks([]) + ax.set_aspect('equal') + fig.add_axes(ax) + return fig, ax + +def select_one_collision(collisions, collision_id): + ds = collisions.sel(collision_id=collision_id) + bgood = ds.sel(stage='before')['name'].where(ds.sel(stage='before')['particle_type'] != 'nan',drop=True) + agood = ds.sel(stage='after')['name'].where(ds.sel(stage='after')['particle_type'] != 'nan',drop=True) + goodname=np.unique(np.concatenate((bgood,agood))) + ds = ds.sel(name=goodname) + + return ds + +def collisions(collisions, collision_id): + fig, ax = _square_plot() + + ds = select_one_collision(collisions, collision_id) + + + return \ No newline at end of file diff --git a/src/collision/collision_generate.f90 b/src/collision/collision_generate.f90 index 7bcf7bdbe..4ef6ec215 100644 --- a/src/collision/collision_generate.f90 +++ b/src/collision/collision_generate.f90 @@ -70,14 +70,7 @@ module subroutine collision_generate_bounce(self, nbody_system, param, t) nimp = size(impactors%id(:)) do i = 1, nimp j = impactors%id(i) - rcom(:) = pl%rb(:,j) - impactors%rbcom(:) - vcom(:) = pl%vb(:,j) - impactors%vbcom(:) - rnorm(:) = .unit. rcom(:) - ! Do the reflection - vcom(:) = vcom(:) - 2 * dot_product(vcom(:),rnorm(:)) * rnorm(:) - ! Shift the positions so that the collision doesn't immediately occur again - pl%rb(:,j) = pl%rb(:,j) + 0.5_DP * pl%radius(j) * rnorm(:) - pl%vb(:,j) = impactors%vbcom(:) + vcom(:) + call collision_util_bounce_one(pl%rb(:,j),pl%vb(:,j),impactors%rbcom(:),impactors%vbcom(:),pl%radius(j)) self%status = DISRUPTED pl%status(j) = ACTIVE pl%ldiscard(j) = .false. diff --git a/src/collision/collision_module.f90 b/src/collision/collision_module.f90 index 21b606202..d6323f522 100644 --- a/src/collision/collision_module.f90 +++ b/src/collision/collision_module.f90 @@ -84,7 +84,7 @@ module collision real(DP), dimension(NDIM) :: v_unit !! velocity direction unit vector of collisional system real(DP), dimension(NDIM) :: rbcom !! Center of mass position vector of the collider nbody_system in nbody_system barycentric coordinates real(DP), dimension(NDIM) :: vbcom !! Velocity vector of the center of mass of the collider nbody_system in nbody_system barycentric coordinates - real(DP), dimension(NDIM) :: rbimp !! Impact point position vector of the collider nbody_system in nbody_system barycentric coordinates + real(DP), dimension(NDIM) :: rcimp !! Impact point position vector of the collider nbody_system in nbody_system barycentric coordinates real(DP), dimension(NDIM) :: bounce_unit !! The impact point velocity vector is the component of the velocity in the distance vector direction contains @@ -397,6 +397,13 @@ module subroutine collision_util_add_fragments_to_collider(self, nbody_system, p class(base_parameters), intent(in) :: param !! Current Swiftest run configuration parameters end subroutine collision_util_add_fragments_to_collider + module subroutine collision_util_bounce_one(r,v,rcom,vcom,radius) + implicit none + real(DP), dimension(:), intent(inout) :: r,v + real(DP), dimension(:), intent(in) :: rcom,vcom + real(DP), intent(in) :: radius + end subroutine collision_util_bounce_one + module subroutine collision_util_construct_constraint_system(collider, nbody_system, param, constraint_system, phase) implicit none class(collision_basic), intent(inout) :: collider !! Collision system object diff --git a/src/collision/collision_resolve.f90 b/src/collision/collision_resolve.f90 index 5e0f2296a..bafbf7de8 100644 --- a/src/collision/collision_resolve.f90 +++ b/src/collision/collision_resolve.f90 @@ -359,6 +359,8 @@ module subroutine collision_resolve_mergeaddsub(nbody_system, param, t, status) plnew%rot(:, 1:nfrag) = fragments%rot(:, 1:nfrag) end if + call plnew%set_rhill(cb) + ! if (param%ltides) then ! plnew%Q = pl%Q(ibiggest) ! plnew%k2 = pl%k2(ibiggest) @@ -369,7 +371,6 @@ module subroutine collision_resolve_mergeaddsub(nbody_system, param, t, status) volume = 4.0_DP/3.0_DP * PI * plnew%radius(i)**3 plnew%density(i) = fragments%mass(i) / volume end do - call plnew%set_rhill(cb) select case(status) case(SUPERCATASTROPHIC) diff --git a/src/collision/collision_util.f90 b/src/collision/collision_util.f90 index a73238ef7..63d623de4 100644 --- a/src/collision/collision_util.f90 +++ b/src/collision/collision_util.f90 @@ -60,6 +60,29 @@ module subroutine collision_util_add_fragments_to_collider(self, nbody_system, p return end subroutine collision_util_add_fragments_to_collider + module subroutine collision_util_bounce_one(r,v,rcom,vcom,radius) + !! Author: David A. Minton + !! + !! Performs a "bounce" operation on a single body by reversing its velocity in a center of mass frame. + implicit none + ! Arguments + real(DP), dimension(:), intent(inout) :: r,v + real(DP), dimension(:), intent(in) :: rcom,vcom + real(DP), intent(in) :: radius + ! Internals + real(DP), dimension(NDIM) :: rrel, vrel, rnorm + + rrel(:) = r(:) - rcom(:) + vrel(:) = v(:) - vcom(:) + rnorm(:) = .unit. rrel(:) + ! Do the reflection + vrel(:) = vrel(:) - 2 * dot_product(vrel(:),rnorm(:)) * rnorm(:) + ! Shift the positions so that the collision doesn't immediately occur again + r(:) = r(:) + 0.5_DP * radius * rnorm(:) + v(:) = vcom(:) + vrel(:) + + return + end subroutine collision_util_bounce_one module subroutine collision_util_construct_constraint_system(collider, nbody_system, param, constraint_system, phase) !! Author: David A. Minton @@ -336,7 +359,7 @@ module subroutine collision_util_dealloc_impactors(self) self%v_unit(:) = 0.0_DP self%rbcom(:) = 0.0_DP self%vbcom(:) = 0.0_DP - self%rbimp(:) = 0.0_DP + self%rcimp(:) = 0.0_DP return end subroutine collision_util_dealloc_impactors @@ -667,8 +690,8 @@ module subroutine collision_util_set_coordinate_impactors(self) impactors%vc(:,1) = impactors%vb(:,1) - impactors%vbcom(:) impactors%vc(:,2) = impactors%vb(:,2) - impactors%vbcom(:) - ! Find the point of impact between the two bodies - impactors%rbimp(:) = impactors%rb(:,1) + impactors%radius(1) * impactors%y_unit(:) - impactors%rbcom(:) + ! Find the point of impact between the two bodies, defined as the location (in the collisional coordinate system) at the surface of body 1 along the line connecting the two bodies. + impactors%rcimp(:) = impactors%rb(:,1) + impactors%radius(1) * impactors%y_unit(:) - impactors%rbcom(:) ! Set the velocity direction as the "bounce" direction" for disruptions, and body 2's direction for hit and runs if (impactors%regime == COLLRESOLVE_REGIME_HIT_AND_RUN) then @@ -898,7 +921,7 @@ module subroutine collision_util_set_natural_scale_factors(self) ! Scale all dimensioned quantities of impactors and fragments impactors%rbcom(:) = impactors%rbcom(:) / collider%dscale impactors%vbcom(:) = impactors%vbcom(:) / collider%vscale - impactors%rbimp(:) = impactors%rbimp(:) / collider%dscale + impactors%rcimp(:) = impactors%rcimp(:) / collider%dscale impactors%rb(:,:) = impactors%rb(:,:) / collider%dscale impactors%vb(:,:) = impactors%vb(:,:) / collider%vscale impactors%rc(:,:) = impactors%rc(:,:) / collider%dscale @@ -950,7 +973,7 @@ module subroutine collision_util_set_original_scale_factors(self) ! Restore scale factors impactors%rbcom(:) = impactors%rbcom(:) * collider%dscale impactors%vbcom(:) = impactors%vbcom(:) * collider%vscale - impactors%rbimp(:) = impactors%rbimp(:) * collider%dscale + impactors%rcimp(:) = impactors%rcimp(:) * collider%dscale impactors%mass = impactors%mass * collider%mscale impactors%Gmass(:) = impactors%Gmass(:) * (collider%dscale**3/collider%tscale**2) diff --git a/src/fraggle/fraggle_generate.f90 b/src/fraggle/fraggle_generate.f90 index b28480e59..a5151d2d3 100644 --- a/src/fraggle/fraggle_generate.f90 +++ b/src/fraggle/fraggle_generate.f90 @@ -54,9 +54,34 @@ module subroutine fraggle_generate(self, nbody_system, param, t) end select call self%set_mass_dist(param) call self%disrupt(nbody_system, param, t, lfailure) + if (lfailure) then - call swiftest_io_log_one_message(COLLISION_LOG_OUT, "Fraggle failed to find an energy-losing solution. Treating this as a hit and run.") - call self%hitandrun(nbody_system, param, t) + call swiftest_io_log_one_message(COLLISION_LOG_OUT, "Fraggle failed to find an energy-losing solution. Simplifying the collisional model.") + + impactors%mass_dist(1) = impactors%mass(1) + impactors%mass_dist(2) = max(0.5_DP * impactors%mass(2), self%min_mfrag) + impactors%mass_dist(3) = impactors%mass(2) - impactors%mass_dist(2) + impactors%regime = COLLRESOLVE_REGIME_DISRUPTION + call self%set_mass_dist(param) + call self%disrupt(nbody_system, param, t, lfailure) + + if (lfailure) then + call swiftest_io_log_one_message(COLLISION_LOG_OUT, "Fraggle failed to find an energy-losing solution. Treating this as a bounce.") + call collision_util_bounce_one(impactors%rb(:,1),impactors%vb(:,1),impactors%rbcom(:),impactors%vbcom(:),impactors%radius(1)) + call collision_util_bounce_one(impactors%rb(:,2),impactors%vb(:,2),impactors%rbcom(:),impactors%vbcom(:),0.0_DP) + call impactors%set_coordinate_system() + call self%setup_fragments(2) + associate (fragments => self%fragments) + fragments%mass(1:2) = impactors%mass(1:2) + fragments%Gmass(1:2) = impactors%Gmass(1:2) + fragments%radius(1:2) = impactors%radius(1:2) + fragments%rb(:,1:2) = impactors%rb(:,1:2) + fragments%vb(:,1:2) = impactors%vb(:,1:2) + fragments%Ip(:,1:2) = impactors%Ip(:,1:2) + fragments%rot(:,1:2) = impactors%rot(:,1:2) + end associate + + end if end if associate (fragments => self%fragments) @@ -134,9 +159,11 @@ module subroutine fraggle_generate_disrupt(self, nbody_system, param, t, lfailur lfailure = .false. call self%get_energy_and_momentum(nbody_system, param, phase="before") self%fail_scale = fail_scale_initial - call fraggle_generate_pos_vec(self) - call fraggle_generate_rot_vec(self, nbody_system, param) - call fraggle_generate_vel_vec(self, nbody_system, param, lfailure) + call fraggle_generate_pos_vec(self, nbody_system, param, lfailure) + if (.not.lfailure) then + call fraggle_generate_rot_vec(self, nbody_system, param) + call fraggle_generate_vel_vec(self, nbody_system, param, lfailure) + end if if (.not.lfailure) then if (self%fragments%nbody /= nfrag_start) then @@ -255,7 +282,7 @@ module subroutine fraggle_generate_hitandrun(self, nbody_system, param, t) end subroutine fraggle_generate_hitandrun - module subroutine fraggle_generate_pos_vec(collider) + module subroutine fraggle_generate_pos_vec(collider, nbody_system, param, lfailure) !! Author: Jennifer L.L. Pouplin, Carlisle A. Wishard, and David A. Minton !! !! Initializes the position vectors of the fragments around the center of mass based on the collision style. @@ -264,7 +291,10 @@ module subroutine fraggle_generate_pos_vec(collider) !! regenerated if they overlap. implicit none ! Arguments - class(collision_fraggle), intent(inout) :: collider !! Fraggle collision system object + class(collision_fraggle), intent(inout) :: collider !! Fraggle collision system object + class(swiftest_nbody_system), intent(inout) :: nbody_system !! Swiftest nbody system object + class(swiftest_parameters), intent(inout) :: param !! Current run configuration parameters + logical, intent(out) :: lfailure !! Did the velocity computation fail? ! Internals real(DP) :: dis, direction, rdistance real(DP), dimension(NDIM,2) :: fragment_cloud_center @@ -273,14 +303,16 @@ module subroutine fraggle_generate_pos_vec(collider) real(DP), dimension(collider%fragments%nbody) :: mass_rscale, phi, theta, u integer(I4B) :: i, j, loop, istart logical :: lsupercat, lhitandrun - integer(I4B), parameter :: MAXLOOP = 20000 + integer(I4B), parameter :: MAXLOOP = 100 real(DP), parameter :: rdistance_scale_factor = 1.0_DP ! Scale factor to apply to distance scaling of cloud centers in the event of overlap ! The distance is chosen to be close to the original locations of the impactors ! but far enough apart to prevent a collisional cascade between fragments real(DP), parameter :: cloud_size_scale_factor = 3.0_DP ! Scale factor to apply to the size of the cloud relative to the distance from the impact point. ! A larger value puts more space between fragments initially + - associate(fragments => collider%fragments, impactors => collider%impactors, nfrag => collider%fragments%nbody) + associate(fragments => collider%fragments, impactors => collider%impactors, nfrag => collider%fragments%nbody, & + pl => nbody_system%pl, tp => nbody_system%tp, npl => nbody_system%pl%nbody, ntp => nbody_system%tp%nbody) lsupercat = (impactors%regime == COLLRESOLVE_REGIME_SUPERCATASTROPHIC) lhitandrun = (impactors%regime == COLLRESOLVE_REGIME_HIT_AND_RUN) @@ -300,30 +332,23 @@ module subroutine fraggle_generate_pos_vec(collider) mass_rscale(:) = (mass_rscale(:) + 1.0_DP) / 2 mass_rscale(:) = mass_rscale(:) * (fragments%mtot / fragments%mass(1:nfrag))**(0.125_DP) ! The power is arbitrary. It just gives the velocity a small mass dependence mass_rscale(:) = mass_rscale(:) / maxval(mass_rscale(:)) + istart = 3 do loop = 1, MAXLOOP if (.not.any(loverlap(:))) exit - if (lhitandrun) then + if (lhitandrun) then ! Keep the target unchanged and place the largest fragment at rdistance away from the projectile along its trajectory fragment_cloud_radius(:) = impactors%radius(:) fragment_cloud_center(:,1) = impactors%rc(:,1) fragment_cloud_center(:,2) = impactors%rc(:,2) + rdistance * impactors%bounce_unit(:) - else if (lsupercat) then - fragment_cloud_center(:,1) = impactors%rc(:,1) - fragment_cloud_center(:,2) = impactors%rc(:,2) - fragment_cloud_radius(:) = cloud_size_scale_factor * rdistance - else - fragment_cloud_center(:,1) = impactors%rbimp(:) - impactors%radius(1) * impactors%y_unit(:) - fragment_cloud_center(:,2) = impactors%rbimp(:) + impactors%radius(2) * impactors%y_unit(:) + else ! Keep the first and second bodies at approximately their original location, but so as not to be overlapping + fragment_cloud_center(:,1) = impactors%rc(:,1) + fragment_cloud_center(:,2) = impactors%rcimp(:) + 1.001_DP * max(fragments%radius(2),impactors%radius(2)) * impactors%y_unit(:) fragment_cloud_radius(:) = cloud_size_scale_factor * rdistance end if - if (lsupercat) then - istart = 1 - else - fragments%rc(:,1) = fragment_cloud_center(:,1) - istart = 2 - end if + fragments%rc(:,1) = fragment_cloud_center(:,1) + fragments%rc(:,2) = fragment_cloud_center(:,2) - do i = istart, nfrag + do i = 1, nfrag if (loverlap(i)) then call random_number(phi(i)) call random_number(theta(i)) @@ -331,6 +356,7 @@ module subroutine fraggle_generate_pos_vec(collider) end if end do + ! Randomly place the n>2 fragments inside their cloud until none are overlapping do concurrent(i = istart:nfrag, loverlap(i)) j = fragments%origin_body(i) @@ -348,27 +374,48 @@ module subroutine fraggle_generate_pos_vec(collider) fragments%rc(:,i) = fragments%rc(:,i) + fragment_cloud_center(:,j) ! Make sure that the fragments are positioned away from the impact point - direction = dot_product(fragments%rc(:,i) - impactors%rbimp(:), fragment_cloud_center(:,j) - impactors%rbimp(:)) + direction = dot_product(fragments%rc(:,i) - impactors%rcimp(:), fragment_cloud_center(:,j) - impactors%rcimp(:)) if (direction < 0.0_DP) then fragments%rc(:,i) = fragments%rc(:,i) - fragment_cloud_center(:,j) fragments%rc(:,i) = -fragments%rc(:,i) + fragment_cloud_center(:,j) end if - end do + + ! Because body 1 and 2 are initialized near the original impactor positions, then if these bodies are still overlapping + ! when the rest are not, we will randomly walk their position in space so as not to move them too far from their starting position + if (all(.not.loverlap(istart:nfrag)) .and. any(loverlap(1:istart-1))) then + do concurrent(i = 1:istart-1,loverlap(i)) + dis = 0.1_DP * fragments%radius(i) + fragments%rmag(i) = dis * u(i)**(THIRD) + fragments%rc(1,i) = fragments%rmag(i) * sin(theta(i)) * cos(phi(i)) + fragments%rc(2,i) = fragments%rmag(i) * sin(theta(i)) * sin(phi(i)) + fragments%rc(3,i) = fragments%rmag(i) * cos(theta(i)) + end do + end if + fragments%rmag(:) = .mag. fragments%rc(:,:) ! Check for any overlapping bodies. loverlap(:) = .false. do j = 1, nfrag + ! Check for overlaps between fragments do i = j + 1, nfrag dis = .mag.(fragments%rc(:,j) - fragments%rc(:,i)) loverlap(i) = loverlap(i) .or. (dis <= (fragments%radius(i) + fragments%radius(j))) loverlap(j) = loverlap(j) .or. (dis <= (fragments%radius(i) + fragments%radius(j))) end do + ! Check for overlaps with existing bodies that are not involved in the collision + do i = 1, npl + if (any(impactors%id(:) == i)) cycle + dis = .mag. (fragments%rc(:,j) - (pl%rb(:,i) / collider%dscale - impactors%rbcom(:))) + loverlap(j) = loverlap(j) .or. (dis <= (pl%radius(i) / collider%dscale + fragments%radius(j))) + end do end do rdistance = rdistance * collider%fail_scale end do + lfailure = any(loverlap(:)) + call collision_util_shift_vector_to_origin(fragments%mass, fragments%rc) call collider%set_coordinate_system() @@ -392,43 +439,16 @@ module subroutine fraggle_generate_rot_vec(collider, nbody_system, param) class(collision_fraggle), intent(inout) :: collider !! Fraggle collision system object class(swiftest_nbody_system), intent(inout) :: nbody_system !! Swiftest nbody system object class(swiftest_parameters), intent(inout) :: param !! Current run configuration parameters - ! Internals - real(DP), dimension(NDIM) :: Lbefore, Lafter, L_spin, rotdir - real(DP) :: v_init, v_final, mass_init, mass_final, rotmag, dKE, KE_init, KE_final - real(DP), parameter :: random_scale_factor = 0.01_DP !! The relative scale factor to apply to the random component of the rotation vector - integer(I4B) :: i - logical :: lhitandrun associate(fragments => collider%fragments, impactors => collider%impactors, nfrag => collider%fragments%nbody) - lhitandrun = (impactors%regime == COLLRESOLVE_REGIME_HIT_AND_RUN) - - ! We will start by assuming that kinetic energy gets partitioned such that the change in kinetic energy of body 1 is equal to the - ! change in kinetic energy between bodies 2 and all fragments. This will then be used to compute a torque on body/fragment 1. - ! All other fragments will be given a random velocity with a magnitude scaled by the change in the orbital system angular momentum - mass_init = impactors%mass(2) - mass_final = sum(fragments%mass(2:nfrag)) - v_init = .mag.(impactors%vb(:,2) - impactors%vb(:,1)) - KE_init = 0.5_DP * mass_init * v_init**2 - - ! Initialize fragment rotations and velocities to be pre-impact rotations in order to compute the energy. This will get adjusted later + ! Initialize fragment rotations and velocities to be pre-impact rotation for body 1, and randomized for bodies >1 and scaled to the original rotation. + ! This will get updated later when conserving angular momentum fragments%rot(:,1) = impactors%rot(:,1) - fragments%vc(:,1) = impactors%vc(:,1) - do concurrent(i = 2:nfrag) - fragments%rot(:,i) = impactors%rot(:,2) - fragments%vc(:,i) = impactors%vc(:,2) - end do - - ! Initialize the largest body with the combined spin angular momentum of the imapactors - fragments%rot(:,1) = (impactors%L_spin(:,1) + impactors%L_spin(:,2)) / (fragments%mass(1) * fragments%Ip(3,1) * fragments%radius(1)) - fragments%rot(:,2:nfrag) = 0.0_DP + call random_number(fragments%rot(:,2:nfrag)) + fragments%rot(:,2:nfrag) = fragments%rot(:,2:nfrag) * .mag.impactors%rot(:,2) fragments%rotmag(:) = .mag.fragments%rot(:,:) - do i = 1,nfrag - if (fragments%rotmag(i) > collider%max_rot) then - fragments%rotmag(i) = 0.5_DP * collider%max_rot - fragments%rot(:,i) = fragments%rotmag(i) * .unit.fragments%rot(:,i) - end if - end do + end associate return @@ -448,22 +468,22 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu class(swiftest_parameters), intent(inout) :: param !! Current run configuration parameters logical, intent(out) :: lfailure !! Did the velocity computation fail? ! Internals - real(DP), parameter :: ENERGY_SUCCESS_METRIC = 1.0e-2_DP !! Relative energy error to accept as a success (success also must be energy-losing in addition to being within the metric amount) - real(DP), parameter :: MOMENTUM_SUCCESS_METRIC = 1.0e-12_DP !! Relative angular momentum error to accept as a success (should be *much* stricter than energy) - integer(I4B) :: i, j, loop, try, istart, nfrag, nsteps + real(DP), parameter :: ENERGY_SUCCESS_METRIC = 1.0e-4_DP !! Relative energy error to accept as a success (success also must be energy-losing in addition to being within the metric amount) + real(DP) :: MOMENTUM_SUCCESS_METRIC = 10*epsilon(1.0_DP) !! Relative angular momentum error to accept as a success (should be *much* stricter than energy) + integer(I4B) :: i, j, loop, try, istart, nfrag, nsteps, nsteps_best logical :: lhitandrun, lsupercat real(DP), dimension(NDIM) :: vimp_unit, rimp, vrot, L_residual, L_residual_unit, dL, drot, rot_new - real(DP) :: vimp, vmag, vesc, dE, E_residual, ke_min, ke_avail, ke_remove, dE_best, E_residual_best, fscale, dE_metric, dM, mfrag, dL_metric, dL_best + real(DP) :: vimp, vmag, vesc, dE, E_residual, E_residual_best, E_residual_last, ke_min, ke_avail, ke_remove, dE_best, fscale, dE_metric, mfrag, dL_metric, dL_best, rn integer(I4B), dimension(collider%fragments%nbody) :: vsign - real(DP), dimension(collider%fragments%nbody) :: vscale, volume + real(DP), dimension(collider%fragments%nbody) :: vscale + real(DP), parameter :: L_ROT_VEL_RATIO = 0.2_DP ! Ratio of angular momentum to put into rotation relative to velocity shear of fragments ! For the initial "guess" of fragment velocities, this is the minimum and maximum velocity relative to escape velocity that the fragments will have real(DP) :: vmin_guess = 1.01_DP real(DP) :: vmax_guess real(DP) :: delta_v, GC - integer(I4B), parameter :: MAXLOOP = 50 - integer(I4B), parameter :: MAXTRY = 50 - real(DP), parameter :: MAX_REDUCTION_RATIO = 0.1_DP ! Ratio of difference between first and second fragment mass to remove from the largest fragment in case of a failure - real(DP), parameter :: ROT_MAX_FRAC = 0.001_DP !! Fraction of difference between current rotation and maximum to add when angular momentum budget gets too high + integer(I4B), parameter :: MAXINNER = 50 + integer(I4B), parameter :: MAXOUTER = 10 + integer(I4B), parameter :: MAXANGMTM = 10000 class(collision_fraggle), allocatable :: collider_local character(len=STRMAX) :: message @@ -477,9 +497,6 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu allocate(collider_local, source=collider) associate(fragments => collider_local%fragments) - volume(:) = 4.0_DP / 3.0_DP * PI * (fragments%radius(:))**3 - fragments%density(:) = fragments%mass(:) / volume(:) - ! The fragments will be divided into two "clouds" based on identified origin body. ! These clouds will collectively travel like two impactors bouncing off of each other. where(fragments%origin_body(:) == 1) @@ -510,12 +527,13 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu dE_metric = huge(1.0_DP) dE_best = huge(1.0_DP) dL_best = huge(1.0_DP) - - outer: do try = 1, MAXTRY + nsteps_best = 0 + nsteps = 0 + outer: do try = 1, MAXOUTER ! Scale the magnitude of the velocity by the distance from the impact point ! This will reduce the chances of fragments colliding with each other immediately, and is more physically correct do concurrent(i = 2:nfrag) - rimp(:) = fragments%rc(:,i) - impactors%rbimp(:) + rimp(:) = fragments%rc(:,i) - impactors%rcimp(:) vscale(i) = .mag. rimp(:) / (.mag. (impactors%rb(:,2) - impactors%rb(:,1))) end do vscale(1) = 0.9_DP * minval(vscale(2:nfrag)) @@ -537,7 +555,7 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu end if else vmag = vscale(i) - rimp(:) = fragments%rc(:,i) - impactors%rbimp(:) + rimp(:) = fragments%rc(:,i) - impactors%rcimp(:) vimp_unit(:) = .unit. (rimp(:) + vsign(i) * impactors%bounce_unit(:)) fragments%vc(:,i) = vmag * vimp_unit(:) + vrot(:) end if @@ -549,45 +567,51 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu call fragments%set_coordinate_system() ke_min = 0.5_DP * fragments%mtot * vesc**2 - inner: do loop = 1, MAXLOOP - nsteps = loop * try + E_residual = huge(1.0_DP) + inner: do loop = 1, MAXINNER + nsteps = nsteps + 1 ! Try to put residual angular momentum into the spin, but if this would go past the spin barrier, then put it into velocity shear instead - angmtm: do j = 1, MAXTRY + angmtm: do j = 1, MAXANGMTM call collider_local%get_energy_and_momentum(nbody_system, param, phase="after") L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1)) - dL_metric = .mag.L_residual(:) / .mag.collider_local%L_total(:,1) + dL_metric = .mag.L_residual(:) / .mag.(collider_local%L_total(:,1)) - if (dL_metric <= MOMENTUM_SUCCESS_METRIC) exit angmtm + if (dL_metric / MOMENTUM_SUCCESS_METRIC <= 1.0_DP) exit angmtm L_residual_unit(:) = .unit. L_residual(:) do i = 1, fragments%nbody mfrag = sum(fragments%mass(i:fragments%nbody)) - drot(:) = -L_residual(:) / (fragments%mtot * fragments%Ip(3,i) * fragments%radius(i)**2) + drot(:) = -L_ROT_VEL_RATIO * L_residual(:) / (fragments%mtot * fragments%Ip(3,i) * fragments%radius(i)**2) rot_new(:) = fragments%rot(:,i) + drot(:) if (.mag.rot_new(:) < collider_local%max_rot) then fragments%rot(:,i) = rot_new(:) fragments%rotmag(i) = .mag.fragments%rot(:,i) else ! We would break the spin barrier here. Put less into spin and more into velocity shear. - dL(:) = -L_residual(:) * fragments%mass(i) / fragments%mtot + drot(:) * fragments%Ip(3,i) * fragments%mass(i) * fragments%radius(i)**2 - call fraggle_generate_velocity_torque(dL, fragments%mass(i), fragments%rc(:,i), fragments%vc(:,i)) - call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc) - fragments%vmag(i) = .mag.fragments%vc(:,i) - if (i >= istart) then call random_number(drot) - drot(:) = (0.5_DP * collider_local%max_rot - fragments%rotmag(i)) * 2 * (drot(:) - 0.5_DP) + call random_number(rn) + drot(:) = (rn * collider_local%max_rot - fragments%rotmag(i)) * 2 * (drot(:) - 0.5_DP) fragments%rot(:,i) = fragments%rot(:,i) + drot(:) fragments%rotmag(i) = .mag.fragments%rot(:,i) if (fragments%rotmag(i) > collider%max_rot) then fragments%rotmag(i) = 0.5_DP * collider%max_rot fragments%rot(:,i) = fragments%rotmag(i) * .unit. fragments%rot(:,i) end if + else + drot(:) = 0.0_DP end if - end if + + dL(:) = -L_residual(:) * fragments%mass(i) / fragments%mtot - drot(:) * fragments%Ip(3,i) * fragments%mass(i) * fragments%radius(i)**2 + call fraggle_generate_velocity_torque(dL, fragments%mass(i), fragments%rc(:,i), fragments%vc(:,i)) + call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc) + fragments%vmag(i) = .mag.fragments%vc(:,i) + end do end do angmtm + call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc) + call fragments%set_coordinate_system() call collider_local%get_energy_and_momentum(nbody_system, param, phase="after") ke_avail = 0.0_DP do i = 1, fragments%nbody @@ -595,6 +619,7 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu end do dE = collider_local%te(2) - collider_local%te(1) + E_residual_last = E_residual E_residual = dE + impactors%Qloss L_residual(:) = (collider_local%L_total(:,2) - collider_local%L_total(:,1)) @@ -606,6 +631,7 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu E_residual_best = E_residual dE_best = dE dL_best = dL_metric + nsteps_best = nsteps do concurrent(i = 1:fragments%nbody) fragments%vb(:,i) = fragments%vc(:,i) + impactors%vbcom(:) @@ -614,6 +640,8 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu if (allocated(collider%fragments)) deallocate(collider%fragments) allocate(collider%fragments, source=fragments) dE_metric = abs(E_residual) / impactors%Qloss + else if (abs(E_residual) >= abs(E_residual_last)) then + exit inner ! We are no longer converging on a solution. At this point, it's best to give up end if if ((dE_best < 0.0_DP) .and. (dE_metric <= ENERGY_SUCCESS_METRIC * try)) exit outer ! As the tries increase, we relax the success metric. What was once a failure might become a success end if @@ -623,43 +651,26 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu fscale = sqrt((max(fragments%ke_orbit_tot - ke_remove, 0.0_DP))/fragments%ke_orbit_tot) fragments%vc(:,:) = fscale * fragments%vc(:,:) fragments%vmag(:) = .mag.fragments%vc(:,:) + fragments%rc(:,:) = 1.0_DP / fscale * fragments%rc(:,:) + fragments%rmag(:) = .mag.fragments%rc(:,:) ! Update the unit vectors and magnitudes for the fragments based on their new orbits and rotations call collision_util_shift_vector_to_origin(fragments%mass, fragments%vc) call fragments%set_coordinate_system() - end do inner - ! We didn't converge. Reset the fragment positions and velocities and try a new configuration with some slightly different parameters - ! Reduce the fragment masses and add it to the largest remenant and try again - if (any(fragments%mass(2:nfrag) > collider%min_mfrag)) then - do i = 2, nfrag - if (fragments%mass(i) > collider%min_mfrag) then - dM = min(MAX_REDUCTION_RATIO * fragments%mass(i), fragments%mass(i) - collider%min_mfrag) - fragments%mass(i) = fragments%mass(i) - dM - fragments%mass(1) = fragments%mass(1) + dM - end if - end do - else - exit outer - end if - fragments%Gmass(:) = GC * fragments%mass(:) - - volume(:) = fragments%mass(:) / fragments%density(:) - fragments%radius(:) = (3._DP * volume(:) / (4._DP * PI))**(THIRD) - call fragments%reset() - call fraggle_generate_pos_vec(collider_local) + call fraggle_generate_pos_vec(collider_local, nbody_system, param, lfailure) + if (lfailure) exit call fraggle_generate_rot_vec(collider_local, nbody_system, param) ! Increase the spatial size factor to get a less dense cloud - collider_local%fail_scale = collider_local%fail_scale*1.001_DP + collider_local%fail_scale = collider_local%fail_scale * 1.001_DP ! Bring the minimum and maximum velocities closer together delta_v = 0.125_DP * (vmax_guess - vmin_guess) vmin_guess = vmin_guess + delta_v vmax_guess = vmax_guess - delta_v - end do outer lfailure = (dE_best > 0.0_DP) .or. (dL_best > MOMENTUM_SUCCESS_METRIC) @@ -673,6 +684,8 @@ module subroutine fraggle_generate_vel_vec(collider, nbody_system, param, lfailu else call swiftest_io_log_one_message(COLLISION_LOG_OUT,"Fraggle velocity calculation converged after " // trim(adjustl(message)) // " steps.") end if + write(message,*) nsteps_best + call swiftest_io_log_one_message(COLLISION_LOG_OUT,"Best solution came after " // trim(adjustl(message)) // " steps.") end associate end associate diff --git a/src/fraggle/fraggle_module.f90 b/src/fraggle/fraggle_module.f90 index b113dc2b0..01412154d 100644 --- a/src/fraggle/fraggle_module.f90 +++ b/src/fraggle/fraggle_module.f90 @@ -16,18 +16,18 @@ module fraggle public type, extends(collision_basic) :: collision_fraggle - real(DP) :: fail_scale !! Scale factor to apply to distance values in the position model when overlaps occur. + real(DP) :: fail_scale !! Scale factor to apply to distance values in the position model when overlaps occur. contains - procedure :: disrupt => fraggle_generate_disrupt !! Generates a system of fragments in barycentric coordinates that conserves energy and momentum. - procedure :: generate => fraggle_generate !! A simple disruption models that does not constrain energy loss in collisions - procedure :: hitandrun => fraggle_generate_hitandrun !! Generates either a pure hit and run, or one in which the runner is disrupted - procedure :: set_mass_dist => fraggle_util_set_mass_dist !! Sets the distribution of mass among the fragments depending on the regime type + procedure :: disrupt => fraggle_generate_disrupt !! Generates a system of fragments in barycentric coordinates that conserves energy and momentum. + procedure :: generate => fraggle_generate !! A simple disruption models that does not constrain energy loss in collisions + procedure :: hitandrun => fraggle_generate_hitandrun !! Generates either a pure hit and run, or one in which the runner is disrupted + procedure :: set_mass_dist => fraggle_util_set_mass_dist !! Sets the distribution of mass among the fragments depending on the regime type end type collision_fraggle interface module subroutine fraggle_generate(self, nbody_system, param, t) implicit none - class(collision_fraggle), intent(inout) :: self !! Fraggle fragment system object + class(collision_fraggle), intent(inout) :: self !! Fraggle fragment system object class(base_nbody_system), intent(inout) :: nbody_system !! Swiftest nbody system object class(base_parameters), intent(inout) :: param !! Current run configuration parameters real(DP), intent(in) :: t !! The time of the collision @@ -50,9 +50,12 @@ module subroutine fraggle_generate_hitandrun(self, nbody_system, param, t) real(DP), intent(in) :: t !! Time of collision end subroutine fraggle_generate_hitandrun - module subroutine fraggle_generate_pos_vec(collider) + module subroutine fraggle_generate_pos_vec(collider, nbody_system, param, lfailure) implicit none - class(collision_fraggle), intent(inout) :: collider !! Fraggle ollision system object + class(collision_fraggle), intent(inout) :: collider !! Fraggle collision system object + class(swiftest_nbody_system), intent(inout) :: nbody_system !! Swiftest nbody system object + class(swiftest_parameters), intent(inout) :: param !! Current run configuration parameters + logical, intent(out) :: lfailure !! Did the velocity computation fail? end subroutine fraggle_generate_pos_vec module subroutine fraggle_generate_rot_vec(collider, nbody_system, param) diff --git a/src/helio/helio_util.f90 b/src/helio/helio_util.f90 index 7d887a6b2..94b464502 100644 --- a/src/helio/helio_util.f90 +++ b/src/helio/helio_util.f90 @@ -22,9 +22,9 @@ module subroutine helio_util_setup_initialize_system(self, param) class(swiftest_parameters), intent(inout) :: param !! Current run configuration parameters call swiftest_util_setup_initialize_system(self, param) - call self%pl%h2b(self%cb) - call self%tp%h2b(self%cb) call self%pl%sort("mass", ascending=.false.) + call self%pl%vh2vb(self%cb) + call self%tp%h2b(self%cb) ! Make sure that the discard list gets allocated initially call self%tp_discards%setup(0, param) diff --git a/src/swiftest/swiftest_driver.f90 b/src/swiftest/swiftest_driver.f90 index eed4fe446..8cc7f9a2a 100644 --- a/src/swiftest/swiftest_driver.f90 +++ b/src/swiftest/swiftest_driver.f90 @@ -106,7 +106,6 @@ program swiftest_driver if (idump == dump_cadence) then idump = 0 call nbody_system%dump(param) - end if call integration_timer%report(message="Integration steps:", unit=display_unit, nsubsteps=istep_out) diff --git a/src/swiftest/swiftest_io.f90 b/src/swiftest/swiftest_io.f90 index a2b604357..af434b2c5 100644 --- a/src/swiftest/swiftest_io.f90 +++ b/src/swiftest/swiftest_io.f90 @@ -129,7 +129,11 @@ module subroutine swiftest_io_conservation_report(self, param, lterminal) associate(nbody_system => self, pl => self%pl, cb => self%cb, npl => self%pl%nbody, display_unit => param%display_unit, nc => self%system_history%nc) - call pl%vb2vh(cb) + select type(self) + class is (helio_nbody_system) ! Don't convert vh to vb for Helio-based integrators, because they are already have that calculated + class is (whm_nbody_system) + call pl%vh2vb(cb) + end select call pl%rh2rb(cb) call nbody_system%get_energy_and_momentum(param) diff --git a/src/swiftest/swiftest_util.f90 b/src/swiftest/swiftest_util.f90 index c1efe888f..05f9383ee 100644 --- a/src/swiftest/swiftest_util.f90 +++ b/src/swiftest/swiftest_util.f90 @@ -1808,7 +1808,7 @@ module subroutine swiftest_util_rearray_pl(self, nbody_system, param) class(encounter_list), allocatable :: plplenc_old logical :: lencounter - associate(pl => self, tp => nbody_system%tp, pl_adds => nbody_system%pl_adds) + associate(pl => self, tp => nbody_system%tp, cb => nbody_system%cb, pl_adds => nbody_system%pl_adds) npl = pl%nbody nadd = pl_adds%nbody @@ -1858,6 +1858,8 @@ module subroutine swiftest_util_rearray_pl(self, nbody_system, param) call pl%sort("mass", ascending=.false.) call pl%flatten(param) + call pl%set_rhill(cb) + ! Reset the kinship trackers call pl%reset_kinship([(i, i=1, npl)]) @@ -3038,6 +3040,24 @@ module subroutine swiftest_util_snapshot_system(self, param, nbody_system, t, ar ! Internals class(swiftest_nbody_system), allocatable :: snapshot + ! To allow for runs to be restarted in a bit-identical way, we'll need to run the same coordinate conversion routines we would run upon restarting + select type(pl => nbody_system%pl) + class is (whm_pl) + call pl%h2j(nbody_system%cb) + class is (helio_pl) + call pl%vh2vb(nbody_system%cb) + end select + + select type(tp => nbody_system%tp) + class is (helio_tp) + select type(cb => nbody_system%cb) + class is (helio_cb) + call tp%vh2vb(vbcb = -cb%ptbeg) + end select + end select + + call nbody_system%pl%set_rhill(nbody_system%cb) + ! Take a minimal snapshot wihout all of the extra storage objects allocate(snapshot, mold=nbody_system) allocate(snapshot%cb, source=nbody_system%cb )