From 5d112e1f661a0a0c932808c8437b065291fefa33 Mon Sep 17 00:00:00 2001
From: Noe Brucy <noe.brucy@cea.fr>
Date: Fri, 21 May 2021 13:52:51 +0200
Subject: [PATCH] [galaxy] add a file with external routine useful for galaxies

---
 galaxy.py | 131 ++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 131 insertions(+)
 create mode 100644 galaxy.py

diff --git a/galaxy.py b/galaxy.py
new file mode 100644
index 0000000..94ce663
--- /dev/null
+++ b/galaxy.py
@@ -0,0 +1,131 @@
+# coding: utf-8
+
+import numpy as np
+from plotter import U
+
+
+def get_dispersion(dset, name):
+    """
+    Compute dispersion from dset["name"]
+    """
+    vel = dset[name]
+    mass = dset["mass_kg"]
+    mass_tot = np.sum(mass)
+    mean = np.sum(mass * vel) / mass_tot
+    return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot)
+
+
+def get_polar_sigma(dset):
+    """
+    Get speed dispersion in polar coordinates
+    """
+    return {
+        velname: get_dispersion(dset, velname) for velname in ["velphi", "velr", "velz"]
+    }
+
+
+def get_gas_dm_stars(pp):
+    # Load  arrays
+    pp.load_parts()
+    pp.load_cells()
+    cells = pp.cells
+    parts = pp.parts
+
+    # Compute extra fields and convert units
+    for dset in (cells, parts):
+        dset["pos_kpc"] = dset["pos"] - np.array([0.5, 0.5, 0.5])
+        dset["pos_kpc"] *= pp.info["unit_length"].express(U.kpc)
+        dset["r"] = np.sqrt(np.sum(dset["pos_kpc"][:, :2] ** 2, axis=1))
+        dset["phi"] = np.angle(dset["pos_kpc"][:, 0] + dset["pos_kpc"][:, 1] * 1j)
+        dset["velphi"] = pp.getter_vect_phi(dset, "vel") * pp.info[
+            "unit_velocity"
+        ].express(U.km_s)
+        dset["velr"] = pp.getter_vect_r(dset, "vel") * pp.info["unit_velocity"].express(
+            U.km_s
+        )
+        dset["velz"] = dset["vel"][:, 2] * pp.info["unit_velocity"].express(U.km_s)
+
+    cells["mass_kg"] = (
+        cells["rho"]
+        * cells["dx"] ** 3
+        * (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.kg)
+    )
+
+    parts["mass_kg"] = parts["mass"] * pp.info["unit_mass"].express(U.kg)
+
+    for dset in (cells, parts):
+        dset["ek"] = dset["mass_kg"] * np.sum(dset["vel"] ** 2, axis=1)
+        dset["ek"] *= (U.kg * pp.info["unit_velocity"] ** 2).express(U.J)
+
+    # Separate DM from stars
+    mass_dm = np.max(parts["mass_kg"])
+    mask_dm = parts["mass_kg"] == mass_dm
+    mask_star = parts["mass_kg"] < mass_dm
+
+    # Create separated arrays
+    gas = cells
+    dm = {key: parts[key][mask_dm] for key in parts}
+    stars = {key: parts[key][mask_star] for key in parts}
+
+    # Store arrays and return them
+    pp.gas = gas
+    pp.dm = dm
+    pp.stars = stars
+    return gas, dm, stars
+
+
+def extract_polar_region(dset, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
+    """
+    Returns a mask for dset of the polar volume within polar coordinates
+    [r - dr, r + dr] x [phi - dphi, phi + dphi] x [z -dz, z + dz]
+    """
+    mask_box = (
+        (np.abs(dset["r"] - r) < dr)
+        & (np.abs(dset["phi"] - phi) < dphi)
+        & (np.abs(dset["pos"][:, 2] - z) < dz)
+    )
+    return {key: dset[key][mask_box] for key in dset}
+
+
+def sector_analysis(pp, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
+    """
+    Analyze box at given coordinates
+    """
+    gds = [pp.gas, pp.dm, pp.stars]
+    gds_box = [extract_polar_region(dset, r, dr, phi, dphi, z, dz) for dset in gds]
+    result = {}
+    result["ek"] = np.array([np.sum(dset["ek"]) for dset in gds_box])  # J
+    result["mass"] = np.array([np.sum(dset["mass_kg"]) for dset in gds_box])  # kg
+    result["ek_spe"] = result["ek"] / result["mass"]  # J.kg^-1
+    for dset in gds_box:
+        sigma = get_polar_sigma(dset)
+        for key in sigma:
+            keyres = "sigma_" + key
+            if keyres in result:
+                result[keyres].append(sigma[key])
+            else:
+                result[keyres] = [sigma[key]]
+    return result
+
+
+def ring_analysis(pp, r=4, dr=0.5, dphi=0.125, z=0, dz=0.5, do_mean=True):
+    """
+    Compute the average at a given of quantities computed in polar sectors.
+    """
+
+    phi_sectors = np.arange(-np.pi + dphi, np.pi - dphi, 2 * dphi)
+    data_sec = [sector_analysis(pp, r, dr, phi, dphi, z, dz) for phi in phi_sectors]
+
+    data = {}
+    mwavg = {}
+    tot_mass = [np.sum([d["mass"][i] for d in data_sec]) for i in range(3)]
+    for key in data_sec[0]:
+        mwavg[key] = []
+        for i in range(3):
+            mwavg[key].append(
+                np.sum([d["mass"][i] * d[key][i] for d in data_sec]) / tot_mass[i]
+            )
+
+        data[key] = np.array([d[key] for d in data_sec])
+
+    return mwavg, phi_sectors, data