"""
Provides the :class:`~martini.sources.tng_source.TNGSource` class.
Facilitates working with IllustrisTNG simulations as input.
"""
import io
import os
import numpy as np
from scipy.spatial.transform import Rotation
from astropy import units as U, constants as C
from astropy.coordinates import ICRS
from .sph_source import SPHSource
from ..sph_kernels import _CubicSplineKernel, find_fwhm
from ..L_coords import L_coords
from typing import Any, TYPE_CHECKING
from requests.models import Response
if TYPE_CHECKING:
from astropy.coordinates.builtin_frames.baseradec import BaseRADecFrame
[docs]
def api_get(
path: str, params: dict | None = None, api_key: str | None = None
) -> Response | dict:
"""
Make a request to the TNG web API service.
Parameters
----------
path : str
The request to submit to the API.
params : dict, optional
Additional options for the API request.
api_key : str
API key to authenticate to the TNG web API service.
Returns
-------
requests.models.Response
Response from the API, a JSON-encoded string.
"""
import requests
baseUrl = "http://www.tng-project.org/api/"
r = requests.get(f"{baseUrl}/{path}", params=params, headers={"api-key": api_key})
r.raise_for_status()
if r.headers["content-type"] == "application/json":
return r.json()
return r
[docs]
def cutout_file(simulation: str, snapNum: int, haloID: int) -> str:
"""
Generate a string identifying a cutout file.
Parameters
----------
simulation : str
Identifier of the simulation.
snapNum : int
Snapshot identifier.
haloID : int
Halo identifier.
Returns
-------
str
A string to use for a cutout file.
"""
return f"martini-cutout-{simulation}-{snapNum}-{haloID}.hdf5"
[docs]
class TNGSource(SPHSource):
"""
Class abstracting HI sources for use with IllustrisTNG simulations.
If used in the IllustrisTNG JupyterLab environment
(https://www.tng-project.org/data/lab/), files on disk are accessed directly.
Otherwise, particles for the galaxy of interest are automatically retrieved using
the TNG web API (https://www.tng-project.org/data/docs/api/).
Use of the TNG web API requires an API key: login at
https://www.tng-project.org/users/login/ or register at
https://www.tng-project.org/users/register/ then obtain your API
key from https://www.tng-project.org/users/profile/ and pass to TNGSource as the
``api_key`` parameter. An API key is not required if logged into the TNG JupyterLab.
Parameters
----------
simulation : str
Simulation identifier string, for example ``"TNG100-1"``, see
https://www.tng-project.org/data/docs/background/.
snapNum : int
Snapshot number. In TNG, snapshot 99 is the final output. Note that
if a 'mini' snapshot (see
http://www.tng-project.org/data/docs/specifications/#sec1a) is selected then
some additional approximations are used.
subID : int
Subhalo ID of the target object. Note that all particles in the FOF
group to which the subhalo belongs are used to construct the data cube.
This avoids strange 'holes' at the locations of other subhaloes in the
same group, and gives a more realistic treatment of foreground and
background emission local to the source. An object of interest could be
found using https://www.tng-project.org/data/search/, for instance. The
"ID" column in the search results on that page is the subID.
api_key : str, optional
Use of the TNG web API requires an API key: login at
https://www.tng-project.org/users/login/ or register at
https://www.tng-project.org/users/register/ then obtain your API
key from https://www.tng-project.org/users/profile/ and provide as a string. An
API key is not required if logged into the TNG JupyterLab.
cutout_dir : str, optional
Ignored if running on the TNG JupyterLab. Directory in which to search for and
save cutout files (see documentation at
https://www.tng-project.org/data/docs/api/ for a description of cutouts). Cutout
filenames are enforced by MARTINI. If `cutout_dir==None` (the default), then the
data will always be downloaded. If a `cutout_dir` is provided, it will first be
searched for the required data. If the data are found, the local copy is used,
otherwise the data are downloaded and a local copy is saved in `cutout_dir` for
future use.
distance : ~astropy.units.Quantity
:class:`~astropy.units.Quantity`, with dimensions of length.
Source distance, also used to set the velocity offset via Hubble's law.
vpeculiar : ~astropy.units.Quantity, optional
:class:`~astropy.units.Quantity`, with dimensions of velocity.
Source peculiar velocity along the direction to the source centre.
rotation : ~scipy.spatial.transform.Rotation, optional
A rotation to apply to the source particles, specified using the
:class:`~scipy.spatial.transform.Rotation` class. That class supports many ways to
specify a rotation (Euler angle, rotation matrices, quaternions, etc.). Refer to
the :mod:`scipy` documentation for details. Note that the ``y-z`` plane will be
the one eventually placed in the plane of the "sky". Cannot be used at the same
time as ``L_coords``.
L_coords : ~martini.L_coords.L_coords, optional
A named tuple specifying 3 angles. Import it as ``from martini import L_coords``.
The angles are used to orient the galaxy relative to its angular momentum vector,
"L". The routine will first identify a preferred plane based on the angular
momenta of the central 1/3 of HI gas. This plane will then be rotated to lie in
the plane of the "sky" (``y-z`` plane), rotated by an angle ``az_rot`` around the
angular momentum vector (rotation around ``x``), then inclined by ``incl`` towards
or away from the line of sight (rotation around ``y``) and finally rotated on the
sky to set the position angle ``pa`` (second rotation around ``x``). All rotations
are extrinsic. The position angle refers to the receding side of the galaxy
measured East of North. The angles should be specified using syntax like:
``L_coords=L_coords(incl=0 * U.deg, pa=270 * U.deg, az_rot=0 * U.deg)``. These
example values are the defaults. Cannot be used at the same time as ``rotation``.
ra : ~astropy.units.Quantity, optional
:class:`~astropy.units.Quantity`, with dimensions of angle.
Right ascension for the source centroid.
dec : ~astropy.units.Quantity, optional
:class:`~astropy.units.Quantity`, with dimensions of angle.
Declination for the source centroid.
coordinate_frame : ~astropy.coordinates.builtin_frames.baseradec.BaseRADecFrame, \
optional
The coordinate frame assumed in converting particle coordinates to RA and Dec, and
for transforming coordinates and velocities to the data cube frame. The frame
needs to have a well-defined velocity as well as spatial origin. Recommended
frames are :class:`~astropy.coordinates.GCRS`, :class:`~astropy.coordinates.ICRS`,
:class:`~astropy.coordinates.HCRS`, :class:`~astropy.coordinates.LSRK`,
:class:`~astropy.coordinates.LSRD` or :class:`~astropy.coordinates.LSR`. The frame
should be passed initialized, e.g. ``ICRS()`` (not just ``ICRS``).
"""
def __init__(
self,
simulation: str,
snapNum: int,
subID: int,
*,
api_key: str | None = None,
cutout_dir: str | None = None,
distance: U.Quantity[U.Mpc],
vpeculiar: U.Quantity[U.km / U.s] = 0 * U.km / U.s,
rotation: Rotation | None = None,
L_coords: L_coords | None = None,
ra: U.Quantity[U.deg] = 0.0 * U.deg,
dec: U.Quantity[U.deg] = 0.0 * U.deg,
coordinate_frame: "BaseRADecFrame" = ICRS(),
) -> None:
# optional dependencies for this source class
import h5py
X_H = 0.76
full_fields_g = [
"Masses",
"Velocities",
"InternalEnergy",
"ElectronAbundance",
"NeutralHydrogenAbundance",
"StarFormationRate",
"Density",
"CenterOfMass",
"GFM_Metals",
]
mdi_full = [None, None, None, None, None, None, 0]
mini_fields_g = [
"Masses",
"Velocities",
"InternalEnergy",
"ElectronAbundance",
"StarFormationRate",
"Density",
"Coordinates",
]
# are we running on the TNG jupyterlab?
jupyterlab = os.path.exists("/home/tnguser/sims.TNG")
if not jupyterlab:
data_header: dict[str, Any] = {}
data_sub = {}
data_g = {}
from requests import HTTPError
if api_key is None:
raise ValueError(
"A TNG API key is required: login at "
"https://www.tng-project.org/users/login/ or register at "
"https://www.tng-project.org/users/register/ then obtain your API "
"key from https://www.tng-project.org/users/profile/"
)
if cutout_dir is not None:
grnr_file = os.path.join(
cutout_dir,
f"martini-cutout-grnr-{simulation}-{snapNum}-{subID}.npy",
)
if os.path.exists(grnr_file):
haloID = np.load(grnr_file)
have_cutout = True
else:
have_cutout = False
else:
print("No cutout_dir provided, cutout will be downloaded.")
have_cutout = False
if have_cutout:
# check for an existing local cutout file
assert cutout_dir is not None # guaranteed
cfname = os.path.join(
cutout_dir, cutout_file(simulation, snapNum, haloID)
)
if not os.path.exists(cfname):
have_cutout = False
else:
print(f"Using local cutout file {os.path.basename(cfname)}")
if not have_cutout: # not else because previous if can modify have_cutout
print("No local cutout found, cutout will be downloaded.")
sub_api_path = f"{simulation}/snapshots/{snapNum}/subhalos/{subID}"
sub = api_get(sub_api_path, api_key=api_key)
assert isinstance(sub, dict)
haloID = sub["grnr"]
np.save(grnr_file, haloID)
data_sub["SubhaloPos"] = np.array([sub[f"pos_{ax}"] for ax in "xyz"])
data_sub["SubhaloVel"] = np.array([sub[f"vel_{ax}"] for ax in "xyz"])
cutout_api_path = (
f"{simulation}/snapshots/{snapNum}/halos/{haloID}/cutout.hdf5"
)
cutout_request = {"gas": ",".join(full_fields_g)}
try:
cutout = api_get(
cutout_api_path, params=cutout_request, api_key=api_key
)
except HTTPError as exc:
if "400 Client Error: Bad Request for url" in exc.args[0]:
cutout_request = {"gas": ",".join(mini_fields_g)}
cutout = api_get(
cutout_api_path, params=cutout_request, api_key=api_key
)
else:
raise
assert isinstance(cutout, Response)
# hold file in memory
cfbuf = io.BytesIO(cutout.content)
if cutout_dir is not None:
# write a copy to disk for later use
ofile = os.path.join(
cutout_dir, cutout_file(simulation, snapNum, haloID)
)
print(f"Writing downloaded cutout to {ofile}")
with open(ofile, "wb") as of:
of.write(cutout.content)
with h5py.File(ofile, "r+") as of:
of.create_group(f"{subID}")
of[f"{subID}"].attrs["pos"] = data_sub["SubhaloPos"]
of[f"{subID}"].attrs["vel"] = data_sub["SubhaloVel"]
with h5py.File(cfbuf if "cfbuf" in locals() else cfname, "r") as cf:
minisnap = "CenterOfMass" not in cf["PartType0"].keys()
fields_g = mini_fields_g if minisnap else full_fields_g
data_g = {field: cf["PartType0"][field][()] for field in fields_g}
if len(data_header) == 0:
data_header["HubbleParam"] = cf["Header"].attrs["HubbleParam"]
data_header["Redshift"] = cf["Header"].attrs["Redshift"]
data_header["Time"] = cf["Header"].attrs["Time"]
if len(data_sub) == 0:
data_sub["SubhaloPos"] = cf[f"{subID}"].attrs["pos"]
data_sub["SubhaloVel"] = cf[f"{subID}"].attrs["vel"]
X_H_g = X_H if minisnap else data_g["GFM_Metals"][:, 0]
else:
from ._illustris_tools import (
loadHeader,
loadSingle,
loadSubset,
getSnapOffsets,
)
if cutout_dir is not None:
print("Running on TNG JupyterLab, cutout_dir ignored.")
basePath = f"/home/tnguser/sims.TNG/{simulation}/output/"
data_header = loadHeader(basePath, snapNum)
data_sub = loadSingle(basePath, snapNum, subhaloID=subID)
haloID = int(data_sub["SubhaloGrNr"])
subset_g = getSnapOffsets(basePath, snapNum, haloID, "Group")
try:
data_g = loadSubset(
basePath,
snapNum,
"gas",
fields=full_fields_g,
subset=subset_g,
mdi=mdi_full,
)
minisnap = False
except Exception as exc:
if ("Particle type" in exc.args[0]) and (
"does not have field" in exc.args[0]
):
data_g.update(
loadSubset(
basePath,
snapNum,
"gas",
fields=["CenterOfMass"],
subset=subset_g,
sq=False,
)
)
minisnap = True
X_H_g = X_H
else:
raise
X_H_g = (
X_H if minisnap else data_g["GFM_Metals"]
) # only loaded column 0: Hydrogen
a = data_header["Time"]
z = data_header["Redshift"]
h = data_header["HubbleParam"]
xe_g = data_g["ElectronAbundance"]
rho_g = data_g["Density"] << 1e10 / h * U.Msun * np.power(a / h * U.kpc, -3)
u_g = data_g["InternalEnergy"] << (U.km / U.s) ** 2
mu_g = 4 / (1 + 3 * X_H_g + 4 * X_H_g * xe_g)
gamma = 5.0 / 3.0 # see http://www.tng-project.org/data/docs/faq/#gen4
T_g = (gamma - 1) / C.k_B * (u_g * mu_g * C.m_p) << U.K
m_g = data_g["Masses"] << 1e10 / h * U.Msun
nH_g = rho_g * X_H_g / C.m_p << U.cm**-3
if "NeutralHydrogenAbundance" in data_g:
fneutral_g = data_g["NeutralHydrogenAbundance"].copy()
else:
from warnings import warn
from Hdecompose.RahmatiEtal2013 import neutral_frac
warn(
"NeutralHydrogenAbundance not available for mini snapshots, approximating"
" following Rahmati et al. (2013) - to avoid this use a full snapshot"
" instead.",
UserWarning,
)
fneutral_g = neutral_frac(
z,
nH_g,
T_g,
onlyA1=True,
mu=mu_g,
# explicitly blank unused:
SFR=None,
TNG_corrections=False, # Diemer et al. 2018 corrections applied below
gamma=None,
fH=None,
Habundance=None,
T0=None,
rho=None,
)
gamma = 5.0 / 3.0
# cold
mu_c = 4 / (1 + 3 * X_H_g) << C.m_p
u_c = C.k_B * (1e3 << U.K) / (mu_c * (gamma - 1.0)) << (U.km / U.s) ** 2
del mu_c
# hot
mu_h = 4 / (3 + 5 * X_H_g) << C.m_p # He fully ionized
T_h = (
1e3
+ 5.73e7
/ (1 + 573 * np.maximum(1.0, nH_g.to_value(U.cm**-3) / 0.13) ** -0.8)
) << U.K # Springel & Hernquist 03; Stevens 19
u_h = C.k_B * T_h / (mu_h * (gamma - 1.0)) << (U.km / U.s) ** 2
del mu_h
sfr_g = data_g["StarFormationRate"] << U.Msun / U.yr
possfr_mask = sfr_g > 0
fneutral_g[possfr_mask] = (1.0 / (u_h - u_c) * (u_h - u_g))[
possfr_mask
].to_value(U.dimensionless_unscaled)
del sfr_g, possfr_mask, u_c, u_h
# partial pressure is used; see Marinacci 17 & Diemer 18 eq. 6
P_g = (gamma - 1.0) / C.k_B * u_g * fneutral_g * rho_g << U.K / U.cm**3
fatomic_g = 1.0 / (
1.0 + (1.0 / (1.7e4 << U.K / U.cm**3) * P_g) ** 0.8 # Leroy 08
# (1. / (4.3e4 << U.K / U.cm**3) * P_g) ** 0.92 # Blitz & Rosolowsky 06
)
del P_g, rho_g, u_g
mHI_g = m_g * X_H_g * fatomic_g * fneutral_g << U.Msun
del m_g, X_H_g, fatomic_g, fneutral_g
try:
xyz_g = data_g["CenterOfMass"] * (a / h) << U.kpc
except KeyError:
xyz_g = data_g["Coordinates"] * (a / h) << U.kpc
vxyz_g = data_g["Velocities"] * np.sqrt(a) << U.km / U.s
V_cell = (
data_g["Masses"] / data_g["Density"] * np.power(a / h * U.kpc, 3)
) # Voronoi cell volume
r_cell = np.power(3.0 * V_cell / 4.0 / np.pi, 1.0 / 3.0).to(U.kpc)
# hsm_g has in mind a cubic spline that =0 at r=h, I think
hsm_g = 2.5 * r_cell * find_fwhm(_CubicSplineKernel().kernel)
xyz_centre = data_sub["SubhaloPos"] * (a / h) << U.kpc
xyz_g -= xyz_centre
vxyz_centre = data_sub["SubhaloVel"] * np.sqrt(a) << U.km / U.s
vxyz_g -= vxyz_centre
super().__init__(
distance=distance,
vpeculiar=vpeculiar,
rotation=rotation,
L_coords=L_coords,
ra=ra,
dec=dec,
h=h,
T_g=T_g,
mHI_g=mHI_g,
xyz_g=xyz_g,
vxyz_g=vxyz_g,
hsm_g=hsm_g,
)
return