add new utils

utils/__init__.py

@@ -5,3 +5,4 @@ from .os_tools import *
 from .plotting import *
 from .simulation import *
 from .radius import *
+from .data import *

utils/astronomical.py

@@ -1,4 +1,7 @@
+from dataclasses import dataclass
+
 from rebound import Particle
+from scipy.constants import kilo, astronomical_unit
 
 habitable_zone_inner = 0.75
 habitable_zone_outer = 1.5
@@ -6,3 +9,24 @@ habitable_zone_outer = 1.5
 
 def is_potentially_habitable(planet: Particle):
     return habitable_zone_inner <= planet.a <= habitable_zone_outer
+
+
+@dataclass
+class ReferencePlanet:
+    a: float  # km
+    e: float
+    mass: float
+
+    @property
+    def a_au(self):
+        return self.a * kilo / astronomical_unit
+
+
+# from NASA HORIZONS
+# https://ssd.jpl.nasa.gov/api/horizons.api?format=text&COMMAND=%27Mercury+Barycenter%27&OBJ_DATA=%27YES%27&MAKE_EPHEM=%27YES%27&EPHEM_TYPE=ELEMENTS&CENTER=%400&START_TIME=1961-08-06&STOP_TIME=1961-08-07&STEP_SIZE=%272+days%27
+inner_solar_system_data = {
+    "mercury": ReferencePlanet(mass=3.302e23, a=6.132591901350513E+07, e=2.330561345288357E-01),
+    "venus": ReferencePlanet(mass=48.685e23, a=1.084671627283828E+08, e=3.297503114845414E-03),
+    "earth": ReferencePlanet(mass=5.97219e24, a=1.473393857505501E+08, e=2.365126062626348E-02),
+    "mars": ReferencePlanet(mass=6.4171e23, a=2.277188818140753E+08, e=9.685171612265037E-02),
+}

utils/data.py

@@ -0,0 +1,22 @@
+import numpy as np
+
+
+def reorder(data, water_lost_median: float, std=False):
+    return [
+        *data[:9], water_lost_median,
+        np.nan, np.nan,
+        0 if std else data[9], np.nan,
+        data[13]
+    ]
+
+
+def get_cb_data(pm=False):
+    ncols = 14
+    data = np.loadtxt("cb_data2.txt" if pm else "cb_data1.txt", usecols=range(1, ncols + 1))
+    means = data.mean(axis=0).tolist()
+    stds = data.std(axis=0).tolist()
+
+    print(means)
+    water_lost_median = 210 if pm else 257
+    return reorder(means, water_lost_median=water_lost_median), \
+           reorder(stds, water_lost_median=water_lost_median, std=True)

utils/os_tools.py

@@ -8,7 +8,7 @@ from setproctitle import setproctitle
 
 
 def filename_from_argv(argument: str = None) -> Path:
-    if len(argv) < 2:
+    if len(argv) < 2 and not argument:
         raise ValueError("specify filename")
     if argument:
         fn = argument

utils/plotting.py

@@ -1,8 +1,24 @@
+import re
+from pathlib import Path
 from typing import Tuple
 
+import matplotlib
+import numpy as np
 from matplotlib import pyplot as plt
 from matplotlib.axes import Axes
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import Colormap, LinearSegmentedColormap, Normalize
 from matplotlib.figure import Figure
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+christoph_cmap: Colormap = LinearSegmentedColormap.from_list("mycmap", ["#FF4500", "blue"])
+
+scenario_colors = {
+    "rbf": "C3",
+    "nn": "C1",
+    "lz": "C2",
+    "pm": "C0",
+}
 
 
 def create_figure() -> Tuple[Figure, Axes]:
@@ -15,4 +31,32 @@ def create_figure() -> Tuple[Figure, Axes]:
 
 
 def plot_settings() -> None:
-    plt.style.use("dark_background")
+    ...
+    # plt.style.use("dark_background")
+
+
+def get_water_cmap() -> Colormap:
+    # return christoph_cmap
+    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))
+    cmap: Colormap = LinearSegmentedColormap.from_list("mycmap", colors)
+    return cmap
+
+
+def add_water_colormap(fig: Figure, ax: Axes, cmap: Colormap) -> None:
+    divider = make_axes_locatable(ax)
+    cax = divider.append_axes("bottom", size="10%", pad=0.5)
+
+    fig.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap),
+                 label="log(water mass fraction)", cax=cax, orientation="horizontal")
+
+
+def add_au_e_label(ax: Axes) -> None:
+    ax.set_xlabel("semi-major axis [AU]")
+    ax.set_ylabel("eccentricity")
+
+
+def mode_from_fn(fn: Path) -> str:
+    return re.search(r"final_(\w+)_", str(fn)).group(1)