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