From 31bc3cc2092eb140315a9a6d9163878afbab928d Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Mon, 26 Apr 2021 15:33:24 +0200
Subject: [PATCH] [postprocessor] add rule for radial sigma and 2D pwspec

---
 plotter.py       | 11 +++++++
 postprocessor.py | 79 +++++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/plotter.py b/plotter.py
index 3ab4592..0398489 100644
--- a/plotter.py
+++ b/plotter.py
@@ -1689,6 +1689,17 @@ class Plotter(Aggregator, BaseProcessor):
                 dependencies=["radial_centers", "rad_mwavg_" + name],
             )
 
+            self.rules["dispersion_rad_" + name] = PlotRule(
+                self,
+                partial(
+                    self._plot,
+                    "/radial/radial_centers",
+                    "/radial/dispersion_rad_" + name,
+                ),
+                "Radial dispersion of {}".format(name),
+                dependencies=["radial_centers", "dispersion_rad_" + name],
+            )
+
             self.rules["avg_map_" + name] = PlotRule(
                 self,
                 partial(self._plot_map, "avg_map_" + name),
diff --git a/postprocessor.py b/postprocessor.py
index 69e6386..af3d080 100644
--- a/postprocessor.py
+++ b/postprocessor.py
@@ -28,6 +28,7 @@ from pymses.analysis import (
 from pymses.filters import CellsToPoints, RegionFilter
 
 from fil_finder import FilFinder2D
+from scipy import fft
 import pspec_new
 
 
@@ -163,7 +164,39 @@ def find_center(distance, skeleton, i_center, j_center, i, j):
         return i_center[i, j], j_center[i, j]
 
 
-# PostProcessor class
+# Power spectrum helpers
+
+
+def pspec(map2D):
+    """
+    Computes the 2D powerspectrum of a 2D map
+
+    Parameters
+    ----------
+    map2D : square map to compute the fft from
+    """
+
+    # Resolution of the map
+    n = map2D.shape[0]
+    assert map2D.shape[1] == n
+
+    # Create bin array for the wavenumber k
+    # (Take into account the symmetry)
+    kbins = np.linspace(0, n // 2, n // 4)
+    k_alias = np.arange(n, dtype=np.float64)
+    k = np.where(k_alias >= n // 2, k_alias - n, k_alias)
+    kx, ky = np.meshgrid(k, k, indexing="ij")
+    # Compute map of k
+    kmap = np.sqrt(kx ** 2 + ky ** 2)
+    # Compute fft
+    fmap = fft.fft2(map2D)
+    # Compute the power map from the fft
+    pmap = pspec_new.pcube(fmap)
+    # Use the power map and the fft to compute the powerspectrum
+    # This is typically an histogram of k weighted by the fourier transform value
+    pspec, kbins, pspec2, fbins = pspec_new.pspectrum(pmap, kmap, kbins, 1, 0)
+    # Return bin center and power spectrum
+    return 0.5 * (kbins[1:] + kbins[:-1]), pspec
 
 
 class PostProcessor(HDF5Container):
@@ -1054,6 +1087,42 @@ class PostProcessor(HDF5Container):
             avg_map = self.get_value("/maps/avg_map_" + name + "_" + ax_los)
         return dmap / avg_map
 
+    def _dispersion_rad(self, name, ax_los="z"):
+        """
+        Computes the dispersion in radial bins of the quantity `name`
+
+        Parameters
+        ----------
+
+        name : name of 2D map of a scalar quantity
+        ax_los : axis of the line of sight
+        """
+
+        # 1. Get full storage names for the map and its azimuthal avererage
+        map_name = f"/maps/{name}_{ax_los}"
+        avg_map_name = f"/maps/avg_map_{name}_{ax_los}"
+
+        # 2. Get maps from the storage
+        dmap = self.get_value(map_name)
+        avgmap = self.get_value(avg_map_name)
+
+        # 3. Get radial bins and centers
+        bins_on_map = self.get_value(f"/maps/bins_on_map_{ax_los}")
+        centers = self.get_value(f"/radial/radial_centers_{ax_los}")
+
+        # 4. Initialize dispersion array
+        sigma = np.zeros(len(centers))
+
+        # 5. Compute RMS in each bin
+        for j in range(len(centers)):
+            mask_bin = bins_on_map == j
+            sigma[j] = np.sqrt(
+                np.sum((dmap[mask_bin] - avgmap[mask_bin]) ** 2) / np.sum(mask_bin)
+            )
+
+        # 6. Returns computed RMS
+        return sigma
+
     def _rad_fluct_pdf(self, name, ax_los="z", mass_weighted=False):
         if mass_weighted:
             fluct_map = self.get_value("/maps/mwfluct_" + name + "_" + ax_los)
@@ -1784,6 +1853,14 @@ class PostProcessor(HDF5Container):
                 "/maps",
                 dependencies=[name, "avg_map_" + name],
             )
+            self.rules["dispersion_rad_" + name] = Rule(
+                self,
+                partial(self._dispersion_rad, name),
+                "radial RMS of {}".format(name),
+                "/radial",
+                dependencies=[name, "avg_map_" + name],
+                unit=unit,
+            )
             self.rules["mwfluct_" + name] = Rule(
                 self,
                 partial(self._fluct_map, name, mass_weigthed=True),