Noe Brucy
225 lines
6.5 KiB
Python
225 lines
6.5 KiB
Python
# 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_kpc"][:, 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):
|
|
"""
|
|
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
|
|
|
|
|
|
def get_time_from_relax(pp):
|
|
pp.load_parts()
|
|
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 get_last_sfr(pp, r=4, dr=0.5, z=0, dz=0.5):
|
|
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
|
|
res["num"] = pp.num
|
|
res["coldens"] = colmean_ring(pp, r)
|
|
res["sfr"] = get_last_sfr(pp, r)
|
|
res["time"] = get_time_from_relax(pp)
|
|
ring = ring_analysis(pp, r, dphi=1.0 / (2 * r))[0]
|
|
for key in ring:
|
|
for i, fluid in enumerate(["gas", "dm", "stars"]):
|
|
res[f"{key}_{fluid}"] = ring[key][i]
|
|
return res
|
|
|
|
|
|
def analyse_disk(pp, rmax=6.0, zmax=1.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["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"]):
|
|
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)
|
|
res = fun(pp)
|
|
pp.unload_cells()
|
|
pp.unload_parts()
|
|
del pp.dm
|
|
del pp.gas
|
|
del pp.stars
|
|
return res
|
|
|
|
|
|
def allinone(pp):
|
|
def fun(pp):
|
|
return analyse_disk(pp), analyse_ring(pp, 4), analyse_ring(pp, 6)
|
|
|
|
return load_wrapper(pp, fun)
|