mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
update final scripts
This commit is contained in:
parent
db5f4e88d7
commit
f2dc99bfb6
GPG key ID:
54DE4D798D244853
2 changed files with 121 additions and 55 deletions
132
final_results.py
132
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
|
||||
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)
|
||||
|
||||
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}")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue