Add rules to compute $\alpha$

aggregator.py

@@ -8,7 +8,7 @@ def _map_rule(rule, arg, overwrite, path, path_out, pp_params, run_num):
         )
     except Exception as e:
         print(e)
-    return pp.process(rule, arg, overwrite)
+    return pp.process(rule, arg, overwrite, overwrite)
 
 
 class Aggregator:

plotter.py

@@ -78,6 +78,9 @@ class Plotter(Aggregator, BaseProcessor):
         "coldens0": "$\Sigma_0$",
         "sfr_avg_window": "window",
         "bx_bound": "$B_0$",
+        "levelmax": "$l_{\max}$",
+        "levelmin": "$l_{\min}$",
+        "comp_frac": "$1 - \\zeta$",
     }
 
     # Conversion table from namelist values (from amses config file)  into LaTex strings
@@ -96,7 +99,7 @@ class Plotter(Aggregator, BaseProcessor):
         pp_params=None,
         selector=None,
         tag=None,
-        **kwargs
+        **kwargs,
     ):
 
         """
@@ -147,7 +150,9 @@ class Plotter(Aggregator, BaseProcessor):
         Check if the dependency belongs to the plotter object or to another one (comp, pp, ..)
         """
         if dep in self.comp.rules:
-            done = self.comp.process(dep, dep_arg, overwrite, overwrite)
+            done = self.comp.process(
+                dep, dep_arg, overwrite, overwrite_dep=self.overwrite_dep
+            )
             self.just_done.extend(done)
         else:
             super(Plotter, self)._not_self_dep(name, dep, dep_arg, overwrite, **kwargs)
@@ -171,10 +176,10 @@ class Plotter(Aggregator, BaseProcessor):
         ax=None,
         movie=False,
         from_cells=False,
-        **kwargs
+        **kwargs,
     ):
         """
-        Generic method to process a rule, with its dependencies
+        Open storage and figure if needed before processing a rule
         """
         if not arg is None:
             name_full = name + "_" + str(arg)
@@ -217,7 +222,7 @@ class Plotter(Aggregator, BaseProcessor):
                                 overwrite,
                                 ax=ax[i],
                                 run=run,
-                                **kwargs
+                                **kwargs,
                             )
                         except TypeError as e:
                             if str(e) in [
@@ -234,7 +239,7 @@ class Plotter(Aggregator, BaseProcessor):
                                     overwrite,
                                     ax=ax,
                                     run=run,
-                                    **kwargs
+                                    **kwargs,
                                 )
                             elif ax is None:
                                 fig = P.figure()
@@ -246,7 +251,7 @@ class Plotter(Aggregator, BaseProcessor):
                                     overwrite,
                                     ax=P.gca(),
                                     run=run,
-                                    **kwargs
+                                    **kwargs,
                                 )
                             else:
                                 raise
@@ -414,7 +419,7 @@ class Plotter(Aggregator, BaseProcessor):
         norm="log",
         put_cbar=True,
         autoscale=True,
-        **kwargs
+        **kwargs,
     ):
         """
         Plot data on a map
@@ -477,39 +482,81 @@ class Plotter(Aggregator, BaseProcessor):
             P.title(title)
 
         for i, plot_overlay in enumerate(overlays):
-            try:
+            if plot_overlay in self.overlays:
-                plot_overlay(ax_los, **overlays_kwargs[i])
+                plot_overlay = self.overlays[plot_overlay]
-            except:
-                plot_overlay(ax_los)
 
-    def _overlay_levels(self, ax_los):
+            try:
+                plot_overlay(ax_los, im_extent, **overlays_kwargs[i])
+            except:
+                plot_overlay(ax_los, im_extent)
+
+    def _overlay_contour(
+        self,
+        ax_los,
+        im_extent,
+        map_name,
+        log=False,
+        lvl_array=None,
+        lw=None,
+        lvl_th=None,
+        lvl_max_lbl=np.inf,
+        lbl_fmt="%g",
+        **kwargs,
+    ):
+        """
+        Add an overlay : contour of other map
+        """
+        map_contour = self.save.get_node("/maps/{}_{}".format(map_name, ax_los)).read()
+        if log:
+            map_contour = np.log10(map_contour)
+        # Computing linewidths
+        mask_fin = np.isfinite(map_contour)
+        if lvl_array is None:
+            lvl_array = np.arange(
+                np.min(map_contour[mask_fin]), np.max(map_contour[mask_fin]) + 1
+            )
+
+        if lw is None:
+            lw = np.ones(lvl_array.size) * 2
+        if lvl_th:
+            lw[lvl_array >= lvl_th] = lw[lvl_array >= lvl_th] ** (
+                lvl_th - lvl_array[lvl_array >= lvl_th]
+            )
+            lw[lvl_array < lvl_th] = 1.0
+
+        cont = P.contour(
+            map_contour,
+            extent=im_extent,
+            origin="lower",
+            linewidths=lw,
+            levels=lvl_array,
+            **kwargs,
+        )
+
+        P.clabel(
+            cont,
+            lvl_array[lvl_array < lvl_max_lbl],
+            inline=1,
+            fontsize=8.0,
+            fmt=lbl_fmt,
+        )
+
+    def _overlay_levels(self, ax_los, im_extent, **kwargs):
         """
         Add an overlay : AMR levels
         """
-        map_level = self.save.get_node("/maps/{}_{}".format("levels", ax_los)).read()
+        return self._overlay_contour(
-        # Computing linewidths
+            ax_los,
-        levels_ar = np.arange(np.min(map_level), np.max(map_level) + 1)
+            im_extent,
-        lw = np.ones(levels_ar.size) * 2
+            "levels",
-        lvl_th = 8  # Level threeshold for reducing linewidths
+            lbl_fmt="%1d",
-        lw[levels_ar >= lvl_th] = lw[levels_ar >= lvl_th] ** (
+            lvl_th=8,
-            lvl_th - levels_ar[levels_ar >= lvl_th]
+            lvl_max_lbl=11,
+            **kwargs,
         )
-        lw[levels_ar < lvl_th] = 1.0
-
-        cont = P.contour(
-            map_level,
-            extent=self.save.root.maps._v_attrs.im_extent,
-            origin="lower",
-            colors="grey",
-            linewidths=lw,
-            levels=levels_ar,
-        )
-        cont.levels = cont.levels + 1
-
-        P.clabel(cont, cont.levels[cont.levels < 11], inline=1, fontsize=8.0, fmt="%1d")
 
     def _overlay_speed(
-        self, ax_los, unit=cst.km_s, unit_coeff=1.0, key_v=None, **kwargs
+        self, ax_los, im_extent, unit=cst.km_s, unit_coeff=1.0, key_v=None, **kwargs
     ):
         """
         Add an overlay : velocity vector field
@@ -532,6 +579,8 @@ class Plotter(Aggregator, BaseProcessor):
         )  # take only a subset of velocities
         map_vv_red = dmap_vv[::vel_red, ::vel_red] * unit_old.express(unit)
 
+        # TODO : redo this with im_extent
+
         nh = map_vh_red.shape[0]
         nv = map_vv_red.shape[1]
         vec_h = (
@@ -559,7 +608,7 @@ class Plotter(Aggregator, BaseProcessor):
             coordinates="figure",
         )
 
-    def _overlay_B(self, ax_los, **kwargs):
+    def _overlay_B(self, ax_los, im_extent, **kwargs):
         """
         Add an overlay : magnetic streamlines
         """
@@ -569,6 +618,8 @@ class Plotter(Aggregator, BaseProcessor):
         dmap_Bh = dmap_Bh_node.read()
         dmap_Bv = self.save.get_node("/maps/B_v_{}".format(ax_los)).read()
 
+        # TODO : redo this with im_extent
+
         vel_red = self.pp_params.plot.vel_red
         radius = self.save.root.maps._v_attrs.radius
         center = self.save.root.maps._v_attrs.center
@@ -590,14 +641,33 @@ class Plotter(Aggregator, BaseProcessor):
 
         P.streamplot(hh, vv, map_Bh_red, map_Bv_red)
 
-    def _plot_radial(self, name, ax_los, label=None, xlog=False, ylog=False):
+    def _plot_radial(
+        self,
+        name,
+        ax_los,
+        run,
+        ylabel=None,
+        xlog=False,
+        ylog=False,
+        ytransform=None,
+        title=None,
+        nml_key=None,
+        put_title=True,
+        put_time=True,
+        time_unit=cst.Myr,
+        **kwargs,
+    ):
         """
         Plot a radial profile (for disks, OUTDATED)
         """
 
         radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
         bin_centers = 0.5 * (radial_bins[1:] + radial_bins[:-1])
-        mean_bin = self.save.get_node("/radial/{}_{}".format(name, ax_los)).read()
+        node = self.save.get_node("/radial/{}_{}".format(name, ax_los))
+        mean_bin = node.read()
+
+        if ytransform is not None:
+            mean_bin = ytransform(mean_bin)
 
         P.grid()
         P.xlabel(r"$r$")
@@ -606,20 +676,37 @@ class Plotter(Aggregator, BaseProcessor):
             P.xscale("log")
         if ylog:
             P.yscale("log")
-        P.plot(bin_centers, mean_bin)
+        P.plot(bin_centers, mean_bin, **kwargs)
 
-        if not label is None:
+        if not ylabel is None:
-            P.ylabel(label)
+            P.ylabel(ylabel)
+
+        if put_title:
+            title = self._label_run(run, node, title, nml_key)
+
+            if put_time:
+                time = self.save.root._v_attrs.time * self.comp.info["unit_time"]
+                time_str = self.pp_params.plot.time_fmt.format(
+                    time.express(time_unit), time_unit.latex
+                )
+                if len(title) > 0:
+                    title = title + " | " + time_str
+                else:
+                    title = time_str
+
+            P.title(title)
 
     def _plot_hist(
         self,
         name,
-        ax_los,
+        ax_los=None,
-        run,
+        run=None,
         group="/hist/",
-        label=None,
+        xlabel=None,
         unit=None,
         unit_coeff=1.0,
+        ytransform=None,
+        label=None,
         title=None,
         nml_key=None,
         put_time=True,
@@ -633,7 +720,7 @@ class Plotter(Aggregator, BaseProcessor):
         nml_color=None,
         fit=None,
         fitlabel=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Plot an histogram (PDF, etc ...)
@@ -649,7 +736,7 @@ class Plotter(Aggregator, BaseProcessor):
             except:
                 xlog = False
 
-        label, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff)
+        xlabel, unit_old, unit = self._ax_label_unit(node, label, unit, unit_coeff)
 
         if "mean" in node:
             index = node["runs"].read().index(run)
@@ -662,6 +749,9 @@ class Plotter(Aggregator, BaseProcessor):
         else:
             centers = centers * unit_old.express(unit)
 
+        if ytransform is not None:
+            values = ytransform(values)
+
         title = self._label_run(run, node, title, nml_key)
 
         if put_time:
@@ -686,19 +776,22 @@ class Plotter(Aggregator, BaseProcessor):
                 except:
                     color = colors(nml)
 
+        if label == None:
+            label = title
+
         if kind == "bar":
             width = centers[1] - centers[0]
-            P.bar(centers, values, width, log=ylog, color=color, label=title, **kwargs)
+            P.bar(centers, values, width, log=ylog, color=color, label=label, **kwargs)
         elif kind == "step":
             if ylog:
                 P.yscale("log")
-            P.step(centers, values, where="mid", color=color, label=title, **kwargs)
+            P.step(centers, values, where="mid", color=color, label=label, **kwargs)
         else:
             raise ValueError("kind must be 'bar' or 'step'")
         P.grid()
 
         if not label is None:
-            P.xlabel(label)
+            P.xlabel(xlabel)
         if not ylabel is None:
             P.ylabel(ylabel)
 
@@ -749,7 +842,7 @@ class Plotter(Aggregator, BaseProcessor):
         legend=None,
         subname_x=None,
         subname_y=None,
-        **kwargs
+        **kwargs,
     ):
         """
         Generic plot routine, with name_x and name_y two path in the hdf5 file
@@ -991,7 +1084,7 @@ class Plotter(Aggregator, BaseProcessor):
 
     def def_rules(self):
         """
-        That is where rules are defined
+        This is where rules are defined
         """
         self.rules = {
             # Generic rules
@@ -1000,13 +1093,69 @@ class Plotter(Aggregator, BaseProcessor):
             ),
             "coldens": PlotRule(
                 self,
-                partial(self._plot_map, "coldens", label=r"$\Sigma$", unit=cst.coldens),
+                partial(
+                    self._plot_map,
+                    "coldens",
+                    label=r"$\Sigma$",
+                    # unit=cst.coldens
+                ),
                 "Column density map",
-                dependencies=["coldens", "speed_h", "speed_v"],
+                dependencies=["coldens"],
+            ),
+            "alpha_disk": PlotRule(
+                self,
+                partial(self._plot_map, "alpha_disk", label=r"$\alpha$"),
+                "Map of the Shakura&Sunaev alpha parameter for disks",
+                dependencies=["alpha_disk"],
+            ),
+            "alpha_grav": PlotRule(
+                self,
+                partial(self._plot_map, "alpha_grav", label=r"$\alpha_g$"),
+                "Map of the grav Shakura&Sunaev alpha parameter for disks",
+                dependencies=["alpha_grav"],
+            ),
+            "vphi": PlotRule(
+                self,
+                partial(
+                    self._plot_map,
+                    "vphi",
+                    label=r"$v_\phi$",
+                    # unit=cst.km_s
+                ),
+                "Azimuthal speed",
+                dependencies=["vphi"],
+            ),
+            "vr": PlotRule(
+                self,
+                partial(
+                    self._plot_map,
+                    "vr",
+                    label=r"$v_r$",
+                    # unit=cst.km_s
+                ),
+                "Radial speed",
+                dependencies=["vr"],
+            ),
+            "P_avg": PlotRule(
+                self,
+                partial(self._plot_map, "P_avg", label=r"$P$"),
+                "Pressure (average)",
+                dependencies=["P_avg"],
+            ),
+            "rho_avg": PlotRule(
+                self,
+                partial(self._plot_map, "rho_avg", label=r"$\rho$"),
+                "Density (average)",
+                dependencies=["rho_avg"],
             ),
             "rho": PlotRule(
                 self,
-                partial(self._plot_map, "rho", label=r"$\rho$", unit=cst.Msun_pc3),
+                partial(
+                    self._plot_map,
+                    "rho",
+                    label=r"$\rho$",
+                    # unit=cst.Msun_pc3
+                ),
                 "Density slice at s = 0, with s = x, y or z.",
                 dependencies=["rho"],
             ),
@@ -1058,18 +1207,6 @@ class Plotter(Aggregator, BaseProcessor):
                 "Density slice with magnetic field and velocity  overlay",
                 dependencies=["rho", "B_h", "B_v", "speed_h", "speed_v"],
             ),
-            "B_int_B_vel": PlotRule(
-                self,
-                partial(
-                    self._plot_map,
-                    "B_int",
-                    label="B",
-                    unit=cst.T,
-                    overlays=[self._overlay_B, self._overlay_speed],
-                ),
-                "Magnetic slice with magnetic field and velocity  overlay",
-                dependencies=["B_int", "B_h", "B_v", "speed_h", "speed_v"],
-            ),
             "jeans_ratio": PlotRule(
                 self,
                 partial(
@@ -1303,7 +1440,19 @@ class Plotter(Aggregator, BaseProcessor):
             ),
         }
 
-        averageables = ["coldens", "rho", "T", "Q"]
+        averageables = [
+            "coldens",
+            "rho",
+            "T",
+            "Q",
+            "vr",
+            "vphi",
+            "rho_avg",
+            "P_avg",
+            "T_avg",
+            "alpha_disk",
+            "alpha_grav",
+        ]
         for name in averageables:
             self.rules["rad_" + name] = PlotRule(
                 self,
@@ -1331,6 +1480,7 @@ class Plotter(Aggregator, BaseProcessor):
                 "Fluctuation of {}".format(name),
                 dependencies=["fluct_" + name],
             )
+
             self.rules["pdf_" + name] = PlotRule(
                 self,
                 partial(
@@ -1343,4 +1493,22 @@ class Plotter(Aggregator, BaseProcessor):
                 dependencies=["fit_pdf_" + name],
             )
 
+            for name_bin in averageables:
+                if name_bin is not name:
+                    group = "mbb_{}_{}".format(name, name_bin)
+                    self.rules["mbb_" + name + "_" + name_bin] = PlotRule(
+                        self,
+                        partial(self._plot_hist, group, ylabel=r"$\alpha$"),
+                        "Mean of {} by bins of {}".format(name, name_bin),
+                        dependencies=[group],
+                    )
+
+        # Dict of overlays
+        self.overlays = {
+            "B": self._overlay_B,
+            "speed": self._overlay_speed,
+            "levels": self._overlay_levels,
+            "contour": self._overlay_contour,
+        }
+
         super(Plotter, self).def_rules()

postprocessor.py

@@ -2,10 +2,39 @@
 import pandas as pd
 import pspec_new
 from baseprocessor import *
+import pymses.utils.regions as reg
+from pymses.filters import RegionFilter
 
-mass_func = lambda dset: dset["rho"] * dset["dx"] ** 3  # Mass function
+
-vol_func = lambda dset: dset["dx"] ** 3  # Volume function
+# Getters
-getter_T = lambda dset: dset["P"] / dset["rho"]  # Temperature
+
+
+def mass_func(dset):
+    try:
+        dx = dset["dx"]
+    except:
+        dx = dset.get_sizes()
+    return dset["rho"] * dx ** 3  # Mass function
+
+
+def vol_func(dset):
+    return dset["dx"] ** 3  # Volume function
+
+
+def getter_T(dset):
+    return dset["P"] / dset["rho"]  # Temperature
+
+
+def getter_P(dset):
+    return dset["P"]
+
+
+def getter_abs_cos_vB(dset):
+    B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
+    v_norm = np.sqrt(np.sum(dset["vel"] ** 2, axis=1))
+    # Compute the dot product in each cell
+    dot_prod = np.einsum("ij,ij->i", dset["vel"], dset["Br"])
+    return np.abs(dot_prod) / (v_norm * B_norm)
 
 
 def getter_B_int(dset):
@@ -22,6 +51,55 @@ def getter_v_norm(dset):
     return v_norm
 
 
+# Helpers
+
+
+def mean_by_bins(
+    x,
+    y,
+    bins=100,
+    logbins=False,
+):
+    """
+    Compute the mean of y in bins of x
+
+    Parameters
+    ----------
+    x, y : np.array of same dimensions
+    bins : int, number of bins
+    logbins : bool, if true, the bins will be logaritmically distributed
+    """
+    mask = np.isfinite(x) & np.isfinite(y)
+    x = x[mask].flatten()
+    y = y[mask].flatten()
+
+    if logbins:
+        minvalue = np.min(x[x > 0])
+        x_bins = np.logspace(np.log10(minvalue), np.log10(np.max(x)), bins, base=10)
+    else:
+        x_bins = np.linspace(np.min(x), np.max(x), bins)
+
+    # For each cell, bin_number contains the number of the bins it belongs to
+    bin_number = np.zeros(len(y))
+
+    # Go through the min value of x of each bin
+    for x_min in x_bins[1:-1]:
+        bin_number = bin_number + (x > x_min).astype(int)
+
+    # Compute the mean in each bin
+    y_mean = np.zeros(len(x_bins) - 1)
+    for i in range(len(y_mean)):
+        y_mean[i] = np.mean(y[bin_number == i])
+
+    # Get the center of each bin
+    if logbins:
+        centers = 10 ** (0.5 * (np.log10(x_bins[1:]) + np.log10(x_bins[:-1])))
+    else:
+        centers = 0.5 * (x_bins[1:] + x_bins[:-1])
+
+    return centers, y_mean
+
+
 class PostProcessor(HDF5Container):
     """
     This class enable to compute and save derived quantities from the raw output
@@ -73,6 +151,15 @@ class PostProcessor(HDF5Container):
             verbose=self.pp_params.pymses.verbose,
         )
         self._amr = self._ro.amr_source(self.pp_params.pymses.variables)
+
+        if self.pp_params.pymses.filter:
+            self.min_coords = np.array(self.pp_params.pymses.min_coords)
+            self.max_coords = np.array(self.pp_params.pymses.max_coords)
+            box = reg.Box((self.min_coords, self.max_coords))
+
+            amr_filt = RegionFilter(box, self._amr)
+            self._amr = amr_filt
+
         self.info = self._ro.info.copy()
 
         # Density operator
@@ -88,16 +175,31 @@ class PostProcessor(HDF5Container):
         else:
             self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op)
 
+        if not self.pp_params.pymses.multiprocessing:
+            self._rt.disable_multiprocessing()
+
         # Set the extend of the image
         self._radius = 0.5 / self.pp_params.pymses.zoom
         self.lbox = self.info["boxlen"]
-        center = self.pp_params.pymses.center
+
-        im_extent = [
+        if self.pp_params.pymses.filter:
-            (-self._radius + center[0]),
+            center = (self.max_coords + self.min_coords) / 2.0
-            (self._radius + center[0]),
+            im_extent = [
-            (-self._radius + center[1]),
+                self.min_coords[0],
-            (self._radius + center[1]),
+                self.max_coords[0],
-        ]
+                self.min_coords[1],
+                self.max_coords[1],
+            ]
+            distance = (self.max_coords[2] - self.min_coords[2]) / 2.0
+        else:
+            center = self.pp_params.pymses.center
+            im_extent = [
+                (-self._radius + center[0]),
+                (self._radius + center[0]),
+                (-self._radius + center[1]),
+                (self._radius + center[1]),
+            ]
+            distance = self._radius
 
         # Get time
         time = self._ro.info["time"]  # time in codeunits
@@ -125,9 +227,9 @@ class PostProcessor(HDF5Container):
             self._cam[ax_los] = Camera(
                 center=self.pp_params.pymses.center,
                 line_of_sight_axis=ax_los,
-                region_size=[2.0 * self._radius, 2.0 * self._radius],
+                region_size=[im_extent[1] - im_extent[0], im_extent[3] - im_extent[2]],
-                distance=self._radius,
+                distance=distance,
-                far_cut_depth=self._radius,
+                far_cut_depth=distance,
                 up_vector=ax_v,
                 map_max_size=self.pp_params.pymses.map_size,
             )
@@ -184,6 +286,40 @@ class PostProcessor(HDF5Container):
             del self.cells
             self.cells_loaded = False
 
+    def getter_pos_disk(self, dset):
+        """
+        Returns the position in normalized units centered on the position of the star
+        """
+        pos = dset.get_cell_centers()
+        pos = pos - (np.array(self.pp_params.disk.pos_star) / self.lbox)
+        return pos
+
+    def getter_vect_r(self, dset, name_vect):
+        """ Radial component of a vector """
+        r = self.getter_pos_disk(dset)[:, :, :2]
+        ur = np.transpose(
+            (np.transpose(r, (2, 0, 1)) / np.sqrt(np.sum(r ** 2, axis=2))), (1, 2, 0)
+        )
+        return np.einsum("ikj, ikj -> ik", dset[name_vect][:, :, :2], ur)
+
+    def getter_vect_phi(self, dset, name_vect):
+        """ Azimuthal  component of a vector """
+
+        r = self.getter_pos_disk(dset)[:, :, :2]
+        r_norm = np.sqrt(np.sum(r ** 2, axis=2))
+        rot = np.array([[0, -1], [1, 0]])
+        uphi = np.transpose(np.einsum("ij, klj -> ikl", rot, r) / r_norm, (1, 2, 0))
+        vect = dset[name_vect][:, :, :2]
+
+        return np.einsum("ikj,ikj -> ik", vect, uphi)
+
+    def getter_vr(self, dset):
+        return self.getter_vect_r(dset, "vel")
+
+    def getter_vphi(self, dset):
+        """ Azimuthal velocity """
+        return self.getter_vect_phi(dset, "vel")
+
     def _slice(self, getter, ax_los="z", z=0, unit=cst.none):
         """
         Slice process function.
@@ -216,25 +352,42 @@ class PostProcessor(HDF5Container):
     ):
         """
         Map of the average of a quantity (given by getter) along an axis (ax_los)
-        Return 2D array
+        Returns 2D array if getter returns a scalar quantity
+
+        If surf_qty is set (projection mode), mass_weighted is ignored
         """
-        if mass_weighted:
+        if surf_qty:
-
-            def num(cells):
-                value = getter(cells)
-                mass = mass_func(cells)
-                # Transpose (.T) is for vectorial values
-                return (mass * value.T).T
-
-            op = FractionOperator(num, mass_function, unit)
-        else:
             op = ScalarOperator(getter, unit)
+        else:
+            if mass_weighted:
+
+                def num(dset):
+                    value = getter(dset)
+                    rho = getter_rho(dset)
+                    return rho * value
+
+                def denum(dset):
+                    return getter_rho(dset)
+
+            else:  # Volume weighted
+
+                def num(dset):
+                    value = getter(dset)
+                    return value
+
+                def denum(dset):
+                    return 1.0
+
+            op = FractionOperator(num, denum, unit)
 
         if self.pp_params.pymses.fft:
             rt = splatting.SplatterProcessor(self._amr, self._ro.info, op)
         else:
             rt = raytracing.RayTracer(self._amr, self._ro.info, op)
 
+        if not self.pp_params.pymses.multiprocessing:
+            rt.disable_multiprocessing()
+
         datamap = rt.process(self._cam[ax_los], surf_qty=surf_qty)
         return datamap.map.T
 
@@ -250,6 +403,7 @@ class PostProcessor(HDF5Container):
         """
         Profile of the average of a quantity (given by getter) perpendicular to an axis
         WARNING : This version only works on an uniform grid, need of a box version for AMR
+        Returns 1D array if getter returns a scalar quantity
         """
         self.load_cells()
         if isinstance(axis, str):
@@ -273,6 +427,10 @@ class PostProcessor(HDF5Container):
         return df.groupby("axis").mean().values[:, 0]
 
     def _vol_avg(self, getter, mass_weighted=True):
+        """
+        Global volumic (or mass_weighted) average of the quantity returned by getter
+        Returns a scalar (or a vctor if the quantity returned by getter is a getter, eg. speed)
+        """
         self.load_cells()
         value = getter(self.cells)
         if mass_weighted:
@@ -307,35 +465,32 @@ class PostProcessor(HDF5Container):
         B = getter_B_int(self.cells)
         rho = getter_rho(self.cells)
 
-        if logbins:
+        centers, B_mean = mean_by_bins(rho, B, bins, logbins)
-            rho_bins = np.logspace(
-                np.log10(np.min(rho)), np.log10(np.max(rho)), bins, base=10
-            )
-        else:
-            rho_bins = np.linspace(np.min(rho), np.max(rho), bins)
-
-        # For each cell, bin_number contains the number of the bins it belongs to
-        bin_number = np.zeros(len(B))
-
-        # Go through the min value of rho of each bin
-        for rho_min in rho_bins[:-1]:
-            bin_number = bin_number + (rho > rho_min).astype(int)
-
-        # Compute the mean in each bin
-        B_mean = np.zeros(len(rho_bins) - 1)
-        for i in range(len(B_mean)):
-            B_mean[i] = np.mean(B[bin_number == i])
-
-        # Get the center of each bin
-        if logbins:
-            centers = 10 ** (0.5 * (np.log10(rho_bins[1:]) + np.log10(rho_bins[:-1])))
-        else:
-            centers = 0.5 * (rho_bins[1:] + rho_bins[:-1])
 
         if self.pp_params.process.unload_cells:
             self.unload_cells()
         return ({"rho": centers, "B": B_mean}, {"logbins": logbins})
 
+    def _mean_by_bins(
+        self, name_x, name_y, ax_los, group="/maps/", bins=100, logbins=True
+    ):
+        """
+        Compute the mean of y in bins of x, where x, y are to array of the hdf5 file
+
+        Parameters
+        ----------
+        x_name, y_name : str, path of x and y in the hdf5 file
+        bins : int, number of bins
+        logbins : bool, if true, the bins will be logaritmically distributed
+        """
+        x = self.save.get_node(group + name_x + "_" + ax_los).read()
+        y = self.save.get_node(group + name_y + "_" + ax_los).read()
+
+        centers, values = mean_by_bins(x, y, bins, logbins)
+        # return ({os.path.basename(name_x) : centers, os.path.basename(name_y) : y_mean},
+        #        {"logbins" : logbins})
+        return (np.stack([values, centers]), {"logbins": logbins})
+
     def _Ek_Eb_rho(self, bins=100, logbins=True):
         """
         Mean of Ek/Eb in rho bins
@@ -458,14 +613,13 @@ class PostProcessor(HDF5Container):
 
     def _B_int(self, ax_los, z=0.0):
         """
-        Slice ont the intensity of the magnétic field
+        Slice ont the intensity of the magnetic field
         """
         B_op = ScalarOperator(
             lambda dset: np.sqrt(np.sum(dset["Br"] ** 2, axis=1)),
             self._ro.info["unit_mag"],
         )
         dmap_B = (slicing.SliceMap(self._amr, self._cam[ax_los], B_op, z=z)).map.T
-        dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read()
         return dmap_B
 
     def _temperature(self, ax_los, z=0.0):
@@ -478,6 +632,8 @@ class PostProcessor(HDF5Container):
         self._amr.set_read_levelmax(self.pp_params.pymses.levelmax)
         level_op = MaxLevelOperator()
         rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op)
+        if not self.pp_params.pymses.multiprocessing:
+            rt_level.disable_multiprocessing()
         datamap = rt_level.process(self._cam[ax_los], surf_qty=True)
         return datamap.map.T
 
@@ -500,8 +656,7 @@ class PostProcessor(HDF5Container):
 
         # Operator to compute the angular speed times rho
         def omega_rho_func(dset):
-            pos = dset.get_cell_centers()
+            pos = self.getter_pos_disk(dset)
-            pos = pos - (self.pp_params.disk.pos_star / self.lbox)
             xx = pos[:, :, 0]
             yy = pos[:, :, 1]
             rc = np.sqrt(xx ** 2 + yy ** 2)  # cylindrical radius
@@ -537,17 +692,19 @@ class PostProcessor(HDF5Container):
         else:
             rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op)
 
-            dmap_omega = rt_omega.process(self._cam[ax_los])
+        if not self.pp_params.pymses.multiprocessing:
-            dmap_cs = rt_cs.process(self._cam[ax_los])
+            rt_omega.disable_multiprocessing()
-            dmap_col = self.save.root.maps.coldens_z.read()
+            rt_cs.disable_multiprocessing()
-            map_Q = (
+
-                (self.lbox * dmap_cs.map.T)
+        dmap_omega = rt_omega.process(self._cam[ax_los])
-                * dmap_omega.map.T
+        dmap_cs = rt_cs.process(self._cam[ax_los])
-                / (np.pi * self.G * dmap_col)
+        dmap_col = self.save.root.maps.coldens_z.read()
-            )
+        map_Q = (
+            (self.lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * self.G * dmap_col)
+        )
         return map_Q
 
-    def _radial_bins(self, _):
+    def _radial_bins(self, ax_los="z"):
         """
         Computes radial bins (for disk)
         """
@@ -578,7 +735,7 @@ class PostProcessor(HDF5Container):
         rad_bins = np.concatenate(([0.0], rad_bins, [rad_of_box]))
         return rad_bins
 
-    def _rr(self, _):
+    def _rr(self, ax_los="z"):
         """
         Computes the radius from the center
         """
@@ -592,7 +749,7 @@ class PostProcessor(HDF5Container):
         rr = np.sqrt((xx - pos_star[0]) ** 2 + (yy - pos_star[1]) ** 2)
         return rr
 
-    def _bins_on_map(self, ax_los):
+    def _bins_on_map(self, ax_los="z"):
         rad_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
         rr = self.save.get_node("/maps/rr_" + ax_los).read()
 
@@ -602,7 +759,7 @@ class PostProcessor(HDF5Container):
             bins = bins + (rr >= r).astype(int)
         return bins
 
-    def _rad_avg(self, name, ax_los):
+    def _rad_avg(self, name, ax_los="z"):
         radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
         bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read()
         dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read()
@@ -613,7 +770,7 @@ class PostProcessor(HDF5Container):
             mean_bin[j] = np.mean(dmap[bins_on_map == j])
         return mean_bin
 
-    def _rad_avg_map(self, name, ax_los):
+    def _rad_avg_map(self, name, ax_los="z"):
 
         radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
         bins_on_map = self.save.get_node("/maps/bins_on_map_" + ax_los).read()
@@ -621,7 +778,7 @@ class PostProcessor(HDF5Container):
         mean_bin = self.save.get_node("/radial/rad_avg_" + name + "_" + ax_los).read()
 
         # Add value for border
-        mean_bin = np.concatenate(([mean_bin[0]], mean_bin))
+        mean_bin = np.concatenate((mean_bin, [mean_bin[-1]]))
 
         rr_flat = rr.flatten()
         bins_on_map_flat = bins_on_map.flatten()
@@ -642,14 +799,14 @@ class PostProcessor(HDF5Container):
 
         return avg_map
 
-    def _fluct_map(self, name, ax_los):
+    def _fluct_map(self, name, ax_los="z"):
 
         dmap = self.save.get_node("/maps/" + name + "_" + ax_los).read()
         avg_map = self.save.get_node("/maps/avg_map_" + name + "_" + ax_los).read()
 
         return dmap / avg_map
 
-    def _rad_fluct_pdf(self, name, ax_los):
+    def _rad_fluct_pdf(self, name, ax_los="z"):
         fluct_map = self.save.get_node("/maps/fluct_" + name + "_" + ax_los).read()
         rr = self.save.get_node("/maps/rr_" + ax_los).read()
 
@@ -666,7 +823,7 @@ class PostProcessor(HDF5Container):
         centers = 0.5 * (edges[1:] + edges[:-1])
         return np.stack([values, centers])
 
-    def _fit_pdf(self, name, ax_los):
+    def _fit_pdf(self, name, ax_los="z"):
         pdf = self.save.get_node("/hist/pdf_" + name + "_" + ax_los)
         values, centers = pdf.read()
         mask_fit = (
@@ -685,6 +842,72 @@ class PostProcessor(HDF5Container):
         pdf.attrs.var = np.var
         return True
 
+    def _alpha_disk(self, ax_los="z"):
+        "Map of the Rayleigh contribution to the Shakura&Sunaev alpha parameter for disks"
+        assert ax_los == "z"
+
+        # Mean part
+
+        T_avg = self.save.get_node("/maps/avg_map_T_avg_z").read()
+
+        radial_bins = self.save.get_node("/radial/radial_bins_" + ax_los).read()
+        mean_bin_vr = self.save.get_node(
+            "/radial/rad_avg_" + "vr" + "_" + ax_los
+        ).read()
+        mean_bin_vphi = self.save.get_node(
+            "/radial/rad_avg_" + "vphi" + "_" + ax_los
+        ).read()
+        mean_bin_vr = np.concatenate((mean_bin_vr, [mean_bin_vr[-1]]))
+        mean_bin_vphi = np.concatenate((mean_bin_vphi, [mean_bin_vr[-1]]))
+
+        # Fluct part
+        def getter_alpha_num(dset):
+            r = np.sqrt(np.sum((self.lbox * self.getter_pos_disk(dset)) ** 2, axis=2))
+
+            bins = np.zeros(r.shape, dtype=int)
+            for r0 in radial_bins[1:]:
+                bins = bins + (r >= r0).astype(int)
+
+            vr_mean = mean_bin_vr[bins]
+            vphi_mean = mean_bin_vphi[bins]
+
+            vr = self.getter_vr(dset)
+            vphi = self.getter_vphi(dset)
+            alpha = (vphi - vphi_mean) * (vr - vr_mean)
+            return alpha
+
+        alpha_f = (
+            self._ax_avg(getter_alpha_num, "z", unit=cst.none, mass_weighted=True)
+            / T_avg
+        )
+
+        # alpha
+        alpha = (2.0 / 3) * alpha_f
+        return alpha
+
+    def _alpha_grav(self, ax_los="z"):
+        "Map of the gravitational contribution to the Shakura&Sunaev alpha parameter for disks"
+        assert ax_los == "z"
+
+        T_avg = self.save.get_node("/maps/avg_map_T_avg_z").read()
+        coldens = self.save.get_node("/maps/avg_map_coldens_z").read()
+
+        def getter_alpha_grav(dset):
+            r2 = np.sum((self.lbox * self.getter_pos_disk(dset)) ** 2, axis=2)
+            e2 = (1.0 / 256.0) ** 2
+            gstar = -self.G * self.pp_params.disk.mass_star / (e2 + r2)
+            gr = self.getter_vect_r(dset, "g") - gstar
+            gphi = self.getter_vect_phi(dset, "g")
+            return gr * gphi / (4 * np.pi * self.G)
+
+        alpha_g = self._ax_avg(getter_alpha_grav, "z", unit=cst.none, surf_qty=True) / (
+            coldens * T_avg
+        )
+
+        # alpha
+        alpha_g = (2.0 / 3) * alpha_g
+        return alpha_g
+
     def _sinks(self):
         csv_name = (
             self.path
@@ -740,6 +963,88 @@ class PostProcessor(HDF5Container):
                 "/maps",
                 unit=self.info["unit_density"] * self.info["unit_length"],
             ),
+            "vr": Rule(
+                self,
+                partial(
+                    self._ax_avg,
+                    self.getter_vr,
+                    mass_weighted=True,
+                    unit=self.info["unit_velocity"],
+                ),
+                "Vertically mass-weighted averaged radial velocity",
+                "/maps",
+                unit=self.info["unit_velocity"],
+            ),
+            "vphi": Rule(
+                self,
+                partial(
+                    self._ax_avg,
+                    self.getter_vphi,
+                    mass_weighted=True,
+                    unit=self.info["unit_velocity"],
+                ),
+                "Vertically mass-weighted averaged azimuthal velocity",
+                "/maps",
+                unit=self.info["unit_velocity"],
+            ),
+            "rho_avg": Rule(
+                self,
+                partial(
+                    self._ax_avg,
+                    getter_rho,
+                    mass_weighted=False,
+                    unit=self.info["unit_density"],
+                ),
+                "Ax mass-weighted averaged azimuthal density",
+                "/maps",
+                unit=self.info["unit_density"],
+            ),
+            "P_avg": Rule(
+                self,
+                partial(
+                    self._ax_avg,
+                    getter_P,
+                    mass_weighted=True,
+                    unit=self.info["unit_pressure"],
+                ),
+                "Ax mass-weighted averaged azimuthal pressure",
+                "/maps",
+                unit=self.info["unit_pressure"],
+            ),
+            "T_avg": Rule(
+                self,
+                partial(
+                    self._ax_avg,
+                    getter_T,
+                    mass_weighted=True,
+                    unit=self.info["unit_pressure"] / self.info["unit_density"],
+                ),
+                "Ax mass-weighted averaged azimuthal temperature",
+                "/maps",
+                unit=self.info["unit_pressure"] / self.info["unit_density"],
+            ),
+            "alpha_disk": Rule(
+                self,
+                self._alpha_disk,
+                "Map of the Shakura&Sunaev alpha parameter for disks",
+                "/maps",
+                unit=cst.none,
+                dependencies=[
+                    "avg_map_rho_avg",
+                    "avg_map_T_avg",
+                    "avg_map_vr",
+                    "avg_map_vphi",
+                ],
+            ),
+            "alpha_grav": Rule(
+                self,
+                self._alpha_grav,
+                "Map of the graviational contrib to\
+                                Shakura&Sunaev alpha parameter for disks",
+                "/maps",
+                unit=cst.none,
+                dependencies=["avg_map_coldens", "avg_map_T_avg"],
+            ),
             "rho": Rule(
                 self,
                 self._rho,
@@ -781,7 +1086,7 @@ class PostProcessor(HDF5Container):
                 "Temperature slice",
                 "/maps",
                 dependencies=["rho"],
-                unit=self.info["unit_temperature"],
+                unit=self.info["unit_pressure"] / self.info["unit_density"],
             ),
             "levels": Rule(
                 self, self._levels, "Max level within line of sight", "/maps"
@@ -806,7 +1111,7 @@ class PostProcessor(HDF5Container):
                 "Toomre Q parameter for a Keplerian disk",
                 "/maps",
                 dependencies=["coldens"],
-                is_valid=lambda _: self.pp_params.disk.on,
+                is_valid=lambda _: self.pp_params.disk.enable,
             ),
             "sinks": Rule(
                 self,
@@ -865,7 +1170,7 @@ class PostProcessor(HDF5Container):
                 partial(self._vol_pdf, getter_T, logbins=True),
                 "Global T-PDF",
                 "/hist",
-                unit=self.info["unit_temperature"],
+                unit=self.info["unit_pressure"] / self.info["unit_density"],
             ),
             "P_pdf": Rule(
                 self,
@@ -946,7 +1251,19 @@ class PostProcessor(HDF5Container):
         }
 
         # Average and other
-        averageables = ["coldens", "rho", "T", "Q"]
+        averageables = [
+            "coldens",
+            "rho",
+            "T",
+            "Q",
+            "vphi",
+            "vr",
+            "rho_avg",
+            "P_avg",
+            "T_avg",
+            "alpha_disk",
+            "alpha_grav",
+        ]
         for name in averageables:
             self.rules["rad_avg_" + name] = Rule(
                 self,
@@ -985,6 +1302,15 @@ class PostProcessor(HDF5Container):
                 "/hist",
                 dependencies=["pdf_" + name],
             )
+            for name_bin in averageables:
+                if name_bin is not name:
+                    self.rules["mbb_" + name + "_" + name_bin] = Rule(
+                        self,
+                        partial(self._mean_by_bins, name_bin, name),
+                        "Mean of {} by bins of {}".format(name, name_bin),
+                        "/hist",
+                        dependencies=[name, name_bin],
+                    )
 
         self._gen_rule_transform("fluct_coldens", np.max, "max", group="/globals")
 

pp_params.yml

@@ -13,6 +13,7 @@ disk: # Disk speficic parameters
     enable    : False                 # Enable specific disk analysis
     pos_star  : [1., 1., 1.]          # Position of the central star
     binning   : "log"                 # Kind of binning (lin : linear, log : logarithmic)
+    mass_star : 1.                    # Mass of the central star
     nb_bin    : 100                   # Number of bins for averaged quantities
     bin_in    : 1e-3                  # Outer radius of the inner bin
     bin_out   : 0.25                  # Inner radius of the outer bin
@@ -35,12 +36,18 @@ pymses: # Parameters for Pymses reader
     levelmax : 20                      # Maximal AMR level visited when looking levels
     fft      : False                   # Quick and dirty rendering using FFT
     verbose  : False                   # Let pymses write on standart output
+    multiprocessing : True             # Whether to use multiprocessing
 
     # Camera settings
     center   :  [0.5, 0.5, 0.5]        # Center of the image
     zoom     : 1.                      # Zoom  of the image
     map_size : 1024                    # Size of the computed maps in pixel
 
+    # Filter parameters
+    filter     : False                 # Enable filtering
+    min_coords : [0.35, 0.35, 0.45]
+    max_coords : [0.65, 0.65, 0.55]
+
 input: # Parameters on how to look for input files (= output from Ramses)
 
     log_prefix     : "run.log"         # Prefix of the log file