[postprocessor] Add Q computation with real epicyclic frequency

postprocessor.py

@@ -876,19 +876,46 @@ class PostProcessor(HDF5Container):
         dmap_omega = rt_omega.process(self._cam[ax_los]).map.T
         return dmap_omega
 
-    def _toomreQ_disk(self, ax_los):
+    def _toomreQ_disk(self, ax_los, omega_approx=False):
         """
-        Compute the Toomre Q parameter in a Keplerian disk
+        Compute the Toomre Q parameter
         """
 
         # Get maps
 
-        dmap_cs2 = self.get_value("/maps/T_mwavg_z")
+        cs2 = self.get_value("/maps/T_mwavg_z")
-        dmap_col = self.get_value("/maps/coldens_z")
+        coldens = self.get_value("/maps/coldens_z")
-        dmap_omega = self.get_value("/maps/omega_mwavg_z")
+        omega = self.get_value("/maps/omega_mwavg_z")
 
         # Compute Q
-        map_Q = (np.sqrt(dmap_cs2) * dmap_omega) / (np.pi * self.G * dmap_col)
+        if omega_approx:
+            # Use angular frequency as epiciclic frequency (true if the disk is Keplerian)
+            kappa = omega
+        else:
+            # Get coordinates
+            im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox
+            map_size = self.pp_params.pymses.map_size
+            pos_star = self.pp_params.disk.pos_star
+
+            # Physical size of cells
+            dx = (im_extent[1] - im_extent[0]) / map_size
+            dy = (im_extent[3] - im_extent[2]) / map_size
+
+            # Physical coordinates of the center of the cells
+            x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx
+            y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy
+
+            xx, yy = np.meshgrid(x, y)
+            rr = np.sqrt((xx - pos_star[0]) ** 2 + (yy - pos_star[1]) ** 2)
+
+            # Compute kappa**2 = (2omega/R)*d(R**2*omega)/dR
+            Gx, Gy = np.gradient(rr ** 2 * omega)
+            Gr = (xx * Gx + yy * Gy) / rr
+            kappa_square = (2 * omega / rr) * Gr
+            kappa = np.square(kappa_square)
+
+        map_Q = (np.sqrt(cs2) * kappa) / (np.pi * self.G * coldens)
+
         return map_Q
 
     def _radial_bins(self, ax_los="z"):