"""Convert between SWIFT internal values and :mod:`astropy` cosmologies."""
from swiftsimio.optional_packages import ASTROPY_AVAILABLE
from unyt import unyt_quantity
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from swiftsimio.metadata.objects import SWIFTUnits
if ASTROPY_AVAILABLE:
from astropy.cosmology import w0waCDM
from astropy.cosmology import Cosmology
import astropy.constants as const
import astropy.units as astropy_units
import numpy as np
[docs]
def swift_neutrinos_to_astropy(
N_eff: float, N_ur: float, M_nu_eV: np.ndarray[float], deg_nu: np.ndarray[float]
) -> astropy_units.Quantity:
"""
Convert SWIFT metadata to the neutrino information needed by :mod:`astropy`.
Parameters
----------
N_eff : float
Fractional number of effective massless neutrinos at high redshift.
N_ur : float
Fractional number of massless neutrino species.
M_nu_eV : array of floats
Masses in eV of massive species only, up to degeneracy.
deg_nu : array of floats
Fractional degeneracies of the massive neutrino species.
Returns
-------
out : ap_m_nu
Array of neutrino masses in eV, replicated according to degeneracy,
including massless species, as desired by astropy.
"""
if np.isscalar(deg_nu):
deg_nu = np.array([deg_nu])
if np.isscalar(M_nu_eV):
M_nu_eV = np.array([M_nu_eV])
if not (deg_nu == deg_nu.astype(int)).all():
raise AttributeError(
"SWIFTsimIO uses astropy, which cannot handle this cosmological model."
)
if not int(N_eff) == deg_nu.astype(int).sum() + int(np.squeeze(N_ur)):
raise AttributeError(
"SWIFTSimIO uses astropy, which cannot handle this cosmological model."
)
ap_m_nu = [[m] * int(d) for m, d in zip(M_nu_eV, deg_nu)] # replicate
ap_m_nu = sum(ap_m_nu, []) + [0.0] * int(
np.squeeze(N_ur)
) # flatten + add massless
ap_m_nu = np.array(ap_m_nu) * astropy_units.eV
return ap_m_nu
def swift_cosmology_to_astropy(cosmo: dict, units: "SWIFTUnits") -> Cosmology:
"""
Convert SWIFT metadata to an :mod:`astropy` cosmology.
Parameters
----------
cosmo: dict
Cosmology dictionary ready straight out of the SWIFT snapshot.
units: SWIFTUnits
The SWIFT Units instance associated with this snapshot.
Returns
-------
out : Cosmology
An instance of ``astropy.cosmology.Cosmology`` filled with the
correct parameters.
"""
H0 = unyt_quantity(cosmo["H0 [internal units]"][0], units=1.0 / units.time)
Omega_b = cosmo["Omega_b"][0]
Omega_lambda = cosmo["Omega_lambda"][0]
Omega_r = cosmo["Omega_r"][0]
Omega_m = cosmo["Omega_m"][0]
w_0 = cosmo["w_0"][0]
w_a = cosmo["w_a"][0]
# For backwards compatibility with previous cosmology constructs
# in snapshots
Tcmb0 = None
Neff = None
N_ur = None
M_nu_eV = None
deg_nu = None
try:
Tcmb0 = cosmo["T_CMB_0 [K]"][0]
assert Tcmb0 != 0
except (IndexError, KeyError, AttributeError, AssertionError):
# expressions taken directly from astropy, since they do no longer
# allow access to these attributes (since version 5.1+)
critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value
a_B_c2 = (4.0 * const.sigma_sb / const.c**3).cgs.value
# SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy.
# This is an exact inversion of the procedure performed in astropy.
critical_density_0 = astropy_units.Quantity(
critdens_const * H0.to_value("1/s") ** 2,
astropy_units.g / astropy_units.cm**3,
)
Tcmb0 = (Omega_r * critical_density_0.to_value("g/cm**3") / a_B_c2) ** (
1.0 / 4.0
)
try:
Neff = cosmo["N_eff"][0]
except (IndexError, KeyError, AttributeError):
Neff = 3.04 # Astropy default
try:
M_nu_eV = cosmo["M_nu_eV"]
except (IndexError, KeyError, AttributeError):
M_nu_eV = 0.0
try:
deg_nu = cosmo["deg_nu"]
except (IndexError, KeyError, AttributeError):
deg_nu = 0.0
try:
N_ur = cosmo["N_ur"]
except (IndexError, KeyError, AttributeError):
N_ur = 3.04 # Astropy default
ap_m_nu = swift_neutrinos_to_astropy(Neff, N_ur, M_nu_eV, deg_nu)
return w0waCDM(
H0=H0.to_astropy(),
Om0=Omega_m,
Ode0=Omega_lambda,
w0=w_0,
wa=w_a,
Tcmb0=Tcmb0,
Ob0=Omega_b,
Neff=Neff,
m_nu=ap_m_nu,
)
else:
[docs]
def swift_cosmology_to_astropy(cosmo: dict, units: "SWIFTUnits") -> dict:
"""
Mock converting to an astropy cosmology if unavailable.
Parameters
----------
cosmo : dict
Cosmology dictionary ready straight out of the SWIFT snapshot.
units : SWIFTUnits
The SWIFT Units instance associated with this snapshot.
Returns
-------
out : dict
The same metadata dictionary that was input.
"""
return cosmo