"""Sub-Halo Abundance Matching (SHAM) stellar-mass–halo-mass relation in JAX."""
import jax
import jax.numpy as jnp
from jax.scipy.special import erfc
from functools import partial
_GIRELLI20_NO_SCATTER = dict(
B=11.79, mu=0.20, C=0.046, nu=-0.38, D=0.709, eta=-0.18, F=0.043, E=0.96
)
_GIRELLI20_SCATTER = dict(
B=11.83, mu=0.18, C=0.047, nu=-0.40, D=0.728, eta=-0.16, F=0.052, E=0.92
)
[docs]
@jax.jit
def smhm_moster13(
log10mhalo: jnp.ndarray,
z: float,
m10: float = 11.590,
m11: float = 1.195,
n10: float = 0.0351,
n11: float = -0.0247,
beta10: float = 1.376,
beta11: float = -0.826,
gamma10: float = 0.608,
gamma11: float = 0.329,
) -> jnp.ndarray:
"""Stellar mass fraction M_star / M_halo — Moster+2013 parametrisation.
Redshift evolution follows the Moster+2013 log-linear prescription.
Parameters
----------
log10mhalo : jnp.ndarray
log10 of halo mass in M_sun/h.
z : float
Redshift.
Accuracy
--------
M_*/M_h < 1 everywhere (physical constraint; verified over [10, 15] dex).
Peak of M_*/M_h at log10(M_h) ≈ 11.5 ± 0.5 (Moster+2013 Fig. 1, z=0);
verified to < 0.5 dex (2026-04-23).
Timing
------
~ 21 µs / call (JIT-compiled, N=200 masses, CPU x86-64, 2026-04-23).
"""
log10m1 = m10 + m11 * z / (1.0 + z)
n = n10 + n11 * z / (1.0 + z)
beta = beta10 + beta11 * z / (1.0 + z)
gamma = gamma10 + gamma11 * z / (1.0 + z)
log10ratio = log10mhalo - log10m1
ratio = jnp.power(10.0, log10ratio)
mstar_over_mhalo = 2.0 * n / (ratio ** (-beta) + ratio**gamma)
log10mstar = log10mhalo + jnp.log10(mstar_over_mhalo)
return log10mstar
[docs]
@jax.jit
def smhm_behroozi13(
log10mhalo: jnp.ndarray,
z: float,
eps0: float = -1.777,
eps_a: float = -0.006,
eps_z: float = 0.0,
eps_a2: float = -0.119,
m0: float = 11.514,
m_a: float = -1.793,
m_z: float = -0.251,
alpha0: float = -1.412,
alpha_a: float = 0.731,
delta0: float = 3.508,
delta_a: float = 2.608,
delta_z: float = -0.043,
gamma0: float = 0.316,
gamma_a: float = 1.319,
gamma_z: float = 0.279,
) -> jnp.ndarray:
"""Stellar mass log10(M_star / M_sun) — Behroozi+2013 parametrisation.
Implements the full redshift evolution of Behroozi, Wechsler & Conroy 2013
(ApJ 770, 57), Eq. 3-4. Every redshift correction is damped by the factor
``nu(a) = exp(-4 a^2)`` (with ``a = 1/(1+z)``), except the ``eps_a2 (a-1)``
term which is not. Omitting ``nu`` left the curve correct at z=0 but off by
~0.25 dex at z~0.13 and ~0.4 dex at z~0.26.
Parameters
----------
log10mhalo : jnp.ndarray
log10 of halo mass in M_sun/h (h=0.7 convention inside Behroozi+2013).
z : float
Redshift.
Accuracy
--------
M_*/M_h < 1 everywhere (physical constraint; verified over [10, 15] dex).
Reproduces Behroozi+2013 Fig. 5 characteristic mass M_*(z=0) to < 0.2 dex;
z=0 output is unchanged from the previous (z=0-only correct) implementation.
Timing
------
~ 25 µs / call (JIT-compiled, N=200 masses, CPU x86-64, 2026-04-23).
"""
a = 1.0 / (1.0 + z)
nu = jnp.exp(-4.0 * a * a)
log10eps = eps0 + (eps_a * (a - 1.0) + eps_z * z) * nu + eps_a2 * (a - 1.0)
log10m1 = m0 + (m_a * (a - 1.0) + m_z * z) * nu
alpha = alpha0 + (alpha_a * (a - 1.0)) * nu
delta = delta0 + (delta_a * (a - 1.0) + delta_z * z) * nu
gamma = gamma0 + (gamma_a * (a - 1.0) + gamma_z * z) * nu
x = log10mhalo - log10m1
f_x = -jnp.log10(jnp.power(10.0, alpha * x) + 1.0) + delta * (
jnp.log10(1.0 + jnp.exp(x))
) ** gamma / (1.0 + jnp.exp(jnp.power(10.0, -x)))
# f_0 == f(x=0): the denominator is 1 + exp(10**0) = 1 + e, NOT 2. (Using
# 2.0 here left the whole relation ~0.55 dex too low at every mass/redshift.)
f_0 = -jnp.log10(2.0) + delta * jnp.log10(2.0) ** gamma / (1.0 + jnp.exp(1.0))
log10mstar = log10eps + log10m1 + f_x - f_0
return log10mstar
[docs]
@jax.jit
def smhm_girelli20(
log10mhalo: jnp.ndarray,
z: float,
B: float = 11.79,
mu: float = 0.20,
C: float = 0.046,
nu: float = -0.38,
D: float = 0.709,
eta: float = -0.18,
F: float = 0.043,
E: float = 0.96,
) -> jnp.ndarray:
"""Stellar mass :math:`\\log_{10}(M_*/M_\\odot)` — Girelli+2020 parametrisation.
Double power-law SHMR with redshift-evolving parameters (Eq. 6 of
Girelli et al. 2020, A&A 634, A135):
.. math::
\\frac{M_*}{M_h}(z) = \\frac{2A(z)}{(M_h/M_A)^{-\\beta} + (M_h/M_A)^{\\gamma}}
with :math:`\\log_{10} M_A = B + z\\mu`, :math:`A = C(1+z)^\\nu`,
:math:`\\gamma = D(1+z)^\\eta`, :math:`\\beta = Fz + E`.
Default parameters are from Table 3 of Girelli+2020 (best-fit without
intrinsic scatter). Pass ``_GIRELLI20_SCATTER`` values for the σ=0.2 dex
scatter fit (Table 4).
Parameters
----------
log10mhalo : jnp.ndarray
:math:`\\log_{10}(M_h / (M_\\odot/h))`.
z : float
Redshift.
B, mu : float
:math:`\\log_{10}(M_A/M_\\odot)` pivot and linear-redshift slope.
C, nu : float
Normalisation amplitude and power-law redshift index.
D, eta : float
High-mass slope amplitude and power-law redshift index.
F, E : float
Linear-redshift slope and zero-point of the low-mass slope :math:`\\beta`.
Returns
-------
jnp.ndarray
:math:`\\log_{10}(M_* / (M_\\odot/h))`.
Accuracy
--------
M_*/M_h < 1 everywhere (physical constraint; verified over [10, 15] dex).
Reproduces Girelli+2020 Fig. 4 at z=0 to < 0.2 dex rms for the default
(no-scatter) parameters (2026-04-23).
Timing
------
~ 21 µs / call (JIT-compiled, N=200 masses, CPU x86-64, 2026-04-23).
"""
log10_MA = B + z * mu
A = C * (1.0 + z) ** nu
gamma = D * (1.0 + z) ** eta
beta = F * z + E
log10ratio = log10mhalo - log10_MA
ratio = jnp.power(10.0, log10ratio)
mstar_over_mhalo = 2.0 * A / (ratio ** (-beta) + ratio**gamma)
return log10mhalo + jnp.log10(mstar_over_mhalo)
[docs]
class SHAMModel:
"""Stellar-mass–halo-mass relation with log-normal scatter.
Parameters
----------
parametrisation : {"moster13", "behroozi13", "girelli20"}
scatter_dex : float
Log-normal scatter in M_star at fixed M_halo [dex].
"""
_SMHM_MAP = {
"moster13": smhm_moster13,
"behroozi13": smhm_behroozi13,
"girelli20": smhm_girelli20,
}
def __init__(
self,
parametrisation: str = "moster13",
scatter_dex: float = 0.2,
):
if parametrisation not in self._SMHM_MAP:
raise ValueError(
f"parametrisation must be one of {list(self._SMHM_MAP)}"
)
self.parametrisation = parametrisation
self.scatter_dex = scatter_dex
self._smhm = self._SMHM_MAP[parametrisation]
[docs]
@partial(jax.jit, static_argnums=(0,))
def log10mstar(self, log10mhalo: jnp.ndarray, z: float) -> jnp.ndarray:
"""Mean log10 M_star [M_sun] at given halo mass and redshift."""
return self._smhm(log10mhalo, z)
[docs]
@partial(jax.jit, static_argnums=(0,))
def sample(
self,
log10mhalo: jnp.ndarray,
z: float,
key: jax.random.PRNGKey,
) -> jnp.ndarray:
"""Draw log10 M_star with log-normal scatter around the mean."""
mu = self.log10mstar(log10mhalo, z)
noise = jax.random.normal(key, shape=mu.shape) * self.scatter_dex
return mu + noise