Emulators¶

Base class¶

class_e.core.emulator ¶

BaseEmulator ¶

Bases: ABC

Interface shared by all emulators (sklearn, JAX and VAE backends, wrappers).

fit `abstractmethod` ¶

fit(X: ndarray, Y: ndarray) -> FitResult

Fit the emulator.

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Input parameters, shape (N, D).	required
`Y`	`ndarray`	Outputs (e.g. spectra), shape (N, K).	required

Returns:

Type	Description
`FitResult`	Success flag and backend diagnostics.

predict `abstractmethod` ¶

predict(X: ndarray) -> np.ndarray

Predict the outputs at new parameters.

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Input parameters, shape (M, D).	required

Returns:

Type	Description
`ndarray`	Predicted outputs, shape (M, K).

save ¶

save(path: str) -> None

Pickle the emulator to path with joblib.

Prefer class_e.core.bundle.EmulatorBundle, which stores no pickles; loading a pickle executes code from the file.

load `staticmethod` ¶

load(path: str) -> 'BaseEmulator'

Load an emulator saved with save (only from a source you trust).

FitResult `dataclass` ¶

FitResult(ok: bool, details: dict)

Outcome of BaseEmulator.fit.

Attributes:

Name	Type	Description
`ok`	`bool`	Whether the fit succeeded.
`details`	`dict`	Backend-specific diagnostics (e.g. number of components, optimizer settings).

Backends¶

class_e.stats.gp ¶

GaussianProcessPCA ¶

GaussianProcessPCA(config: Config | None = None)

Bases: BaseEmulator

Gaussian Process emulator on PCA-compressed outputs.

With n_components=None no compression is applied: a single GP (shared kernel, per-output normalization) is fitted on the full output vector.

set_bounds ¶

set_bounds(bounds)

Set parameter bounds for prior-based standardization.

Parameters:

Name	Type	Description	Default
`bounds`	`list of tuples`	[(min, max), ...] for each parameter	required

Examples:

>>> emulator.set_bounds([(0.02, 0.025), (0.1, 0.2), (0.6, 0.8)])

set_baseline ¶

set_baseline(Y: ndarray) -> None

Set baseline Y_baseline for residual training (log-ratio), log(Y/Y_baseline).

fit ¶

fit(X: ndarray, Y: ndarray) -> FitResult

Fit PCA (if configured) and the GP to the training set.

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Parameters, shape (N, D).	required
`Y`	`ndarray`	Outputs, shape (N, K).	required

Returns:

Type	Description
`FitResult`	`ok` and the fit diagnostics (kernel, PCA details).

hyperparameter_notes ¶

hyperparameter_notes(
    names: Optional[Sequence[str]] = None,
    rtol: float = 0.01,
) -> list[str]

Short notes on the fitted kernel's hyperparameters that sit at a bound.

names label the input dimensions of ARD length scales. A length scale at its upper bound means that input barely affects the output; the noise at its lower bound is expected for noiseless data.

predict ¶

predict(X: ndarray) -> np.ndarray

Predict outputs at parameters X, shape (M, D); returns shape (M, K).

Raises:

Type	Description
`RuntimeError`	If the emulator is not fitted.

reconstruct ¶

reconstruct(Y: ndarray) -> np.ndarray

Y compressed and decompressed by the fitted output transform and PCA, no GP.

The difference to Y is the compression floor (zero without PCA).

Config ¶

Bases: BaseModel

Configuration for PCA-based GaussianProcess emulators with validation.

Ensures all parameters are physically valid and kernel names are recognized.

validate_kernel `classmethod` ¶

validate_kernel(v: str) -> str

Ensure kernel name is recognized.

class_e.stats.jax_gp ¶

JAX-based Gaussian Process emulator using tinygp for fast training and inference.

This module provides a JAX implementation of the PCA-based GP emulator that matches the sklearn implementation but with significantly better performance for large datasets.

GaussianProcessPCAJAX ¶

GaussianProcessPCAJAX(config: Optional[JAXGPConfig] = None)

JAX/tinygp implementation mirroring the sklearn GaussianProcessPCA API.

This emulator uses JAX and tinygp for fast Gaussian Process regression on PCA-compressed outputs. It provides the same interface as GaussianProcessPCA but with significant performance improvements, especially for large datasets.

Examples:

>>> from class_e.stats.jax_gp import GaussianProcessPCAJAX, JAXGPConfig
>>> import numpy as np
>>>
>>> # Create configuration
>>> config = JAXGPConfig(n_components=10, optimize_hyperparams=True, steps=100)
>>> emulator = GaussianProcessPCAJAX(config)
>>>
>>> # Generate training data
>>> X_train = np.random.randn(100, 5)
>>> Y_train = np.random.randn(100, 50)
>>>
>>> # Fit the emulator
>>> emulator.fit(X_train, Y_train)
>>>
>>> # Make predictions
>>> X_test = np.random.randn(20, 5)
>>> Y_pred = emulator.predict(X_test)

set_bounds ¶

set_bounds(bounds: list) -> None

Set parameter bounds for prior-based standardization.

Parameters:

Name	Type	Description	Default
`bounds`	`list of tuples`	[(min, max), ...] for each parameter	required

Examples:

>>> emulator.set_bounds([(0.02, 0.025), (0.1, 0.2), (0.6, 0.8)])

reconstruct ¶

reconstruct(Y: ndarray) -> np.ndarray

Y compressed and decompressed by the fitted output transform and PCA, no GP.

The difference to Y is the compression floor: the error an exact GP in the principal components would still make.

predict ¶

predict(X: ndarray) -> np.ndarray

Make predictions on new inputs.

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Input parameters, shape (N, D)	required

Returns:

Type	Description
`ndarray`	Predicted outputs, shape (N, K)

predict_jax ¶

predict_jax(X: Any) -> Any

JAX-native prediction (fully differentiable).

This method keeps everything as JAX arrays and is fully differentiable with respect to the inputs, enabling gradient-based inference (HMC, NUTS).

Parameters:

Name	Type	Description	Default
`X`	`ndarray`	Input parameters, shape (N, D) or (D,) for single point	required

Returns:

Type	Description
`ndarray`	Predicted outputs, shape (N, K) or (K,) for single point

Examples:

>>> # Single point prediction with gradients
>>> x = jnp.array([0.67, 0.022, 0.12, 2.1e-9, 0.96, 0.055])
>>> y_pred = emulator.predict_jax(x)
>>>
>>> # Compute gradient of prediction w.r.t. parameters
>>> grad_fn = jax.grad(lambda params: jnp.sum(emulator.predict_jax(params)))
>>> gradient = grad_fn(x)
>>>
>>> # Batch prediction
>>> X_batch = jnp.array([...])  # Shape: (N, D)
>>> Y_pred = emulator.predict_jax(X_batch)

log_likelihood_jax ¶

log_likelihood_jax(
    params: Any, data: Any, sigma: Optional[Any] = None
) -> Any

Compute log-likelihood for MCMC sampling (fully differentiable).

This function is designed for use with gradient-based MCMC samplers like HMC and NUTS. It computes the log-likelihood of observing data given cosmological params using the emulator predictions.

Parameters:

Name	Type	Description	Default
`params`	`ndarray`	Cosmological parameters, shape (D,)	required
`data`	`ndarray`	Observed data, shape (K,)	required
`sigma`	`ndarray`	Measurement uncertainties, shape (K,). If None, assumes sigma=1.0	`None`

Returns:

Type	Description
`float`	Log-likelihood value

Examples:

>>> # Define log-likelihood function
>>> data_obs = jnp.array([...])  # Observed CMB spectrum
>>> sigma_obs = jnp.array([...])  # Uncertainties
>>>
>>> def log_prob(params):
...     return emulator.log_likelihood_jax(params, data_obs, sigma_obs)
>>>
>>> # Use with JAX gradient
>>> grad_log_prob = jax.grad(log_prob)
>>> gradient = grad_log_prob(params_test)
>>>
>>> # Use with numpyro HMC
>>> import numpyro
>>> from numpyro.infer import MCMC, NUTS
>>>
>>> def model():
...     # Prior
...     h = numpyro.sample("h", numpyro.distributions.Uniform(0.6, 0.8))
...     omega_b = numpyro.sample("omega_b", numpyro.distributions.Uniform(0.02, 0.025))
...     # ... more parameters
...     params = jnp.array([h, omega_b, ...])
...
...     # Likelihood using emulator
...     numpyro.factor("obs", log_prob(params))
>>>
>>> nuts_kernel = NUTS(model)
>>> mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000)
>>> mcmc.run(jax.random.PRNGKey(0))

gradient_wrt_params ¶

gradient_wrt_params(
    params: Any, output_idx: Optional[int] = None
) -> Any

Compute gradient of predictions with respect to input parameters.

This is useful for understanding parameter sensitivities and for gradient-based optimization/sampling.

Parameters:

Name	Type	Description	Default
`params`	`ndarray`	Input parameters, shape (D,)	required
`output_idx`	`int`	If provided, compute gradient for specific output dimension. Otherwise, returns full Jacobian.	`None`

Returns:

Type	Description
`ndarray`	If output_idx is None: Jacobian matrix, shape (K, D) If output_idx is int: Gradient vector, shape (D,)

Examples:

>>> # Gradient of first output w.r.t. all parameters
>>> grad = emulator.gradient_wrt_params(params, output_idx=0)
>>>
>>> # Full Jacobian
>>> jacobian = emulator.gradient_wrt_params(params)
>>> # jacobian[i, j] = d(output_i) / d(param_j)

diagnose_standardization ¶

diagnose_standardization(
    X_test: ndarray, param_names: Optional[list] = None
) -> dict

Diagnose potential standardization bias.

This method helps identify when standardization may introduce systematic bias due to mismatch between training distribution and inference distribution.

Parameters:

Name	Type	Description	Default
`X_test`	`ndarray`	Test/inference points (e.g., MCMC samples or fiducial parameters), shape (N, D) or (D,)	required
`param_names`	`list of str`	Parameter names for more readable output	`None`

Returns:

Type	Description
`dict`	Dictionary with diagnostic information: - 'scaler_mean': Standardization center, shape (D,) - 'scaler_scale': Standardization width, shape (D,) - 'test_mean': Mean of test points, shape (D,) - 'offset': Offset between scaler and test (scaler_mean - test_mean), shape (D,) - 'relative_offset': Offset in units of scale, shape (D,) - 'param_names': Parameter names (if provided)

Examples:

>>> # Diagnose potential bias for MCMC samples
>>> mcmc_samples = np.array([...])  # Shape: (N, D)
>>> diagnostics = emulator.diagnose_standardization(mcmc_samples, param_names)
>>>
>>> # Check for significant offsets
>>> for i, name in enumerate(param_names):
...     if abs(diagnostics['relative_offset'][i]) > 0.5:
...         print(f"Warning: {name} has large offset: {diagnostics['relative_offset'][i]:.2f} sigma")
>>>
>>> # Diagnose at fiducial point
>>> fiducial = np.array([0.67, 0.022, 0.12, 2.1e-9, 0.965, 0.055])
>>> diagnostics = emulator.diagnose_standardization(fiducial, param_names)

Notes

Large relative offsets (|relative_offset| > 0.5) indicate potential bias. Consider using standardization_mode='prior' or 'manual' to reduce bias.

JAXGPConfig `dataclass` ¶

JAXGPConfig(
    n_components: int = 12,
    length_scale: float = 1.0,
    log_space: bool = False,
    log_eps: float = 1e-30,
    standardize_X: bool = True,
    standardization_mode: str = "training",
    standardization_center: Optional[ndarray] = None,
    standardization_scale: Optional[ndarray] = None,
    standardize_Y: bool = True,
    residual_mode: str = "none",
    kernel: str = "rbf",
    kernel_params: dict = None,
    amplitude: float = 1.0,
    noise: float = 1e-06,
    ard: bool = True,
    optimize_hyperparams: bool = True,
    optimizer: str = "auto",
    steps: int = 200,
    tol: float = 1e-05,
    learning_rate: float = 0.05,
    n_restarts: int = 3,
    random_state: int = 0,
    amplitude_bounds: Tuple[float, float] = (0.01, 100.0),
    length_scale_bounds: Tuple[float, float] = (
        0.01,
        100.0,
    ),
    noise_bounds: Tuple[float, float] = (1e-08, 0.1),
    use_derivatives: bool = False,
    derivative_noise: float = 0.0001,
    derivative_noise_bounds: Tuple[float, float] = (
        1e-10,
        1.0,
    ),
)

Configuration for JAX-based PCA GP emulator.

Attributes:

Name	Type	Description
`n_components`	`int`	Number of PCA components to use (default: 12)
`length_scale`	`float`	Initial kernel length scale (default: 1.0)
`log_space`	`bool`	Whether to train in log space (default: False)
`log_eps`	`float`	Small epsilon for log operations to avoid $\log 0$ (default: 1e-30)
`standardize_X`	`bool`	Whether to standardize input features (default: True)
`standardization_mode`	`str`	How to standardize inputs: 'training' (use training data statistics), 'prior' (use prior bounds), or 'manual' (user-provided) (default: 'training')
`standardization_center`	`Optional[ndarray]`	Manual standardization center (required if standardization_mode='manual')
`standardization_scale`	`Optional[ndarray]`	Manual standardization scale (optional, auto-computed from bounds if None)
`standardize_Y`	`bool`	Standardize each output feature (after the log/residual transform) before PCA, so components are not dominated by the largest-amplitude features (default: True)
`residual_mode`	`str`	Residual mode: 'none' or 'log_ratio' (default: 'none')
`kernel`	`str`	Kernel name: 'rbf', 'matern32', 'matern52', etc. (default: 'rbf')
`kernel_params`	`dict`	Additional kernel parameters (default: None)
`amplitude`	`float`	Initial kernel amplitude (default: 1.0)
`noise`	`float`	Initial observation noise level (default: 1e-6)
`ard`	`bool`	Automatic Relevance Determination: one length scale per input parameter (default: True). Cosmological parameters affect the outputs on very different scales, which a single isotropic length scale cannot capture.
`optimize_hyperparams`	`bool`	Whether to optimize hyperparameters by maximizing the marginal likelihood (default: True)
`optimizer`	`str`	'lbfgs' (L-BFGS with line search, run to convergence), 'adam' (fixed number of steps) or 'auto': L-BFGS in float64, Adam in float32, where the line search cannot resolve the loss accurately enough (default: 'auto')
`steps`	`int`	Maximum L-BFGS iterations, or number of Adam steps (default: 200)
`tol`	`float`	L-BFGS stops when the gradient norm of the loss falls below this (default: 1e-5)
`learning_rate`	`float`	Learning rate for the Adam optimizer (default: 0.05)
`n_restarts`	`int`	Extra optimizer runs from random starting points, all inside one compiled program; the best marginal likelihood wins (default: 3)
`random_state`	`int`	Seed for the random starting points (default: 0)
`amplitude_bounds`	`Tuple[float, float]`	(min, max) bounds for amplitude during optimization (default: (0.01, 100.0))
`length_scale_bounds`	`Tuple[float, float]`	(min, max) bounds for length_scale during optimization (default: (0.01, 100.0)); with ARD, inputs the outputs barely depend on need length scales well beyond the standardized prior range
`noise_bounds`	`Tuple[float, float]`	(min, max) bounds for noise during optimization (default: (1e-8, 0.1))

Notes

With standardization_mode='prior' (and bounds set with set_bounds) each input is centred on the middle of its prior and scaled so that the prior box spans $\pm 3$:

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

Unlike 'training', this does not depend on where the training points fall, which avoids a systematic bias when a uniform design is used to infer parameters away from its centre.

class_e.stats.vae_gp ¶

VAE+GP hybrid emulator for cosmological parameter estimation.

This module implements a Variational Autoencoder (VAE) combined with Gaussian Process regression for emulating high-dimensional cosmological observables. The VAE learns a nonlinear compression of the output space, while GPs model the mapping from input parameters to the learned latent space.

The VAE is trained on the standardized outputs by maximizing the $\beta$-weighted evidence lower bound (ELBO), with the latent sample drawn by the reparameterization trick:

\[ \mathcal{L}(Y) = \mathbb{E}_{q(z|Y)}\big[\log p(Y|z)\big] - \beta\, \mathrm{KL}\big(q(z|Y)\,\|\,\mathcal{N}(0, I)\big), \qquad z = \mu + \sigma \odot \epsilon,\quad \epsilon \sim \mathcal{N}(0, I). \]

$\beta = 1$ is the ELBO itself.

Classes:

Name	Description
- ``VAEGPConfig``: configuration of the VAE+GP emulator.
- ``VAEEncoder``: VAE encoder network
- ``VAEDecoder``: VAE decoder network
- ``VAE``: complete VAE model.
- ``GaussianProcessVAEJAX``: main emulator class combining VAE and GP.

Functions:

Name	Description
- ``train_vae``: train the VAE with $\beta$-annealing and early stopping.

GaussianProcessVAEJAX ¶

GaussianProcessVAEJAX(config: Optional[VAEGPConfig] = None)

Bases: BaseEmulator

VAE+GP hybrid emulator using JAX.

A variational autoencoder compresses the (transformed, standardized) outputs nonlinearly to n_latent dimensions; one Gaussian process with a kernel shared by all latents (as in the JAX PCA+GP backend) maps inputs to the latent means; the decoder maps them back. Prediction is a single compiled, differentiable function. diagnose separates the error of the compression from that of the latent regression.

Examples:

>>> config = VAEGPConfig(n_latent=10, vae_epochs=300, standardization_mode="prior")
>>> emulator = GaussianProcessVAEJAX(config)
>>> emulator.set_bounds(bounds)
>>> result = emulator.fit(X_train, Y_train)
>>> Y_pred = emulator.predict_numpy(X_test)
>>> emulator.diagnose(X_test, Y_test)

set_bounds ¶

set_bounds(bounds: list) -> None

Set parameter bounds for prior-based standardization.

set_baseline ¶

set_baseline(Y: ndarray) -> None

Set the baseline for residual_mode='log_ratio'.

set_random_key ¶

set_random_key(key) -> None

Set the JAX random key (VAE initialization, batching, validation split).

fit ¶

fit(X: ndarray, Y: ndarray) -> FitResult

Fit the VAE on the outputs, then the latent GP on the inputs.

Returns:

Type	Description
`FitResult`	`details`: `vae_recon_error` (best validation per-feature MSE of `decode(mu)`, standardized units), `vae_final_val_loss` (last epoch's), `vae_best_epoch`, `vae_epochs_trained`, `n_latent`, and the latent GP's `optimized_amplitude/length_scale/noise`.

predict ¶

predict(X)

Predict outputs (JAX array; differentiable in X, e.g. for HMC/NUTS).

Parameters:

Name	Type	Description	Default
`X`	`(array - like, shape(M, D_in))`		required

Returns:

Type	Description
`(Array, shape(M, D_out))`

predict_jax ¶

predict_jax(X)

Alias of predict (JAX-native), with (D,) inputs giving (D_out,).

predict_numpy ¶

predict_numpy(X: ndarray) -> np.ndarray

Predict outputs as a NumPy array.

diagnose ¶

diagnose(X: ndarray, Y: ndarray) -> dict

Split the emulator's error by stage on held-out data.

Errors are RMS over output features in standardized units (1 = the training scatter of each feature), as median and 95th percentile over samples:

reconstruction: decode(encode(y)) against y (compression);
latent: GP-predicted latents against encode(y), in units of the latents' training scale (regression);
total: prediction against y.

save ¶

save(path: str) -> None

Save to a single .npz file (arrays plus a JSON config; no pickle).

load `staticmethod` ¶

load(path: str) -> GaussianProcessVAEJAX

Load an emulator written by save.

VAEGPConfig `dataclass` ¶

VAEGPConfig(
    n_latent: int = 10,
    encoder_dims: Tuple[int, ...] = (1024, 512, 256),
    decoder_dims: Tuple[int, ...] = (256, 512, 1024),
    activation: str = "silu",
    use_layernorm: bool = True,
    vae_learning_rate: float = 0.001,
    vae_epochs: int = 500,
    vae_batch_size: int = 64,
    beta_schedule: str = "linear",
    beta_max: float = 1.0,
    beta_warmup_epochs: int = 100,
    early_stopping_patience: int = 50,
    validation_split: float = 0.1,
    recon_loss: str = "mse",
    huber_delta: float = 1.0,
    standardize_X: bool = True,
    standardization_mode: str = "prior",
    standardization_center: Optional[ndarray] = None,
    standardization_scale: Optional[ndarray] = None,
    kernel: str = "rbf",
    length_scale: float = 1.0,
    amplitude: float = 1.0,
    noise: float = 1e-06,
    kernel_params: Optional[dict] = None,
    ard: bool = True,
    amplitude_bounds: Tuple[float, float] = (0.01, 100.0),
    length_scale_bounds: Tuple[float, float] = (
        0.01,
        100.0,
    ),
    noise_bounds: Tuple[float, float] = (1e-08, 0.1),
    optimize_hyperparams: bool = True,
    optimizer: str = "auto",
    steps: int = 200,
    tol: float = 1e-05,
    learning_rate: float = 0.05,
    n_restarts: int = 3,
    random_state: int = 0,
    log_space: bool = False,
    log_eps: float = 1e-30,
    residual_mode: str = "none",
)

Configuration for VAE+GP hybrid emulator.

This configuration combines VAE architecture settings, VAE training parameters, GP settings, and output transformation options.

VAE Architecture

n_latent : int Number of latent dimensions (default: 10). Typically 2-5x smaller than PCA components for same accuracy. encoder_dims : Tuple[int, ...] Hidden layer sizes for encoder (default: (1024, 512, 256)). Adjust based on output dimensionality. decoder_dims : Tuple[int, ...] Hidden layer sizes for decoder (default: (256, 512, 1024)). activation : str Activation function: 'silu', 'gelu', 'relu', 'tanh' (default: 'silu'). use_layernorm : bool Apply LayerNorm after each hidden layer (default: True).

VAE Training

vae_learning_rate : float Initial learning rate for VAE (default: 1e-3). vae_epochs : int Number of training epochs (default: 500). vae_batch_size : int Batch size for VAE training (default: 64). beta_schedule : str KL weight schedule: 'constant', 'linear', 'cyclical' (default: 'linear'). beta_max : float Maximum $\beta$ weight for KL term (default: 1.0). beta_warmup_epochs : int Epochs before starting $\beta$ increase (default: 100). early_stopping_patience : int Stop if the validation reconstruction doesn't improve for N epochs (default: 50). validation_split : float Fraction of data for validation (default: 0.1). recon_loss : str Reconstruction term on the standardized outputs: 'mse' (Gaussian, summed over features, so $\beta = 1$ is the ELBO) or 'huber' (default: 'mse'). huber_delta : float Transition of the Huber loss, in standardized units (default: 1.0).

GP Configuration

standardize_X : bool Whether to standardize input features (default: True). standardization_mode : str 'training', 'prior', or 'manual' (default: 'prior'). standardization_center : Optional[np.ndarray] Manual standardization center (for mode='manual'). standardization_scale : Optional[np.ndarray] Manual standardization scale (for mode='manual'). kernel : str GP kernel: 'rbf', 'matern32', 'matern52' (default: 'rbf'). length_scale : float Initial kernel length scale (default: 1.0). amplitude : float Initial kernel amplitude (default: 1.0). noise : float Initial observation noise (default: 1e-6). kernel_params : dict Extra kernel parameters (default: none). ard : bool Per-dimension length scales (default: True). amplitude_bounds, length_scale_bounds, noise_bounds : Tuple[float, float] Bounds for the hyperparameter optimization (as in JAXGPConfig). optimize_hyperparams : bool Optimize the latent GP's hyperparameters (one kernel shared by all latents, by maximizing the summed marginal likelihood; default: True). optimizer, steps, tol, learning_rate, n_restarts, random_state Hyperparameter optimizer settings (as in JAXGPConfig; default: L-BFGS in float64 / Adam in float32, at most 200 iterations, 3 extra random starts).

Output Transforms

log_space : bool Train VAE in log-space (default: False). log_eps : float Epsilon for log transform (default: 1e-30). residual_mode : str 'none' or 'log_ratio' (default: 'none').

Input transforms and segmentation¶

class_e.stats.inputs.InputTransform ¶

InputTransform(
    emulator: Any,
    param_names: Sequence[str],
    inputs: Mapping[str, str],
)

Bases: BaseEmulator

An emulator trained on expressions of the physical parameters.

Parameters:

Name	Type	Description	Default
`emulator`	`BaseEmulator`	The wrapped emulator (any backend); it sees the inputs, in the order of `inputs`.	required
`param_names`	`sequence of str`	Physical parameter names, in column order (the interface of this emulator).	required
`inputs`	`mapping of str to str`	The wrapped emulator's inputs, `{name: expression}`; parameters not listed are not seen (list them as `{"h": "h"}` to keep them).	required

input_names `property` ¶

input_names: list[str]

Names of the wrapped emulator's inputs.

info `property` ¶

info: Any

The wrapped emulator's configuration.

transform ¶

transform(X: Any, xp: Any = np) -> Any

The wrapped emulator's inputs at physical inputs X, shape (M, D) -> (M, len(inputs)).

set_bounds ¶

set_bounds(bounds: Sequence[Sequence[float]]) -> None

Physical prior box; the wrapped emulator gets the box of the inputs.

fit ¶

fit(
    X: ndarray, Y: ndarray, *args: Any, **kwargs: Any
) -> FitResult

Fit the wrapped emulator on the transformed inputs; extra arguments are passed on.

X holds the physical parameters, shape (N, D).

predict_jax ¶

predict_jax(X: Any) -> Any

Differentiable prediction; (D,) gives (K,), (M, D) gives (M, K).

predict ¶

predict(X: ndarray) -> np.ndarray

Predict at physical parameters X, shape (M, D) or (D,).

reconstruct ¶

reconstruct(Y: ndarray) -> np.ndarray

The wrapped emulator's compression floor (see its reconstruct).

class_e.stats.segmented ¶

One emulator per multipole range, stitched into a single spectrum emulator.

Different parts of a CMB spectrum respond to different physics: the reionization bump of EE at low ell is set by tau, the damping tail at high ell by A_s e^{-2 tau}, Silk damping and lensing, and very high ell needs CLASS settings (non-linear lensing, large l_max) that are expensive to run everywhere. A SegmentedEmulator combines emulators trained on separate ell ranges -- possibly on different datasets, backends or input parametrizations -- into one emulator on the full grid.

Adjacent segments overlap. Inside each blend window the prediction moves from the lower segment to the upper one with a cosine taper,

P(ell) = (1 - s) P_lower(ell) + s P_upper(ell),   s = (1 - cos(pi t)) / 2,

with t going from 0 to 1 across the window, so the result is smooth (no jump at the seam) and equals each child exactly outside the windows.

A segment may see transformed inputs, e.g. {"logA_eff": "logA - 2 * tau_reio", ...} for the high-ell part, where tau only enters through A_s e^{-2 tau}. Transforms are arithmetic expressions of the physical parameters (see class_e.utils.expressions), evaluated in JAX so the composition stays differentiable.

SegmentedEmulator ¶

SegmentedEmulator(
    param_names: Sequence[str],
    segments: Sequence[Segment],
    overlap: int = 100,
)

Bases: BaseEmulator

Emulators on overlapping ell ranges, blended into one spectrum.

Parameters:

Name	Type	Description	Default
`param_names`	`sequence of str`	Physical input names, in column order (shared by all segments).	required
`segments`	`sequence of Segment`	Ordered by increasing ell; each must overlap the next one.	required
`overlap`	`int`	Width (in multipoles) of each blend window: the window between segments i and i+1 is the top `overlap` multipoles of segment i, which segment i+1 must cover.	`100`

segment_inputs ¶

segment_inputs(X: Any, i: int, xp: Any = np) -> Any

Input columns of segment i for physical inputs X of shape (M, D).

segment_bounds ¶

segment_bounds(
    bounds: Sequence[Sequence[float]], i: int
) -> list[tuple[float, float]]

Box of segment i's inputs: the range of each transform over the corners of bounds.

Exact for transforms that are monotonic in each parameter (e.g. linear ones).

fit ¶

fit(X: ndarray, Y: ndarray) -> FitResult

Fit every child on its own ell range of Y (given on self.ell).

Configure the children first (e.g. set_bounds(self.segment_bounds(bounds, i))). Children trained elsewhere can be combined without calling this.

predict_jax ¶

predict_jax(X: Any) -> Any

Differentiable prediction; (D,) gives (L,), (M, D) gives (M, L).

predict ¶

predict(X: ndarray) -> np.ndarray

Predict the stitched spectrum at parameters X, shape (M, D); returns (M, L) on self.ell.

reconstruct ¶

reconstruct(Y: ndarray) -> np.ndarray

Compression floor: each child's reconstruct on its ell range, blended as predictions.

Y is given on self.ell; every child must implement reconstruct and lie within self.ell (true for stitched bundles whose children cover only their own ranges).

stitch ¶

stitch(
    bundles: Sequence[Any],
    overlap: int = 100,
    meta: Optional[Mapping[str, Any]] = None,
) -> Any

Combine bundles covering different ell ranges into one bundle.

Each spectrum is stitched from the bundles that contain it, ordered by their largest multipole (e.g. a 2-40 EE bundle below a 2-2508 TT/TE/EE one, or a 2-2508 bundle below a 2000-5000 one). Scalar observables (e.g. h) are taken from the first bundle that has them.

Parameters:

Name	Type	Description	Default
`bundles`	`sequence of EmulatorBundle or paths`	Same parameter names, in the same order. The stitched prior box is the intersection of their boxes.	required
`overlap`	`int`	Width of each blend window (the top `overlap` multipoles of the lower bundle).	`100`
`meta`	`mapping`	Metadata of the result; the sources (name, file checksum, ell range, metrics) and the union of the sampling constraints are added.	`None`

Distances and expansion history¶

class_e.stats.distances ¶

Expansion history and distances from one emulator: $H(z)$ and $\chi(z) = c\int_0^z dz'/H(z')$.

The emulator learns the integrand of the comoving distance in $x = \ln(1+z)$,

\[ g(x) = \frac{d\chi}{dx} = \frac{c\,(1+z)}{H(z)} = r(x;\theta)\, b(x;\theta), \]

as the ratio $r$ to a background $b$, on the nodes of a background grid (the Hz observable, grid/z_bg). Between the nodes $r$ is the cubic spline through them, so

\[ H(z) = \frac{c\,(1+z)}{r(x)\,b(x)}, \qquad \chi(z) = \int_0^{x} r\,b\,dx' \]

follow from the same prediction at any $z$ of the grid, and $\chi' = c/H$ holds exactly. Spline values are fixed matrices applied to the emulated $r$ (spline_operators) and $\chi$ is a composite Gauss-Legendre sum, so both are linear in the emulator's output, as the integral of a GP is a GP.

The background $b$ is either

the training mean of $g$ at each node (default; splined between them), or
CPLBackground, the analytic expansion of CPL dark energy, matter, radiation and curvature. The emulator then learns only what it leaves out (massive neutrinos, details of the radiation era, early dark energy, ...). For w0waCDM with a free neutrino mass (2048 training points, 6 parameters, neutrinos through m_ncdm) the error on $\chi$ drops from $7\times10^{-5}$ (median) and $9\times10^{-3}$ (worst) with the training mean to $5\times10^{-9}$ and $6\times10^{-7}$, on $H(z)$ to $3\times10^{-6}$ at worst at any redshift, and the quadrature floor from $8\times10^{-5}$ to $10^{-9}$.

emu = DistanceEmulator(GaussianProcessPCAJAX(config), z=ds.grids["z_bg"], ... baseline=CPLBackground(ds.names, m_ncdm="m_nu", w0="w0", wa="wa", N_ur=2.0308)) emu.set_bounds(ds.bounds) emu.fit(ds.theta, ds.observables["Hz"]) # H(z) in km/s/Mpc on the nodes emu.comoving_distance(theta, z=[0.3, 0.5, 1100.0]) # Mpc emu.hubble(theta, z=[0.3, 0.5]) # km/s/Mpc emu.distances(theta, z=[0.3, 0.5])["D_V"] # D_M, D_A, D_L, D_H, D_V too

distances turns $\chi$ into the transverse distance $D_M$ with the curvature of the background (transverse_comoving_distance).

The scalars of the recombination era (rs_drag, z_drag, rs_star, z_star, theta_star_100) are emulated the same way, as a ratio to an analytic value (RecombinationBaseline: sound horizons on the CPL background, fitting formulas for the redshifts) with RatioEmulator. Near $w_0 + w_a = 0$ early dark energy shrinks $r_s$ by up to 10%; the baseline follows it, and the ratio varies by < 1% across a w0waCDM box.

rs = RatioEmulator(GaussianProcessPCAJAX(JAXGPConfig(n_components=1)), ... RecombinationBaseline(background, "rs_drag")) rs.fit(ds.theta, ds.observables["rs_drag"])

DistanceEmulator ¶

DistanceEmulator(
    emulator: Any,
    z: Sequence[float],
    baseline: Optional[CPLBackground] = None,
    aux: Optional[Mapping[str, Any]] = None,
)

Bases: BaseEmulator

$H(z)$ and $\chi(z)$ from one emulator of $d\chi/d\ln(1+z)$ (see the module docstring).

Parameters:

Name	Type	Description	Default
`emulator`	`BaseEmulator`	The wrapped emulator (any backend). It learns the ratio of $c(1+z)/H(z)$ to the background at each node: values of order one at every redshift.	required
`z`	`sequence of float`	The nodes, the grid of the `Hz` observable: from 0, increasing.	required
`baseline`	`CPLBackground`	Analytic background the ratio is taken to; default the training mean.	`None`
`aux`	`mapping of str to emulator`	Fitted scalar emulators (one output each) whose predictions the background reads as extra parameters, after the emulator's own: e.g. `{"H0": emu_H0}` with a `CPLBackground` over `param_names + ["H0"]` and `h="H0 / 100"`, for a model sampled in the acoustic angle rather than $h$. The background then uses the emulated value both in training and in prediction, so the ratio stays consistent.	`None`

info `property` ¶

info: Any

The wrapped emulator's configuration.

z_max `property` ¶

z_max: float

Largest redshift of the grid.

set_bounds ¶

set_bounds(bounds: Sequence[Sequence[float]]) -> None

Prior box, passed to the wrapped emulator.

integrand ¶

integrand(H: ndarray) -> np.ndarray

$c\,(1+z)/H(z)$ in Mpc at the nodes, for H in km/s/Mpc of shape (N, Z).

with_aux ¶

with_aux(X: Any, xp: Any = np) -> Any

The parameters the background reads: X (M, D) and the aux predictions after it.

ratio ¶

ratio(X: ndarray, H: ndarray) -> np.ndarray

What the wrapped emulator learns: the integrand over the background at the nodes.

fit ¶

fit(
    X: ndarray, H: ndarray, *args: Any, **kwargs: Any
) -> FitResult

Fit on H(z) in km/s/Mpc at the nodes, shape (N, Z); extra arguments are passed on.

operators ¶

operators(
    z: Optional[Sequence[float]] = None,
) -> _Operators

Spline and quadrature operators from the nodes to z (default: the nodes), cached per grid.

predict ¶

predict(X: ndarray) -> np.ndarray

H(z) in km/s/Mpc at the nodes, shape (M, Z): what fit was given.

hubble ¶

hubble(
    X: ndarray, z: Optional[Sequence[float]] = None
) -> np.ndarray

H(z) in km/s/Mpc, shape (M, len(z)); z defaults to the nodes.

comoving_distance ¶

comoving_distance(
    X: ndarray, z: Optional[Sequence[float]] = None
) -> np.ndarray

Line-of-sight comoving distance in Mpc, shape (M, len(z)); z defaults to the nodes.

predict_jax ¶

predict_jax(X: Any, z: Optional[Any] = None) -> Any

Differentiable H(z) in km/s/Mpc; (D,) gives (len(z),), (M, D) gives (M, len(z)).

z may be a traced array (redshifts known only inside jax.jit).

comoving_distance_jax ¶

comoving_distance_jax(
    X: Any, z: Optional[Any] = None
) -> Any

Differentiable comoving distance in Mpc; shapes as predict_jax.

distances ¶

distances(
    X: ndarray,
    z: Optional[Sequence[float]] = None,
    Omega_k: Optional[Any] = None,
) -> dict

$H$ (km/s/Mpc) and the distances (Mpc) at z (default: the nodes), each (M, len(z)).

Keys: H, chi (line of sight), D_M (transverse), D_A = D_M/(1+z), D_L = (1+z) D_M, D_H = c/H and D_V = (z D_M^2 D_H)^{1/3}. Omega_k (scalar or (M,)) defaults to the CPLBackground's, or flat with the training mean; $H_0$ is the emulated $H(0)$.

distances_jax ¶

distances_jax(
    X: Any,
    z: Optional[Any] = None,
    Omega_k: Optional[Any] = None,
) -> dict

Differentiable distances; (D,) gives (len(z),) arrays, (M, D) gives (M, len(z)).

z may be a traced array: the spline and quadrature are evaluated in JAX, with the same nodes and Gauss-Legendre rule as distances (identical results).

predict_ratio_jax ¶

predict_ratio_jax(X: Any) -> tuple[Any, Any]

The emulated ratio at the nodes, for distances_from_ratio: (X_background, r).

Predict once per parameter point, then ask for distances at any redshifts: the two steps of distances_jax, for a sampler whose likelihoods query several sets of redshifts. X is (M, D); X_background adds the aux predictions the background reads.

distances_from_ratio ¶

distances_from_ratio(
    X_background: Any,
    r: Any,
    z: Any,
    Omega_k: Optional[Any] = None,
) -> dict

distances (JAX, (M, len(z)) arrays) from predict_ratio_jax; z may be traced.

reconstruct ¶

reconstruct(X: ndarray, H: ndarray) -> np.ndarray

H(z) at the nodes through the wrapped emulator's compression only (its floor).

CPLBackground ¶

CPLBackground(
    param_names: Sequence[str],
    h: str = "h",
    omega_m: str = "omega_b + omega_cdm",
    w0: str = "-1",
    wa: str = "0",
    omega_r: Optional[str] = None,
    N_ur: float = 3.044,
    T_cmb: float | str = 2.7255,
    Omega_k: str = "0",
    m_ncdm: str = "0",
    N_ncdm: Optional[int | str] = None,
    T_ncdm: float | str = T_NCDM,
)

Analytic $c(1+z)/H(z)$ in Mpc: CPL dark energy, matter, radiation, massive neutrinos, curvature.

\[ \frac{H^2}{H_0^2} = \Omega_m a^{-3} + \Omega_r a^{-4} + \Omega_\nu(a) + \Omega_k a^{-2} + \Omega_{\rm DE}\, a^{-3(1+w_0+w_a)} e^{-3 w_a (1-a)} \]

with $\Omega_i = \omega_i/h^2$, $\Omega_{\rm DE}$ closing the budget today and, for one massive species of mass m_ncdm, $\Omega_\nu(a) = \Omega_\nu^{m=0} a^{-4} f(m a/k_B T_\nu)$ (neutrino_density_ratio): radiation early, matter late, as in CLASS. Each argument is an expression of the parameters (see class_e.utils.expressions) or a number. Without m_ncdm massive neutrinos can be counted as matter (m_nu / 93.14 in omega_m), which misses their radiation era: H(z) off by ~1% above z ~ 100.

Parameters:

Name	Type	Description	Default
`param_names`	`sequence of str`	Parameter names, in the column order of the emulator's inputs.	required
`h`	`str`	Expressions of $h$, $\omega_m = \Omega_m h^2$ (baryons and CDM; neutrinos with `m_ncdm` apart), $w_0$ and $w_a$.	`'h'`
`omega_m`	`str`	Expressions of $h$, $\omega_m = \Omega_m h^2$ (baryons and CDM; neutrinos with `m_ncdm` apart), $w_0$ and $w_a$.	`'h'`
`w0`	`str`	Expressions of $h$, $\omega_m = \Omega_m h^2$ (baryons and CDM; neutrinos with `m_ncdm` apart), $w_0$ and $w_a$.	`'h'`
`wa`	`str`	Expressions of $h$, $\omega_m = \Omega_m h^2$ (baryons and CDM; neutrinos with `m_ncdm` apart), $w_0$ and $w_a$.	`'h'`
`m_ncdm`	`str`	Expression of the mass in eV of each massive neutrino species (CLASS's `m_ncdm`).	`'0'`
`N_ncdm`	`int`	Number of degenerate massive species: default 1 with `m_ncdm`, else 0. A species of zero mass stays as radiation, as in CLASS (no jump at `m_nu = 0`).	`None`
`Omega_k`	`str`	Expression of the curvature density $\Omega_k$ (CLASS's `Omega_k`; default flat).	`'0'`
`omega_r`	`str`	Expression of $\omega_r$; default photons at `T_cmb` plus `N_ur` massless species.	`None`
`N_ur`	`float`	Massless neutrino species (3.044 without massive ones; CLASS's 2.0308 with one).	`3.044`
`T_cmb`	`float`	CMB temperature in K.	`2.7255`
`T_ncdm`	`float`	Temperature of the massive species over `T_cmb` (CLASS's default 0.71611).	`T_NCDM`

omega_nu ¶

omega_nu(X: Any, xp: Any = np) -> Any

$\omega_\nu = \Omega_\nu h^2$ of the massive species today, shape (M,): $\approx m_\nu/93.14$ eV.

curvature ¶

curvature(X: Any, xp: Any = np) -> Any

$\Omega_k$ at parameters X (M, D): shape (M,).

RatioEmulator ¶

RatioEmulator(
    emulator: Any,
    baseline: RecombinationBaseline,
    aux: Optional[Mapping[str, Any]] = None,
)

Bases: BaseEmulator

A scalar observable as an emulated ratio to its analytic RecombinationBaseline.

The wrapped emulator learns $y/b(\theta)$, values of order one varying by < 1%, and predicts $y = b\,(y/b)$.

Parameters:

Name	Type	Description	Default
`emulator`	`BaseEmulator`	The wrapped emulator (any backend), one output.	required
`baseline`	`RecombinationBaseline`	The analytic value.	required
`aux`	`mapping of str to emulator`	Fitted scalar emulators the baseline reads as parameters after the emulator's own (see `DistanceEmulator`), e.g. `H0` for a model sampled in the acoustic angle.	`None`

info `property` ¶

info: Any

The wrapped emulator's configuration.

set_bounds ¶

set_bounds(bounds: Sequence[Sequence[float]]) -> None

Prior box, passed to the wrapped emulator.

base ¶

base(X: Any, xp: Any = np) -> Any

The baseline at X (M, D), as a column (M, 1).

fit ¶

fit(
    X: ndarray, Y: ndarray, *args: Any, **kwargs: Any
) -> FitResult

Fit on the observable Y, shape (N, 1); extra arguments are passed on.

predict ¶

predict(X: ndarray) -> np.ndarray

The observable at X (M, D): shape (M, 1).

predict_jax ¶

predict_jax(X: Any) -> Any

Differentiable prediction; (D,) gives (1,), (M, D) gives (M, 1).

reconstruct ¶

reconstruct(X: ndarray, Y: ndarray) -> np.ndarray

Y through the wrapped emulator's compression only (its floor).

RecombinationBaseline ¶

RecombinationBaseline(
    background: CPLBackground,
    quantity: str,
    omega_b: str = "omega_b",
    T_cmb: float = 2.7255,
    theta_s: Optional[str] = None,
)

Analytic value of a scalar of the recombination era, from a CPLBackground.

z_drag, z_star: the fitting formulas z_drag_fit, z_star_fit;
rs_drag, rs_star: the comoving sound horizon $r_s = \int_{z}^{\infty} c_s\,dz'/H$ at those redshifts, with $c_s = c/\sqrt{3(1+R)}$ and $R = 3\omega_b/(4\omega_\gamma(1+z))$, on the background's $H(z)$;
theta_star_100: $100\,r_s(z_*)/D_M(z_*)$;
H0: $100\,h$; with theta_s (a model sampled in the acoustic angle), the $h$ at which the analytic $100\,\theta$ equals it, so the dependence of $H_0$ on $w_0, w_a$ along the geometric degeneracy is analytic (CLASS's $H_0$ over it: 0.987 $\pm$ 0.4% across a w0waCDM box where $H_0$ spans 47-140).

It carries the dependence an emulator finds hard (e.g. early dark energy as $w_0 + w_a \to 0$ shrinks $r_s$ by 10%): CLASS's value over it varies by < 1% across a w0waCDM box, the smooth ratio RatioEmulator learns.

Parameters:

Name	Type	Description	Default
`background`	`CPLBackground`	Its parameters, $h$, $\omega_m$ and curvature are used too.	required
`quantity`	`str`	One of `z_drag`, `z_star`, `rs_drag`, `rs_star`, `theta_star_100`, `H0`.	required
`omega_b`	`str`	Expression of $\omega_b$.	`'omega_b'`
`T_cmb`	`float`	CMB temperature in K (photon density of the baryon loading $R$).	`2.7255`
`theta_s`	`str`	For `H0` only: expression of $100\,\theta_s$ to solve $h$ from (the background's `h` is then unused). The solution is differentiable (implicit function theorem).	`None`

param_names `property` ¶

param_names: list[str]

The parameters it reads (the background's).

transverse_comoving_distance ¶

transverse_comoving_distance(
    chi: Any, H0: Any, Omega_k: Any, xp: Any = np
) -> Any

$D_M$ from the line-of-sight distance $\chi$ (Mpc), $H_0$ (km/s/Mpc) and $\Omega_k$.

\[ D_M = \chi\, S\!\left(\Omega_k (H_0\chi/c)^2\right), \qquad S(u) = \frac{\sinh\sqrt{u}}{\sqrt{u}}\ (u > 0), \quad \frac{\sin\sqrt{-u}}{\sqrt{-u}}\ (u < 0), \]

with the series $1 + u/6 + u^2/120$ near flatness, so it is exact, smooth (and differentiable with jax.numpy as xp) across $\Omega_k = 0$. chi is (M, Z); H0 and Omega_k are scalars or (M,).

neutrino_density_ratio ¶

neutrino_density_ratio(y: Any) -> Any

$\rho_\nu/\rho_\nu^{m=0}$ of one massive species with $y = m_\nu a/(k_B T_\nu)$.

Komatsu et al. (2011, eq. 26): $f(y) = [1 + (A y)^p]^{1/p}$, $A = 180\zeta(3)/(7\pi^4)$, $p = 1.83$; exact in both limits ($f \to 1$ relativistic, $f \to A y$ as matter), 0.35% in between. Smooth and differentiable (numpy or jax arrays).

spline_operators ¶

spline_operators(
    z_nodes: Sequence[float], z: Sequence[float]
) -> tuple[np.ndarray, np.ndarray]

Matrices of the cubic spline in $x = \ln(1+z)$ through nodal values, and of its integral.

For values $g$ at z_nodes (the first one $z = 0$), A @ g is the not-a-knot cubic spline through them at z, and B @ g its integral $\int_0^{x(z)} g\,dx$. Both are of shape (len(z), len(z_nodes)).

z_drag_fit ¶

z_drag_fit(omega_b: Any, omega_m: Any) -> Any

Redshift of the baryon drag epoch, Eisenstein & Hu (1998, eq. 4); ~1% from CLASS's.

z_star_fit ¶

z_star_fit(omega_b: Any, omega_m: Any) -> Any

Redshift of photon decoupling, Hu & Sugiyama (1996, eq. E1); ~0.1% from CLASS's.

Name	Type	Description	Default
`background`	`CPLBackground`	Its parameters, \(h\), \(\omega_m\) and curvature are used too.	required
`quantity`	`str`	One of `z_drag`, `z_star`, `rs_drag`, `rs_star`, `theta_star_100`, `H0`.	required
`omega_b`	`str`	Expression of \(\omega_b\).	`'omega_b'`
`T_cmb`	`float`	CMB temperature in K (photon density of the baryon loading \(R\)).	`2.7255`
`theta_s`	`str`	For `H0` only: expression of \(100\,\theta_s\) to solve \(h\) from (the background's `h` is then unused). The solution is differentiable (implicit function theorem).	`None`

Emulators¶

Base class¶

class_e.core.emulator ¶

BaseEmulator ¶

fit abstractmethod ¶

predict abstractmethod ¶

save ¶

load staticmethod ¶

FitResult dataclass ¶

Backends¶

class_e.stats.gp ¶

GaussianProcessPCA ¶

set_bounds ¶

set_baseline ¶

fit ¶

hyperparameter_notes ¶

predict ¶

reconstruct ¶

Config ¶

validate_kernel classmethod ¶

class_e.stats.jax_gp ¶

GaussianProcessPCAJAX ¶

set_bounds ¶

reconstruct ¶

predict ¶

predict_jax ¶

log_likelihood_jax ¶

gradient_wrt_params ¶

diagnose_standardization ¶

JAXGPConfig dataclass ¶

class_e.stats.vae_gp ¶

GaussianProcessVAEJAX ¶

set_bounds ¶

set_baseline ¶

set_random_key ¶

fit ¶

predict ¶

predict_jax ¶

predict_numpy ¶

diagnose ¶

save ¶

load staticmethod ¶

VAEGPConfig dataclass ¶

Input transforms and segmentation¶

class_e.stats.inputs.InputTransform ¶

input_names property ¶

info property ¶

transform ¶

set_bounds ¶

fit ¶

predict_jax ¶

predict ¶

reconstruct ¶

class_e.stats.segmented ¶

SegmentedEmulator ¶

segment_inputs ¶

segment_bounds ¶

fit ¶

predict_jax ¶

predict ¶

reconstruct ¶

stitch ¶

Distances and expansion history¶

class_e.stats.distances ¶

DistanceEmulator ¶

info property ¶

z_max property ¶

set_bounds ¶

integrand ¶

with_aux ¶

ratio ¶

fit ¶

operators ¶

predict ¶

hubble ¶

comoving_distance ¶

predict_jax ¶

comoving_distance_jax ¶

distances ¶

distances_jax ¶

fit `abstractmethod` ¶

predict `abstractmethod` ¶

load `staticmethod` ¶

FitResult `dataclass` ¶

validate_kernel `classmethod` ¶

JAXGPConfig `dataclass` ¶

load `staticmethod` ¶

VAEGPConfig `dataclass` ¶

input_names `property` ¶

info `property` ¶

info `property` ¶

z_max `property` ¶

info `property` ¶

param_names `property` ¶