pipeline/postprocessor.py

# coding: utf-8
import pandas as pd
import pspec_new
from baseprocessor import *
import pymses.utils.regions as reg
from pymses.filters import RegionFilter


# Getters


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):
    B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
    return B_norm


def getter_rho(dset):
    return dset["rho"]


def getter_v_norm(dset):
    v_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
    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
    """

    # Axes information
    _ax_nb = {"x": 0, "y": 1, "z": 2}  # Number of each axes
    _axes_h = {"x": "y", "y": "x", "z": "x"}  # Associated horizontal axe
    _axes_v = {"x": "z", "y": "z", "z": "y"}  # Associated vertical axe

    G = 1.0  # Gravitational constant

    cells_loaded = False

    def __init__(self, path=None, num=None, path_out=None, pp_params=None, tag=None):
        """
        Creates the basic structures needed for the outputs
        """

        super(PostProcessor, self).__init__(path, path_out, pp_params, tag)

        # Open outfile
        if not self.pp_params.out.tag == "":
            tag_name = self.pp_params.out.tag + "_"
        else:
            tag_name = ""

        self.filename = (
            self.path_out + "/postproc_" + tag_name + format(num, "05") + ".h5"
        )

        self.cells_filename = (
            self.path_out + "/cells_" + tag_name + format(num, "05") + ".h5"
        )

        if not os.path.exists(self.path_out):
            os.makedirs(self.path_out)

        self.open()

        # Ramses Output
        self.path = path
        self.run = os.path.basename(path)
        self.num = num
        self._ro = pymses.RamsesOutput(
            path,
            num,
            order=self.pp_params.pymses.order,
            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
        self._rho_op = ScalarOperator(
            lambda dset: dset["rho"], self._ro.info["unit_density"]
        )

        # Density ray tracer
        if self.pp_params.pymses.fft:
            self._rt = splatting.SplatterProcessor(
                self._amr, self._ro.info, self._rho_op
            )
        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"]

        if self.pp_params.pymses.filter:
            center = (self.max_coords + self.min_coords) / 2.0
            im_extent = [
                self.min_coords[0],
                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

        # Set post processing attributes
        self.save.root._v_attrs.dir = os.path.dirname(path)
        self.save.root._v_attrs.run = os.path.basename(path)
        self.save.root._v_attrs.num = num
        self.save.root._v_attrs.lbox = self.lbox
        self.save.root._v_attrs.unit_length = self.info["unit_length"]
        self.save.root._v_attrs.time = time

        if not "/maps" in self.save:
            self.save.create_group("/", "maps", "2D maps")
        self.save.root.maps._v_attrs.center = center
        self.save.root.maps._v_attrs.radius = self._radius
        self.save.root.maps._v_attrs.im_extent = im_extent

        # Initialize cameras
        self._cam = {}
        for ax_los in self._ax_nb:  # los = line of sight
            ax_h = self._axes_h[ax_los]
            ax_v = self._axes_v[ax_los]

            self._cam[ax_los] = Camera(
                center=self.pp_params.pymses.center,
                line_of_sight_axis=ax_los,
                region_size=[im_extent[1] - im_extent[0], im_extent[3] - im_extent[2]],
                distance=distance,
                far_cut_depth=distance,
                up_vector=ax_v,
                map_max_size=self.pp_params.pymses.map_size,
            )

        self.close()

        self.log_id = "[{}, {}] ".format(self.run, self.num)

        self.def_rules()

    def load_cells(self):
        """
        Load all cells from the source file in the memory.
        Cells will be accessible trough self.cells
        (/!\ Long and memory heavy)
        """
        if not self.cells_loaded:
            if os.path.exists(self.cells_filename):
                cells_hdf5 = tables.open_file(self.cells_filename, mode="r")
                try:
                    node = cells_hdf5.get_node("/cells")
                    self.cells = {}
                    for key in node._v_children:
                        self.cells[key] = cells_hdf5.get_node("/cells/" + key).read()
                finally:
                    cells_hdf5.close()
            else:
                cell_source = CellsToPoints(self._amr)
                cells_pymses = cell_source.flatten()
                self.cells = {}
                for key in cells_pymses.fields:
                    self.cells[key] = cells_pymses[key]
                self.cells["dx"] = cells_pymses.get_sizes()
                self.cells["pos"] = cells_pymses.points

                if self.pp_params.process.save_cells:
                    cells_hdf5 = tables.open_file(self.cells_filename, mode="w")
                    try:
                        for key in self.cells:
                            cells_hdf5.create_array(
                                "/cells", key, self.cells[key], "", createparents=True
                            )
                    finally:
                        cells_hdf5.close()

        self.cells_loaded = True

    def unload_cells(self):
        """
        Free space in the memory by telling the garbage collectors that
        self.cells is not needed
        """
        if self.cells_loaded:
            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.
        Return a slice of the source box.

        Parameters
        ----------
        getter : callable
             A callable that extract the wanted data from a pymses dataset

        ax_los : string
             The axis perpendicular to the slice plane

        z : float
             Coordinate of the slice on the ax_los axis

        unit : cst.Unit
             Unit of the resulting dataset

        Returns
        -------
        A numpy array containing the slice
        """
        op = ScalarOperator(getter, unit)
        datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z)
        return datamap.map.T

    def _ax_avg(
        self, getter, ax_los, unit=cst.none, mass_weighted=True, surf_qty=False
    ):
        """
        Map of the average of a quantity (given by getter) along an axis (ax_los)
        Returns 2D array if getter returns a scalar quantity

        If surf_qty is set (projection mode), mass_weighted is ignored
        """
        if surf_qty:
            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

    def _get_axis(self, axis):

        if isinstance(axis, str):
            axis = self._ax_nb[axis]

        self.load_cells()
        return np.sort(np.unique(self.cells["pos"][:, axis]))

    def _plane_avg_uniform(self, getter, axis, unit=cst.none, mass_weighted=False):
        """
        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):
            axis = self._ax_nb[axis]
        axis_data = self.cells["pos"][:, axis]
        value = getter(self.cells)

        df = pd.DataFrame({"axis": axis_data})
        if mass_weighted:
            mass = mass_func(self.cells)
            tot_mass = np.sum(mass)
            df["value"] = value * mass / tot_mass
        else:
            df["value"] = value

        if self.pp_params.process.unload_cells:
            self.unload_cells()

        df.sort_values("axis", inplace=True)

        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:
            mass = mass_func(self.cells)
            # Transpose (.T) is for vectorial values
            data = np.sum((mass * value.T).T, axis=0) / np.sum(mass)
        else:
            data = np.sum(value, axis=0)

        if self.pp_params.process.unload_cells:
            self.unload_cells()
        return data

    def _vol_pdf(self, getter, bins=100, logbins=False, weight_func=vol_func):
        self.load_cells()
        data = getter(self.cells)
        if logbins:
            data = np.log10(data)
        weights = weight_func(self.cells)
        if self.pp_params.process.unload_cells:
            self.unload_cells()

        values, edges = np.histogram(data, bins, weights=weights)
        centers = 0.5 * (edges[1:] + edges[:-1])
        return (np.stack([values, centers]), {"logbins": logbins})

    def _Brho(self, bins=100, logbins=True):
        """
        Mean of B in rho bins
        """
        self.load_cells()
        B = getter_B_int(self.cells)
        rho = getter_rho(self.cells)

        centers, B_mean = mean_by_bins(rho, B, bins, logbins)

        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
        """
        self.load_cells()
        mean_speed = self.save.get_node("/globals/mwa_speed").read()
        mean_speed = mean_speed * self.info["unit_velocity"].express(cst.km_s)
        vel_fluct = (self.cells)["vel"] * self.info["unit_velocity"].express(
            cst.km_s
        ) - mean_speed
        B_norm = getter_B_int(self.cells)
        B_norm = B_norm * self.info["unit_mag"].express(cst.T)
        v_norm = np.sqrt(
            np.sum((vel_fluct * 10 ** (3)) ** 2, axis=1)
        )  # v_norm [m/s] et vel_fluct [km/s]
        rho = getter_rho(self.cells)
        rho_kg_m3 = rho * self.info["unit_density"].express(cst.kg_m3)
        eb = 0.5 * (B_norm) ** 2 / (4 * np.pi * 10 ** (-7))  # mettre le bon mu
        ek = 0.5 * v_norm ** 2 * rho_kg_m3
        rapport = ek / eb

        if 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_norm))

        # 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
        ek_eb = np.zeros(len(rho_bins) - 1)
        for i in range(len(ek_eb)):
            ek_eb[i] = np.mean(rapport[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, "Ek_Eb_rho": ek_eb}, {"logbins": logbins})

    def cos_vfluct_B(self):
        "return the cos of the angle between the magnetic field and the velocity fluctuation field"
        mean_speed = self.save.get_node("/globals/mwa_speed").read()

        def getter_cos_vfluct_B(dset):
            vel_fluct = dset["vel"] - mean_speed
            B_norm = np.sqrt(np.sum(dset["Br"] ** 2, axis=1))
            v_norm = np.sqrt(np.sum(vel_fluct ** 2, axis=1))
            # Compute the dot product in each cell
            dot_prod = np.einsum("ij,ij->i", vel_fluct, dset["Br"])
            return np.abs(dot_prod) / (v_norm * B_norm)

        return self._vol_pdf(getter_cos_vfluct_B)

    def _mwa_sigma(self, axes=["x", "y", "z"]):
        mw_speed = self.save.get_node("/globals/mwa_speed").read()

        if axes == ["x", "y", "z"]:

            def getter(dset):
                return np.sum((dset["vel"] - mw_speed) ** 2, axis=1)

        else:

            def getter(dset):
                sigma_squared = 0.0
                for ax in axes:
                    ax_nb = self._ax_nb[ax]
                    sigma_sq_ax = (dset["vel"][:, ax_nb] - mw_speed[ax_nb]) ** 2
                    sigma_squared = sigma_squared + sigma_sq_ax
                return sigma_squared

        return np.sqrt(self._vol_avg(getter, mass_weighted=True))

    def _coldens(self, ax_los):
        datamap = self._rt.process(self._cam[ax_los], surf_qty=True)
        return datamap.map.T

    def _rho(self, ax_los, z=0.0):
        datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=z)
        return (datamap_rho.map).T

    def _vector_h(self, name, unit, ax_los, z=0.0):
        h_op = ScalarOperator(
            lambda dset: dset[name][:, self._ax_nb[self._axes_h[ax_los]]],
            self._ro.info[unit],
        )
        dmap_h = slicing.SliceMap(self._amr, self._cam[ax_los], h_op, z=z).map.T
        return dmap_h

    def _vector_v(self, name, unit, ax_los, z=0.0):
        v_op = ScalarOperator(
            lambda dset: dset[name][:, self._ax_nb[self._axes_v[ax_los]]],
            self._ro.info[unit],
        )
        dmap_v = slicing.SliceMap(self._amr, self._cam[ax_los], v_op, z=z).map.T
        return dmap_v

    def _speed_h(self, ax_los, z=0.0):
        return self._vector_h("vel", "unit_velocity", ax_los, z)

    def _speed_v(self, ax_los, z=0.0):
        return self._vector_v("vel", "unit_velocity", ax_los, z)

    def _B_h(self, ax_los, z=0.0):
        return self._vector_h("Br", "unit_mag", ax_los, z)

    def _B_v(self, ax_los, z=0.0):
        return self._vector_v("Br", "unit_mag", ax_los, z)

    def _B_int(self, ax_los, z=0.0):
        """
        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
        return dmap_B

    def _temperature(self, ax_los, z=0.0):
        P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"])
        dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=z)).map.T
        dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read()
        return dmap_P / dmap_rho

    def _levels(self, ax_los):
        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

    def _jeans(self, ax_los):
        dmap_T = self.save.get_node("/maps/T_" + ax_los).read()
        dmap_rho = self.save.get_node("/maps/rho_" + ax_los).read()
        dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho)
        return dmap_jeans

    def _jeans_ratio(self, ax_los):
        dmap_jeans = self.save.get_node("/maps/jeans_" + ax_los).read()
        dmap_levels = self.save.get_node("/maps/levels_" + ax_los).read()
        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
        """

        # Operator to compute the angular speed times rho
        def omega_rho_func(dset):
            pos = self.getter_pos_disk(dset)
            xx = pos[:, :, 0]
            yy = pos[:, :, 1]
            rc = np.sqrt(xx ** 2 + yy ** 2)  # cylindrical radius
            vx = dset["vel"][:, :, 0]
            vy = dset["vel"][:, :, 1]
            omega_rho = 1.0 / rc ** 2
            omega_rho = omega_rho * dset["rho"]
            vyx = vy * xx
            vxy = vx * yy
            omega_rho = omega_rho * (vyx - vxy)
            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"]
        )

        # Operator to compute the sound speed
        cs_op = FractionOperator(
            lambda dset: dset["P"],
            lambda dset: dset["rho"],
            self._ro.info["unit_velocity"],
        )

        # 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_cs = splatting.SplatterProcessor(
                self._amr, self._ro.info, cs_op, surf_qty=False
            )
        else:
            rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op)

        if not self.pp_params.pymses.multiprocessing:
            rt_omega.disable_multiprocessing()
            rt_cs.disable_multiprocessing()

        dmap_omega = rt_omega.process(self._cam[ax_los])
        dmap_cs = rt_cs.process(self._cam[ax_los])
        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, ax_los="z"):
        """
        Computes radial bins (for disk)
        """
        pos_star = self.pp_params.disk.pos_star
        im_extent = np.array(self.save.root.maps._v_attrs.im_extent) * self.lbox

        # radius of the corner of the box plus a margin
        rad_of_box = (
            np.sqrt(
                (im_extent[1] - pos_star[0]) ** 2 + (im_extent[3] - pos_star[1]) ** 2
            )
            + 0.1
        )

        bin_in = self.pp_params.disk.bin_in
        bin_out = self.pp_params.disk.bin_out
        nb_bin = self.pp_params.disk.nb_bin

        # radial bins
        if self.pp_params.disk.binning == "log":
            lrad_in = np.log10(bin_in)
            lrad_ext = np.log10(bin_out)
            rad_bins = np.logspace(lrad_in, lrad_ext, num=nb_bin)
        elif self.pp_params.disk.binning == "lin":
            rad_bins = np.linspace(bin_in, bin_out, num=nb_bin)

        # Add boundaries
        rad_bins = np.concatenate(([0.0], rad_bins, [rad_of_box]))
        return rad_bins

    def _rr(self, ax_los="z"):
        """
        Computes the radius from the center
        """
        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

        x = np.linspace(im_extent[0], im_extent[1], map_size)
        y = np.linspace(im_extent[2], im_extent[3], map_size)
        xx, yy = np.meshgrid(x, y)
        rr = np.sqrt((xx - pos_star[0]) ** 2 + (yy - pos_star[1]) ** 2)
        return rr

    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()

        # Find appropriate bin for each coordinate set
        bins = np.zeros(rr.shape, dtype=int)
        for r in rad_bins[1:]:
            bins = bins + (rr >= r).astype(int)
        return bins

    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()

        # mean of all the cells in the bin
        mean_bin = np.zeros(len(radial_bins) - 1)
        for j in range(len(radial_bins) - 1):
            mean_bin[j] = np.mean(dmap[bins_on_map == j])
        return mean_bin

    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()
        rr = self.save.get_node("/maps/rr_" + ax_los).read()
        mean_bin = self.save.get_node("/radial/rad_avg_" + name + "_" + ax_los).read()

        # Add value for border
        mean_bin = np.concatenate((mean_bin, [mean_bin[-1]]))

        rr_flat = rr.flatten()
        bins_on_map_flat = bins_on_map.flatten()

        # Compute the map azimuthally averaged
        # use linear interpolation to improve accuracy
        avg_flat = (radial_bins[bins_on_map_flat + 1] - rr_flat) * mean_bin[
            bins_on_map_flat
        ]
        avg_flat = (
            avg_flat
            + (rr_flat - radial_bins[bins_on_map_flat]) * mean_bin[bins_on_map_flat + 1]
        )
        avg_flat = avg_flat / (
            radial_bins[bins_on_map_flat + 1] - radial_bins[bins_on_map_flat]
        )
        avg_map = np.reshape(avg_flat, rr.shape)

        return avg_map

    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="z"):
        fluct_map = self.save.get_node("/maps/fluct_" + name + "_" + ax_los).read()
        rr = self.save.get_node("/maps/rr_" + ax_los).read()

        mask_pdf = (rr > self.pp_params.disk.rmin_pdf) & (
            rr < self.pp_params.disk.rmax_pdf
        )

        nb_cells = np.sum(mask_pdf.flatten())
        values, edges = np.histogram(
            np.log10(fluct_map[mask_pdf].flatten()),
            self.pp_params.pdf.nb_bin,
            weights=np.ones(nb_cells) / nb_cells,
        )
        centers = 0.5 * (edges[1:] + edges[:-1])
        return np.stack([values, centers])

    def _fit_pdf(self, name, ax_los="z"):
        pdf = self.save.get_node("/hist/pdf_" + name + "_" + ax_los)
        values, centers = pdf.read()
        mask_fit = (
            (centers > self.pp_params.pdf.xmin_fit)
            & (centers < self.pp_params.pdf.xmax_fit)
            & (values > 0)
        )
        (slope, origin, correlation, _, stderr) = linregress(
            centers[mask_fit], np.log10(values[mask_fit])
        )

        pdf.attrs.slope = slope
        pdf.attrs.origin = origin
        pdf.attrs.correlation = correlation
        pdf.attrs.stderr = stderr
        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
            + "/output_"
            + str(self.num).zfill(5)
            + "/sink_"
            + str(self.num).zfill(5)
            + ".csv"
        )
        sinks = np.loadtxt(csv_name, delimiter=",")
        header = [
            "Id",
            "M",
            "dmf",
            "x",
            "y",
            "z",
            "vx",
            "vy",
            "vz",
            "rot_period",
            "lx",
            "ly",
            "lz",
            "acc_rate",
            "acc_lum",
            "age",
            "int_lum",
            "Teff",
        ]
        if len(sinks) == 0:
            sinks = np.zeros(len(header))

        sinks_dict = {}
        for key, a in zip(header, sinks):
            sinks_dict[key] = a

        return sinks_dict

    def _pspec(self):
        outfile = self.path_out + "/pspec.h5"
        pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile)
        return outfile

    def def_rules(self):

        self.rules = {
            # Base rules
            "coldens": Rule(
                self,
                self._coldens,
                "Column density",
                "/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,
                "Density slice",
                "/maps",
                unit=self.info["unit_density"],
            ),
            "speed_h": Rule(
                self,
                self._speed_h,
                "Horizontal speed slice wrt the line of sight",
                "/maps",
                unit=self.info["unit_velocity"],
            ),
            "speed_v": Rule(
                self,
                self._speed_v,
                "Vertical speed slice wrt the line of sight",
                "/maps",
                unit=self.info["unit_velocity"],
            ),
            "B_h": Rule(
                self,
                self._speed_h,
                "Horizontal slice of the magnetic field wrt the line of sight",
                "/maps",
                unit=self.info["unit_mag"],
            ),
            "B_v": Rule(
                self,
                self._speed_v,
                "Vertical  slice  of the magnetic field wrt the line of sight",
                "/maps",
                unit=self.info["unit_mag"],
            ),
            "T": Rule(
                self,
                self._temperature,
                "Temperature slice",
                "/maps",
                dependencies=["rho"],
                unit=self.info["unit_pressure"] / self.info["unit_density"],
            ),
            "levels": Rule(
                self, self._levels, "Max level within line of sight", "/maps"
            ),
            "jeans": Rule(
                self,
                self._jeans,
                "Jeans lenght slice",
                "/maps",
                dependencies=["rho", "T"],
            ),
            "jeans_ratio": Rule(
                self,
                self._jeans_ratio,
                "Jeans' lenght divided by the max resolution",
                "/maps",
                dependencies=["jeans", "levels"],
            ),
            "Q": Rule(
                self,
                self._toomreQ_disk,
                "Toomre Q parameter for a Keplerian disk",
                "/maps",
                dependencies=["coldens"],
                is_valid=lambda _: self.pp_params.disk.enable,
            ),
            "sinks": Rule(
                self,
                self._sinks,
                group="/datasets",
                unit={
                    "Id": cst.none,
                    "M": cst.Msun,
                    "dmf": cst.Msun,
                    "x": "",
                    "y": "",
                    "z": "",
                    "vx": "",
                    "vy": "",
                    "vz": "",
                    "rot_period": "[y]",
                    "lx": "|l|",
                    "ly": "|l|",
                    "lz": "|l|",
                    "acc_rate": "[Msol/y]",
                    "acc_lum": "[Lsol]",
                    "age": cst.year,
                    "int_lum": "[Lsol]",
                    "Teff": cst.K,
                },
            ),
            "pspec": Rule(self, self._pspec, "Power spectrum", "/hdf5"),
            # Helpers
            "radial_bins": Rule(self, self._radial_bins, "Radial bins", "/radial"),
            "rr": Rule(self, self._rr, "Coordinate map", "/maps"),
            "bins_on_map": Rule(
                self,
                self._bins_on_map,
                "Convert map coordinates to bins",
                "/maps",
                dependencies=["radial_bins", "rr"],
            ),
            "B_int": Rule(
                self,
                self._B_int,
                "Magnetic intensity slice",
                "/maps",
                dependencies=["rho"],
                unit=self.info["unit_mag"],
            ),
            # PDF
            "rho_pdf": Rule(
                self,
                partial(self._vol_pdf, partial(simple_getter, "rho"), logbins=True),
                "Global rho-PDF",
                "/hist",
                unit=self.info["unit_density"],
            ),
            "T_pdf": Rule(
                self,
                partial(self._vol_pdf, getter_T, logbins=True),
                "Global T-PDF",
                "/hist",
                unit=self.info["unit_pressure"] / self.info["unit_density"],
            ),
            "P_pdf": Rule(
                self,
                partial(self._vol_pdf, getter_T, logbins=True),
                "Global P-PDF",
                "/hist",
                unit=self.info["unit_pressure"],
            ),
            "cos_pdf": Rule(
                self,
                partial(self.cos_vfluct_B),
                "Global cos fluctuation-PDF",
                "/hist",
                dependencies=["mwa_speed"],
                unit=cst.none,
            ),
            "Brho": Rule(
                self,
                self._Brho,
                "Average of B as a function of rho",
                "/datasets",
                unit={"rho": self.info["unit_density"], "B": self.info["unit_mag"]},
            ),
            "Ek_Eb_rho": Rule(
                self,
                self._Ek_Eb_rho,
                "Average of Ek/Eb as a function of rho",
                "/datasets",
                dependencies=["mwa_speed"],
                unit={"rho": self.info["unit_density"], "Ek_Eb_rho": cst.none},
            ),
            # Profiles
            "axis": Rule(
                self,
                partial(self._get_axis),
                "Axis",
                "/profile",
                unit=self.info["unit_length"],
            ),
            "rho_prof": Rule(
                self,
                partial(self._plane_avg_uniform, partial(simple_getter, "rho")),
                "Rho profile",
                "/profile",
                unit=self.info["unit_density"],
                dependencies=["axis"],
            ),
            # globals
            "time_num": Rule(
                self,
                lambda: self.info["time"],
                "Time",
                "/globals",
                unit=self.info["unit_time"],
            ),
            "mwa_speed": Rule(
                self,
                partial(self._vol_avg, partial(simple_getter, "vel")),
                "Mass weighted speed average",
                "/globals",
                unit=self.info["unit_velocity"],
            ),
            "mwa_sigma": Rule(
                self,
                self._mwa_sigma,
                "Mass weighted speed average",
                "/globals",
                dependencies={"mwa_speed": None},
                unit=self.info["unit_velocity"],
            ),
            "mwa_B_int": Rule(
                self,
                partial(self._vol_avg, getter_B_int),
                "Mass weighted Magnetic intensity average",
                "/globals",
                unit=self.info["unit_mag"],
            ),
        }

        # Average and other
        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,
                partial(self._rad_avg, name),
                "Azimuthal average of {}".format(name),
                "/radial",
                dependencies=["radial_bins", "bins_on_map", name],
            )

            self.rules["avg_map_" + name] = Rule(
                self,
                partial(self._rad_avg_map, name),
                "Interpolated map of azimuthal average of {}".format(name),
                "/maps",
                dependencies=["radial_bins", "bins_on_map", "rr", "rad_avg_" + name],
            )
            self.rules["fluct_" + name] = Rule(
                self,
                partial(self._fluct_map, name),
                "Fluctuation wrt to average of {}".format(name),
                "/maps",
                dependencies=[name, "avg_map_" + name],
            )
            self.rules["pdf_" + name] = Rule(
                self,
                partial(self._rad_fluct_pdf, name),
                "Probability density function of {} fluctuations".format(name),
                "/hist",
                dependencies=["rr", "fluct_" + name],
            )

            self.rules["fit_pdf_" + name] = Rule(
                self,
                partial(self._fit_pdf, name),
                "Fit the PDF of {} fluctuations".format(name),
                "/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")

        super(PostProcessor, self).def_rules()


def get_time(path, num):
    """
    Return the time of the output (code units)

    Parameters
    ----------
    num       output number
    path_out  path of the pipeline output

    Returns
    -------
    time      the time of the output (code units)
    """
    try:
        f = open(
            path
            + "/output_"
            + str(num).zfill(5)
            + "/info_"
            + str(num).zfill(5)
            + ".txt"
        )
        for line in f:
            ls = line.split()
            if len(ls) > 1 and ls[0] == "time":
                time = float(ls[2])
                break
        f.close()
        return time
    except IOError as e:
        print(e)
        return np.nan