Datasets and bundles

Datasets

class_e.core.dataset

Training datasets on disk: reading (both formats) and the v2 layout.

A dataset is an .npz (plus a .meta.json sidecar, see class_e.generation.engine.save_dataset). Format v2 holds plain arrays only, so it loads with allow_pickle=False:

Key	Content
schema	"class-e-dataset/2"
theta	(N, D) sampled parameters
names	(D,) parameter names
bounds	(D, 2) prior box
obs/<name>	(N, K) one observable (obs/TT, obs/P(z=0.5), obs/Tk, obs/Hz)
grid/<axis>	a grid (grid/ell, grid/k, grid/z, grid/tk_k, grid/z_bg, ...)
deriv/<name>	(N, D, K) finite-difference derivatives, if computed
axes	JSON {name: [axis, ...]}: the grids a row of each observable spans

Format v1 (class-e <= 0.0.4) stored the observables as a pickled dict; read_dataset still reads it, and class_e.generation.engine.upgrade_dataset rewrites a v1 file as v2.

Dataset dataclass

Dataset(
    theta: ndarray,
    names: List[str],
    bounds: Optional[List[Tuple[float, float]]],
    observables: Dict[str, ndarray],
    grids: Dict[str, ndarray] = dict(),
    axes: Dict[str, Tuple[str, ...]] = dict(),
    derivatives: Dict[str, ndarray] = dict(),
    meta: Dict[str, Any] = dict(),
    path: Optional[Path] = None,
)

A training (or test) set: parameters, observables and their grids.

n property

n: int

Number of points.

ell property

ell: Optional[ndarray]

The multipole grid, if the dataset has one.

k property

k: Optional[ndarray]

The wavenumber grid [h/Mpc], if the dataset has one.

z property

z: Optional[ndarray]

The redshift grid, if the dataset has one.

read_dataset

read_dataset(path: str | Path) -> Dataset

Read a dataset (format v2, or v1) and its .meta.json sidecar.

Emulator bundles

class_e.core.bundle.EmulatorBundle

EmulatorBundle(
    emulators: Mapping[str, Any],
    param_names: Sequence[str],
    bounds: Sequence[Sequence[float]],
    grids: Mapping[str, Any] | None = None,
    meta: Mapping[str, Any] | None = None,
)

Emulators for several observables sharing one parameter space.

Parameters:

Name	Type	Description	Default
emulators	mapping of str to emulator	Fitted emulators keyed by observable name (e.g. "TT").	required
param_names	sequence of str	Input parameter names, in the column order the emulators were trained on.	required
bounds	sequence of (min, max)	Prior box of each parameter, same order as param_names.	required
grids	mapping of str to array	Output grids, e.g. {"ell": ell} or {"k": k, "z": z}.	None
meta	mapping	JSON-serializable metadata (name, description, dataset, metrics, ...).	None

observables property

observables: list[str]

Names of the observables the bundle holds an emulator for.

predict

predict(
    params: Mapping[str, Any] | ndarray,
    observables: Sequence[str] | None = None,
    out_of_bounds: str = "warn",
) -> dict[str, np.ndarray]

Predict observables for one or several cosmologies.

Parameters:

Name	Type	Description	Default
params	mapping or array	Either a mapping from every parameter name to a scalar or a 1-D array (one value per cosmology), or an array of shape (D,) or (N, D) with columns in param_names order.	required
observables	sequence of str	Subset of observables to predict (default: all).	None
out_of_bounds	('warn', 'raise', 'ignore')	What to do when a cosmology lies outside the prior box the emulators were trained on (emulators do not extrapolate reliably).	"warn"

Returns:

Type	Description
dict	Observable name -> prediction, shape (L,) for a single cosmology or (N, L) otherwise.

save

save(path: str | PathLike, overwrite: bool = False) -> Path

Write the bundle to path (conventionally *.cemu).

read_manifest staticmethod

read_manifest(path: str | PathLike) -> dict

Return the manifest of a bundle without loading or checking its arrays.

load classmethod

load(
    path: str | PathLike, verify: bool = True
) -> "EmulatorBundle"

Load a bundle written by :meth:save.

Parameters:

Name	Type	Description	Default
path	path - like	The .cemu file.	required
verify	bool	Recompute the stored reference predictions and raise :class:BundleError if they are not reproduced.	True

describe

describe() -> str

Human-readable summary of the bundle.