diff --git a/postprocessor.py b/postprocessor.py index 995f2d1..627359c 100644 --- a/postprocessor.py +++ b/postprocessor.py @@ -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""" - self.load_parts() - self.unload_parts() + 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,