improve io checks

This commit is contained in:
Noe Brucy
2023-04-05 16:50:43 +02:00
parent 301d08dce7
commit bd520d7189
5 changed files with 105 additions and 83 deletions
+39 -28
View File
@@ -287,6 +287,9 @@ class SnapshotProcessor(HDF5Container):
"phi": "unit_gravpot",
"mass": "unit_mass",
"epoch": "unit_time",
"size": "unit_length",
"dx": "unit_length",
"pos": "unit_length",
"id": U.none,
"level": U.none,
}
@@ -422,20 +425,7 @@ class SnapshotProcessor(HDF5Container):
self.def_rules()
def init_pymses(self):
# If ratarmount was used
if os.path.exists(f"{self.path}/output_{self.num:05}/output_{self.num:05}"):
path = f"{self.path}/output_{self.num:05}"
else:
path = self.path
self._ro = pymses.RamsesOutput(
path,
self.num,
order=self.params.pymses.order,
verbose=self.params.pymses.verbose,
check_endianness=False,
)
def check_variables(self, path):
# Check if variables are in output
name_conv = {
"rho": "density",
@@ -479,16 +469,16 @@ class SnapshotProcessor(HDF5Container):
def is_available(available_vars, pymsesrc, var):
if var in ["g", "phi"]:
if not has_grav:
self.logger.warning(f"Variable {var} not in output")
self.logger.debug(f"Variable {var} not in output")
if var not in pymsesrc:
self.logger.warning(f"Variable {var} not in pymsesrc")
self.logger.debug(f"Variable {var} not in pymsesrc")
return has_grav and var in pymsesrc
else:
if var in name_conv:
if name_conv[var] not in available_vars:
self.logger.warning(f"Variable {var} not in output")
self.logger.debug(f"Variable {var} not in output")
if var not in pymsesrc:
self.logger.warning(f"Variable {var} not in pymsesrc")
self.logger.debug(f"Variable {var} not in pymsesrc")
return name_conv[var] in available_vars and var in pymsesrc
else:
self.logger.warning(f"Variable {var} is unknown")
@@ -528,6 +518,24 @@ class SnapshotProcessor(HDF5Container):
pymses.rcConfig.Ramses.amr_fields.remove_field("g")
pymses.rcConfig.Ramses.amr_fields.remove_field("phi")
def init_pymses(self):
# If ratarmount was used
if os.path.exists(f"{self.path}/output_{self.num:05}/output_{self.num:05}"):
path = f"{self.path}/output_{self.num:05}"
else:
path = self.path
self._ro = pymses.RamsesOutput(
path,
self.num,
order=self.params.pymses.order,
verbose=self.params.pymses.verbose,
check_endianness=False,
)
if self.params.pymses.check_variables:
self.check_variables(path)
self._amr = self._ro.amr_source(self.params.pymses.variables)
self._part = self._ro.particle_source(self.params.pymses.part_variables)
@@ -676,12 +684,12 @@ class SnapshotProcessor(HDF5Container):
if len(data[key] > 0):
if f"/{group}/{key}" in hdf5 and overwrite:
hdf5.remove_node(f"/{group}/{key}")
hdf5.create_array(
f"/{group}", key, data[key], "", createparents=True
)
unit = self._get_units(self.unit_key[key])
hdf5.get_node("/{group}/{key}").unit = unit
nb_written += 1
hdf5.create_array(
f"/{group}", key, data[key], "", createparents=True
)
unit = self._get_units(self.unit_key[key])
hdf5.get_node(f"/{group}/{key}").unit = unit
nb_written += 1
else:
self.logger.warning("Empty key")
if "namelist" not in hdf5.root._v_attrs:
@@ -739,12 +747,12 @@ class SnapshotProcessor(HDF5Container):
self.cells = self.load_data(
cells_src,
filename,
self.params.process.save_cells and save,
True,
keys=keys,
group="cells",
)
self.cells_loaded = True
self.logger.info("Cells loaded")
self.logger.info("Cells loaded")
def unload_cells(self):
"""
@@ -754,7 +762,7 @@ class SnapshotProcessor(HDF5Container):
if self.cells_loaded:
del self.cells
self.cells_loaded = False
self.logger.info("Cells unloaded")
self.logger.info("Cells unloaded")
def load_destructured(self, save=True):
self.load_cells(save=save)
@@ -1032,6 +1040,7 @@ class SnapshotProcessor(HDF5Container):
unit=None,
logbins=False,
weight_func=vol_func,
**kwargs,
):
self.load_cells()
data = getter(self.cells)
@@ -1043,7 +1052,9 @@ class SnapshotProcessor(HDF5Container):
if self.params.process.unload_cells:
self.unload_cells()
values, edges = np.histogram(data, bins, weights=weights, density=True)
values, edges = np.histogram(
data, bins, weights=weights, density=True, **kwargs
)
centers = 0.5 * (edges[1:] + edges[:-1])
return (np.stack([values, centers]), {"logbins": logbins})