From a286c08b72f1d91f77861b47d867693d4cb5fb48 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Wed, 23 Jun 2021 17:14:09 +0200
Subject: [PATCH] [postprocessor] add option to degrade maps

---
 postprocessor.py | 57 ++++++++++++++++++++++++++++++++++++++++++++----
 pspec_new.py     |  2 ++
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/postprocessor.py b/postprocessor.py
index ef1f2fa..3472017 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -215,6 +215,46 @@ def pspec(map2D):
     return 0.5 * (kbins[1:] + kbins[:-1]), pspec
 
 
+def degrade_map(dmap, nnew, integrate=False):
+    """Degrade a data dmap to a coarser resolution
+
+    Parameters:
+    -----------
+    cube
+        (numpy.ndarray, ndim=2) input data cube, lengths in each direction must
+
+    integrate
+        (bool, default False) if True, values are added instead of averaged
+
+    Return value:
+    -------------
+    degraded_cube
+        (numpy.ndarray, ndim=3) output data cube, 2**lvl values in each
+        direction
+    """
+
+    assert dmap.ndim == 2
+
+    nold = dmap.shape[0]
+    assert nnew <= nold
+    nsum, rem = divmod(nold, nnew)
+    assert rem == 0
+
+    # For each direction, we split the corresponding axis (length nold) into
+    # 2 axes, the first one being the axis for the output dmap (length nnew),
+    # and the second one the axis to average or integrate over (fine cells).
+
+    # reshape creates a view, no copy is involved
+    v = dmap.reshape((nnew, nsum, nnew, nsum), order="C")
+    # Even indexes are kept, odd ones are to be summed
+    dmap_new = np.einsum("iajb->ij", v)
+
+    if not integrate:
+        dmap_new /= nsum
+
+    return dmap_new
+
+
 class PostProcessor(HDF5Container):
     """
     This class enable to compute and save derived quantities from the raw output
@@ -944,7 +984,7 @@ class PostProcessor(HDF5Container):
         dmap_omega = rt_omega.process(self._cam[ax_los]).map.T
         return dmap_omega
 
-    def _toomreQ_disk(self, ax_los, omega_approx=True, G1_units=True):
+    def _toomreQ_disk(self, ax_los, omega_approx=True, G1_units=True, coarsen_factor=1):
         """
         Compute the Toomre Q parameter
         """
@@ -952,17 +992,26 @@ class PostProcessor(HDF5Container):
         # Get maps
 
         if G1_units:
-            G = 1.
+            G = 1.0
             cs2 = self.get_value("/maps/T_mwavg_z")
             coldens = self.get_value("/maps/coldens_z")
             omega = self.get_value("/maps/omega_mwavg_z")
             space_coeff = self.lbox
         else:
-            G = U.G.express(U.kg**-1 * U.m ** 3 * U.s ** -2)
+            G = U.G.express(U.kg ** -1 * U.m ** 3 * U.s ** -2)
             cs2 = self.get_value("/maps/T_mwavg_z", unit=U.m ** 2 * U.s ** -2)
             coldens = self.get_value("/maps/coldens_z", unit=U.kg * U.m ** -2)
             omega = self.get_value("/maps/omega_mwavg_z", unit=U.s ** -1)
             space_coeff = self.info["unit_length"].express(U.m)
+
+        map_size = self.pp_params.pymses.map_size
+
+        if coarsen_factor > 1:
+            map_size = map_size // coarsen_factor
+            omega = degrade_map(omega, map_size)
+            coldens = degrade_map(coldens, map_size)
+            cs2 = degrade_map(cs2, map_size)
+
         # Compute Q
         if omega_approx:
             # Use angular frequency as epiciclic frequency (true if the disk is Keplerian)
@@ -970,7 +1019,7 @@ class PostProcessor(HDF5Container):
         else:
             # Get coordinates
             im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * space_coeff
-            map_size = self.pp_params.pymses.map_size
+
             center = np.array(self.pp_params.disk.center) * space_coeff
 
             # Physical size of cells
diff --git a/pspec_new.py b/pspec_new.py
index 6299922..009505c 100644
--- a/pspec_new.py
+++ b/pspec_new.py
@@ -62,6 +62,8 @@ def degrade_cube(cube, lvl, integrate=False):
     return cube_new
 
 
+
+
 def calc_k(n, nbinsk, nbig, dkbig, dim=3, saxis=2):
     """Make cubes containing the wave vectors, a list of bins and the
     associated normalization factors