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_arepo.py
+++ b/sectors_extraction_arepo.py
@@ -0,0 +1,84 @@
+# coding: utf-8
+
+import numpy as np
+import h5py
+from astropy import units as u
+
+
+def load_fields(path):
+    """
+    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
+    """
+
+    snap = h5py.File(path)
+    info = snap["Header"].attrs
+    gas = snap["PartType0"]
+    stars = snap["PartType4"]
+    time = info["Time"]
+    box_size = info["BoxSize"]
+    UnitLength_in_kpc = info["UnitLength_in_cm"] * u.cm.to(u.kpc)
+    UnitVelocity_in_km_per_s = info["UnitVelocity_in_cm_per_s"] * u.cm.to(u.km)
+    UnitTime_in_Myr = (
+        info["UnitLength_in_cm"] / info["UnitVelocity_in_cm_per_s"]
+    ) * u.s.to(u.Myr)
+    UnitMass_in_Msun = info["UnitMass_in_g"] * u.g.to(u.Msun)
+
+    # Create dataset
+    data = {
+        "header": {"time": time * UnitTime_in_Myr, "fluids": ["gas", "stars"]},
+        "gas": {
+            "position": (
+                np.asarray(gas["CenterOfMass"]) - np.array([0.5, 0.5, 0.5]) * box_size
+            )
+            * UnitLength_in_kpc,
+            "volume": (np.asarray(gas["Masses"]) / np.asarray(gas["Density"]))
+            * UnitLength_in_kpc**3,
+            "velocity": np.asarray(gas["Velocities"]) * UnitVelocity_in_km_per_s,
+            "mass": np.asarray(gas["Masses"]) * UnitMass_in_Msun,
+        },
+        "stars": {
+            "position": (
+                np.asarray(stars["Coordinates"]) - np.array([0.5, 0.5, 0.5]) * box_size
+            )
+            * UnitLength_in_kpc,
+            "velocity": np.asarray(stars["Velocities"]) * UnitVelocity_in_km_per_s,
+            "mass": np.asarray(stars["Masses"]) * UnitMass_in_Msun,
+            "birth_time": np.asarray(stars["StellarFormationTime"]) * UnitTime_in_Myr,
+        },
+    }
+    return data
+
+
+if __name__ == "__main__":
+    from snapshotprocessor import SnapshotProcessor
+
+    data = load_fields(
+        "/home/nbrucy/Travail/Postdoc/Ecogal/MW/junia2_0.25kpc/OUTPUT_SN/snap_150.hdf5"
+    )