from math import isclose 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, \ habitable_zone_inner, habitable_zone_outer, get_water_cmap, create_figure, add_au_e_label, \ inner_solar_system_data, is_ci, get_cb_data, initial_saturn_a, initial_jupiter_a # pd.set_option('display.max_columns', None) pd.options.display.max_columns = None pd.options.display.width = 0 minmax = True plot = True def chunks(lst, lst2): """Yield successive n-sized chunks from lst.""" for i in range(0, len(lst), 2): yield lst[i:i + 2] + lst2[i:i + 2] def get_masses(ed: ExtraData, planets: List[Particle]): total = 0 water = 0 for p in planets: data = ed.pd(p) # print(data.water_mass) total += data.total_mass / earth_mass water += data.water_mass / earth_water_mass return total, water methods = { "rbf": "data/final_rbf_NUM.bin", "rbf_sm": "data/final_rbf_NUM.bin", "nn": "data/final_nn_NUM.bin", "lz": "data/final_lz_correct_NUM.bin", "pm": "data/final_pm_NUM.bin", } columns = [ "$N_\\text{planets}$ [1]", "$N_\\text{planets,pot}$ [1]", "$M_\\text{planets}$ [$M_\\oplus$]", "$M_\\text{planets,pot}$ [$M_\\oplus$]", "$M_\\text{water}$ [$M_{w,\\oplus}$]", "$M_\\text{water,pot}$ [$M_{w,\\oplus}$]", "$M_\\text{esc}$ [$M_\\oplus$]", "$M_\\text{esc,water}$ [$M_{w,\\oplus}$]", "$M_\\text{sun}$ [$M_\\oplus$]", "$M_\\text{sun,water}$ [$M_{w,\\oplus}$]", "$M_\\text{gas-giant}$ [$M_\\oplus$]", "$M_\\text{gas-giant,water}$ [$M_{w,\\oplus}$]", "$M_\\text{col}$ [$M_\\oplus$]", "$M_\\text{col,water}$ [$M_{w,\\oplus}$]", "$t_\\text{last-col}$ [Myr]" ] maintable = [] final_bodies = {} for name, filepath in methods.items(): table = [] rows = [] fin_as = [] fin_es = [] fin_mass = [] fin_core_mass = [] 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 sa = SimulationArchive(str(fn.with_suffix(".bin"))) last_sim: Simulation = sa[-1] planets = [] a_values_per_planet = {} e_values_per_planet = {} for sim in sa: if sim.t < ed.meta.tmax - 10 * mega: continue 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) gas_giants = [] for particle in last_sim.particles: particle_data = ed.pd(particle) if particle_data.type in ["sun", "planetesimal"]: continue if particle_data.type == "gas giant": gas_giants.append(particle) 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) if len(gas_giants) != 2: print(f"it seems like gas giants got lost ({len(gas_giants)} left)") elif not ( isclose(gas_giants[0].a, initial_saturn_a, rel_tol=0.05) and isclose(gas_giants[1].a, initial_jupiter_a, rel_tol=0.05) ): print("gas giants moved") pothab_planets = [p for p in planets if is_potentially_habitable(p)] num_planets = len(planets) num_planets_pot = len(pothab_planets) M_planets, M_water = get_masses(ed, planets) M_planets_pot, M_water_pot = get_masses(ed, pothab_planets) # mass of particles thrown into Sun/escaped escaped_mass = 0 escaped_water_mass = 0 sun_mass = 0 sun_water_mass = 0 for particle in ed.pdata.values(): if particle.escaped: if particle.type in ["sun", "gas giant"]: print(fn, particle, "escaped", particle.type) continue escaped_mass += particle.total_mass / earth_mass escaped_water_mass += particle.water_mass / earth_water_mass elif particle.collided_with_sun: sun_mass += particle.total_mass / earth_mass sun_water_mass += particle.water_mass / earth_water_mass # count mass lost to gas giants gas_giant_mass = 0 gas_giant_water_mass = 0 for col_id, collision in ed.tree.get_tree().items(): gas_parent = None other_parent = None gas_parent_counter = 0 for parent in collision["parents"]: if ed.pdata[parent].type == "gas giant": gas_parent_counter += 1 gas_parent = parent else: other_parent = parent if gas_parent_counter > 1: print(f"it seems like {gas_parent_counter} gas giants collided with each other in {col_id}") continue if gas_parent: previous_body = ed.pdata[other_parent] gas_giant_mass += previous_body.total_mass / earth_mass gas_giant_water_mass += previous_body.water_mass / earth_water_mass # count mass lost in collisions collision_mass = 0 collision_water_mass = 0 for col_id, collision in ed.tree.get_tree().items(): parent_mass = 0 parent_water_mass = 0 for parent in collision["parents"]: parent_body = ed.pdata[parent] parent_mass += parent_body.total_mass / earth_mass parent_water_mass += parent_body.water_mass / earth_water_mass child = ed.pdata[col_id] diff = parent_mass - child.total_mass / earth_mass diff_water = parent_water_mass - child.water_mass / earth_water_mass collision_mass += diff collision_water_mass += diff_water if collision_water_mass < 1e-10: collision_water_mass = 0 if collision_mass < 1e-10: collision_mass = 0 # last collision time last_collision_time = 0 for col_id, collision in ed.tree.get_tree().items(): meta: CollisionMeta = collision["meta"] if meta.time > last_collision_time: 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, 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) if minmax: maintable.append(list(zip(df.min(), df.max()))) else: maintable.append(list(zip(df.mean(), df.std()))) maintable.append(list(zip(*get_cb_data(minmax)))) maintable.append(list(zip(*get_cb_data(minmax, pm=True)))) transposed = list(map(list, zip(*maintable))) for row, name in zip(transposed, columns): n, unit = name.split() strings = [" ".join([n, unit])] for value, pm in row: if np.isnan(value): strings.append("{-}") elif minmax: # in that case (value,pm) = (min,max) strings.append(f"{value:.1f} -- {pm:.1f}") else: strings.append(f"{value:.1f}\\pm {pm:.1f}") print(" & ".join(strings) + r" \\")