111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
# coding: utf-8
|
|
|
|
import numpy as np
|
|
from snapshotprocessor import U
|
|
from tables import NoSuchNodeError
|
|
|
|
|
|
def load_fields(pp):
|
|
"""
|
|
Parameters
|
|
----------
|
|
path : _type_
|
|
_description_
|
|
output : _type_
|
|
_description_
|
|
|
|
Returns
|
|
-------
|
|
dict
|
|
A dataset of cells in the following format:
|
|
gas:
|
|
position (Ngas, 3) [kpc], centered
|
|
volume (Ngas) [pc^3]
|
|
velocity (Ngas, 3) [km/s]
|
|
mass (Ngas) [Msun]
|
|
stars:
|
|
position (Nstar, 3) [kpc], centered
|
|
velocity (Nstar, 3) [km/s]
|
|
mass (Nstar) [Msun]
|
|
birth_time (Nstar) [Myr]
|
|
dm:
|
|
position (Ngas, 3) [kpc], centered
|
|
velocity (Ngas, 3) [km/s]
|
|
mass (Ngas) [Msun]
|
|
maps:
|
|
extent (xmin, xmax, ymin, ymax) Coordinates of the edges of the map, centered
|
|
gas_coldens (Nx, Ny) [Msun/pc^2], map of column density
|
|
"""
|
|
|
|
# Load arrays
|
|
pp.load_cells(keys=["pos", "vel", "dx", "rho"])
|
|
|
|
try:
|
|
pp.load_parts(keys=["pos", "vel", "mass", "epoch"])
|
|
except (KeyError, NoSuchNodeError):
|
|
pp.load_parts(keys=["pos", "vel", "mass"])
|
|
|
|
cells = pp.cells
|
|
parts = pp.parts
|
|
|
|
if "epoch" not in parts:
|
|
parts["epoch"] = np.zeros(len(pp.parts["mass"]))
|
|
|
|
# Compute extra fields and convert units
|
|
for dset in (cells, parts):
|
|
dset["position"] = dset["pos"] - np.array([0.5, 0.5, 0.5])
|
|
|
|
cells["mass"] = (
|
|
cells["rho"]
|
|
* cells["dx"] ** 3
|
|
* (pp.info["unit_density"] * pp.info["unit_length"] ** 3).express(U.Msun)
|
|
)
|
|
|
|
cells["volume"] = cells["dx"] ** 3 * (pp.info["unit_length"] ** 3).express(
|
|
U.kpc**3
|
|
)
|
|
|
|
# Separate DM from stars
|
|
mass_dm = np.max(parts["mass"])
|
|
mask_dm = parts["mass"] == mass_dm
|
|
mask_star = parts["mass"] < mass_dm
|
|
|
|
# Create dataset
|
|
data = {
|
|
"header": {"time": pp.time * pp.info["unit_time"].express(U.Myr),
|
|
"fluids": ["gas", "stars", "dm"],
|
|
"box_size": pp.info["unit_length"].express(U.kpc)},
|
|
"gas": {
|
|
"position": cells["position"] * pp.info["unit_length"].express(U.kpc),
|
|
"volume": cells["volume"],
|
|
"velocity": cells["vel"] * pp.info["unit_velocity"].express(U.km_s),
|
|
"mass": cells["mass"],
|
|
},
|
|
"stars": {
|
|
"position": parts["position"][mask_star]
|
|
* pp.info["unit_length"].express(U.kpc),
|
|
"velocity": parts["vel"][mask_star]
|
|
* pp.info["unit_velocity"].express(U.km_s),
|
|
"mass": parts["mass"][mask_star] * pp.info["unit_mass"].express(U.Msun),
|
|
"birth_time": parts["epoch"][mask_star]
|
|
* pp.info["unit_time"].express(U.Myr),
|
|
},
|
|
"dm": {
|
|
"position": parts["position"][mask_dm]
|
|
* pp.info["unit_length"].express(U.kpc),
|
|
"velocity": parts["vel"][mask_dm]
|
|
* pp.info["unit_velocity"].express(U.km_s),
|
|
"mass": parts["mass"][mask_dm] * pp.info["unit_mass"].express(U.Msun),
|
|
},
|
|
}
|
|
return data
|
|
|
|
|
|
if __name__ == "__main__" and False:
|
|
from snapshotprocessor import SnapshotProcessor
|
|
|
|
pp = SnapshotProcessor(
|
|
"/home/nbrucy/ecogal/galturb/F20H_15_4pc_frig_from_relax", num=50, params="params_gal.yml"
|
|
)
|
|
data = load_fields(pp)
|