finish collisionhistory.py

collisionhistory.py

@@ -1,28 +1,58 @@
+import random
+from sys import argv
+
 from matplotlib import pyplot as plt
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from rebound import SimulationArchive, Simulation
+from scipy.constants import mega
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv, earth_mass, earth_water_mass, plot_settings, is_ci
+from utils import filename_from_argv, earth_mass, earth_water_mass, plot_settings, is_ci, is_potentially_habitable
+
+files = argv[1:]
+multifile = len(files) > 1
 
 plot_settings()
-
-fn = filename_from_argv()
-ed = ExtraData.load(fn)
-sa = SimulationArchive(str(fn.with_suffix(".bin")))
-
-last_sim: Simulation = sa[-1]
-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)
 
-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, 200 * mega)
+    ax.set_xlabel("time [yr]")
+    ax.set_xscale("log")
+    ax.set_yscale("log")
+
+earth_mass_ax = ax_masses.secondary_yaxis("right",
+                                          functions=(lambda x: x / earth_mass, lambda x: x * earth_mass))
+earth_mass_ax.set_ylabel('masses [$M_\\oplus$]')
+ax_wmfs.axhline(earth_water_mass / earth_mass, linestyle="dotted")
+
+num_formed_planets = 0
+num_large_planets = 0
+num_habitable_planets = 0
+num_water_rich_planets = 0
+
+random.seed(1)
+random.shuffle(files)
+for file in files:
+    fn = filename_from_argv(file)
+
+    ed = ExtraData.load(fn)
+    sa = SimulationArchive(str(fn.with_suffix(".bin")))
+
+    last_sim: Simulation = sa[-1]
+    print([p.hash.value for p in last_sim.particles])
+    print(last_sim.t)
+
+    for particle in last_sim.particles:
         if ed.pd(particle).type in ["sun", "gas giant"]:
             continue
+        # if not is_potentially_habitable(particle):
+        #     continue
         masses = []
         objects = []
         times = []
@@ -30,6 +60,14 @@ for particle in last_sim.particles:
         objects.append(ed.pdata[hash])
         times.append(ed.meta.current_time)
 
+        num_formed_planets += 1
+        if particle.m > .6 * earth_mass:
+            num_large_planets += 1
+            if is_potentially_habitable(particle):
+                num_habitable_planets += 1
+                if ed.pd(particle).water_mass_fraction > 1e-4:
+                    num_water_rich_planets += 1
+
         while True:
             print(f"looking at {hash}")
             try:
@@ -49,29 +87,34 @@ for particle in last_sim.particles:
             else:
                 hash = parents[1]
         objects.append(ed.pdata[hash])
-    times.append(0)  # TODO: check log-x
+        times.append(0)
         if len(times) < 3:
             continue
         masses = [p.total_mass for p in objects]
         wmfs = [p.water_mass_fraction for p in objects]
         figs = []
-    ax_masses.step(times, masses, label=particle.hash.value)
-    ax_wmfs.step(times, wmfs, label=particle.hash.value)
-    ax_wmfs.set_ylabel("water mass fraction")
-    ax_masses.set_ylabel("masses [kg]")
+        if multifile:
+            args = {
+                "linewidth": 1,
+                "color": "black",
+                "alpha": .3
+            }
+        else:
+            args = {}
+        ax_masses.step(times, masses, label=particle.hash.value, **args)
+        ax_wmfs.step(times, wmfs, label=particle.hash.value, **args)
+
+if not multifile:
     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")
+
+print(num_large_planets / num_formed_planets)
+habitable_large_planet_fraction = num_habitable_planets / num_large_planets
+water_rich_planet_fraction = num_water_rich_planets / num_habitable_planets
+print(habitable_large_planet_fraction)
+print(water_rich_planet_fraction)
+
 fig.tight_layout()
 if not is_ci():
-    fig.savefig("/home/lukas/tmp/collisionhistory.pdf", transparent=True)
+    fig.savefig(f"/home/lukas/tmp/collisionhistory_{fn.name}.pdf", transparent=True)
 plt.show()