Updated Chambers 2013 initial conditions

examples/Chambers2013/init_cond.py

@@ -2,6 +2,7 @@
 import swiftest
 import numpy as np
 from numpy.random import default_rng
+import matplotlib.pyplot as plt
 
 # Initialize simulation object
 sim = swiftest.Simulation(compute_conservation_values=True,  rotation=True, init_cond_format="EL",collision_model="fraggle",encounter_save="none")
@@ -11,20 +12,65 @@
 Ns = 140
 Mb = 2.8e-7 * 14 / Nb
 Ms = 2.8e-8 * 140 / Ns
-dens = 3000.0 / (sim.param['MU2KG'] / sim.param['DU2M']**3)
-Rb = (3 * Mb / (4 * np.pi * dens) )**(1.0 / 3.0)
-Rs = (3 * Ms / (4 * np.pi * dens) )**(1.0 / 3.0)
+dens = 3000.0 / (sim.MU2KG / sim.DU2M**3)
+
 mtiny = 1e-2 * Ms
 mininum_fragment_mass = 1e-4 * Ms
-rng = default_rng(seed=170834)
+rng = default_rng(seed=3031179)
+
+runname = "Chambers (2013)"
+def a_profile(n_bodies):
+    """
+    Generates the profile described in Sec. 2.1 of Chambers:  
+    
+    *In all cases, the surface density R = 8 g/cm2 at 1 AU, varying as a**(-3/2), where a is the orbital semi-major axis. 
+    The region with a < 0.7 AU deviates from this law, declining linearly with decreasing distance until R = 0 at 0.3 AU. 
+    The outer edge of the disk is 2 AU in all cases.
+    """
+    def sample(r_inner, r_break, r_outer, slope1, slope2):
+       """ 
+       Define the functions to pull random semi-major axes from a distribution using a rejection sampling method
+       This defines a 2-slope model with a break at r_break
+       Based on (https://stackoverflow.com/questions/66874819/random-numbers-with-user-defined-continuous-probability-distribution)
+       """
+       while True:
+          x = rng.uniform(r_inner, r_outer, size=1)
+
+          # The proportionality factor A ensures that the PDF approaches the same value on either side of the break point
+          # Assumes the break point is the max of the PDF
+          if x < r_break:
+              slope = slope1 + 1
+              A = 1.0
+          else:
+              slope = slope2 + 1
+              A = r_break**(slope1-slope2) 
+          y = rng.uniform(0, A*r_break**slope, size=1)
+          pdf = A*x**(slope)
+          if (y < pdf):
+                return x
+
+    a_inner = 0.3
+    a_break = 0.7
+    a_outer = 2.0
+    slope1 = 1.0
+    slope2 = -1.5
+
+    a_vals = np.zeros(n_bodies)
+    for k in range(n_bodies):
+        a_vals[k] = sample(a_inner, a_break, a_outer, slope1, slope2)
+    return a_vals
 
 # Define the initial orbital elements of the big and small bodies
-avalb = rng.uniform(0.3, 2.0, Nb)
-avals = rng.uniform(0.3, 2.0, Ns)
-evalb = rng.uniform(0.0, 0.01, Nb)
-evals = rng.uniform(0.0, 0.01, Ns)
-incvalb = rng.uniform(0.0, 0.005 * 180 / np.pi, Nb)
-incvals = rng.uniform(0.0, 0.005 * 180 / np.pi, Ns)
+avalb = a_profile(Nb)
+avals = a_profile(Ns)
+
+esigma = 0.01
+isigma = np.rad2deg(0.5 * esigma)
+evalb = rng.rayleigh(scale=esigma, size=Nb)
+evals = rng.rayleigh(scale=esigma, size=Ns)
+incvalb = rng.rayleigh(scale=isigma, size=Nb)
+incvals = rng.rayleigh(scale=isigma, size=Ns)
+
 capomvalb = rng.uniform(0.0, 360.0, Nb)
 capomvals = rng.uniform(0.0, 360.0, Ns)
 omegavalb = rng.uniform(0.0, 360.0, Nb)
@@ -35,20 +81,63 @@
 Ipvals = np.full((Ns,3), 0.4)
 rotvalb = np.zeros_like(Ipvalb)
 rotvals = np.zeros_like(Ipvals)
-GMvalb = np.full(Nb, Mb * sim.GU)
-GMvals = np.full(Ns, Ms * sim.GU)
-Rvalb = np.full(Nb, Rb)
-Rvals = np.full(Ns, Rs)
+
+noise_digits = 4 # Approximately the number of digits of precision to vary the mass values to avoid duplicate masses
+epsilon = np.finfo(float).eps
+Mnoiseb = 1.0 + 10**noise_digits * rng.uniform(-epsilon,epsilon, Nb)
+Mnoises = 1.0 + 10**noise_digits * rng.uniform(-epsilon,epsilon, Ns)
+
+Mvalb = Mb * Mnoiseb
+Mvals = Ms * Mnoises
+
+Rvalb = (3 * Mvalb / (4 * np.pi * dens) )**(1.0 / 3.0)
+Rvals = (3 * Mvals / (4 * np.pi * dens) )**(1.0 / 3.0)
 
 # Give the bodies unique names
 nameb = [f"Big{i:03}" for i in range(Nb)]
 names = [f"Small{i:03}" for i in range(Ns)]
 
 # Add the modern planets and the Sun using the JPL Horizons Database.
 sim.add_solar_system_body(["Sun","Jupiter","Saturn","Uranus","Neptune"])
-sim.add_body(name=nameb, a=avalb, e=evalb, inc=incvalb, capom=capomvalb, omega=omegavalb, capm=capmvalb, Gmass=GMvalb, radius=Rvalb, rot=rotvalb, Ip=Ipvalb)
-sim.add_body(name=names, a=avals, e=evals, inc=incvals, capom=capomvals, omega=omegavals, capm=capmvals, Gmass=GMvals, radius=Rvals, rot=rotvals, Ip=Ipvals)
+sim.add_body(name=nameb, a=avalb, e=evalb, inc=incvalb, capom=capomvalb, omega=omegavalb, capm=capmvalb, mass=Mvalb, radius=Rvalb, rot=rotvalb, Ip=Ipvalb)
+sim.add_body(name=names, a=avals, e=evals, inc=incvals, capom=capomvals, omega=omegavals, capm=capmvals, mass=Mvals, radius=Rvals, rot=rotvals, Ip=Ipvals)
 sim.set_parameter(mtiny=mtiny, minimum_fragment_mass=mininum_fragment_mass)
 
-sim.run(tstop=3e8, dt=6.0875/365.25, istep_out=60000, dump_cadence=1)
+sim.set_parameter(tstop=3e8, dt=6.0875/365.25, istep_out=60000, dump_cadence=10)
+sim.clean()
+sim.write_param()
+
+# Plot the initial conditions
+fig = plt.figure(figsize=(8.5, 11))
+ax1 = plt.subplot(2, 1, 1)
+ax2 = plt.subplot(2, 1, 2)
+fig.suptitle(runname)
+
+ic = sim.init_cond.isel(time=0)
+radius = ic['radius'].values
+markersize = radius * 4e5
+markercolor = np.where(radius < 2.0e-5, np.full_like(markersize, "black",dtype=object), np.full_like(markersize,"red",dtype=object))
+a = ic['a'].values
+e = ic['e'].values
+inc = ic['inc'].values
+
+ax1.scatter(a, e, s=markersize, color=markercolor, label=None)
+ax1.set_xlabel("Semimajor Axis (AU)")
+ax1.set_ylabel("Eccentricity")
+ax1.set_xlim([0.0, 2.5])
+ax1.set_ylim([0.0,4*esigma])
+
+ax2.scatter(a, inc, s=markersize, color=markercolor, label=None)
+ax2.set_xlabel("Semimajor Axis (AU)")
+ax2.set_ylabel("Inclination (deg)")
+ax2.set_xlim([0.0, 2.5])
+ax2.set_ylim([0.0,8*isigma])
+
+ax1.scatter(-1, -1, label='planetesimal', color='black')
+ax1.scatter(-1, -1, label='embryo', color='red') 
+ax1.legend(loc='upper right', frameon=False)
+
+plt.tight_layout()
+plt.show()
+fig.savefig('initial_conditions.png',dpi=300)
 

examples/Chambers2013/scattermovie.py

@@ -4,11 +4,12 @@
 import matplotlib.pyplot as plt
 from matplotlib import animation
 import matplotlib.colors as mcolors
+plt.switch_backend('agg')
 
 titletext = "Chambers (2013)"
 valid_plot_styles = ["aescatter", "aiscatter"]
-xlim={"aescatter" : (0.0, 2.25),
-      "aiscatter" : (0.0, 2.25)}
+xlim={"aescatter" : (0.0, 2.5),
+      "aiscatter" : (0.0, 2.5)}
 ylim={"aescatter" : (0.0, 1.0),
       "aiscatter" : (0.0, 40.0)}
 xlabel={"aescatter": "Semimajor axis (AU)",
@@ -17,7 +18,7 @@
         "aiscatter": "Inclination (deg)"}
 
 
-plot_style = valid_plot_styles[1]
+plot_style = valid_plot_styles[0]
 framejump = 1
 animation_file = f"Chambers2013-{plot_style}.mp4"
 
@@ -45,19 +46,7 @@ def __init__(self, ds, param):
         self.ax.set_xlim(xlim[plot_style])
         self.ax.set_ylim(ylim[plot_style])
         fig.add_axes(self.ax)
-        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}')
-
-    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):
+
         # First frame
         """Initial drawing of the scatter plot."""
         t, npl, pl, radmarker, origin = next(self.data_stream(0))
@@ -73,16 +62,31 @@ def setup_plot(self):
         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)
+
         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}')
+
+    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)
-            name_good = d['name'].where(d['name'] != "Sun", drop=True)
-            d = d.sel(name=name_good)
+            n=len(d['name'])
+            d = d.isel(name=range(1,n))
             d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale
 
             t = d['time'].values