various update

snapshotprocessor.py

@@ -293,6 +293,37 @@ class SnapshotProcessor(HDF5Container):
         "id": U.none,
         "level": U.none,
     }
+    # Units sinks
+    unit_sinks = {
+        "id": U.none,
+        "msink": U.Msun,
+        "dmfsink": U.Msun,
+        "x": "unit_density",
+        "y": "unit_density",
+        "z": "unit_density",
+        "vx": "unit_velocity",
+        "vy": "unit_velocity",
+        "vz": "unit_velocity",
+        "rot_period": "unit_time",
+        "lx": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
+        "ly": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
+        "lz": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
+        "acc_rate": {"unit_mass": 1, "unit_time": -1},
+        "acc_lum": U.none,
+        "age": U.year,
+        "int_lum": U.none,
+        "Teff": U.K,
+        "level": U.none,
+        "mbh": "unit_mass",
+        "tform": "unit_time",
+        "del_mass": U.Msun,
+        "rho_gas": "unit_density",
+        "cs**2": {"unit_velocity": 2},
+        "etherm": {"unit_mass": 1, "unit_velocity": 2},
+        "vx_gas": "unit_velocity",
+        "vy_gas": "unit_velocity",
+        "vz_gas": "unit_velocity",
+    }
 
     G = 1.0  # Gravitational constant
 
@@ -675,9 +706,10 @@ class SnapshotProcessor(HDF5Container):
                 self.save_data(data, filename, group)
         return data
 
-    def save_data(self, data, filename, group, overwrite=False):
+    def save_data(self, data, filename, group, overwrite=False, units=None):
         self.logger.debug(f"Writing {filename}")
         hdf5 = tables.open_file(filename, mode="a")
+        units = units or self.unit_key
         try:
             nb_written = 0
             for key in data:
@@ -688,7 +720,7 @@ class SnapshotProcessor(HDF5Container):
                         hdf5.create_array(
                             f"/{group}", key, data[key], "", createparents=True
                         )
-                        unit = self._get_units(self.unit_key[key])
+                        unit = self._get_units(units[key])
                         hdf5.get_node(f"/{group}/{key}").unit = unit
                         nb_written += 1
                 else:
@@ -775,10 +807,13 @@ class SnapshotProcessor(HDF5Container):
 
     def convert_hdf5(self, filename=None):
         self.load_destructured(save=False)
+        self.load_sinks()
+        sinks = {key: self.sinks[key].values for key in self.sinks}
         if filename is None:
             filename = self.snap_filename
         self.save_data(self.cells, filename, "cells")
         self.save_data(self.parts, filename, "parts")
+        self.save_data(sinks, filename, "sinks", units=self.unit_sinks)
         self.unload_destructured()
 
     def get_nml(self, nml_key):
@@ -1600,12 +1635,15 @@ class SnapshotProcessor(HDF5Container):
         alpha_g = (2.0 / 3) * alpha_g
         return alpha_g
 
-    def _sinks(self):
+    def load_sinks(self):
 
         csv_name = f"{self.path}/output_{self.num:05}/sink_{self.num:05}.csv"
         if not os.path.exists(csv_name):  # If ratarmount was used
             csv_name = f"{self.path}/output_{self.num:05}/output_{self.num:05}/sink_{self.num:05}.csv"
 
+        if not os.path.exists(csv_name):
+            self.sinks = {}
+            return {}
         f = open(csv_name)
         first_line = f.readline()
         second_line = f.readline()
@@ -1642,6 +1680,7 @@ class SnapshotProcessor(HDF5Container):
                 "Teff",
             ]
             df = pd.read_csv(csv_name, header=None, names=header)
+        self.sinks = df
 
         return {key: df[key].values for key in df}
 
@@ -1935,39 +1974,10 @@ class SnapshotProcessor(HDF5Container):
                 "/maps",
                 dependencies=["coldens", "T_mwavg", "omega_mwavg"],
             ),
-            "sinks": Rule(
-                self._sinks,
+            "load_sinks_rule": Rule(
+                self.load_sinks,
                 group="/datasets",
-                unit={
-                    "id": U.none,
-                    "msink": U.Msun,
-                    "dmfsink": U.Msun,
-                    "x": "unit_density",
-                    "y": "unit_density",
-                    "z": "unit_density",
-                    "vx": "unit_velocity",
-                    "vy": "unit_velocity",
-                    "vz": "unit_velocity",
-                    "rot_period": "unit_time",
-                    "lx": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
-                    "ly": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
-                    "lz": {"unit_mass": 1, "unit_length": 2, "unit_time": -1},
-                    "acc_rate": {"unit_mass": 1, "unit_time": -1},
-                    "acc_lum": U.none,
-                    "age": U.year,
-                    "int_lum": U.none,
-                    "Teff": U.K,
-                    "level": U.none,
-                    "mbh": "unit_mass",
-                    "tform": "unit_time",
-                    "del_mass": U.Msun,
-                    "rho_gas": "unit_density",
-                    "cs**2": {"unit_velocity": 2},
-                    "etherm": {"unit_mass": 1, "unit_velocity": 2},
-                    "vx_gas": "unit_velocity",
-                    "vy_gas": "unit_velocity",
-                    "vz_gas": "unit_velocity",
-                },
+                unit=self.unit_sinks,
             ),
             "pspec": Rule(self.pspec, "Power spectrum", "/hdf5"),
             "write_particles": Rule(self._write_particles, "Particles file", "/hdf5"),