mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
add new scripts
This commit is contained in:
parent
f2dc99bfb6
commit
aef8bda87b
3 changed files with 147 additions and 0 deletions
22
compress.py
Normal file
22
compress.py
Normal file
|
@ -0,0 +1,22 @@
|
||||||
|
"""
|
||||||
|
The .energylog.csv files logged the energy every 100 years which resulted in huge files.
|
||||||
|
This script converts the data to a compressed hdf5 archive.
|
||||||
|
"""
|
||||||
|
import h5py
|
||||||
|
|
||||||
|
from utils import filename_from_argv
|
||||||
|
|
||||||
|
times = []
|
||||||
|
values = []
|
||||||
|
fn = filename_from_argv()
|
||||||
|
with fn.with_suffix(".energylog.csv").open() as f:
|
||||||
|
for line in f:
|
||||||
|
time_str, val_str = line.split(",")
|
||||||
|
times.append(int(float(time_str)))
|
||||||
|
values.append(float(val_str))
|
||||||
|
print("creating hdf5 file")
|
||||||
|
with h5py.File(fn.with_suffix(".energylog.hdf5"), "w") as f:
|
||||||
|
times_dset = f.create_dataset("times", data=times, compression="gzip", shuffle=True, fletcher32=True)
|
||||||
|
print(times_dset.dtype)
|
||||||
|
values_dset = f.create_dataset("values", data=values, compression="gzip", shuffle=True, fletcher32=True)
|
||||||
|
print(values_dset.dtype)
|
60
show_initial_conditions.py
Normal file
60
show_initial_conditions.py
Normal file
|
@ -0,0 +1,60 @@
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
import matplotlib
|
||||||
|
import numpy as np
|
||||||
|
from matplotlib import pyplot as plt
|
||||||
|
from matplotlib.cm import ScalarMappable
|
||||||
|
from matplotlib.colors import Colormap, Normalize, LinearSegmentedColormap
|
||||||
|
from matplotlib.figure import Figure
|
||||||
|
from mpl_toolkits.axes_grid1 import make_axes_locatable
|
||||||
|
from rebound import Simulation
|
||||||
|
|
||||||
|
from extradata import ExtraData
|
||||||
|
from water_sim import add_particles_from_conditions_file
|
||||||
|
|
||||||
|
|
||||||
|
def cm_to_inch(cm: float) -> float:
|
||||||
|
return cm / 2.54
|
||||||
|
|
||||||
|
|
||||||
|
sim = Simulation()
|
||||||
|
ed = ExtraData()
|
||||||
|
add_particles_from_conditions_file(sim, ed, "initcon/conditions_final.input")
|
||||||
|
|
||||||
|
size_factor = 2
|
||||||
|
min_val, max_val = 0.2, 1.0
|
||||||
|
orig_cmap: Colormap = matplotlib.cm.get_cmap('plasma_r')
|
||||||
|
colors = orig_cmap(np.linspace(min_val, max_val, 20))
|
||||||
|
# colors = ["black", "lightblue"]
|
||||||
|
cmap: Colormap = LinearSegmentedColormap.from_list("mycmap", colors)
|
||||||
|
a_list = []
|
||||||
|
e_list = []
|
||||||
|
m_list = []
|
||||||
|
wf_list = []
|
||||||
|
for p in sim.particles[3:]:
|
||||||
|
orbit = p.calculate_orbit(primary=sim.particles[0])
|
||||||
|
a_list.append(orbit.a)
|
||||||
|
e_list.append(orbit.e)
|
||||||
|
m_list.append(p.m)
|
||||||
|
wf_list.append(ed.pd(p).water_mass_fraction)
|
||||||
|
m_list = np.asarray(m_list)
|
||||||
|
with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
|
||||||
|
color_val = (np.log10(wf_list) + 5) / 5
|
||||||
|
colors = cmap(color_val)
|
||||||
|
|
||||||
|
plt.figure(figsize=(cm_to_inch(16.5), cm_to_inch(6)))
|
||||||
|
fig: Figure = plt.gcf()
|
||||||
|
print(fig.get_size_inches())
|
||||||
|
|
||||||
|
plt.scatter(a_list, e_list, s=size_factor * m_list / m_list.mean(), c=colors)
|
||||||
|
plt.xlabel("a [AU]")
|
||||||
|
plt.ylabel("e")
|
||||||
|
divider = make_axes_locatable(plt.gca())
|
||||||
|
cax = divider.append_axes("bottom", size="10%", pad=0.5)
|
||||||
|
|
||||||
|
plt.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap),
|
||||||
|
label="log(water mass fraction)", cax=cax, orientation="horizontal")
|
||||||
|
|
||||||
|
plt.tight_layout()
|
||||||
|
plt.savefig(Path("~/tmp/out.pdf").expanduser())
|
||||||
|
plt.show()
|
65
timeplot_static.py
Normal file
65
timeplot_static.py
Normal file
|
@ -0,0 +1,65 @@
|
||||||
|
from os.path import expanduser
|
||||||
|
|
||||||
|
import matplotlib.pyplot as plt
|
||||||
|
import numpy as np
|
||||||
|
from matplotlib.axes import Axes
|
||||||
|
from matplotlib.cm import ScalarMappable
|
||||||
|
from matplotlib.colors import Normalize
|
||||||
|
from matplotlib.figure import Figure
|
||||||
|
from rebound import SimulationArchive
|
||||||
|
from scipy.constants import mega
|
||||||
|
|
||||||
|
from extradata import ExtraData
|
||||||
|
from utils import filename_from_argv, get_water_cmap
|
||||||
|
|
||||||
|
mean_mass = 5.208403167890638e+24 # constant between plots
|
||||||
|
size_mult = 100
|
||||||
|
|
||||||
|
|
||||||
|
def plot_file(file, time, ax: Axes, mode):
|
||||||
|
fn = filename_from_argv(file)
|
||||||
|
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
||||||
|
ed = ExtraData.load(fn)
|
||||||
|
sim = sa.getSimulation(t=time)
|
||||||
|
|
||||||
|
water_fractions = [ed.pd(p).water_mass_fraction for p in sim.particles[1:]]
|
||||||
|
a = [p.a for p in sim.particles[1:]]
|
||||||
|
e = [p.e for p in sim.particles[1:]]
|
||||||
|
m = np.array([p.m for p in sim.particles[1:]])
|
||||||
|
# m[:2] /= 1e2 # reduce size of gas giants
|
||||||
|
sizes = (np.array(m) / mean_mass) ** (2 / 3) * size_mult
|
||||||
|
with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
|
||||||
|
color_val = (np.log10(water_fractions) + 5) / 5
|
||||||
|
cmap = get_water_cmap()
|
||||||
|
colors = cmap(color_val)
|
||||||
|
ax.set_xlim(0, 5.5)
|
||||||
|
ax.set_ylim(-0.05, .8)
|
||||||
|
ax.text(0.05, 0.93, f"{time / mega:.0f} Myr", size=10,
|
||||||
|
horizontalalignment='left', verticalalignment='top', transform=ax.transAxes)
|
||||||
|
ax.text(0.95, 0.93, mode, size=10,
|
||||||
|
horizontalalignment='right', verticalalignment='top', transform=ax.transAxes)
|
||||||
|
ax.scatter(a, e, s=sizes, zorder=10, c=colors)
|
||||||
|
|
||||||
|
|
||||||
|
def main():
|
||||||
|
fig: Figure
|
||||||
|
modes = ["RBF", "NN", "PM", "LZ"]
|
||||||
|
fig, axes = plt.subplots(5, 4, sharex=True, sharey=True, figsize=(10, 13))
|
||||||
|
for col_nr, file in enumerate(
|
||||||
|
["data/final_rbf_1.bin", "data/final_nn_1.bin", "data/final_pm_1.bin", "data/final_lz_1.bin"]):
|
||||||
|
for row_nr, time in enumerate([1 * mega, 5 * mega, 20 * mega, 50 * mega, 100 * mega]):
|
||||||
|
plot_file(file, time, axes[row_nr][col_nr], mode=modes[col_nr])
|
||||||
|
fig_colorbar = plt.figure()
|
||||||
|
fig_colorbar.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=get_water_cmap()), aspect=20,
|
||||||
|
label="log(water mass fraction)", orientation="horizontal")
|
||||||
|
fig.supxlabel("semi-major axis [AU]")
|
||||||
|
fig.supylabel("eccentricity")
|
||||||
|
fig.tight_layout()
|
||||||
|
fig_colorbar.gca().remove()
|
||||||
|
fig.savefig(expanduser(f"~/tmp/timeplot.pdf"), dpi=300)
|
||||||
|
fig_colorbar.savefig(expanduser(f"~/tmp/timeplot_colorbar.pdf"), bbox_inches='tight')
|
||||||
|
plt.show()
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
main()
|
Loading…
Reference in a new issue