plotting improvements (for presentation)

analyze.py

@@ -2,7 +2,9 @@ import matplotlib.pyplot as plt
 from rebound import SimulationArchive, Particle, Simulation
 
 from extradata import ExtraData
-from utils import filename_from_argv
+from utils import filename_from_argv, plot_settings
+
+plot_settings()
 
 fn = filename_from_argv()
 sa = SimulationArchive(str(fn.with_suffix(".bin")))
@@ -28,7 +30,9 @@ for name, d in data.items():
     if False:
         plt.scatter(times, values, label=name, s=.9)
     else:
-        plt.plot(times, values, label=name)
+        plt.plot(times, values, label=name, linewidth=0.6)
 # plt.legend()
 # OrbitPlot(sim, slices=1)
+plt.tight_layout()
+plt.savefig("/home/lukas/tmp/time.pdf", transparent=True)
 plt.show()

collisionhistory.py

@@ -4,7 +4,9 @@ from matplotlib.figure import Figure
 from rebound import SimulationArchive, Simulation
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv
+from utils import filename_from_argv, earth_mass, earth_water_mass, plot_settings
+
+plot_settings()
 
 fn = filename_from_argv()
 ed = ExtraData.load(fn)
@@ -15,8 +17,8 @@ print([p.hash.value for p in last_sim.particles])
 print(last_sim.t)
 
 fig: Figure = plt.figure()
-ax_masses: Axes = fig.add_subplot(2,1, 1)
-ax_wmfs: Axes = fig.add_subplot(2,1, 2)
+ax_masses: Axes = fig.add_subplot(2, 1, 1)
+ax_wmfs: Axes = fig.add_subplot(2, 1, 2)
 
 for particle in last_sim.particles:
     if ed.pd(particle).type in ["sun", "gas giant"]:
@@ -37,6 +39,7 @@ for particle in last_sim.particles:
             break
         meta: CollisionMeta = collision["meta"]
         parents = collision["parents"]
+        print("mass:", ed.pdata[hash].total_mass / earth_mass)
         masses.append(ed.pdata[hash].total_mass)
         objects.append(ed.pdata[hash])
         times.append(meta.time)
@@ -57,9 +60,17 @@ for particle in last_sim.particles:
     ax_wmfs.set_ylabel("water mass fraction")
     ax_masses.set_ylabel("masses [kg]")
     for ax in [ax_wmfs, ax_masses]:
+        ax.set_xlim(1e4, ed.meta.current_time)
         ax.set_xlabel("time [yr]")
         ax.set_xscale("log")
         ax.set_yscale("log")
         ax.legend()
+    twin_ax = ax_masses.twinx()
+    mn, mx = ax_masses.get_ylim()
+    twin_ax.set_ylim(mn / earth_mass, mx / earth_mass)
+    twin_ax.set_ylabel('[$M_\\oplus$]')
+    twin_ax.set_yscale("log")
+    ax_wmfs.axhline(earth_water_mass/earth_mass,linestyle="dotted")
 fig.tight_layout()
+fig.savefig("/home/lukas/tmp/collisionhistory.pdf", transparent=True)
 plt.show()

collisionstats.py

@@ -5,7 +5,9 @@ from matplotlib import pyplot as plt
 from scipy.constants import mega
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv, create_figure
+from utils import filename_from_argv, create_figure, plot_settings
+
+plot_settings()
 
 dotsize = 1.5
 
@@ -36,7 +38,7 @@ ax4.set_yscale("log")
 
 ax4.set_xlabel(time_label)
 ax4.set_ylabel("water loss")
-
+sum = 0
 for file in argv[1:]:
     fn = filename_from_argv(file)
     print(fn)
@@ -51,6 +53,7 @@ for file in argv[1:]:
     times = []
 
     for collision in ed.tree.get_tree().values():
+        sum += 1
         meta: CollisionMeta = collision["meta"]
         vs.append(meta.input.velocity_esc)
         angles.append(meta.input.alpha)
@@ -60,14 +63,15 @@ for file in argv[1:]:
         water_loss.append(loss)
         times.append(meta.time / mega)
 
-    ax1.scatter(angles, vs,s=dotsize)
-    ax2.scatter(times, angles,s=dotsize)
-    ax3.scatter(times, vs,s=dotsize)
-    ax4.scatter(times, water_loss,s=dotsize)
+    ax1.scatter(angles, vs, s=dotsize)
+    ax2.scatter(times, angles, s=dotsize)
+    ax3.scatter(times, vs, s=dotsize)
+    ax4.scatter(times, water_loss, s=dotsize)
 
 ax4.autoscale(enable=True, axis='y')
 
 for i, fig in enumerate([fig1, fig2, fig3, fig4]):
-    fig.savefig(Path("plots") / fn.with_suffix(f".collision{i}.pdf").name)
-
+    fig.tight_layout()
+    fig.savefig(Path("plots") / fn.with_suffix(f".collision{i}.pdf").name, transparent=True)
+print(sum)
 plt.show()

particle_numbers.py

@@ -6,7 +6,9 @@ from matplotlib.figure import Figure
 from rebound import SimulationArchive, Simulation
 
 from extradata import ExtraData
-from utils import filename_from_argv
+from utils import filename_from_argv, plot_settings
+
+plot_settings()
 
 fig: Figure = plt.figure()
 ax: Axes = plt.gca()
@@ -29,11 +31,13 @@ for file in argv[1:]:
         N = num_embryos + num_planetesimals
         Ns.append(N)
         ts.append(sim.t)
-
-    ax.step(ts, Ns, label=fn, where="post")
+    perfect_merging = "pm" in str(fn)
+    ax.step(ts, Ns, label=fn, where="post", linestyle="dashed" if perfect_merging else "solid",linewidth=0.7)
 ax.set_xlabel("time [yr]")
 ax.set_ylabel("number of objects")
 # ax.set_xscale("log")
 
 plt.legend()
+plt.tight_layout()
+plt.savefig("/home/lukas/tmp/particle_numbers.pdf", transparent=True)
 plt.show()

timeplot.py

@@ -11,10 +11,19 @@ from matplotlib.collections import PathCollection
 from matplotlib.colors import Normalize, Colormap
 from matplotlib.text import Text
 from rebound import SimulationArchive, Particle
+from scipy.constants import mega
 
 from extradata import ExtraData, ParticleData
 from utils import filename_from_argv
 
+output_plots = False
+if output_plots:
+    size_factor = 1
+    figsize = None
+else:
+    size_factor = 3
+    figsize = (15, 10)
+
 
 class MyProgramArgs(argparse.Namespace):
     save_video: bool
@@ -38,8 +47,9 @@ def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraD
         time = 10 ** ((num + log10(50000) / log_timestep) * log_timestep)
     else:
         timestep = ed.meta.current_time / total_frames
+        print(num / total_frames)
         time = num * timestep
-    print(time)
+    print(f"{num / total_frames:.2f}, {time:.0f}")
     sim = sa.getSimulation(t=time)
     if time < 1e3:
         timestr = f"{time:.0f}"
@@ -71,14 +81,15 @@ def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraD
     else:
         raise ValueError("invalid y-axis")
     dots.set_offsets(bla.T)
-    water_fractions = np.array(water_fractions)
+    with np.errstate(divide='ignore'):  # allow 0 water (becomes -inf)
         color_val = (np.log10(water_fractions) + 5) / 5
     colors = cmap(color_val)
     if not mean_mass:
         mean_mass = np.mean(m[3:])
-    dots.set_sizes(3 * m / mean_mass)
+    dots.set_sizes(size_factor * m / mean_mass)
     dots.set_color(colors)
-    # plt.savefig("tmp/" + str(num) + ".pdf",transparent=True)
+    if output_plots:
+        plt.savefig(f"tmp/{num}_{time / mega}.pdf", transparent=True)
     return dots, title
 
 
@@ -88,7 +99,7 @@ def main(args: MyProgramArgs):
     logtime = False
     cmap: Colormap = matplotlib.cm.get_cmap("Blues")
 
-    fig1 = plt.figure(figsize=(15, 10))
+    fig1 = plt.figure(figsize=figsize)
 
     l: PathCollection = plt.scatter([1], [1])
 

utils/plotting.py

@@ -12,3 +12,7 @@ def create_figure() -> Tuple[Figure, Axes]:
     fig: Figure = plt.figure()
     ax: Axes = fig.gca()
     return fig, ax
+
+
+def plot_settings() -> None:
+    plt.style.use("dark_background")

visualize_collision.py

@@ -5,11 +5,13 @@ from matplotlib.patches import FancyArrowPatch
 from mpl_toolkits.mplot3d import Axes3D, proj3d
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv
+from utils import filename_from_argv, plot_settings
 
 fn = filename_from_argv()
 ed = ExtraData.load(fn)
 
+plot_settings()
+
 
 def get_circle(dx, dy, dz, r):
     u, v = np.mgrid[0:2 * np.pi:40j, 0:np.pi:20j]
@@ -64,5 +66,5 @@ for collision in ed.tree.get_tree().values():
     ax.set_zlim3d((zmin, zmin - diff))
     # ax.set_ylim3d(XYZlim)
     # ax.set_zlim3d(XYZlim * 3/4)
-
+    plt.savefig("/home/lukas/tmp/3d.pdf")
     plt.show()