Generation¶

Engine¶

class_e.generation.engine ¶

GenerationConfig `dataclass` ¶

GenerationConfig(
    yaml_path: Optional[Path],
    bounds: Dict[str, Tuple[float, float]],
    observables: List[str],
    lmax: int,
    lensed: bool,
    k_values: Optional[ndarray],
    z_values: Optional[ndarray],
    n_samples: int,
    method: str,
    sampling_kwargs: Dict[str, Any],
    output_path: Path,
    model_name: str,
    compute_derivatives: bool = False,
    derivative_epsilon: float = 0.0001,
    tk_k_values: Optional[ndarray] = None,
    tk_z_values: Optional[ndarray] = None,
    bg_z_values: Optional[ndarray] = None,
    constraints: Tuple[str, ...] = (),
    max_failure_fraction: float = 0.01,
    seed: Optional[int] = None,
    design: Optional[str] = None,
    model_data: Dict[str, Any] = dict(),
)

Settings of one design's data generation, parsed by parse_config.

Attributes:

Name	Type	Description
`yaml_path`	`(Path, optional)`	Where the configuration came from (provenance only).
`bounds`	`dict of str to tuple of float`	Prior box of each sampled parameter.
`observables`	`list of str`	Observables to compute and store.
`lmax`	`int`	Maximum stored multipole.
`lensed`	`bool`	Whether CMB spectra are lensed.
`k_values, z_values`	`(ndarray, optional)`	Wavenumber and redshift grids of the matter power spectrum.
`n_samples`	`int`	Number of points of the design.
`method`	`str`	Sampling method (`sobol`, `lhc`, `lhs`, `cov`).
`sampling_kwargs`	`dict`	Extra arguments of the sampler (e.g. the covariance matrix).
`output_path`	`Path`	Default dataset path.
`model_name`	`str`	Name of the model, used in file names.
`compute_derivatives`	`bool, default False`	Also compute central finite-difference derivatives of the observables.
`derivative_epsilon`	`float, default 1e-4`	Step of the finite differences.
`tk_k_values, tk_z_values`	`(ndarray, optional)`	Grids of the `Tk` observable.
`bg_z_values`	`(ndarray, optional)`	Redshift grid of the background observables (`Hz`, `chi`).
`constraints`	`tuple of str`	Expressions the sampled points must satisfy.
`max_failure_fraction`	`float, default 0.01`	Fraction of CLASS failures tolerated (the points are dropped) before erroring.
`seed`	`(int, optional)`	Design seed; `None` draws a new design every run.
`design`	`(str, optional)`	Name of the design in `samples.designs`, if any.
`model_data`	`dict`	The model section of the YAML (parameters, derived, theory, ...), overrides applied.

parse_config ¶

parse_config(
    source: str | Path | Mapping[str, Any],
    design: Optional[str] = None,
) -> GenerationConfig

The GenerationConfig of one design of a configuration.

source is a YAML path, or the configuration itself (a dict, e.g. built in code or with overrides applied; see class_e.config.PipelineConfig).

compute_dataset ¶

compute_dataset(
    conf: GenerationConfig,
    seed: Optional[int] = None,
    *,
    rows: Optional[Tuple[int, int]] = None,
    checkpoint: Optional[str | Path] = None,
    checkpoint_every: int = 10,
    progress: Optional[ProgressCallback] = None,
    cancel: Optional[Event] = None,
    classy: Optional[str | Path] = None,
    point_retries: int = None,
) -> TrainingData

Compute a training dataset (serial only).

Parameters:

Name	Type	Description	Default
`seed`	`int`	Design seed; defaults to `conf.seed` (YAML `samples.seed`). With neither, every run draws a different design.	`None`
`rows`	`(start, stop)`	Compute only these rows of the design (the full design of `conf.n_samples` points is drawn, so chunks run separately merge into the same dataset, see `merge_datasets`). The failure limit (`conf.max_failure_fraction`) then applies to the whole design, when the chunks are merged, not to each chunk.	`None`
`checkpoint`	`path`	File where finished rows are saved every `checkpoint_every` points. If it exists (from an interrupted run with the same settings), its rows are reused and only the missing ones are computed.	`None`
`progress`	`callable`	Receives the `class_e.core.events` of stage `"generate"` (and `"generate/derivatives"`); nothing is printed.	`None`
`cancel`	`Event`	Once set, the run stops before the next point (after saving the checkpoint) and raises `class_e.core.exceptions.Cancelled`; rerunning resumes from the checkpoint.	`None`
`classy`	`path`	The CLASS build to use (see `class_e.core.boltzmann.resolve_classy_path`: this, else `$CLASS_E_CLASSY`, else the YAML's `theory.classy.path`, else the installed `classy`). Recorded in the meta (`classy`: version and file).	`None`
`point_retries`	`int`	Extra attempts at a point where CLASS fails (default `CLASS_RETRIES`); with a checkpoint, also how often a point may crash the whole process before it is recorded as failed (see `_CrashGuard`).	`None`

save_dataset ¶

save_dataset(
    result: TrainingData, out_path: str | Path
) -> Path

Write result as a v2 dataset (see class_e.core.dataset), then its sidecar.

merge_datasets ¶

merge_datasets(
    parts: Sequence[str | Path],
    out_path: str | Path,
    max_failure_fraction: Optional[float] = None,
) -> Path

Merge datasets computed as row chunks of one design (compute_dataset(rows=...)).

The chunks must be complete (see dataset_problem), come from the same settings and design (fingerprint and design_sha256 in their sidecars; not the YAML path, so chunks computed elsewhere merge), and together cover consecutive rows. CLASS failures are checked here, over the whole merged design: more than max_failure_fraction (default: the YAML's samples.max_failure_fraction, recorded by each chunk) raises, since dropping that many points biases the design. Writes out_path and its .meta.json sidecar.

upgrade_dataset ¶

upgrade_dataset(path: str | Path) -> Path

Rewrite a dataset in the v1 format (pickled observables) as v2, in place.

Atomic, and the sidecar then records the new file; a v2 dataset is left untouched.

Runner¶

class_e.generation.runner ¶

Running generation: chunks of a design, crash-isolated workers, merging.

A design of N points is split into chunks (rows [i*N/K, (i+1)*N/K)); every chunk redraws the full seeded design, computes its rows with a checkpoint, and writes a part <parts>/<prefix>.part<iii>of<KKK>.npz. merge_chunks joins the parts.

With run.isolate (the default) each chunk runs in its own process (python -m class_e.generation.worker), supervised here: a CLASS segfault costs that worker, which is restarted from its checkpoint (run.chunk_retries); a point that keeps killing the process is recorded as failed (run.point_retries, see class_e.generation.engine._CrashGuard). Progress comes back as events.

run_design / generate_all are the library behind class-e generate CFG; run_chunk_supervised behind one SLURM array task (--chunk i --n-chunks K).

run_design ¶

run_design(
    config,
    design: Optional[str] = None,
    *,
    out: Optional[str | Path] = None,
    flags: Optional[dict] = None,
    n_chunks: Optional[int] = None,
    workers: Optional[int] = None,
    classy: Optional[str] = None,
    overwrite: bool = False,
    progress: Optional[ProgressCallback] = None,
    cancel: Optional[Event] = None,
)

Generate the dataset of one design of a class_e.config.PipelineConfig.

The design is split into chunks (run.chunk_size, n_chunks, or one per worker) computed by run.workers workers in parallel (in-process and serially if not run.isolate), then merged into dataset_path(design, out). A complete dataset with the same settings is kept unless overwrite, which recomputes everything; otherwise parts already computed (same settings and CLASS version) are reused. Returns a class_e.generation.dataset.DatasetResult.

generate_all ¶

generate_all(
    config,
    designs: Optional[Sequence[Optional[str]]] = None,
    **kwargs,
) -> Dict[Any, Any]

run_design for every design (default: all of the configuration), one after the other.

run_chunk ¶

run_chunk(
    conf,
    *,
    chunk: int,
    n_chunks: int,
    out_dir: str | Path,
    checkpoint_every: int = 8,
    point_retries: Optional[int] = None,
    classy: Optional[str] = None,
    progress: Optional[ProgressCallback] = None,
    cancel: Optional[Event] = None,
) -> Path

Compute one chunk of a design in this process, resumable from its checkpoint.

A complete part computed with the same settings (and CLASS version) is kept; an incomplete or stale part is recomputed, and a checkpoint of other settings discarded.

merge_chunks ¶

merge_chunks(
    prefix: str | Path,
    out: str | Path,
    max_failure_fraction: Optional[float] = None,
    n_chunks: Optional[int] = None,
) -> Path

Merge every part of a design into one dataset (see generation.engine.merge_datasets).

All K parts must exist; CLASS failures are checked over the whole design.

chunk_status ¶

chunk_status(out_dir: str | Path) -> dict

Progress of chunked generation in out_dir, per design prefix.

Returns {prefix: {"n_chunks", "finished", "incomplete", "points", "failed", "running": {chunk: [done, size]}, "crashes": {chunk: n}}}; incomplete lists parts whose write did not finish, crashes the worker deaths recorded for running chunks.