Update sectors extractions scripts

2023-05-05 18:11:31 +02:00
parent 85cfd71b66
commit 946f11725e
3 changed files with 197 additions and 2 deletions
--- a/sectors_extraction_ramses.py
+++ b/sectors_extraction_ramses.py
@@ -0,0 +1,108 @@
+# coding: utf-8
+
+import numpy as np
+from snapshotprocessor import U
+from tables import NoSuchNodeError
+
+
+def load_fields(pp):
+    """
+    Parameters
+    ----------
+    path : _type_
+        _description_
+    output : _type_
+        _description_
+
+    Returns
+    -------
+    dict
+        A dataset of cells in the following format:
+        gas:
+            position (Ngas, 3) [kpc], centered
+            volume   (Ngas)    [pc^3]
+            velocity (Ngas, 3) [km/s]
+            mass     (Ngas)    [Msun]
+        stars:
+            position (Nstar, 3) [kpc], centered
+            velocity (Nstar, 3) [km/s]
+            mass     (Nstar)    [Msun]
+            birth_time      (Nstar)    [Myr]
+        dm:
+            position (Ngas, 3) [kpc], centered
+            velocity (Ngas, 3) [km/s]
+            mass     (Ngas)    [Msun]
+        maps:
+            extent      (xmin, xmax, ymin, ymax) Coordinates of the edges of the map, centered
+            gas_coldens (Nx, Ny) [Msun/pc^2], map of column density
+    """
+
+    # Load  arrays
+    pp.load_cells(keys=["pos", "vel", "dx", "rho"])
+
+    try:
+        pp.load_parts(keys=["pos", "vel", "mass", "epoch"])
+    except (KeyError, NoSuchNodeError):
+        pp.load_parts(keys=["pos", "vel", "mass"])
+
+    cells = pp.cells
+    parts = pp.parts
+
+    if "epoch" not in parts:
+        parts["epoch"] = np.zeros(len(pp.parts["mass"]))
+
+    # Compute extra fields and convert units
+    for dset in (cells, parts):
+        dset["position"] = dset["pos"] - np.array([0.5, 0.5, 0.5])
+
+    cells["mass"] = (
+        cells["rho"]
+        * cells["dx"] ** 3
+        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun)
+    )
+
+    cells["volume"] = cells["dx"] ** 3 * (pp.info["unit_length"] ** 3).express(
+        U.kpc**3
+    )
+
+    # Separate DM from stars
+    mass_dm = np.max(parts["mass"])
+    mask_dm = parts["mass"] == mass_dm
+    mask_star = parts["mass"] < mass_dm
+
+    # Create dataset
+    data = {
+        "header": {"time": pp.time * pp.info["unit_time"].express(U.Myr)},
+        "gas": {
+            "position": cells["position"] * pp.info["unit_length"].express(U.kpc),
+            "volume": cells["volume"],
+            "velocity": cells["vel"] * pp.info["unit_velocity"].express(U.km_s),
+            "mass": cells["mass"],
+        },
+        "stars": {
+            "position": parts["position"][mask_star]
+            * pp.info["unit_length"].express(U.kpc),
+            "velocity": parts["vel"][mask_star]
+            * pp.info["unit_velocity"].express(U.km_s),
+            "mass": parts["mass"][mask_star] * pp.info["unit_mass"].express(U.Msun),
+            "birth_time": parts["epoch"][mask_star]
+            * pp.info["unit_time"].express(U.Myr),
+        },
+        "dm": {
+            "position": parts["position"][mask_dm]
+            * pp.info["unit_length"].express(U.kpc),
+            "velocity": parts["vel"][mask_dm]
+            * pp.info["unit_velocity"].express(U.km_s),
+            "mass": parts["mass"][mask_dm] * pp.info["unit_mass"].express(U.Msun),
+        },
+    }
+    return data
+
+
+if __name__ == "__main__":
+    from snapshotprocessor import SnapshotProcessor
+
+    pp = SnapshotProcessor(
+        "/home/nbrucy/ecogal/galturb/F20H_7_1pc_frig", num=80, params="params_gal.yml"
+    )
+    data = load_fields(pp)