From f2dc99bfb6d411e630e3c0e2a0090d9f3ff75a0e Mon Sep 17 00:00:00 2001 From: Lukas Winkler Date: Sun, 30 Jan 2022 19:53:25 +0100 Subject: [PATCH] update final scripts --- final_results.py | 136 ++++++++++++++++++++++++++++++++++---------- particle_numbers.py | 40 +++++-------- 2 files changed, 121 insertions(+), 55 deletions(-) diff --git a/final_results.py b/final_results.py index e3084e8..e9c8847 100644 --- a/final_results.py +++ b/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", - "nn": "data/final_nn_*.bin", - "lz": "data/final_lz_*.bin", - "pm": "data/final_pm_*.bin", + "rbf": "data/final_rbf_NUM.bin", + "rbf_sm": "data/final_rbf_NUM.bin", + "nn": "data/final_nn_NUM.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", - "escaped_water_mass", "sun_mass", "sun_water_mass", "gas_giant_mass", "gas_giant_water_mass", - "collision_mass", "collision_water_mass","last_collision_time"] +columns = ["num_planets [1]", "num_planets_pot [1]", "M_planets [$M_\\oplus$]", "M_planets_pot [$M_\\oplus$]", + "M_water [$M_{w,\\oplus}$]", "M_water_pot [$M_{w,\\oplus}$]", "escaped_mass [$M_\\oplus$]", + "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)): - fn = filename_from_argv(file) - # print(fn) - - ed = ExtraData.load(fn) + fin_as = [] + fin_es = [] + fin_mass = [] + fin_core_mass = [] + 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: - if ed.pd(particle).type in ["sun", "gas giant"]: + a_values_per_planet = {} + 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, - escaped_water_mass, sun_mass, sun_water_mass, gas_giant_mass, gas_giant_water_mass, + values = [num_planets, num_planets_pot, M_planets, M_planets_pot, M_water, M_water_pot, + 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}") diff --git a/particle_numbers.py b/particle_numbers.py index a02b824..edc068a 100644 --- a/particle_numbers.py +++ b/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) - for a in [ax, axins2]: - a.step(ts, Ns, label=mode, where="post", color=colors[mode], linewidth=0.7) - # break + mode = mode_from_fn(fn) + ax.step(ts, Ns, label=mode, where="post", color=scenario_colors[mode], linewidth=0.7,alpha=.5) 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") - -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") +ax.set_xscale("log") handles, labels = ax.get_legend_handles_labels() by_label = dict(zip(labels, handles)) -ax.legend(by_label.values(), by_label.keys()) -# plt.tight_layout() +new_labels = { + "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()