Ax-aware plotting, More PDFs, improve format

postprocessor.py

@@ -2,8 +2,9 @@
 
 from baseprocessor import *
 
-mass_func = lambda dset: dset["rho"] * dset.get_sizes() ** 3  # Mass function
-vol_func = lambda dset: dset.get_sizes() ** 3  # Volume function
+mass_func = lambda dset: dset["rho"] * dset["dx"] ** 3  # Mass function
+vol_func = lambda dset: dset["dx"] ** 3  # Volume function
+getter_T = lambda dset: dset["P"] / dset["rho"]  # Temperature
 
 
 class PostProcessor(HDF5Container):
@@ -35,6 +36,10 @@ class PostProcessor(HDF5Container):
 
         self.filename = path_out + "/postproc_" + tag_name + format(num, "05") + ".h5"
 
+        self.cells_filename = (
+            path_out + "/cells_" + tag_name + format(num, "05") + ".h5"
+        )
+
         if not os.path.exists(path_out):
             os.makedirs(path_out)
 
@@ -68,10 +73,10 @@ class PostProcessor(HDF5Container):
         self._lbox = self.info["boxlen"]
         center = pp_params.pymses.center
         im_extent = [
-            (-self._radius + center[0]) * self._lbox,
-            (self._radius + center[0]) * self._lbox,
-            (-self._radius + center[1]) * self._lbox,
-            (self._radius + center[1]) * self._lbox,
+            (-self._radius + center[0]),
+            (self._radius + center[0]),
+            (-self._radius + center[1]),
+            (self._radius + center[1]),
         ]
 
         # Get time
@@ -82,13 +87,14 @@ class PostProcessor(HDF5Container):
         self.save.root._v_attrs.run = os.path.basename(path)
         self.save.root._v_attrs.num = num
         self.save.root._v_attrs.lbox = self._lbox
+        self.save.root._v_attrs.unit_length = self.info["unit_length"]
         self.save.root._v_attrs.time = time
 
         if not "/maps" in self.save:
             self.save.create_group("/", "maps", "2D maps")
-            self.save.root.maps._v_attrs.center = center
-            self.save.root.maps._v_attrs.radius = self._radius
-            self.save.root.maps._v_attrs.im_extent = im_extent
+        self.save.root.maps._v_attrs.center = center
+        self.save.root.maps._v_attrs.radius = self._radius
+        self.save.root.maps._v_attrs.im_extent = im_extent
 
         # Initialize cameras
         self._cam = {}
@@ -119,9 +125,34 @@ class PostProcessor(HDF5Container):
         (/!\ Long and memory heavy)
         """
         if not self.cells_loaded:
-            cell_source = CellsToPoints(self._amr)
-            self.cells = cell_source.flatten()
-            self.cells_loaded = True
+            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()
+
+                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
 
     def unload_cells(self):
         """
@@ -195,16 +226,16 @@ class PostProcessor(HDF5Container):
         else:
             return np.sum(value, axis=0)
 
-    def _vol_pdf(self, getter, log=False, weight_func=vol_func):
+    def _vol_pdf(self, getter, bins=100, logbins=False, weight_func=vol_func):
         self.load_cells()
         data = getter(self.cells)
         if logbins:
             data = np.log10(data)
         weights = weight_func(self.cells)
 
-        values, edges = np.histogram(data, weights=weights)
+        values, edges = np.histogram(data, bins, weights=weights)
         centers = 0.5 * (edges[1:] + edges[:-1])
-        return np.stack([values, centers])
+        return (np.stack([values, centers]), {"logbins": logbins})
 
     def _mwa_sigma(self, axes=["x", "y", "z"]):
         mw_speed = self.save.get_node("/globals/mwa_speed").read()
@@ -635,9 +666,24 @@ class PostProcessor(HDF5Container):
             # PDF
             "rho_pdf": Rule(
                 self,
-                partial(self._vol_pdf, partial(simple_getter, "rho")),
+                partial(self._vol_pdf, partial(simple_getter, "rho"), logbins=True),
                 "Global rho-PDF",
                 "/hist",
+                unit=self.info["unit_density"],
+            ),
+            "T_pdf": Rule(
+                self,
+                partial(self._vol_pdf, getter_T, logbins=True),
+                "Global T-PDF",
+                "/hist",
+                unit=self.info["unit_temperature"],
+            ),
+            "P_pdf": Rule(
+                self,
+                partial(self._vol_pdf, getter_T, logbins=True),
+                "Global P-PDF",
+                "/hist",
+                unit=self.info["unit_pressure"],
             ),
             # globals
             "time_num": Rule(