add new scripts

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)

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()

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()