API Reference¶

Core Module¶

jaxcapse.jaxcapse ¶

IdentityInterpolation ¶

IdentityInterpolation(prediction_ell_grid)

Return emulator outputs unchanged on their training grid.

MLP ¶

MLP(emulator: FlaxEmulator, in_MinMax: ndarray, out_MinMax: ndarray, postprocessing: Callable, emulator_description: Dict[str, Any], training_ell_grid: ndarray, interpolation: str = 'auto', max_spline_knots: int = 2048, endpoint_tolerance: float = 0.1, InterpolationMethod: Optional[Callable] = None)

CAPSE MLP emulator using jaxace infrastructure.

This class wraps a jaxace FlaxEmulator with CAPSE-specific functionality for CMB power spectrum computation.

Initialize MLP with jaxace emulator and CAPSE-specific components.

Parameters:

Name	Type	Description	Default
`emulator`	`FlaxEmulator`	jaxace FlaxEmulator instance	required
`in_MinMax`	`ndarray`	Input normalization parameters	required
`out_MinMax`	`ndarray`	Output normalization parameters	required
`postprocessing`	`Callable`	Postprocessing function (must be JAX-compatible)	required
`emulator_description`	`Dict[str, Any]`	Emulator metadata	required
`training_ell_grid`	`ndarray`	Multipoles used while training the emulator	required
`interpolation`	`str`	Automatic, disabled, or forced cubic interpolation	`'auto'`
`max_spline_knots`	`int`	Largest source grid automatically interpolated	`2048`
`endpoint_tolerance`	`float`	Distance used to snap source bounds to integers	`0.1`
`InterpolationMethod`	`Optional[Callable]`	Optional preconstructed interpolation method	`None`

get_Cl ¶

get_Cl(input_data: ndarray) -> jnp.ndarray

Compute CMB power spectrum Cl values with JIT compilation.

Parameters:

Name	Type	Description	Default
`input_data`	`ndarray`	Cosmological parameters as JAX array	required

Returns:

Type	Description
`ndarray`	Processed Cl values

get_Cl_batch ¶

get_Cl_batch(input_batch: ndarray) -> np.ndarray

Compute CMB power spectrum Cl values for a batch of inputs using vectorization.

Parameters:

Name	Type	Description	Default
`input_batch`	`ndarray`	Array of cosmological parameters, shape (n_samples, n_params)	required

Returns:

Type	Description
`ndarray`	Array of processed Cl values, shape (n_samples, n_cls)

get_ell_grid ¶

get_ell_grid() -> jnp.ndarray

Return the multipole grid corresponding to get_Cl output.

get_training_ell_grid ¶

get_training_ell_grid() -> jnp.ndarray

Return the original emulator training multipoles.

predict ¶

predict(input_data: ndarray) -> jnp.ndarray

Alias for get_Cl to maintain compatibility with tests and examples.

Parameters:

Name	Type	Description	Default
`input_data`	`ndarray`	Cosmological parameters as JAX array	required

Returns:

Type	Description
`ndarray`	Processed Cl values

SplinePlan ¶

SplinePlan(training_ell_grid, plan_type=CubicSplinePlan, endpoint_tolerance=0.1)

Interpolate fixed training multipoles onto a dense integer grid.

load_emulator ¶

load_emulator(folder_path: str, interpolation: str = 'auto', max_spline_knots: int = 2048, endpoint_tolerance: float = 0.1) -> MLP

Load a CAPSE emulator using jaxace infrastructure.

Parameters:

Name	Type	Description	Default
`folder_path`	`str`	Path to the emulator folder containing: - nn_setup.json: Neural network specification - weights.npy: Trained weights - inminmax.npy: Input normalization parameters - outminmax.npy: Output normalization parameters - l.npy: Multipoles used during training - postprocessing.py: Postprocessing function	required
`interpolation`	`str`	`"auto"` (default), `"none"`, or `"cubic"`.	`'auto'`
`max_spline_knots`	`int`	Maximum source-grid size selected for automatic cubic interpolation. Defaults to 2048.	`2048`
`endpoint_tolerance`	`float`	Maximum distance used to snap source bounds to an integer. Defaults to 0.1; bounds farther away are moved inward.	`0.1`

Returns:

Type	Description
`MLP`	MLP instance ready for inference

load_preprocessing ¶

load_preprocessing(root_path: str, filename: str) -> callable

Load postprocessing function from Python file.

Parameters:

Name	Type	Description	Default
`root_path`	`str`	Directory containing the postprocessing file	required
`filename`	`str`	Name of the postprocessing file (without .py extension)	required

Returns:

Type	Description
`callable`	The postprocessing function

prepare_interpolation_method ¶

prepare_interpolation_method(training_ell_grid, interpolation='auto', max_spline_knots=2048, endpoint_tolerance=0.1)

Choose identity or cubic interpolation for an emulator training grid.