update final scripts

final_results.py

@@ -1,18 +1,25 @@
-import glob
+from os.path import expanduser
+from statistics import mean
 from typing import List
 
+import matplotlib.pyplot as plt
+import numpy as np
 import pandas as pd
 from rebound import SimulationArchive, Simulation
 from scipy.constants import mega
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv, is_potentially_habitable, Particle, earth_mass, earth_water_mass
+from utils import filename_from_argv, is_potentially_habitable, Particle, earth_mass, earth_water_mass, \
+    habitable_zone_inner, habitable_zone_outer, get_water_cmap, add_water_colormap, create_figure, add_au_e_label, \
+    inner_solar_system_data, is_ci, get_cb_data
 
 # pd.set_option('display.max_columns', None)
 pd.options.display.max_columns = None
 pd.options.display.width = 0
 
-minmax=True
+minmax = True
+plot = True
+
 
 def chunks(lst, lst2):
     """Yield successive n-sized chunks from lst."""
@@ -20,13 +27,6 @@ def chunks(lst, lst2):
         yield lst[i:i + 2] + lst2[i:i + 2]
 
 
-def print_row(mean, std, name: str):
-    strings = [name]
-    for value, pm in zip(list(mean), list(std)):
-        strings.append(f"{value:.1f}\\pm {pm:.1f}")
-    print(" & ".join(strings) + r" \\")
-
-
 def get_masses(ed: ExtraData, planets: List[Particle]):
     total = 0
     water = 0
@@ -39,25 +39,40 @@ def get_masses(ed: ExtraData, planets: List[Particle]):
 
 
 methods = {
-    "rbf": "data/final_rbf_*.bin",
+    "rbf": "data/final_rbf_NUM.bin",
-    "nn": "data/final_nn_*.bin",
+    "rbf_sm": "data/final_rbf_NUM.bin",
-    "lz": "data/final_lz_*.bin",
+    "nn": "data/final_nn_NUM.bin",
-    "pm": "data/final_pm_*.bin",
+    "lz": "data/final_lz_NUM.bin",
+    "pm": "data/final_pm_NUM.bin",
 }
-columns = ["num_planets", "num_planets_pot", "M_planets", "M_planets_pot", "M_water", "M_water_pot", "escaped_mass",
+columns = ["num_planets [1]", "num_planets_pot [1]", "M_planets [$M_\\oplus$]", "M_planets_pot [$M_\\oplus$]",
-           "escaped_water_mass", "sun_mass", "sun_water_mass", "gas_giant_mass", "gas_giant_water_mass",
+           "M_water [$M_{w,\\oplus}$]", "M_water_pot [$M_{w,\\oplus}$]", "escaped_mass [$M_\\oplus$]",
-           "collision_mass", "collision_water_mass","last_collision_time"]
+           "escaped_water_mass [$M_{w,\\oplus}$]", "sun_mass [$M_\\oplus$]", "sun_water_mass [$M_{w,\\oplus}$]",
+           "gas_giant_mass [$M_\\oplus$]", "gas_giant_water_mass [$M_{w,\\oplus}$]",
+           "collision_mass [$M_\\oplus$]", "collision_water_mass [$M_{w,\\oplus}$]", "last_collision_time [Myr]"]
 
 maintable = []
 
+final_bodies = {}
+
 for name, filepath in methods.items():
     table = []
     rows = []
-    for file in sorted(glob.glob(filepath)):
+    fin_as = []
-        fn = filename_from_argv(file)
+    fin_es = []
-        # print(fn)
+    fin_mass = []
-
+    fin_core_mass = []
-        ed = ExtraData.load(fn)
+    fin_wmf = []
+    fin_wmf = []
+    max_num = 41 if name == "rbf" else 21
+    for i in range(1, max_num):
+        fn = filename_from_argv(filepath.replace("NUM", str(i)))
+        print(fn)
+        try:
+            ed = ExtraData.load(fn)
+        except FileNotFoundError as e:
+            print(e.filename)
+            continue
         if ed.meta.current_time < ed.meta.tmax:
             print("not yet finished")
             continue
@@ -65,13 +80,32 @@ for name, filepath in methods.items():
         sa = SimulationArchive(str(fn.with_suffix(".bin")))
         last_sim: Simulation = sa[-1]
         planets = []
-        for particle in last_sim.particles:
+        a_values_per_planet = {}
-            if ed.pd(particle).type in ["sun", "gas giant"]:
+        e_values_per_planet = {}
+        for sim in sa:
+            if sim.t < ed.meta.tmax - 10 * mega:
                 continue
-            if ed.pd(particle).type == "planetesimal":
+            for particle in sim.particles[1:]:
+                hash = particle.hash.value
+                orbit = particle.calculate_orbit()
+                if hash not in a_values_per_planet:
+                    a_values_per_planet[hash] = []
+                    e_values_per_planet[hash] = []
+                a_values_per_planet[hash].append(orbit.a)
+                e_values_per_planet[hash].append(orbit.e)
+        for particle in last_sim.particles:
+            particle_data = ed.pd(particle)
+            if particle_data.type in ["sun", "gas giant"]:
+                continue
+            if particle_data.type == "planetesimal":
                 continue
             # print(particle.r * astronomical_unit / earth_radius)
             planets.append(particle)
+            fin_as.append(mean(a_values_per_planet[particle.hash.value]))
+            fin_es.append(mean(e_values_per_planet[particle.hash.value]))
+            fin_mass.append(particle_data.total_mass)
+            fin_core_mass.append(particle_data.total_mass * particle_data.core_mass_fraction)
+            fin_wmf.append(particle_data.water_mass_fraction)
 
         pothab_planets = [p for p in planets if is_potentially_habitable(p)]
         num_planets = len(planets)
@@ -146,15 +180,52 @@ for name, filepath in methods.items():
                 last_collision_time = meta.time
         last_collision_time /= mega
 
-        values = [num_planets, num_planets_pot, M_planets, M_planets_pot, M_water, M_water_pot, escaped_mass,
+        values = [num_planets, num_planets_pot, M_planets, M_planets_pot, M_water, M_water_pot,
-                  escaped_water_mass, sun_mass, sun_water_mass, gas_giant_mass, gas_giant_water_mass,
+                  sun_mass, sun_water_mass, escaped_mass, escaped_water_mass,
+                  gas_giant_mass, gas_giant_water_mass,
                   collision_mass, collision_water_mass, last_collision_time]
         # print(values)
         table.append(values)
         rows.append(str(fn))
 
+    if plot:
+        mean_mass = 5.208403167890638e+24  # constant between plots
+        size_mult = 100
+        fin_mass = np.array(fin_mass)
+        fin_core_mass = np.array(fin_core_mass)
+        print("mean", mean_mass)
+        sizes = (np.array(fin_mass) / mean_mass) ** (2 / 3) * size_mult
+        core_sizes = (np.array(fin_core_mass) / mean_mass) ** (2 / 3) * size_mult
+        with np.errstate(divide='ignore'):  # allow 0 water (becomes -inf)
+            color_val = (np.log10(fin_wmf) + 5) / 5
+        cmap = get_water_cmap()
+        print(color_val)
+        print(np.log10(fin_wmf))
+        colors = cmap(color_val)
+        fig, ax = create_figure()
+        ax.scatter(fin_as, fin_es, s=sizes, zorder=10, cmap=(), c=colors)
+        ax.scatter(fin_as, fin_es, s=core_sizes, zorder=12, color="black")
+        for pname, planet in inner_solar_system_data.items():
+            if pname == "earth":
+                earth_wmf = earth_water_mass / earth_mass
+                earth_color_val = (np.log10(earth_wmf) + 5) / 5
+                fill_color = cmap(earth_color_val)
+            else:
+                fill_color = "white"
+            ax.scatter(planet.a_au, planet.e, s=(planet.mass / mean_mass) ** (2 / 3) * size_mult,
+                       color=fill_color, edgecolors="black", zorder=5)
+        ax.axvspan(habitable_zone_inner, habitable_zone_outer, color="#eee")
+        # add_water_colormap(fig, ax, cmap=cmap)
+        add_au_e_label(ax)
+        ax.set_xlim(0.25, 3.6)
+        ax.set_ylim(-0.05, 0.55)
+        fig.tight_layout()
+        if not is_ci():
+            plt.savefig(expanduser(f"~/tmp/final_bodies_{name}.pdf"))
+        # plt.show()
     pd.set_option('display.float_format', lambda x: '%.1f' % x)
     df = pd.DataFrame(table, index=rows, columns=columns)
+    print([a[2] for a in table])
     # print(df)
     # print("\n-----\n")
     # print_row(df.mean(), df.std(), name)
@@ -163,12 +234,19 @@ for name, filepath in methods.items():
     else:
         maintable.append(list(zip(df.mean(), df.std())))
 
+maintable.append(list(zip(*get_cb_data())))
+maintable.append(list(zip(*get_cb_data(pm=True))))
+
 transposed = list(map(list, zip(*maintable)))
 
 for row, name in zip(transposed, columns):
-    strings = [name.replace("_", "\\_")]
+    n, unit = name.split()
+
+    strings = [" ".join([n.replace("_", "\\_"), unit])]
     for value, pm in row:
-        if minmax:
+        if np.isnan(value):
+            strings.append("{-}")
+        elif minmax:
             strings.append(f"{value:.1f} -- {pm:.1f}")
         else:
             strings.append(f"{value:.1f}\\pm {pm:.1f}")

particle_numbers.py

@@ -1,19 +1,16 @@
 import pickle
 import random
-import re
 from os.path import expanduser
 from pathlib import Path
-from random import shuffle
 from sys import argv
 
 from matplotlib import pyplot as plt
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
-from mpl_toolkits.axes_grid1.inset_locator import mark_inset, inset_axes
 from rebound import SimulationArchive, Simulation
 
 from extradata import ExtraData
-from utils import filename_from_argv, plot_settings, is_ci
+from utils import filename_from_argv, plot_settings, is_ci, scenario_colors, mode_from_fn
 
 cache_file = Path("particle_numbers_cache.pickle.xz")
 if cache_file.exists():
@@ -26,18 +23,10 @@ plot_settings()
 
 fig: Figure = plt.figure()
 ax: Axes = plt.gca()
-# axins2 = zoomed_inset_axes(ax, zoom=1, loc="center")
-axins2 = inset_axes(ax, 2,2, loc="center right")  # no zoom
-random.seed(1)
 
-colors = {
-    "rbf": "C0",
-    "nn": "C1",
-    "lz": "C2",
-    "pm": "C3",
-}
 files = argv[1:]
-shuffle(files)
+random.seed(1)
+random.shuffle(files)
 for file in files:
     fn = filename_from_argv(file)
     print(fn)
@@ -62,27 +51,26 @@ for file in files:
             Ns.append(N)
             ts.append(sim.t)
         cache[fn.name] = Ns, ts
-    mode = re.search(r"final_(\w+)_", str(fn)).group(1)
+    mode = mode_from_fn(fn)
-    for a in [ax, axins2]:
+    ax.step(ts, Ns, label=mode, where="post", color=scenario_colors[mode], linewidth=0.7,alpha=.5)
-        a.step(ts, Ns, label=mode, where="post", color=colors[mode], linewidth=0.7)
-    # break
 
 with cache_file.open("wb") as f:
     pickle.dump(cache, f)
 
 ax.set_xlabel("time [yr]")
 ax.set_ylabel("number of objects")
-# ax.set_xscale("log")
+ax.set_xscale("log")
-
-axins2.set_xlim(1e8, 2e8)
-axins2.set_ylim(0, 15)
-axins2.tick_params(labelleft=False, labelbottom=False)
-mark_inset(ax, axins2, loc1=2, loc2=4, fc="none", ec="0.5")
 
 handles, labels = ax.get_legend_handles_labels()
 by_label = dict(zip(labels, handles))
-ax.legend(by_label.values(), by_label.keys())
+new_labels = {
-# plt.tight_layout()
+    "RBF": by_label["rbf"],
+    "NN": by_label["nn"],
+    "LZ": by_label["lz"],
+    "PM": by_label["pm"],
+}
+ax.legend(new_labels.values(), new_labels.keys())
+plt.tight_layout()
 if not is_ci():
     plt.savefig(expanduser(f"~/tmp/particle_numbers.pdf"))
 plt.show()