Emulator backends¶

All emulators compress the output with PCA and fit one Gaussian process per component. They share the BaseEmulator API (fit, predict, save, load), so changing backend is a one-line change.

Backend (`--backend`)	Class	Install	Use when
`sklearn` (default)	`GaussianProcessPCA`	base install	small datasets, no JAX
`jax`	`GaussianProcessPCAJAX`	`uv sync --extra jax`	large N, GPU, gradients, HMC
`vae`	`GaussianProcessVAEJAX`	`uv sync --extra jax`	nonlinear output compression

from class_e.stats.jax_gp import GaussianProcessPCAJAX, JAXGPConfig

emulator = GaussianProcessPCAJAX(JAXGPConfig(n_components=12, steps=200))
emulator.fit(X_train, Y_train)   # X: (N, D) parameters, Y: (N, K) spectra
Y_pred = emulator.predict(X_test)

Persist emulators as bundles (class-e train --out model.cemu), not pickles.

Configuration (JAX backend)¶

JAXGPConfig(
    n_components=12,           # PCA components
    kernel="rbf",              # 'rbf', 'matern32', 'matern52', ...
    ard=True,                  # one length scale per input
    standardize_X=True,
    standardization_mode="prior",
    log_space=False,           # emulate log(Y) (positive observables)
    optimize_hyperparams=True,
    steps=200,
)

GPU training¶

The JAX backends train on an NVIDIA GPU with the jax-cuda extra (Linux, CUDA 12):

uv sync --extra jax-cuda                  # or: pip install "class-e[jax-cuda]"
XLA_PYTHON_CLIENT_PREALLOCATE=false class-e train cfg.yaml

XLA_PYTHON_CLIENT_PREALLOCATE=false makes JAX allocate GPU memory as needed instead of 75% of it at start-up, so the memory in use is the memory the fit needs. Training holds about six N × N float64 matrices: about 13 GB at N = 16384, whatever the number of inputs.

On the CPU, jaxlib factorizes the kernel matrix with scipy's BLAS/LAPACK, so the BLAS thread count matters. Cluster profiles often set OPENBLAS_NUM_THREADS=1 (against an OpenBLAS crash under concurrent XLA calls while sampling), which serializes training; class-e train warns when it trains a JAX backend on more than 1000 points that way. Unset it, or set it to the cores of the job, for training. Training TT on N = 4095 points (9 inputs) takes:

Setup	Time
64 cores, `OPENBLAS_NUM_THREADS=1`	34 min
64 cores, unthrottled	18 min
1 × A100	82 s

The three agree to 1.5e-5 σ_CV.

Input standardization¶

Gaussian processes are sensitive to input scales, and a poor choice biases inference. Three modes are available (all backends):

Mode	Center and scale	Use when
`"training"` (default)	mean and standard deviation of the training inputs	training and inference regions coincide
`"prior"` (recommended)	from the prior bounds, see below	uniform designs (LHC, Sobol) and inference anywhere in the prior
`"manual"`	`standardization_center` / `standardization_scale`	inference around a known fiducial point

Prior mode maps each parameter with

\[ x_d = \frac{\theta_d - (\theta_{d,\min} + \theta_{d,\max})/2}{(\theta_{d,\max} - \theta_{d,\min})/6}, \]

so the prior box spans \(\pm 3\) in \(x\).

Training mode can bias results

With a uniform design but inference away from the training mean, the standardization offset is systematic and does not shrink with more training data.

config = JAXGPConfig(standardize_X=True, standardization_mode="prior")
emulator = GaussianProcessPCAJAX(config)
emulator.set_bounds(bounds)   # [(min, max), ...] per parameter; required for prior mode
emulator.fit(X_train, Y_train)

diag = emulator.diagnose_standardization(fiducial, param_names)
# relative_offset < 0.3 for every parameter: fine; > 0.5: bias risk

Gradients and sampling¶

The JAX and VAE backends expose differentiable predictions:

import jax, jax.numpy as jnp

y = emulator.predict_jax(jnp.asarray(theta))                    # (K,) or (n, K)
grad = jax.grad(lambda p: jnp.sum(emulator.predict_jax(p)))(jnp.asarray(theta))

GaussianProcessPCAJAX.log_likelihood_jax(params, data, sigma) returns a differentiable Gaussian log-likelihood, ready for HMC/NUTS samplers such as numpyro. The first call includes JIT compilation time.

Input transforms and segments¶

InputTransform trains on expressions of the parameters (for example s = w0 + wa) while the interface stays in the physical parameters.
SegmentedEmulator and class-e stitch join emulators covering different \(\ell\) ranges, blending over an overlap window.
DistanceEmulator gives \(H(z)\) and all distances from one emulator of the expansion history, and RatioEmulator the recombination scalars (\(r_d\), \(\theta_*\), ...): see Distances and expansion history.

See the Emulators API.