From 87e6d7da0877606233e0c0c2ec44563c695f11f6 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Tue, 23 Aug 2022 10:13:04 +0200 Subject: [PATCH] add sectors extraction script --- sectors_extraction.py | 266 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 266 insertions(+) create mode 100644 sectors_extraction.py diff --git a/sectors_extraction.py b/sectors_extraction.py new file mode 100644 index 0000000..8d2b7ae --- /dev/null +++ b/sectors_extraction.py @@ -0,0 +1,266 @@ +# coding: utf-8 + +import numpy as np +import pandas as pd +from plotter import U + + +def get_gas_dm_stars(pp): + # Load arrays + try: + pp.load_parts(keys=["pos", "vel", "mass", "epoch"]) + except: + pp.load_parts(keys=["pos", "vel", "mass"]) + + pp.load_cells(keys=["pos", "vel", "dx", "rho"]) + 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 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_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) + mask = epoch > 0 + masstot, time = np.histogram(epoch[mask], weights=mass[mask], bins=200) + 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.params.pymses.map_size + center = np.array(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 + """ + masks_box = [(np.abs(dset["phi"] - phi) < dphi) for i, dset in enumerate(gds_ring)] + gds_box = [ + {key: dset[key][mask] for key in dset if key in keys} + for dset, mask in zip(gds_ring, masks_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 dir in sigma: + res[f"sigma_{dir}_{fluid}"] = sigma[dir] + return res + + +keys = ["epoch", "ek", "mass_kg", "pos_kpc", "velphi", "velr", "velz", "mass", "phi"] + + +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) + 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 + ] + + res = {} + for key in data_sec[0]: + res[key] = [d[key] for d in data_sec] + + res["r"] = [r] * len(phi_sectors) + res["dr"] = [dr] * len(phi_sectors) + return res + + +def get_time_from_relax(pp): + try: + epoch = pp.parts["epoch"].copy() + epoch *= pp.info["unit_time"].express(U.Myr) + trelax = np.min(epoch[epoch > 0]) + tfromrelax = np.max(epoch - trelax) + except KeyError: + tfromrelax = 0.0 + return tfromrelax + + +def analyse_rings(pp, radius=[4]): + res = {} + for i, r in enumerate(radius): + ring = ring_analysis(pp, r, dphi=1.0 / (2 * r)) + + if i == 0: + res = ring + else: + for key in res: + res[key].extend(ring[key]) + + 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=12.0): + + res = {} + res["run"] = pp.run + res["num"] = pp.num + res["coldens"] = np.mean(pp.coldens_map[pp.rr_map < rmax]) + res["sfr"] = get_sfr(pp, pp.stars) + res["time"] = get_time_from_relax(pp) + + for dset, fluid in zip([pp.gas, pp.dm, pp.stars], ["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 + return res + + +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() + del pp.dm + del pp.gas + del pp.stars + return res + + +def allinone(pp, redo=False): + def fun(pp): + return analyse_disk(pp), analyse_rings(pp, [4, 5, 6, 7, 8]) + + + try: + assert(not redo) + sectors = pd.read_csv("{pp.run}/disk_{pp.run}_{pp.num}.csv") + disk = pd.read_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") + + except (AssertionError, FileNotFoundError): + res = load_wrapper(pp, fun) + disk = pd.DataFrame({key: [res[0][key]] for key in res[0]}) + sectors = pd.DataFrame({key: res[1][key] for key in res[1]}) + sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") + disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") + + return disk, sectors