added galsec

.pre-commit-config.yaml

@@ -20,15 +20,14 @@ repos:
     hooks:
       - id: remove-crlf
   - repo: https://github.com/psf/black
-    rev: 20.8b1
+    rev: 22.6.0
     hooks:
       - id: black
-  - repo: https://github.com/asottile/blacken-docs
+#  - repo: https://github.com/asottile/blacken-docs
-    rev: v1.8.0
+#    rev: v1.8.0
-    hooks:
+#    hooks:
-      - id: blacken-docs
+#      - id: blacken-docs
-        additional_dependencies: [ black==20.8b1 ]
+  - repo: https://github.com/pycqa/flake8/
-  - repo: https://gitlab.com/PyCQA/flake8
+    rev: 6.0.0
-    rev: 3.8.4
     hooks:
       - id: flake8

galturb/galsec.py

@@ -0,0 +1,355 @@
+# coding: utf-8
+"""
+    Galaxy kiloparsec extractor
+
+"""
+import numpy as np
+from astropy.table import QTable, hstack
+from astropy import units as u
+from astropy.units.quantity import Quantity
+
+
+def vect_r(position: np.array, vector: np.array) -> np.array:
+    """Radial component of a vector
+
+    Parameters
+    ----------
+    position : np.array (N, 3)
+        (only reads x and y)
+    vector : np.array (N, 3)
+        (only reads x and y components)
+
+    Returns
+    -------
+    np.array (N)
+        Radial component
+    """
+    r = position[:, :2]
+    ur = np.transpose((np.transpose(r) / np.sqrt(np.sum(r**2, axis=1))))
+    return np.einsum("ij, ij -> i", vector[:, :2], ur)
+
+
+def vect_phi(position: np.array, vector: np.array):
+    """Azimuthal  component of a vecto
+
+    Parameters
+    ----------
+    position : np.array (N, 3)
+        (only reads x and y)
+    vector : np.array (N, 3)
+        (only reads x and y components)
+
+    Returns
+    -------
+    np.array (N)
+        Azimuthal  component
+    """
+    r = position[:, :2]
+    r_norm = np.sqrt(np.sum(r**2, axis=1))
+    rot = np.array([[0, -1], [1, 0]])
+    uphi = np.transpose(np.einsum("ij, kj -> ik", rot, r) / r_norm)
+    return np.einsum("ij,ij -> i", vector[:, :2], uphi)
+
+
+def get_sfr(
+    stars: QTable, time: Quantity[u.Myr], average_on: Quantity[u.Myr] = 30 * u.Myr
+) -> Quantity[u.Msun / u.year]:
+    """_summary_
+
+    Parameters
+    ----------
+    stars : QTable
+        _description_
+    time : Quantity[u.Myr],
+        time at wich the SFR should be computed
+    average_on : Quantity, optional
+        compute the sfr on the last `average_on` years, by default 30*u.Myr
+
+    Returns
+    -------
+    float
+        SFR in Msun/yr
+    """
+    try:
+        mask = (stars["birth_time"] > max(time - average_on, 0 * u.Myr)) & (
+            stars["birth_time"] < time
+        )
+        dtime = min(average_on, time).to(u.year)
+        if dtime == 0:
+            sfr = 0.0 * u.Msun / u.year
+        else:
+            sfr = np.sum(stars["mass"][mask]) / dtime
+    except KeyError:
+        sfr = 0.0 * u.Msun / u.year
+    return sfr
+
+
+def aggregate(
+    grouped_data: QTable,
+    weight_field: str = "mass",
+    extensive_fields: str = ["mass", "ek"],
+    weighted_fields: list = ["velr", "velphi", "velz"],
+) -> QTable:
+    """Aggregate wisely from grouped data
+
+    Parameters
+    ----------
+    grouped_data : QTable
+        should already have group information
+    weight_field : str, optional
+        the field used for the  weighting of the non extensive quantities, by default "mass"
+    extensive_fields : str, optional
+        these will be summed. Should include weight_field, by default ["mass", "ek"]
+    weighted_fields : list, optional
+        the field that will be weighted, by default ["velr", "velphi", "velz"]
+
+    Returns
+    -------
+    QTable
+        a tabble with the aggregated value for each group
+    """
+    assert weight_field in extensive_fields
+
+    for field in weighted_fields:
+        grouped_data[field] *= grouped_data[weight_field]
+
+    binned_data = grouped_data[extensive_fields + weighted_fields].groups.aggregate(
+        np.add
+    )
+
+    for field in weighted_fields:
+        binned_data[field] /= binned_data[weight_field]  # weighted mean
+
+        # Veocity dispersion
+        grouped_data[field] /= grouped_data[weight_field]
+        for i in range(len(grouped_data.groups)):
+            slice = grouped_data.groups.indices[i], grouped_data.groups.indices[i + 1]
+            grouped_data[slice[0] : slice[1]][field] -= binned_data[i][field]
+        grouped_data[field] = grouped_data[weight_field] * grouped_data[field] ** 2
+        binned_data[f"sigma_{field}"] = grouped_data[field,].groups.aggregate(
+            np.add
+        )[field]
+        binned_data[f"sigma_{field}"] = np.sqrt(
+            binned_data[f"sigma_{field}"] / binned_data[weight_field]
+        )
+
+    return binned_data
+
+
+class Galsec:
+    """
+    Galactic sector extractor
+    """
+
+    fluids = ["gas", "stars", "dm"]
+    units = {
+        "gas": {
+            "position": u.kpc,
+            "volume": u.kpc**3,
+            "velocity": u.km / u.s,
+            "mass": u.Msun,
+        },
+        "stars": {
+            "position": u.kpc,
+            "velocity": u.km / u.s,
+            "mass": u.Msun,
+            "birth_time": u.Myr,
+        },
+        "dm": {"position": u.kpc, "velocity": u.km / u.s, "mass": u.Msun},
+    }
+
+    def __init__(self, data: dict, copy=True) -> None:
+        """Initiazise a Galaxasec object
+
+        Parameters
+        ----------
+        cell_data : dict
+            A dataset of cells in the following format:
+            header:
+                time     float  [Myr] time since the start of the simulation
+            gas:
+                position float array (Ngas, 3) [kpc], centered
+                volume   float array (Ngas)    [pc^3]
+                velocity float array (Ngas, 3) [km/s]
+                mass     float array (Ngas)    [Msun]
+            stars:
+                position float array (Nstar, 3) [kpc], centered
+                velocity float array (Nstar, 3) [km/s]
+                mass     float array (Nstar)    [Msun]
+                birth_time      float array (Nstar)    [Myr]
+            dm:
+                position float array (Ngas, 3) [kpc], centered
+                velocity float array (Ngas, 3) [km/s]
+                mass     float array (Ngas)    [Msun]
+            maps:
+                extent      float list  [xmin, xmax, ymin, ymax]  Coordinates of the edges of the map, centered
+                gas_coldens float array (Nx, Ny) [Msun/pc^2]      Map of column density
+        copy : bool, default=True
+            Wheter the data should be copied.
+        """
+
+        self.data = {}
+
+        for fluid in self.fluids:
+            self.data[fluid] = QTable(data[fluid], units=self.units[fluid], copy=copy)
+
+        self.time = data["header"]["time"] * u.Myr
+
+        self.compute_derived_values()
+
+    def compute_derived_values(self) -> None:
+        """
+        Helper function to computed values derivated from input data
+        """
+        for fluid in self.fluids:
+            dset = self.data[fluid]
+            dset["r"] = np.sqrt(
+                np.sum(dset["position"][:, :2] ** 2, axis=1)
+            )  # galactic radius
+            dset["phi"] = np.angle(
+                dset["position"][:, 0] + dset["position"][:, 1] * 1j
+            )  # galactic longitude
+            dset["phi"][dset["phi"] < 0] += (
+                2 * np.pi * u.rad
+            )  # rescaling to get only positive values
+            dset["velphi"] = vect_phi(
+                dset["position"], dset["velocity"]
+            )  # azimuthal velocity
+            dset["velr"] = vect_r(dset["position"], dset["velocity"])  # radial velocity
+            dset["velz"] = dset["velocity"][:, 2]  # vertical velocity
+            dset["ek"] = (dset["mass"] * np.sum(dset["velocity"] ** 2, axis=1)).to(
+                u.erg
+            )  # kinetic energy
+
+    def ring_binning(self, delta_r: Quantity[u.kpc]):
+        """Add radial bin informations to the dataset
+
+        Parameters
+        ----------
+        delta_r : Quantity[u.kpc]
+            spacing between two radial bins
+        """
+        for fluid in self.fluids:
+            r_bin = np.trunc(self.data[fluid]["r"] / delta_r)
+            r_bin = (r_bin + 0.5) * delta_r  # Store the middle value
+            self.data[fluid]["r_bin"] = r_bin
+
+    def sector_binning(self, delta_r: Quantity[u.kpc], delta_l: Quantity[u.kpc]):
+        """Add sector bin information to the dataset
+
+        Parameters
+        ----------
+        delta_r : Quantity[u.kpc]
+            spacing between two radial bins
+        delta_l : Quantity[u.kpc]
+            spacing (in spatial units) between two azimuthal bins
+        """
+
+        self.ring_binning(delta_r)
+
+        for fluid in self.fluids:
+            delta_phi = (delta_l / self.data[fluid]["r_bin"]) * u.rad
+            phi_bin = (np.trunc(self.data[fluid]["phi"] / delta_phi) + 0.5) * delta_phi
+            self.data[fluid]["phi_bin"] = phi_bin
+
+    def sector_analysis(
+        self,
+        delta_r: Quantity[u.kpc] = u.kpc,
+        delta_l: Quantity[u.kpc] = u.kpc,
+        rmin: Quantity[u.kpc] = 1 * u.kpc,
+        rmax: Quantity[u.kpc] = 12 * u.kpc,
+    ):
+        """Compute the aggration of quantities in sectors bins
+
+        Parameters
+        ----------
+        delta_r : Quantity[u.kpc], optional
+            spacing between two radial bins, by default u.kpc
+        delta_l : Quantity[u.kpc], optional
+            spacing (in spatial units) between two azimuthal bins, by default u.kpc
+        rmin : Quantity[u.kpc], optional
+            filter out bin below that radius, by default 1*u.kpc
+        rmax : Quantity[u.kpc], optional
+            filter out bin beyond that radius, by default 12*u.kpc
+        """
+
+        self.sector_binning(delta_r, delta_l)
+        grouped_data = {}
+        self.sectors = {}
+
+        for fluid in self.fluids:
+            if fluid == "gas":
+                extensive_fields = ["mass", "ek", "volume"]
+            else:
+                extensive_fields = ["mass", "ek"]
+            filtered_data = self.data[fluid][
+                np.logical_and(
+                    self.data[fluid]["r"] > rmin, self.data[fluid]["r"] < rmax
+                )
+            ]
+            grouped_data[fluid] = filtered_data.group_by(["r_bin", "phi_bin"])
+            self.sectors[fluid] = hstack(
+                [
+                    grouped_data[fluid]["r_bin", "phi_bin"].groups.aggregate(np.fmin),
+                    aggregate(grouped_data[fluid], extensive_fields=extensive_fields),
+                ]
+            )
+            self.sectors[fluid].rename_column("r_bin", "r")
+            self.sectors[fluid].rename_column("phi_bin", "phi")
+
+        self.sectors["stars"]["sfr"] = (
+            np.zeros(len(self.sectors["stars"]["mass"])) * u.Msun / u.year
+        )
+        for i, group in enumerate(grouped_data["stars"].groups):
+            self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
+
+        self.sectors["stars"]["sfr"][i] = get_sfr(group, self.time)
+
+    def ring_analysis(
+        self,
+        delta_r: Quantity[u.kpc] = u.kpc,
+        rmin: Quantity[u.kpc] = 1 * u.kpc,
+        rmax: Quantity[u.kpc] = 12 * u.kpc,
+    ):
+        """Compute the aggration of quantities in radial bins
+
+        Parameters
+        ----------
+        delta_r : Quantity[u.kpc], optional
+            spacing between two radial bins, by default u.kpc
+        rmin : Quantity[u.kpc], optional
+            filter out bin below that radius, by default 1*u.kpc
+        rmax : Quantity[u.kpc], optional
+            filter out bin beyond that radius, by default 12*u.kpc
+        """
+
+        self.ring_binning(delta_r)
+        grouped_data = {}
+        self.rings = {}
+
+        for fluid in self.fluids:
+            if fluid == "gas":
+                extensive_fields = ["mass", "ek", "volume"]
+            else:
+                extensive_fields = ["mass", "ek"]
+            filtered_data = self.data[fluid][
+                np.logical_and(
+                    self.data[fluid]["r"] > rmin, self.data[fluid]["r"] < rmax
+                )
+            ]
+            grouped_data[fluid] = filtered_data.group_by(["r_bin"])
+            self.rings[fluid] = hstack(
+                [
+                    grouped_data[fluid][
+                        "r_bin",
+                    ].groups.aggregate(np.fmin),
+                    aggregate(grouped_data[fluid], extensive_fields=extensive_fields),
+                ]
+            )
+            self.rings[fluid].rename_column("r_bin", "r")
+
+        self.rings["stars"]["sfr"] = (
+            np.zeros(len(self.rings["stars"]["mass"])) * u.Msun / u.year
+        )
+        for i, group in enumerate(grouped_data["stars"].groups):
+            self.rings["stars"]["sfr"][i] = get_sfr(group, self.time)

galturb/galsec_plot.py

@@ -0,0 +1,53 @@
+# coding: utf-8
+"""
+    Galaxy kiloparsec plotter
+"""
+from astropy.table import QTable
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_radial(table: QTable):
+
+    fig = plt.figure(figsize=(6, 5))
+    # setting the axis limits in [left, bottom, width, height]
+    rect = [0.1, 0.1, 0.8, 0.8]
+
+    # the polar axis:
+    ax_polar = fig.add_axes(rect, polar=True, frameon=False)
+    rmax = 10
+    ax_polar.set_rmax(rmax)
+
+    ax_polar.set_xticklabels([])
+
+    sc = ax_polar.scatter(table["phi"], table["r"], c=table["mass"], s=100, alpha=1)
+    plt.colorbar(sc)
+
+    fig, ax = plt.subplots(subplot_kw=dict(projection="polar"), figsize=(6, 5))
+
+    rbins = np.unique(table["r"])
+    delta_r = rbins[1] - rbins[0]
+
+    for r in rbins:
+
+        mask = table["r"] == r
+        phibins = table["phi"][mask].value
+
+        C = np.log10(table["mass"][mask].value)
+        N = len(C)
+        C = C.reshape(1, N)
+        P = np.zeros(shape=(2, N + 1))
+
+        R = np.ones(shape=(2, N + 1)) * (r + delta_r / 2.0)
+        R[0, :] -= delta_r
+        R = R.value
+
+        deltaphi = phibins[1] - phibins[0]
+        P[0, :-1] = phibins - deltaphi / 2
+        P[1, :-1] = phibins - deltaphi / 2
+        P[0, -1] = P[0, 0]
+        P[1, -1] = P[1, 0]
+
+        pc = ax.pcolormesh(P, R, C, cmap="magma_r", vmin=6, vmax=9)
+        print(P, R, C)
+    fig.colorbar(pc)

galturb/params_gal.yml

@@ -0,0 +1,96 @@
+plot : # Plot parameters
+    put_title : False                 # Add a title to plot
+
+    # Maps
+    ntick     : 6                     # Number of ticks for maps
+
+    # Overlays
+    vel_red   : 40                    # Take point each vel_red for velocities
+
+    time_fmt  :  "time = {:.3g} {}"   # Time format string, 1st field is time and 2nd is unit
+
+disk: # Disk specific parameters
+    enable    : True                 # Enable specific disk analysis
+    center    : [0.5, 0.5, 0.5]       # Position of the center
+    binning   : "log"                 # Kind of binning (lin : linear, log : logarithmic)
+    mass_star : 1.                    # Mass of the central star
+    nb_bin    : 256                   # Number of bins for averaged quantities
+    bin_in    : 1e-3                  # Outer radius of the inner bin
+    bin_out   : 18                    # Inner radius of the outer bin
+    rmin_pdf  : 1                     # Inner radius for PDF computation
+    rmax_pdf  : 18                    # Outer radius for PDF computation
+
+
+pdf: # parameters for probability density functions
+    nb_bin    : 100                   # Number of bins for the PDF
+    range     : [-1.5, 2.5]           # Range of the PDF (log of fluctuation)
+    xmin_fit  : 0.3                   # Lower boundary of the fit (log of fluctuation)
+    xmax_fit  : 1.5                   # Upper boundary of the fit (log of fluctuation)
+    fit_cut   : 1e-4                  # Exclude value that are < fit_cut * maximum
+
+
+pymses: # Parameters for Pymses reader
+
+    # Source settings
+    variables :  ["rho", "vel", "P"]  # Read these variables
+    part_variables :   ["vel","mass","id","level","epoch"]  # Read these variables
+    order     : '<'                                       # Bit order
+
+
+    # Processing options
+    levelmax : 20                      # Maximal AMR level visited when looking levels
+    fft      : False                   # Quick and dirty rendering using FFT
+    verbose  : True                   # 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     : 4.                      # Zoom  of the image
+    map_size : 2048                    # 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
+    label_filename : "label.txt"       # Name of the label file
+    nml_filename   : "galaxy.nml"      # name of the namelist file
+    ramses_ism     : False             # If ramses-ism is used
+
+out: # Parameters for post processing
+    tag         : ""                   # Tag for the image
+    interactive : True                 # Interactive mode (keep figures open)
+    save        : True                 # Save the plots on the disk
+    ext         : '.jpeg'              # extension for plots
+    fmt         : ""                   # Format of the output filename for plots
+    # The following keys are accepted
+    # {out} : The output directory (where hdf5  files are also stored)
+    # {run} : Name of the relevant run
+    # {num} : Name of the input file (from Ramses)
+    # {ext} : Extension defined above
+    # {name} : Name of the rule
+    # {tag} : Tag defined above
+    # {nml[nml_key]} : The value of nml_key in the namelist (ex: {nml[amr_params/levelmin]})
+
+
+process: # General setting of the post-processor module
+    verbose : True                     # Give more infos on what is going on
+    num_process : 1                    # Number of forks
+    save_cells : True                  # Save cells structure on disk
+    save_parts : True                  # Save particles on disk
+    unload_cells : True                # Save memory usage
+
+rules: # Specific rules parameters
+    turb_energy_threshold : -1       # Remove invalid data (<0 = no threshold)
+
+
+astrophysix: # Parameters for astrophysix and galactica
+    simu_fmt          : "{tag}_{run}"      # Format of the name of simulation
+    descr_fmt         : "{tag}_{run}"      # Format of the default description
+    # The following keys are accepted
+    # {run} : Name of the relevant run
+    # {tag} : Tag defined above
+    # {nml[nml_key]} : The value of nml_key in the namelist (ex: {nml[amr_params/levelmin]})

galturb/pymsesrc

@@ -0,0 +1,13 @@
+{
+    "Version": 1,
+    "RAMSES": {
+        "ndimensions": 3,
+        "amr_field_descr": [
+            {"__type__": "scalar_field", "__file_type__": "hydro", "name": "rho", "ivar": 0},
+            {"__type__": "vector_field", "__file_type__": "hydro", "name": "vel", "ivars": [1, 2, 3]},
+            {"__type__": "scalar_field", "__file_type__": "hydro", "name": "P", "ivar": 4},
+            {"__type__": "scalar_field", "__file_type__": "grav", "name": "phi", "ivar": 0},
+            {"__type__": "vector_field", "__file_type__": "grav", "name": "g", "ivars": [1, 2, 3]}
+        ]
+    }
+}

galturb/sectors_extraction_ramses.py

@@ -0,0 +1,108 @@
+# coding: utf-8
+
+import numpy as np
+from snapshotprocessor import U
+from tables import NoSuchNodeError
+
+
+def load_fields(pp):
+    """
+    Parameters
+    ----------
+    path : _type_
+        _description_
+    output : _type_
+        _description_
+
+    Returns
+    -------
+    dict
+        A dataset of cells in the following format:
+        gas:
+            position (Ngas, 3) [kpc], centered
+            volume   (Ngas)    [pc^3]
+            velocity (Ngas, 3) [km/s]
+            mass     (Ngas)    [Msun]
+        stars:
+            position (Nstar, 3) [kpc], centered
+            velocity (Nstar, 3) [km/s]
+            mass     (Nstar)    [Msun]
+            birth_time      (Nstar)    [Myr]
+        dm:
+            position (Ngas, 3) [kpc], centered
+            velocity (Ngas, 3) [km/s]
+            mass     (Ngas)    [Msun]
+        maps:
+            extent      (xmin, xmax, ymin, ymax) Coordinates of the edges of the map, centered
+            gas_coldens (Nx, Ny) [Msun/pc^2], map of column density
+    """
+
+    # Load  arrays
+    pp.load_cells(keys=["pos", "vel", "dx", "rho"])
+
+    try:
+        pp.load_parts(keys=["pos", "vel", "mass", "epoch"])
+    except (KeyError, NoSuchNodeError):
+        pp.load_parts(keys=["pos", "vel", "mass"])
+
+    cells = pp.cells
+    parts = pp.parts
+
+    if "epoch" not in parts:
+        parts["epoch"] = np.zeros(len(pp.parts["mass"]))
+
+    # Compute extra fields and convert units
+    for dset in (cells, parts):
+        dset["position"] = dset["pos"] - np.array([0.5, 0.5, 0.5])
+
+    cells["mass"] = (
+        cells["rho"]
+        * cells["dx"] ** 3
+        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun)
+    )
+
+    cells["volume"] = cells["dx"] ** 3 * (pp.info["unit_length"] ** 3).express(
+        U.kpc**3
+    )
+
+    # Separate DM from stars
+    mass_dm = np.max(parts["mass"])
+    mask_dm = parts["mass"] == mass_dm
+    mask_star = parts["mass"] < mass_dm
+
+    # Create dataset
+    data = {
+        "header": {"time": pp.time * pp.info["unit_time"].express(U.Myr)},
+        "gas": {
+            "position": cells["position"] * pp.info["unit_length"].express(U.kpc),
+            "volume": cells["volume"],
+            "velocity": cells["vel"] * pp.info["unit_velocity"].express(U.km_s),
+            "mass": cells["mass"],
+        },
+        "stars": {
+            "position": parts["position"][mask_star]
+            * pp.info["unit_length"].express(U.kpc),
+            "velocity": parts["vel"][mask_star]
+            * pp.info["unit_velocity"].express(U.km_s),
+            "mass": parts["mass"][mask_star] * pp.info["unit_mass"].express(U.Msun),
+            "birth_time": parts["epoch"][mask_star]
+            * pp.info["unit_time"].express(U.Myr),
+        },
+        "dm": {
+            "position": parts["position"][mask_dm]
+            * pp.info["unit_length"].express(U.kpc),
+            "velocity": parts["vel"][mask_dm]
+            * pp.info["unit_velocity"].express(U.km_s),
+            "mass": parts["mass"][mask_dm] * pp.info["unit_mass"].express(U.Msun),
+        },
+    }
+    return data
+
+
+if __name__ == "__main__":
+    from snapshotprocessor import SnapshotProcessor
+
+    pp = SnapshotProcessor(
+        "/home/nbrucy/simus/F20H_alfven_frig", num=1, params="params_gal.yml"
+    )
+    data = load_fields(pp)