Noe Brucy
266 lines
7.6 KiB
Python
266 lines
7.6 KiB
Python
# 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 KeyError:
|
|
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
|