"""1/Vmax estimator for the luminosity function and stellar mass function.
Reference
---------
Schmidt (1968), ApJ 151, 393.
Condon (1989), ApJ 338, 13 — Poisson error formula.
Johnston (2011), arXiv:1106.2039v3, §3, Eq. 51–53.
"""
from __future__ import annotations
import jax
import jax.numpy as jnp
import numpy as np
from astropy.cosmology import FlatLambdaCDM
from ..catalogue import GalaxyCatalogue
@jax.jit
def _phi_vmax_jax(
x: jnp.ndarray,
log_vmax: jnp.ndarray,
bin_edges: jnp.ndarray,
weight: jnp.ndarray,
) -> tuple[jnp.ndarray, jnp.ndarray]:
"""JAX kernel: 1/Vmax sum and Poisson error per bin.
Implements the Schmidt (1968) estimator with the Condon (1989) Poisson
error (Johnston 2011 Eq. 51–53)::
φ_j = Σ_{i∈bin_j} w_i / V_{max,i} / ΔM_j
σ_j = sqrt(Σ_{i∈bin_j} (w_i / V_{max,i})^2) / ΔM_j
Parameters
----------
x : jnp.ndarray, shape (N,)
Observable (absolute magnitude or log10 stellar mass).
log_vmax : jnp.ndarray, shape (N,)
log10(V_max [Mpc^3]) per galaxy (Johnston 2011 Eq. 51).
bin_edges : jnp.ndarray, shape (n_bins+1,)
Bin edges in the same units as x.
weight : jnp.ndarray, shape (N,)
Per-galaxy weights.
Returns
-------
phi : jnp.ndarray, shape (n_bins,)
φ = Σ w_i / V_max,i per bin divided by bin width.
phi_err : jnp.ndarray, shape (n_bins,)
sqrt(Σ (w_i / V_max,i)^2) / bin_width — Condon (1989) Poisson error
(Johnston 2011 Eq. 53).
"""
n_bins = bin_edges.size - 1
vmax = 10.0**log_vmax
def _bin(i):
lo = bin_edges[i]
hi = bin_edges[i + 1]
dx = hi - lo
mask = (x >= lo) & (x < hi)
contrib = jnp.where(mask, weight / vmax, 0.0)
phi_i = jnp.sum(contrib) / dx
phi_err_i = jnp.sqrt(jnp.sum(jnp.where(mask, (weight / vmax) ** 2, 0.0))) / dx
return phi_i, phi_err_i
phi, phi_err = jax.vmap(_bin)(jnp.arange(n_bins))
return phi, phi_err
[docs]
def luminosity_function_vmax(
cat: GalaxyCatalogue,
z_min: float,
z_max: float,
mag_bins: np.ndarray,
area_sr: float,
cosmo: FlatLambdaCDM,
z_max_individual: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""1/Vmax luminosity function estimator.
Non-parametric estimator using the maximum-volume method of Schmidt
(1968). Each galaxy contributes w_i / V_{max,i} to its magnitude bin
(Johnston 2011 §3, Eq. 51–53).
Parameters
----------
cat : GalaxyCatalogue
Galaxy catalogue with abs_mag populated.
z_min, z_max : float
Redshift limits of the survey volume.
mag_bins : ndarray, shape (n_bins+1,)
Absolute magnitude bin edges [mag].
area_sr : float
Survey solid angle [sr].
cosmo : FlatLambdaCDM
Cosmology for V_max computation.
z_max_individual : ndarray, optional
Per-galaxy maximum observable redshift. Defaults to z_max.
Returns
-------
mag_centres : ndarray, shape (n_bins,)
Absolute magnitude bin centres [mag].
phi : ndarray, shape (n_bins,)
Luminosity function [Mpc^-3 mag^-1].
phi_err : ndarray, shape (n_bins,)
Poisson-like uncertainty [Mpc^-3 mag^-1] (Condon 1989).
References
----------
Schmidt (1968), ApJ 151, 393.
Condon (1989), ApJ 338, 13.
Johnston (2011), arXiv:1106.2039v3, §3, Eq. 51–53.
"""
if cat.abs_mag is None:
raise ValueError("GalaxyCatalogue.abs_mag must be set for LF computation.")
abs_mag = cat.abs_mag if cat.abs_mag.ndim == 1 else cat.abs_mag[:, 0]
vmax = cat.vmax(z_min, z_max, area_sr, cosmo, z_max_individual)
log_vmax = np.log10(np.clip(vmax, 1e-30, None))
mag_bins = np.asarray(mag_bins, dtype=np.float64)
mag_centres = 0.5 * (mag_bins[:-1] + mag_bins[1:])
phi, phi_err = _phi_vmax_jax(
jnp.array(abs_mag),
jnp.array(log_vmax),
jnp.array(mag_bins),
jnp.array(cat.weight),
)
return mag_centres, np.array(phi), np.array(phi_err)
[docs]
def stellar_mass_function_vmax(
cat: GalaxyCatalogue,
z_min: float,
z_max: float,
mstar_bins: np.ndarray,
area_sr: float,
cosmo: FlatLambdaCDM,
z_max_individual: np.ndarray | None = None,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""1/Vmax stellar mass function estimator.
Non-parametric estimator using the maximum-volume method of Schmidt
(1968), applied to stellar mass (Johnston 2011 §3, Eq. 51–53).
Parameters
----------
cat : GalaxyCatalogue
Galaxy catalogue with log10_mstar populated.
z_min, z_max : float
Redshift limits of the survey volume.
mstar_bins : ndarray, shape (n_bins+1,)
log10(M_*/M_sun) bin edges [dex].
area_sr : float
Survey solid angle [sr].
cosmo : FlatLambdaCDM
Cosmology for V_max computation.
z_max_individual : ndarray, optional
Per-galaxy maximum observable redshift. Defaults to z_max.
Returns
-------
mstar_centres : ndarray, shape (n_bins,)
log10(M_*/M_sun) bin centres [dex].
phi : ndarray, shape (n_bins,)
Stellar mass function [Mpc^-3 dex^-1].
phi_err : ndarray, shape (n_bins,)
Poisson-like uncertainty [Mpc^-3 dex^-1] (Condon 1989).
References
----------
Schmidt (1968), ApJ 151, 393.
Condon (1989), ApJ 338, 13.
Johnston (2011), arXiv:1106.2039v3, §3, Eq. 51–53.
"""
if cat.log10_mstar is None:
raise ValueError("GalaxyCatalogue.log10_mstar must be set for SMF computation.")
vmax = cat.vmax(z_min, z_max, area_sr, cosmo, z_max_individual)
log_vmax = np.log10(np.clip(vmax, 1e-30, None))
mstar_bins = np.asarray(mstar_bins, dtype=np.float64)
mstar_centres = 0.5 * (mstar_bins[:-1] + mstar_bins[1:])
phi, phi_err = _phi_vmax_jax(
jnp.array(cat.log10_mstar),
jnp.array(log_vmax),
jnp.array(mstar_bins),
jnp.array(cat.weight),
)
return mstar_centres, np.array(phi), np.array(phi_err)