[postprocessor] Ensure units are good in Q computation

postprocessor.py

@@ -840,61 +840,55 @@ class PostProcessor(HDF5Container):
         dmap_jeans_ratio = dmap_jeans * 2 ** (dmap_levels)
         return dmap_jeans_ratio
 
-    def _toomreQ_disk(self, ax_los):
-        """
-        Compute the Toomre Q parameter in a Keplerian disk
-        """
-
+    def _omega_average(self, ax_los):
         # Operator to compute the angular speed times rho
+        lbox = self.lbox
+
         def omega_rho_func(dset):
             pos = self.oct_getter_pos_disk(dset)
-            xx = pos[:, :, 0]
-            yy = pos[:, :, 1]
-            rc = np.sqrt(xx ** 2 + yy ** 2)  # cylindrical radius
+            xx = pos[:, :, 0] * lbox
+            yy = pos[:, :, 1] * lbox
+            rc2 = xx ** 2 + yy ** 2  # square of cylindrical radius
             vx = dset["vel"][:, :, 0]
             vy = dset["vel"][:, :, 1]
-            omega_rho = 1.0 / rc ** 2
-            omega_rho = omega_rho * dset["rho"]
+            omega_rho = dset["rho"]
             vyx = vy * xx
             vxy = vx * yy
-            omega_rho = omega_rho * (vyx - vxy)
+            omega_rho = omega_rho * (vyx - vxy) / rc2
             return omega_rho
 
         # Operator to compute the angular speed
-        omega_op = FractionOperator(
-            omega_rho_func, lambda dset: dset["rho"], 1.0 / self._ro.info["unit_time"]
+        omega_unit = self._ro.info["unit_velocity"] / (
+            self._ro.info["unit_length"] / lbox
         )
-
-        # Operator to compute the square of the sound speed
-        T_op = FractionOperator(
-            lambda dset: dset["P"],
+        omega_op = FractionOperator(
+            omega_rho_func,
             lambda dset: dset["rho"],
-            self._ro.info["unit_velocity"],
+            omega_unit,
         )
 
         # Ray tracer for the angular speed
         rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op)
 
-        # Ray tracer for the sound speed
-        if self.pp_params.pymses.fft:
-            rt_T = splatting.SplatterProcessor(
-                self._amr, self._ro.info, T_op, surf_qty=False
-            )
-        else:
-            rt_T = raytracing.RayTracer(self._amr, self._ro.info, T_op)
-
         if not self.pp_params.pymses.multiprocessing:
             rt_omega.disable_multiprocessing()
-            rt_T.disable_multiprocessing()
 
-        dmap_omega = rt_omega.process(self._cam[ax_los])
-        dmap_T = rt_T.process(self._cam[ax_los])
-        dmap_col = self.save.root.maps.coldens_z.read()
-        map_Q = (
-            (self.lbox * np.sqrt(dmap_T.map.T))
-            * dmap_omega.map.T
-            / (np.pi * self.G * dmap_col)
-        )
+        dmap_omega = rt_omega.process(self._cam[ax_los]).map.T
+        return dmap_omega
+
+    def _toomreQ_disk(self, ax_los):
+        """
+        Compute the Toomre Q parameter in a Keplerian disk
+        """
+
+        # Get maps
+
+        dmap_cs2 = self.get_value("/maps/T_mwavg_z")
+        dmap_col = self.get_value("/maps/coldens_z")
+        dmap_omega = self.get_value("/maps/omega_mwavg_z")
+
+        # Compute Q
+        map_Q = (np.sqrt(dmap_cs2) * dmap_omega) / (np.pi * self.G * dmap_col)
         return map_Q
 
     def _radial_bins(self, ax_los="z"):
@@ -1352,6 +1346,14 @@ class PostProcessor(HDF5Container):
                 "/maps",
                 unit=self.info["unit_pressure"] / self.info["unit_density"],
             ),
+            "omega_mwavg": Rule(
+                self,
+                partial(self._omega_average),
+                "Ax mass-weighted averaged azimuthal of angular frequency",
+                "/maps",
+                unit=self.info["unit_velocity"]
+                / (self.info["unit_length"] / self.lbox),
+            ),
             "alpha_disk": Rule(
                 self,
                 self._alpha_disk,
@@ -1423,14 +1425,14 @@ class PostProcessor(HDF5Container):
             "jeans": Rule(
                 self,
                 self._jeans,
-                "Jeans lenght slice",
+                "Jeans length slice",
                 "/maps",
                 dependencies=["rho", "T"],
             ),
             "jeans_ratio": Rule(
                 self,
                 self._jeans_ratio,
-                "Jeans' lenght divided by the max resolution",
+                "Jeans' length divided by the max resolution",
                 "/maps",
                 dependencies=["jeans", "levels"],
             ),
@@ -1439,7 +1441,7 @@ class PostProcessor(HDF5Container):
                 self._toomreQ_disk,
                 "Toomre Q parameter for a Keplerian disk",
                 "/maps",
-                dependencies=["coldens"],
+                dependencies=["coldens", "T_mwavg", "omega_mwavg"],
             ),
             "sinks": Rule(
                 self,