API Reference¶

Cosmology¶

jaxace.w0waCDMCosmology `dataclass` ¶

w0waCDMCosmology(ln10As: float, ns: float, h: float, omega_b: float, omega_c: float, omega_k: float = 0.0, m_nu: float = 0.0, w0: float = -1.0, wa: float = 0.0)

D_f_z ¶

D_f_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Linear growth factor and growth rate (D(z), f(z)).

D_z ¶

D_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Linear growth factor D(z).

E_a ¶

E_a(a: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Dimensionless Hubble parameter E(a) = H(a)/H0.

E_z ¶

E_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Dimensionless Hubble parameter E(z) = H(z)/H0.

dA_z ¶

dA_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Angular diameter distance in Mpc.

dL_z ¶

dL_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Luminosity distance at redshift z in Mpc.

dM_z ¶

dM_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Transverse comoving distance in Mpc (affected by curvature).

d̃A_z ¶

d̃A_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Dimensionless angular diameter distance d̃A(z).

d̃M_z ¶

d̃M_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Dimensionless transverse comoving distance d̃M(z).

f_z ¶

f_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Growth rate f(z) = d log D / d log a.

r_z ¶

r_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Line-of-sight comoving distance in Mpc.

r̃_z ¶

r̃_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Dimensionless comoving distance r̃(z).

Ωm_a ¶

Ωm_a(a: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Matter density parameter Ωₘ(a) at scale factor a.

Ωtot_z ¶

Ωtot_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Total density parameter at redshift z (always 1.0 for flat universe).

ρc_z ¶

ρc_z(z: Union[float, ndarray]) -> Union[float, jnp.ndarray]

Critical density at redshift z in M☉/Mpc³.

Background Functions¶

Hubble Functions¶

jaxace.E_z ¶

E_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Dimensionless Hubble parameter E(z) = H(z)/H0.

This is equivalent to E(a) with the transformation a = 1/(1+z).

Returns:

Type	Description
`Union[float, ndarray]`	Hubble parameter E(z). Handles NaN/Inf inputs by propagating them appropriately.

jaxace.E_a ¶

E_a(a: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Dimensionless Hubble parameter E(a) = H(a)/H0.

The normalized Hubble parameter is given by:

\[E(a) = \sqrt{\Omega_{\gamma,0} a^{-4} + \Omega_{\mathrm{cb},0} a^{-3} + \Omega_{\Lambda,0} \rho_{\mathrm{DE}}(a) + \Omega_{\nu}(a) + \Omega_{k,0} a^{-2}}\]

where:

\(\Omega_{\gamma,0}\) is the photon density parameter today
\(\Omega_{\mathrm{cb},0}\) is the cold dark matter + baryon density parameter today
\(\Omega_{\Lambda,0}\) is the dark energy density parameter today (from flatness constraint)
\(\rho_{\mathrm{DE}}(a)\) is the normalized dark energy density
\(\Omega_{\nu}(a)\) is the massive neutrino contribution
\(\Omega_{k,0}\) is the curvature density parameter today

Returns:

Type	Description
`Union[float, ndarray]`	Hubble parameter E(a). Handles NaN/Inf inputs by propagating them appropriately.
`Union[float, ndarray]`	Returns NaN for invalid parameter combinations.

jaxace.dlogEdloga ¶

dlogEdloga(a: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Logarithmic derivative of the Hubble parameter.

\[\frac{\mathrm{d} \ln E}{\mathrm{d} \ln a} = \frac{a}{E} \frac{\mathrm{d}E}{\mathrm{d}a}\]

This quantity appears in the growth factor differential equation.

Returns:

Type	Description
`Union[float, ndarray]`	Logarithmic derivative d(ln E)/d(ln a).

Matter Density¶

jaxace.Ωm_a ¶

Ωm_a(a: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Matter density parameter Ωₘ(a) at scale factor a.

\[\Omega_{\mathrm{m}}(a) = \frac{\Omega_{\mathrm{cb},0} a^{-3}}{E(a)^2}\]

where E(a) is the normalized Hubble parameter.

Returns:

Type	Description
`Union[float, ndarray]`	Matter density parameter Ωₘ(a).

Growth Functions¶

jaxace.D_z ¶

D_z(z, Ωcb0, h, mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0) -> Union[float, jnp.ndarray]

Linear growth factor D(z).

The growth factor is normalized such that D(z=0) = 1. It satisfies the differential equation given in growth_solver.

Returns:

Type	Description
`Union[float, ndarray]`	jnp.ndarray: Linear growth factor D(z). Returns NaN for NaN inputs, handles invalid parameters gracefully.

jaxace.f_z ¶

f_z(z, Ωcb0, h, mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0) -> Union[float, jnp.ndarray]

Growth rate f(z) = d log D / d log a.

The growth rate is defined as:

\[f(z) = \frac{\mathrm{d} \ln D}{\mathrm{d} \ln a}\]

where D is the linear growth factor.

Returns:

Type	Description
`Union[float, ndarray]`	jnp.ndarray: Growth rate f(z). Returns NaN for NaN inputs, handles invalid parameters gracefully.

jaxace.D_f_z ¶

D_f_z(z, Ωcb0, h, mν=0.0, w0=-1.0, wa=0.0, Ωk0=0.0)

Distance Functions¶

jaxace.r_z ¶

r_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Line-of-sight comoving distance r(z) in Mpc.

This is the conformal distance scaled to physical units. Independent of curvature (curvature only affects transverse distances).

Returns:

Type	Description
`Union[float, ndarray]`	Line-of-sight comoving distance in Mpc

jaxace.dA_z ¶

dA_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Angular diameter distance dA(z) in Mpc.

For curved universes, uses transverse comoving distance: dA(z) = dM(z) / (1+z)

Returns:

Type	Description
`Union[float, ndarray]`	Angular diameter distance in Mpc

jaxace.dL_z ¶

dL_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Luminosity distance at redshift z.

For curved universes, uses transverse comoving distance: dL(z) = dM(z) * (1 + z)

Parameters:

Name	Type	Description	Default
`z`	`Union[float, ndarray]`	Redshift	required
`Ωcb0`	`Union[float, ndarray]`	Present-day matter density parameter (CDM + baryons)	required
`h`	`Union[float, ndarray]`	Dimensionless Hubble parameter (H0 = 100h km/s/Mpc)	required
`mν`	`Union[float, ndarray]`	Sum of neutrino masses in eV	`0.0`
`w0`	`Union[float, ndarray]`	Dark energy equation of state parameter	`-1.0`
`wa`	`Union[float, ndarray]`	Dark energy equation of state evolution parameter	`0.0`
`Ωk0`	`Union[float, ndarray]`	Curvature density parameter	`0.0`

Returns:

Type	Description
`Union[float, ndarray]`	Luminosity distance in Mpc

Density Functions¶

jaxace.ρc_z ¶

ρc_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0, Ωk0: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

jaxace.Ωtot_z ¶

Ωtot_z(z: Union[float, ndarray], Ωcb0: Union[float, ndarray], h: Union[float, ndarray], mν: Union[float, ndarray] = 0.0, w0: Union[float, ndarray] = -1.0, wa: Union[float, ndarray] = 0.0) -> Union[float, jnp.ndarray]

Utility Functions¶

jaxace.a_z ¶

a_z(z)

Neural Network Emulators¶

jaxace.init_emulator ¶

init_emulator(nn_dict: Dict[str, Any], weight: ndarray, emulator_type: Type[FlaxEmulator] = FlaxEmulator, validate: bool = True, validate_weights: Optional[bool] = None) -> FlaxEmulator

Initialize an emulator from neural network dictionary and weights.

Parameters:

Name	Type	Description	Default
`nn_dict`	`Dict[str, Any]`	Neural network specification dictionary	required
`weight`	`ndarray`	Flattened weight array	required
`emulator_type`	`Type[FlaxEmulator]`	Type of emulator (currently only FlaxEmulator)	`FlaxEmulator`
`validate`	`bool`	Whether to validate nn_dict structure	`True`
`validate_weights`	`Optional[bool]`	Whether to validate weight dimensions	`None`

Returns:

Type	Description
`FlaxEmulator`	Initialized FlaxEmulator instance

jaxace.load_trained_emulator ¶

load_trained_emulator(path: str, backend: Type[FlaxEmulator] = FlaxEmulator, weights_file: str = 'weights.npy', inminmax_file: str = 'inminmax.npy', outminmax_file: str = 'outminmax.npy', nn_setup_file: str = 'nn_setup.json', postprocessing_file: str = 'postprocessing.py', validate: bool = True) -> GenericEmulator

Load a trained emulator from disk.

This function matches Julia's AbstractCosmologicalEmulators.load_trained_emulator.

Parameters:

Name	Type	Description	Default
`path`	`str`	Directory path containing the emulator files	required
`backend`	`Type[FlaxEmulator]`	Emulator backend type (FlaxEmulator)	`FlaxEmulator`
`weights_file`	`str`	Filename for neural network weights (default: "weights.npy")	`'weights.npy'`
`inminmax_file`	`str`	Filename for input normalization (default: "inminmax.npy")	`'inminmax.npy'`
`outminmax_file`	`str`	Filename for output normalization (default: "outminmax.npy")	`'outminmax.npy'`
`nn_setup_file`	`str`	Filename for network architecture (default: "nn_setup.json")	`'nn_setup.json'`
`postprocessing_file`	`str`	Filename for postprocessing function (default: "postprocessing.py")	`'postprocessing.py'`
`validate`	`bool`	Whether to validate the loaded data (default: True)	`True`

Returns:

Type	Description
`GenericEmulator`	GenericEmulator instance ready for evaluation

Example

emu = load_trained_emulator("/path/to/emulator/") output = emu.run_emulator(input_params)

File Structure

The expected directory structure is:

path/
├── weights.npy          # Neural network weights
├── inminmax.npy         # Input normalization (n_params, 2)
├── outminmax.npy        # Output normalization (n_output, 2)
├── nn_setup.json        # Network architecture
└── postprocessing.py    # Postprocessing function

jaxace.load_trained_emulator_from_artifact ¶

load_trained_emulator_from_artifact(artifact_name: str, artifacts_toml: Optional[str] = None, backend: Type[FlaxEmulator] = FlaxEmulator, **kwargs) -> GenericEmulator

Load a trained emulator from an artifact defined in Artifacts.toml.

This function automatically downloads and caches emulators from remote sources (e.g., Zenodo) using the fetch-artifacts system, then loads them using load_trained_emulator.

Parameters:

Name	Type	Description	Default
`artifact_name`	`str`	Name of the artifact as defined in Artifacts.toml	required
`artifacts_toml`	`Optional[str]`	Optional path to Artifacts.toml file. If None, looks for Artifacts.toml in the package root.	`None`
`backend`	`Type[FlaxEmulator]`	Emulator backend type (FlaxEmulator)	`FlaxEmulator`
`**kwargs`	`Any`	Additional arguments passed to load_trained_emulator (e.g., weights_file, validate, etc.)	`{}`

Returns:

Type	Description
`GenericEmulator`	GenericEmulator instance ready for evaluation

Example

Load from default Artifacts.toml¶

emu = load_trained_emulator_from_artifact("ACE_mnuw0wacdm_sigma8_basis") output = emu.run_emulator(input_params)

Load from custom Artifacts.toml¶

emu = load_trained_emulator_from_artifact( ... "ACE_mnuw0wacdm_sigma8_basis", ... artifacts_toml="/path/to/Artifacts.toml" ... )

Note

On first use, this will download the emulator from the URL specified in Artifacts.toml. Subsequent uses will load from the local cache (typically ~/.fetch_artifacts/).

jaxace.FlaxEmulator `dataclass` ¶

FlaxEmulator(model: Module, parameters: Dict[str, Any], states: Optional[Dict[str, Any]] = None, description: Dict[str, Any] = None)

Bases: AbstractTrainedEmulator

Flax-based emulator with automatic JIT compilation.

Key features: 1. Automatic JIT compilation on first use 2. Automatic batch detection and vmap application 3. Cached compiled functions for performance

Attributes:

Name	Type	Description
`model`	`Module`	Flax model (nn.Module)
`parameters`	`Dict[str, Any]`	Model parameters dictionary
`states`	`Optional[Dict[str, Any]]`	Model states (usually empty for standard feedforward networks)
`description`	`Dict[str, Any]`	Emulator description dictionary

run_emulator ¶

run_emulator(input_data: Union[ndarray, ndarray]) -> jnp.ndarray

Run the emulator with automatic JIT compilation and batch detection.

This method automatically: 1. Converts numpy arrays to JAX arrays 2. Detects if input is a batch or single sample 3. Applies JIT compilation 4. Uses vmap for batch processing

Parameters:

Name	Type	Description	Default
`input_data`	`Union[ndarray, ndarray]`	Input array (single sample or batch) Shape: (n_features,) for single or (n_samples, n_features) for batch	required

Returns:

Type	Description
`ndarray`	Output array from the neural network

call ¶

__call__(input_data: Union[ndarray, ndarray]) -> jnp.ndarray

Allow the emulator to be called directly as a function.

jaxace.GenericEmulator `dataclass` ¶

GenericEmulator(trained_emulator: AbstractTrainedEmulator, in_minmax: ndarray, out_minmax: ndarray, postprocessing: callable = None)

Bases: AbstractTrainedEmulator

Generic emulator that wraps a trained neural network with normalization and postprocessing.

This class provides a complete emulator interface that: 1. Normalizes inputs using min-max scaling 2. Runs the underlying neural network 3. Denormalizes outputs 4. Applies optional postprocessing

This matches the Julia AbstractCosmologicalEmulators.jl GenericEmulator struct.

Attributes:

Name	Type	Description
`trained_emulator`	`AbstractTrainedEmulator`	The underlying trained neural network (FlaxEmulator)
`in_minmax`	`ndarray`	Input normalization parameters, shape (n_input_features, 2) Column 0 is min, column 1 is max
`out_minmax`	`ndarray`	Output normalization parameters, shape (n_output_features, 2) Column 0 is min, column 1 is max
`postprocessing`	`callable`	Optional postprocessing function with signature (input_params, output, emulator) -> processed_output

run_emulator ¶

run_emulator(input_params: Union[ndarray, ndarray], auxiliary_params: Union[ndarray, ndarray, None] = None) -> jnp.ndarray

Run the complete emulator pipeline.

Steps: 1. Normalize inputs using in_minmax 2. Run the neural network 3. Denormalize outputs using out_minmax 4. Apply postprocessing function

Parameters:

Name	Type	Description	Default
`input_params`	`Union[ndarray, ndarray]`	Input parameters, shape (n_features,) or (n_samples, n_features)	required
`auxiliary_params`	`Union[ndarray, ndarray, None]`	Deprecated. Present only for compatibility with older jaxace postprocessing functions.	`None`

Returns:

Type	Description
`ndarray`	Processed output array

call ¶

__call__(input_params: Union[ndarray, ndarray], auxiliary_params: Union[ndarray, ndarray, None] = None) -> jnp.ndarray

Allow the emulator to be called directly as a function.

Postprocessing signature

Custom postprocessing functions should use the ACE.jl-compatible signature postprocessing(input_params, output, emulator). Legacy four-argument functions with auxiliary_params are accepted for backward compatibility but are not the canonical 0.6.0 API.

Utilities¶

jaxace.maximin ¶

maximin(input_data: Union[ndarray, ndarray], minmax: Union[ndarray, ndarray]) -> Union[np.ndarray, jnp.ndarray]

Normalize input data using min-max scaling. Matches Julia's maximin function.

Parameters:

Name	Type	Description	Default
`input_data`	`Union[ndarray, ndarray]`	Input array to normalize (shape: (n_features,) or (n_features, n_samples))	required
`minmax`	`Union[ndarray, ndarray]`	Array of shape (n_features, 2) where column 0 is min, column 1 is max	required

Returns:

Type	Description
`Union[ndarray, ndarray]`	Normalized array in range [0, 1]

jaxace.inv_maximin ¶

inv_maximin(output_data: Union[ndarray, ndarray], minmax: Union[ndarray, ndarray]) -> Union[np.ndarray, jnp.ndarray]

Denormalize output data from min-max scaling. Matches Julia's inv_maximin function.

Parameters:

Name	Type	Description	Default
`output_data`	`Union[ndarray, ndarray]`	Normalized array (shape: (n_features,) or (n_features, n_samples))	required
`minmax`	`Union[ndarray, ndarray]`	Array of shape (n_features, 2) where column 0 is min, column 1 is max	required

Returns:

Type	Description
`Union[ndarray, ndarray]`	Denormalized array

Interpolation¶

jaxace.akima_interpolation ¶

akima_interpolation(u, t, t_new) -> Union[float, jnp.ndarray]

Akima spline interpolation for 1D or 2D data.

This is a direct translation of AbstractCosmologicalEmulators.jl's akima_interpolation function. Evaluates the Akima spline that interpolates the data points (t_i, u_i) at new abscissae t_new.

The Akima spline is a piecewise cubic polynomial that uses weighted averaging of local slopes to determine derivatives at each node. This dampens oscillations without explicit shape constraints. The spline is C¹ continuous but generally not C².

This implementation is fully compatible with JAX's jit and automatic differentiation.

Parameters:

Name	Type	Description	Default
`u`	`ndarray`	Ordinates (function values) at data nodes. The 1D case has shape `(n,)`. The 2D case has shape `(n, n_cols)` where each column is interpolated independently.	required
`t`	`ndarray`	Strictly increasing abscissae (x-coordinates), shape `(n,)`.	required
`t_new`	`ndarray`	Query point(s) where spline is evaluated, scalar or array.	required

Returns:

Type	Description
`Union[float, ndarray]`	jnp.ndarray: Interpolated value(s) at `t_new`. For 1D input this is a scalar if `t_new` is scalar, otherwise an array. For 2D input this is a matrix of shape `(len(t_new), n_cols)`.

Example (1D): >>> import jax.numpy as jnp >>> t = jnp.linspace(0, 1, 10) >>> u = jnp.sin(2 * jnp.pi * t) >>> t_new = jnp.linspace(0, 1, 50) >>> u_new = akima_interpolation(u, t, t_new)

Example (2D - multiple columns): >>> # Interpolate Jacobian with 11 parameter columns >>> k_in = jnp.linspace(0.01, 0.3, 50) >>> jacobian = jnp.randn(50, 11) # 11 parameters >>> k_out = jnp.linspace(0.01, 0.3, 100) >>> result = akima_interpolation(jacobian, k_in, k_out) # (100, 11)

# Works with jit
>>> akima_jit = jax.jit(akima_interpolation)
>>> u_new = akima_jit(u, t, t_new)

# Works with autodiff
>>> grad_fn = jax.grad(lambda u: jnp.sum(akima_interpolation(u, t, t_new)))
>>> grad_u = grad_fn(u)

jaxace.cubic_spline_interpolation ¶

cubic_spline_interpolation(u, t, t_new) -> Union[float, jnp.ndarray]

Natural Cubic Spline interpolation for 1D or 2D data.

This is a direct translation of AbstractCosmologicalEmulators.jl's cubic_spline_interpolation function. Evaluates the Natural Cubic Spline that interpolates the data points (t_i, u_i) at new abscissae t_new.

This implementation is fully compatible with JAX's jit and automatic differentiation.

Parameters:

Name	Type	Description	Default
`u`	`ndarray`	Ordinates (function values) at data nodes. The 1D case has shape `(n,)`. The 2D case has shape `(n, n_cols)` where each column is interpolated independently.	required
`t`	`ndarray`	Strictly increasing abscissae (x-coordinates), shape `(n,)`.	required
`t_new`	`ndarray`	Query point(s) where spline is evaluated, scalar or array.	required

Returns:

Type	Description
`Union[float, ndarray]`	jnp.ndarray: Interpolated value(s) at `t_new`. For 1D input this is a scalar if `t_new` is scalar, otherwise an array. For 2D input this is a matrix of shape `(len(t_new), n_cols)`.

Chebyshev¶

jaxace.ChebyshevPlan ¶

Bases: NamedTuple

Plan for computing Chebyshev coefficients of a function evaluated at Chebyshev nodes.

Attributes:

Name	Type	Description
`K`	`Tuple[int, ...]`	Tuple of polynomial degrees (K+1 nodes per dimension)
`nodes`	`Tuple[ndarray, ...]`	Tuple of evaluation nodes arrays
`dim`	`Tuple[int, ...]`	Tuple of target dimensions for decomposition

jaxace.chebpoints ¶

chebpoints(K: int, x_min: float, x_max: float) -> jnp.ndarray

Generate Chebyshev roots mapped to [x_min, x_max].

Matches AbstractCosmologicalEmulators.chebpoints

jaxace.prepare_chebyshev_plan ¶

prepare_chebyshev_plan(x_min: Union[float, Tuple[float, ...]], x_max: Union[float, Tuple[float, ...]], K: Union[int, Tuple[int, ...]], size_nd: Optional[Tuple[int, ...]] = None, dim: Union[int, Tuple[int, ...]] = 0) -> ChebyshevPlan

Precomputes the Chebyshev nodes required to compute coefficients.

K is the polynomial degree (K+1 nodes). For N-dimensional inputs, specify the target dimensions dim.

Parameters:

Name	Type	Description	Default
`x_min`	`Union[float, Tuple[float, ...]]`	Minimum x value(s)	required
`x_max`	`Union[float, Tuple[float, ...]]`	Maximum x value(s)	required
`K`	`Union[int, Tuple[int, ...]]`	Polynomial degree(s)	required
`size_nd`	`Optional[Tuple[int, ...]]`	Tuple representing input array shape (unused parameter kept for API parity)	`None`
`dim`	`Union[int, Tuple[int, ...]]`	Target dimension(s) for Chebyshev decomposition (default 0)	`0`

Returns:

Type	Description
`ChebyshevPlan`	ChebyshevPlan object containing nodes and settings

jaxace.chebyshev_polynomials ¶

chebyshev_polynomials(x_grid: ndarray, x_min: float, x_max: float, K: int) -> jnp.ndarray

Computes the matrix of Chebyshev polynomials evaluated on x_grid, mapped to [-1, 1] from [x_min, x_max].

Matches AbstractCosmologicalEmulators.chebyshev_polynomials functionality.

Parameters:

Name	Type	Description	Default
`x_grid`	`ndarray`	Grid of evaluation points	required
`x_min`	`float`	Input minimum domain	required
`x_max`	`float`	Input maximum domain	required
`K`	`int`	Polynomial degree	required

Returns:

Type	Description
`ndarray`	Matrix of size (len(x_grid), K+1)

jaxace.chebyshev_decomposition ¶

chebyshev_decomposition(plan: ChebyshevPlan, f_vals: ndarray) -> jnp.ndarray

Computes Chebyshev coefficients for a function evaluated at Chebyshev nodes.

Fully supports batched N-dimensional JAX arrays across plan.dim using custom DCT-1.

Parameters:

Name	Type	Description	Default
`plan`	`ChebyshevPlan`	ChebyshevPlan containing interpolation configuration	required
`f_vals`	`ndarray`	Array of function outputs evaluated exactly on plan.nodes along plan.dim (must have size K+1 along plan.dim)	required

Returns:

Type	Description
`ndarray`	N-Dimensional array of Chebyshev coefficients

API Reference¶

Cosmology¶

jaxace.w0waCDMCosmology dataclass ¶

D_f_z ¶

D_z ¶

E_a ¶

E_z ¶

dA_z ¶

dL_z ¶

dM_z ¶

d̃A_z ¶

d̃M_z ¶

f_z ¶

r_z ¶

r̃_z ¶

Ωm_a ¶

Ωtot_z ¶

ρc_z ¶

Background Functions¶

Hubble Functions¶

jaxace.E_z ¶

jaxace.E_a ¶

jaxace.dlogEdloga ¶

Matter Density¶

jaxace.Ωm_a ¶

Growth Functions¶

jaxace.D_z ¶

jaxace.f_z ¶

jaxace.D_f_z ¶

Distance Functions¶

jaxace.r_z ¶

jaxace.dA_z ¶

jaxace.dL_z ¶

Density Functions¶

jaxace.ρc_z ¶

jaxace.Ωtot_z ¶

Utility Functions¶

jaxace.a_z ¶

Neural Network Emulators¶

jaxace.init_emulator ¶

jaxace.load_trained_emulator ¶

jaxace.load_trained_emulator_from_artifact ¶

Load from default Artifacts.toml¶

Load from custom Artifacts.toml¶

jaxace.FlaxEmulator dataclass ¶

run_emulator ¶

__call__ ¶

jaxace.GenericEmulator dataclass ¶

run_emulator ¶

__call__ ¶

Utilities¶

jaxace.maximin ¶

jaxace.inv_maximin ¶

Interpolation¶

jaxace.akima_interpolation ¶

jaxace.cubic_spline_interpolation ¶

Chebyshev¶

jaxace.ChebyshevPlan ¶

jaxace.chebpoints ¶

jaxace.prepare_chebyshev_plan ¶

jaxace.chebyshev_polynomials ¶

jaxace.chebyshev_decomposition ¶

jaxace.w0waCDMCosmology `dataclass` ¶

jaxace.FlaxEmulator `dataclass` ¶

call ¶

jaxace.GenericEmulator `dataclass` ¶

call ¶