halox.emus: Emulators

halox.emus: Emulators#

halox provides JAX-powered neural network emulators to speed up expensive theory computations. In particular, it packs an emulator of the RMS of density fluctuations \(\sigma\) as a function of mass, redshift, and cosmology, which can be used as a backend to compute halo bias and halo mass function. For tutorials on how to use the emulator as a backend to accelerate computations, see Using the σ(M) emulator.

Available trained emulators#

		Training parameter space
File Name	Architecture	Size	\(\log_{10} M\)	\(z\)	\(\Omega_b\)	\(\Omega_c\)	\(h\)	\(\sigma_8\)	\(n_s\)
sigma_mp4.npz (default)	MLP; 3 hidden layers (size 64); SiLU inner activations	104,858	[11, 16]	[-0.04, 5]	[0.01, 0.08]	[0.085, 0.5]	[0.6, 1.0]	[0.4, 1.0]	[0.8, 1.1]

API#

SigmaMEmulator([weight_file])

Neural network emulator for \(\sigma(M, z)\).

class halox.emus.SigmaMEmulator(weight_file='sigma_mp4.npz')[source]#

Neural network emulator for \(\sigma(M, z)\).

Wraps a pre-trained neural network that emulates the RMS variance of density fluctuations \(\sigma(M, z)\) as a function of halo mass, redshift, and cosmological parameters.

Parameters:: weight_file (str, optional) – Name of the weight file to load from the package data, default "sigma_mp4.npz".

build_input(m, z, c)[source]#

Build and normalize the network input array.

Parameters:

m (ArrayLike) – Halo mass [h-1 Msun]
z (ArrayLike) – Redshift
c (jc.Cosmology) – Underlying cosmology

Returns:

Normalized input array of shape (n, 7).

Return type:

Array

__call__(m, z, cosmo)[source]#

Evaluate \(\sigma(M, z)\) using the emulator.