From 3ec31b7fd05323c6f2621071dd7c397335488974 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Fri, 15 Jan 2021 18:28:43 +0100
Subject: [PATCH] [postprocessor] Add load particles function

---
 postprocessor.py | 99 +++++++++++++++++++++++++++++++++---------------
 pp_params.yml    |  3 +-
 2 files changed, 71 insertions(+), 31 deletions(-)

diff --git a/postprocessor.py b/postprocessor.py
index 56bd1c5..99e2ccd 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -237,6 +237,7 @@ class PostProcessor(HDF5Container):
             verbose=self.pp_params.pymses.verbose,
         )
         self._amr = self._ro.amr_source(self.pp_params.pymses.variables)
+        self._part = self._ro.part_source(["vel", "mass"])
 
         if self.pp_params.pymses.filter:
             self.min_coords = np.array(self.pp_params.pymses.min_coords)
@@ -245,6 +246,8 @@ class PostProcessor(HDF5Container):
 
             amr_filt = RegionFilter(box, self._amr)
             self._amr = amr_filt
+            part_filt = RegionFilter(box, self._part)
+            self._part = part_filt
 
         self.info = self._ro.info.copy()
 
@@ -346,42 +349,72 @@ class PostProcessor(HDF5Container):
 
         self.def_rules()
 
+    def load_data(self, points_src, loaded, filename, save):
+        """
+        Load data from the source file in the memory.
+        (Long and memory heavy)
+        """
+        if not self.part_loaded:
+            if os.path.exists(filename):
+                hdf5 = tables.open_file(filename, mode="r")
+                try:
+                    node = hdf5.get_node("/data")
+                    data = {}
+                    for key in node._v_children:
+                        self.part[key] = hdf5.get_node("/data/" + key).read()
+                finally:
+                    hdf5.close()
+            else:
+                data_pymses = points_src.flatten()
+                data = {}
+                for key in data_pymses.fields:
+                    data[key] = data_pymses[key]
+                data["dx"] = data_pymses.get_sizes()
+                data["pos"] = data_pymses.points
+
+                if save:
+                    hdf5 = tables.open_file(filename, mode="w")
+                    try:
+                        for key in data:
+                            hdf5.create_array(
+                                "/data", key, data[key], "", createparents=True
+                            )
+                    finally:
+                        hdf5.close()
+            return data
+
+    def load_parts(self):
+        self.parts = self.load_data(
+            self._part,
+            self.parts_loaded,
+            self.parts_filename,
+            self.pp_params.process.save_parts,
+        )
+        self.parts_loaded = True
+
+    def unload_parts(self):
+        """
+        Free space in the memory by telling the garbage collectors that
+        self.parts is not needed
+        """
+        if self.parts_loaded:
+            del self.parts
+            self.parts_loaded = False
+
     def load_cells(self):
         """
         Load all cells from the source file in the memory.
         Cells will be accessible trough self.cells
         (Long and memory heavy)
         """
-        if not self.cells_loaded:
-            if os.path.exists(self.cells_filename):
-                cells_hdf5 = tables.open_file(self.cells_filename, mode="r")
-                try:
-                    node = cells_hdf5.get_node("/cells")
-                    self.cells = {}
-                    for key in node._v_children:
-                        self.cells[key] = cells_hdf5.get_node("/cells/" + key).read()
-                finally:
-                    cells_hdf5.close()
-            else:
-                cell_source = CellsToPoints(self._amr)
-                cells_pymses = cell_source.flatten()
-                self.cells = {}
-                for key in cells_pymses.fields:
-                    self.cells[key] = cells_pymses[key]
-                self.cells["dx"] = cells_pymses.get_sizes()
-                self.cells["pos"] = cells_pymses.points
-
-                if self.pp_params.process.save_cells:
-                    cells_hdf5 = tables.open_file(self.cells_filename, mode="w")
-                    try:
-                        for key in self.cells:
-                            cells_hdf5.create_array(
-                                "/cells", key, self.cells[key], "", createparents=True
-                            )
-                    finally:
-                        cells_hdf5.close()
-
-        self.cells_loaded = True
+        cells_src = CellsToPoints(self._amr)
+        self.cells = self.load_data(
+            cells_src,
+            self.cells_loaded,
+            self.cells_filename,
+            self.pp_params.process.save_cells,
+        )
+        self.parts_loaded = True
 
     def unload_cells(self):
         """
@@ -1123,6 +1156,11 @@ class PostProcessor(HDF5Container):
         pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile)
         return True
 
+    def _particles(self):
+        """Ensure particles are written in the hdf5 file"""
+        self.load_parts()
+        self.unload_parts()
+
     def _filaments(self):
 
         datamap_name = self.pp_params.filaments.datamap
@@ -1413,6 +1451,7 @@ class PostProcessor(HDF5Container):
                 },
             ),
             "pspec": Rule(self, self._pspec, "Power spectrum", "/hdf5"),
+            "particles": Rule(self, self._particles, "Particles file", "/hdf5"),
             "filaments": Rule(
                 self,
                 self._filaments,
diff --git a/pp_params.yml b/pp_params.yml
index 903c1c7..fd0af68 100644
--- a/pp_params.yml
+++ b/pp_params.yml
@@ -86,9 +86,10 @@ out: # Parameters for post processing
 
 
 process: # General setting of the post-processor module
-    verbose : False                     # Give more infos on what is going on
+    verbose : False                    # Give more infos on what is going on
     num_process : 1                    # Number of forks
     save_cells : True                  # Save cells structure on disk
+    save_particles : True              # Save particles on disk
     unload_cells : True                # Save memory usage
 
 rules: # Specific rules parameters