finish collisionhistory.py

collisionhistory.py

@@ -1,77 +1,120 @@
+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:
-    if ed.pd(particle).type in ["sun", "gas giant"]:
-        continue
-    masses = []
-    objects = []
-    times = []
-    hash = particle.hash.value
-    objects.append(ed.pdata[hash])
-    times.append(ed.meta.current_time)
+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")
 
-    while True:
-        print(f"looking at {hash}")
-        try:
-            collision = ed.tree.get_tree()[hash]
-        except KeyError:
-            print("found end of the tree")
-            break
-        meta: CollisionMeta = collision["meta"]
-        parents = collision["parents"]
-        print("mass:", ed.pdata[hash].total_mass / earth_mass)
-        masses.append(ed.pdata[hash].total_mass)
+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 = []
+        hash = particle.hash.value
         objects.append(ed.pdata[hash])
-        times.append(meta.time)
-        # print(collision)
-        if ed.pdata[parents[0]].total_mass > ed.pdata[parents[1]].total_mass:
-            hash = parents[0]
+        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:
+                collision = ed.tree.get_tree()[hash]
+            except KeyError:
+                print("found end of the tree")
+                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)
+            # print(collision)
+            if ed.pdata[parents[0]].total_mass > ed.pdata[parents[1]].total_mass:
+                hash = parents[0]
+            else:
+                hash = parents[1]
+        objects.append(ed.pdata[hash])
+        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 = []
+        if multifile:
+            args = {
+                "linewidth": 1,
+                "color": "black",
+                "alpha": .3
+            }
         else:
-            hash = parents[1]
-    objects.append(ed.pdata[hash])
-    times.append(0)  # TODO: check log-x
-    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]")
+            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()