[postprocessor] write particles and cells rules

This commit is contained in:
Noe Brucy
2021-06-15 08:33:27 +02:00
parent f47e422b1a
commit b334c89f61
+21 -6
View File
@@ -394,7 +394,7 @@ class PostProcessor(HDF5Container):
self.def_rules()
def load_data(self, points_src, filename, save):
def load_data(self, points_src, filename, save, keys=None):
"""
Load data from the source file in the memory.
(Long and memory heavy)
@@ -404,7 +404,9 @@ class PostProcessor(HDF5Container):
try:
node = hdf5.get_node("/data")
data = {}
for key in node._v_children:
if keys is None:
keys = node._v_children
for key in keys:
data[key] = hdf5.get_node("/data/" + key).read()
finally:
hdf5.close()
@@ -430,12 +432,13 @@ class PostProcessor(HDF5Container):
hdf5.close()
return data
def load_parts(self):
def load_parts(self, keys=None):
if not self.parts_loaded:
self.parts = self.load_data(
self._part,
self.parts_filename,
self.pp_params.process.save_parts,
keys=keys,
)
self.parts_loaded = True
@@ -448,7 +451,7 @@ class PostProcessor(HDF5Container):
del self.parts
self.parts_loaded = False
def load_cells(self):
def load_cells(self, keys=None):
"""
Load all cells from the source file in the memory.
Cells will be accessible trough self.cells
@@ -460,6 +463,7 @@ class PostProcessor(HDF5Container):
cells_src,
self.cells_filename,
self.pp_params.process.save_cells,
keys=keys,
)
self.cells_loaded = True
@@ -1285,12 +1289,20 @@ class PostProcessor(HDF5Container):
pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile, **kwargs)
return np.array([self.pspec_filename])
def _particles(self):
def _write_particles(self):
"""Ensure particles are written in the hdf5 file"""
if not os.path.exists(self.parts_filename) and not self.parts_loaded:
self.load_parts()
self.unload_parts()
return np.array([self.parts_filename])
def _write_cells(self):
"""Ensure cells are written in the hdf5 file"""
if os.path.exists(self.cells_filename) and not self.cells_loaded:
self.load_cells()
self.unload_cells()
return np.array([self.cells_filename])
def _filaments(self):
datamap_name = self.pp_params.filaments.datamap
@@ -1557,7 +1569,10 @@ class PostProcessor(HDF5Container):
},
),
"pspec": Rule(self, self._pspec, "Power spectrum", "/hdf5"),
"particles": Rule(self, self._particles, "Particles file", "/hdf5"),
"write_particles": Rule(
self, self._write_particles, "Particles file", "/hdf5"
),
"write_cells": Rule(self, self._write_cells, "Cells file", "/hdf5"),
"filaments": Rule(
self,
self._filaments,