diff --git a/compress.py b/compress.py new file mode 100644 index 0000000..a5a19b6 --- /dev/null +++ b/compress.py @@ -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) diff --git a/show_initial_conditions.py b/show_initial_conditions.py new file mode 100644 index 0000000..b073945 --- /dev/null +++ b/show_initial_conditions.py @@ -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() diff --git a/timeplot_static.py b/timeplot_static.py new file mode 100644 index 0000000..9fa0027 --- /dev/null +++ b/timeplot_static.py @@ -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()