[galaxy] move towards stats box analysis

galaxy.py

@@ -4,26 +4,6 @@ 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()
@@ -74,65 +54,133 @@ def get_gas_dm_stars(pp):
     return gas, dm, stars
 
 
-def extract_polar_region(dset, r=4, dr=0.5, phi=0, dphi=0.125, z=0, dz=0.5):
+def get_dispersion(dset, name):
     """
-    Returns a mask for dset of the polar volume within polar coordinates
+    Compute dispersion from dset["name"]
-    [r - dr, r + dr] x [phi - dphi, phi + dphi] x [z -dz, z + dz]
     """
-    mask_box = (
+    vel = dset[name]
-        (np.abs(dset["r"] - r) < dr)
+    mass = dset["mass_kg"]
-        & (np.abs(dset["phi"] - phi) < dphi)
+    mass_tot = np.sum(mass)
-        & (np.abs(dset["pos_kpc"][:, 2] - z) < dz)
+    mean = np.sum(mass * vel) / mass_tot
-    )
+    return np.sqrt(np.sum(mass * (vel - mean) ** 2) / mass_tot)
-    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):
+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_sfr(pp, stars):
+    try:
+        epoch = stars["epoch"].copy()
+        epoch *= pp.info["unit_time"].express(U.year)
+        mass = stars["mass"].copy()
+        mass *= pp.info["unit_mass"].express(U.Msun)
+        masstot, time = np.histogram(epoch, weights=mass, bins=100)
+        dtime = np.diff(time)
+        sfr = masstot[-1] / dtime[-1]
+    except KeyError:
+        sfr = 0.0
+    return sfr
+
+
+def get_coldens(pp):
+    """
+    Get mean column density in a sector
+    """
+    pp.coldens("z")
+    pp.coldens_map = pp.get_value("/maps/coldens_z", unit=U.coldens)
+
+
+    im_extent = np.array(pp.get_attribute("/maps", "im_extent"))
+    im_extent *= pp.info["unit_length"].express(U.kpc)
+    map_size = pp.pp_params.pymses.map_size
+    center = np.array(pp.pp_params.disk.center)
+    center *= pp.info["unit_length"].express(U.kpc)
+
+    # Physical size of cells
+    dx = (im_extent[1] - im_extent[0]) / map_size
+    dy = (im_extent[3] - im_extent[2]) / map_size
+
+    # Physical coordinates of the center of the cells
+    x = np.linspace(im_extent[0], im_extent[1], map_size) + 0.5 * dx - center[0]
+    y = np.linspace(im_extent[2], im_extent[3], map_size) + 0.5 * dy - center[1]
+
+    xx, yy = np.meshgrid(x, y)
+
+    # Physical radius
+    pp.rr_map = np.sqrt(xx ** 2 + yy ** 2)
+    pp.phi_map = np.angle(xx + yy * 1j)
+
+
+def sector_analysis(pp, gds_ring, mask_ring, phi=0, dphi=0.125):
     """
     Analyze box at given coordinates
     """
-    gds = [pp.gas, pp.dm, pp.stars]
+    masks_box = [(np.abs(dset["phi"] - phi) < dphi) for i, dset in enumerate(gds_ring)]
-    gds_box = [extract_polar_region(dset, r, dr, phi, dphi, z, dz) for dset in gds]
+    gds_box = [
-    result = {}
+        {key: dset[key][mask] for key in dset if key in keys}
-    result["ek"] = np.array([np.sum(dset["ek"]) for dset in gds_box])  # J
+        for dset, mask in zip(gds_ring, masks_box)
-    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:
+    res = {}
+
+    # Generic Info
+    res["phi"] = phi
+    res["dphi"] = dphi
+
+    res["sfr"] = get_sfr(pp, gds_box[2])
+    res["coldens"] = np.mean(
+        pp.coldens_map[mask_ring & (np.abs(pp.phi_map - phi) < dphi)]
+    )
+
+    for dset, fluid in zip(gds_box, ["gas", "dm", "stars"]):
+        res[f"ek_{fluid}"] = np.sum(dset["ek"])  # J
+        res[f"mass_{fluid}"] = np.sum(dset["mass_kg"])  # kg
+        res[f"ek_spe_{fluid}"] = res[f"ek_{fluid}"] / res[f"mass_{fluid}"]  # J.kg^-1
         sigma = get_polar_sigma(dset)
-        for key in sigma:
+        for dir in sigma:
-            keyres = "sigma_" + key
+            res[f"sigma_{dir}_{fluid}"] = sigma[dir]
-            if keyres in result:
+    return res
-                result[keyres].append(sigma[key])
+
-            else:
+
-                result[keyres] = [sigma[key]]
+keys = ["epoch", "ek", "mass_kg", "pos_kpc", "velphi", "velr", "velz", "mass", "phi"]
-    return result
 
 
 def ring_analysis(pp, r=4, dr=0.5, dphi=0.125, z=0, dz=0.5):
     """
     Compute the average at a given of quantities computed in polar sectors.
     """
-
+    gds = [pp.gas, pp.dm, pp.stars]
     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]
+    masks_ring = [
+        (np.abs(dset["r"] - r) < dr) & (np.abs(dset["pos_kpc"][:, 2] - z) < dz)
+        for dset in gds
+    ]
+    gds_ring = [
+        {key: dset[key][mask] for key in dset if key in keys}
+        for dset, mask in zip(gds, masks_ring)
+    ]
+    mask_ring = np.abs(pp.rr_map - r) < dr
+    data_sec = [
+        sector_analysis(pp, gds_ring, mask_ring, phi, dphi) for phi in phi_sectors
+    ]
 
-    data = {}
+    res = {}
-    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] = []
+        res[key] = [d[key] for d in data_sec]
-        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])
+    res["r"] = [r] * len(phi_sectors)
+    res["dr"] = [dr] * len(phi_sectors)
+    return res
 
-    return mwavg, phi_sectors, data
 
 
 def get_time_from_relax(pp):
-    pp.load_parts()
     try:
         epoch = pp.parts["epoch"].copy()
         epoch *= pp.info["unit_time"].express(U.Myr)
@@ -143,57 +191,32 @@ def get_time_from_relax(pp):
     return tfromrelax
 
 
-def get_last_sfr(pp, r=4, dr=0.5, z=0, dz=0.5):
+def analyse_rings(pp, radius=[4]):
-    pp.load_parts()
-    try:
-        epoch = pp.stars["epoch"].copy()
-        epoch *= pp.info["unit_time"].express(U.year)
-        mass = pp.stars["mass"].copy()
-        mass *= pp.info["unit_mass"].express(U.Msun)
-        mask = (
-            (epoch > 0)
-            & (np.abs(pp.stars["r"] - r) < dr)
-            & (np.abs(pp.stars["pos_kpc"][:, 2] - z) < dz)
-        )
-        masstot, time = np.histogram(epoch[mask], weights=mass[mask], bins=100)
-        dtime = np.diff(time)
-        sfr = masstot[-1] / dtime[-1]
-    except KeyError:
-        sfr = 0.0
-    return sfr
-
-
-def colmean_ring(pp, r=4, dr=0.5):
-    pp.coldens("z")
-    pp.rr("z")
-    col = pp.get_value("/maps/coldens_z", unit=U.coldens)
-    rr = pp.get_value("/maps/rr_z", unit=U.kpc)
-    colmean = np.mean(col[np.abs(rr - r) < dr])
-    return colmean
-
-
-def analyse_ring(pp, r=4):
     res = {}
-    res["run"] = pp.run
+    for i, r in enumerate(radius):
-    res["num"] = pp.num
+        ring = ring_analysis(pp, r, dphi=1.0 / (2 * r))
-    res["coldens"] = colmean_ring(pp, r)
+
-    res["sfr"] = get_last_sfr(pp, r)
+        if i == 0:
-    res["time"] = get_time_from_relax(pp)
+            res = ring
-    ring = ring_analysis(pp, r, dphi=1.0 / (2 * r))[0]
+        else:
-    for key in ring:
+            for key in res:
-        for i, fluid in enumerate(["gas", "dm", "stars"]):
+                res[key].extend(ring[key])
-            res[f"{key}_{fluid}"] = ring[key][i]
+
+    res["run"] = [pp.run] * len(res["r"])
+    res["num"] = [pp.num] * len(res["r"])
+    time = get_time_from_relax(pp)
+    res["time"] = [time] * len(res["r"])
+
     return res
 
 
-def analyse_disk(pp, rmax=6.0, zmax=1.0):
+def analyse_disk(pp, rmax=12.0):
-    get_gas_dm_stars(pp)
+
     res = {}
     res["run"] = pp.run
     res["num"] = pp.num
-    res["coldens"] = colmean_ring(pp, rmax / 2.0, rmax / 2.0)
+    res["coldens"] = np.mean(pp.coldens_map[pp.rr_map < rmax])
-
+    res["sfr"] = get_sfr(pp, pp.stars)
-    res["sfr"] = get_last_sfr(pp, rmax / 2.0, rmax / 2.0, 0, zmax)
     res["time"] = get_time_from_relax(pp)
 
     for dset, fluid in zip([pp.gas, pp.dm, pp.stars], ["gas", "dm", "stars"]):
@@ -208,6 +231,8 @@ def load_wrapper(pp, fun):
     Wrapper to load_unload data and map function
     """
     get_gas_dm_stars(pp)
+    get_coldens(pp)
+
     res = fun(pp)
     pp.unload_cells()
     pp.unload_parts()
@@ -219,6 +244,6 @@ def load_wrapper(pp, fun):
 
 def allinone(pp):
     def fun(pp):
-        return analyse_disk(pp), analyse_ring(pp, 4), analyse_ring(pp, 6)
+        return analyse_disk(pp), analyse_rings(pp, [4, 5, 6, 7, 8])
 
     return load_wrapper(pp, fun)