import os
import jax
import jax.numpy as jnp
from hmfast.tracers.base_tracer import Tracer
from hmfast.halos.profiles import GalaxyHODProfile, Z07GalaxyHODProfile
from hmfast.download import _get_default_data_path
from hmfast.utils import Const
# Ensure high precision for cosmological integrations
jax.config.update("jax_enable_x64", True)
[docs]
class GalaxyTracer(Tracer):
"""
Galaxy counts tracer.
Attributes
----------
profile : GalaxyHODProfile
Halo occupation distribution profile used to model galaxy number counts.
dndz : tuple of jnp.ndarray
Normalized galaxy redshift distribution stored as :math:`(z, dN/dz)`.
"""
_required_profile_type = GalaxyHODProfile
def __init__(self, profile=None, dndz=None):
super().__init__(profile=profile or Z07GalaxyHODProfile())
if dndz is None:
dndz_path = os.path.join(_get_default_data_path(), "auxiliary_files", "normalised_dndz_cosmos_0.txt")
dndz = self._load_dndz_data(dndz_path)
self.dndz = dndz
@property
def dndz(self):
return self._dndz_data
@dndz.setter
def dndz(self, value):
self._dndz_data = self._normalize_dndz(value)
# --- JAX PyTree Registration ---
def _tree_flatten(self):
# The profile IS the leaf. JAX will automatically
# drill down into the profile's own 5 leaves.
leaves = (self.profile, self._dndz_data)
return (leaves, None)
@classmethod
def _tree_unflatten(cls, aux_data, leaves):
profile, dndz_data = leaves
obj = cls.__new__(cls)
obj.profile = profile
obj._dndz_data = dndz_data
return obj
[docs]
def update(self, profile=None, dndz=None):
"""
Return a new GalaxyTracer instance with updated attributes using PyTree logic.
Parameters
----------
profile : GalaxyHODProfile, optional
New HOD profile to use for the tracer. If None, the profile is unchanged.
dndz : array_like, optional
New redshift distribution (z, dN/dz). If None, the distribution is unchanged.
Returns
-------
GalaxyTracer
New tracer instance with updated attributes.
"""
flat, aux = self._tree_flatten()
new_profile = profile if profile is not None else flat[0]
new_dndz = self._normalize_dndz(dndz) if dndz is not None else flat[1]
return self._tree_unflatten(aux, (new_profile, new_dndz))
[docs]
def kernel(self, cosmology, z):
"""
Compute the galaxy kernel :math:`W_g(\\chi)` at redshift :math:`z`.
The kernel is given by:
.. math::
W_g(\\chi) = \\frac{H(z)}{c} \\frac{dN}{dz}
where :math:`dN/dz` is the normalized redshift distribution of galaxies.
Parameters
----------
cosmology : Cosmology
Cosmology object with required methods and parameters.
z : float or array_like
Redshift(s) at which to compute the kernel.
Returns
-------
W_g : array_like
Galaxy kernel evaluated at redshift(s) :math:`z`.
"""
z = jnp.atleast_1d(z)
z_g, phi_prime_g = self.dndz
phi_prime_g_at_z = jnp.interp(z, z_g, phi_prime_g, left=0.0, right=0.0)
H_grid = cosmology.hubble_parameter(z) / (Const._c_ / 1e3)
return H_grid * phi_prime_g_at_z
jax.tree_util.register_pytree_node(
GalaxyTracer,
lambda obj: obj._tree_flatten(),
lambda aux_data, children: GalaxyTracer._tree_unflatten(aux_data, children)
)