petitRADTRANS.sbi.task#

Task definitions for simulation-based inference workflows.

These classes bridge the user-facing retrieval configuration objects and the runtime-native simulation interfaces used by amortized inference. SBITask is the primary object that allows the user to interface with the SBI module.

Attributes#

`_FINGERPRINT_SCHEMA_VERSION`
`_PRIOR_FINGERPRINT_CUBE_GRID`

Classes#

`NoiseModelConfig`	Describe how observational noise is injected during simulation.
`ObservationValueConstraint`	Admissible range for simulated observation values.
`ObservationSchema`	Describes the type of observation being used. For SBI
`SimulationConfig`	Control how an SBITask generates prior-predictive forward models.
`SBITask`	An SBITask is the core unit of the SBI workflow. It is an immutable datastructure

Functions#

_model_function_digest(→ str)

Return a string description of a model-generating function's implementation.

Module Contents#

class petitRADTRANS.sbi.task.NoiseModelConfig#

Describe how observational noise is injected during simulation.

Attributes:

mode:: Short identifier of the noise model implementation. Can be ‘observational’, ‘gaussian’ or ‘none’.
parameters:: Backend-specific parameters used to instantiate the noise model.
seed:: Optional deterministic seed for repeatable simulation pipelines.
resample_during_training:: When True (default) the trainer redraws observational noise per training batch from the stored noise-free spectra instead of replaying the single realization frozen into the dataset — the network then never sees the same noise twice, which prevents overconfident (collapsed) posteriors. Set False to reproduce legacy frozen-noise runs.

mode: str = 'observational'#

parameters: Mapping[str, Any]#

seed: int | None = None#

resample_during_training: bool = True#

class petitRADTRANS.sbi.task.ObservationValueConstraint#

Admissible range for simulated observation values. This is used to reject forward models that produce invalid results and introduce numerical instability. For example, transit depths and photometric fluxes must be non-negative.

Attributes:

min_value:: Optional lower bound for valid simulated values.
max_value:: Optional upper bound for valid simulated values.
min_inclusive:: Whether min_value is included in the valid interval.
max_inclusive:: Whether max_value is included in the valid interval.

min_value: float | None = None#

max_value: float | None = None#

min_inclusive: bool = True#

max_inclusive: bool = True#

to_payload() → dict[str, Any]#

class petitRADTRANS.sbi.task.ObservationSchema#

Describes the type of observation being used. For SBI we’re primarily concerned with spectra, but we include the other options for compatibility with the retrieval module and potential future extensions of the SBI module.

Attributes:

dataset_names:: Dataset identifiers included in the task.
modalities:: Mapping from dataset name to a short modality label such as 'spectrum', 'photometry', or 'time_series'.
metadata:: Additional task-level information required by encoders or benchmarks, such as instrument names or wavelength coverage.

dataset_names: tuple[str, Ellipsis]#

modalities: Mapping[str, str]#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.task.SimulationConfig#

Control how an SBITask generates prior-predictive forward models.

Attributes:

batch_size:: Number of parameter points produced per simulator call.
n_simulations:: Optional target number of simulations for dataset generation jobs.
use_vectorized_runtime:: Whether to prefer the JAX-vectorized runtime path when available.
store_per_datapoint_log_likelihood:: Persist log-likelihood components alongside simulated observations.
noise_model:: Noise model configuration applied after forward-model evaluation.
observation_value_constraints:: Optional per-dataset admissible value ranges enforced on deterministic projected observations before simulations are accepted.

batch_size: int = 256#

n_simulations: int | None = None#

use_vectorized_runtime: bool = True#

store_per_datapoint_log_likelihood: bool = False#

noise_model: NoiseModelConfig#

observation_value_constraints: Mapping[str, ObservationValueConstraint]#

petitRADTRANS.sbi.task._FINGERPRINT_SCHEMA_VERSION = '1.0.0'#

petitRADTRANS.sbi.task._PRIOR_FINGERPRINT_CUBE_GRID = (0.001, 0.25, 0.5, 0.75, 0.999)#

petitRADTRANS.sbi.task._model_function_digest(function: Any) → str#

Return a string description of a model-generating function’s implementation.

Prefers the source text (so editing the function body invalidates a cached dataset), falling back to bytecode and then the qualified name for callables whose source is unavailable. Never raises – fingerprinting must not fail.

class petitRADTRANS.sbi.task.SBITask#

An SBITask is the core unit of the SBI workflow. It is an immutable datastructure that stores the retrieval configuration, canonical parameter layout, observation schema, and simulation policy needed to generate training data and evaluate amortized posteriors.

SBITask.from_retrieval_config(…) is the standard constructor: it infers: dataset type, the forward-model family, and version strings.

name: str#

retrieval_config: petitRADTRANS.retrieval.retrieval_config.RetrievalConfig#

parameter_layout: petitRADTRANS.retrieval.runtime.ParameterLayout#

observation_schema: ObservationSchema#

simulation_config: SimulationConfig#

task_version: str = '0.1.0'#

preprocessing_version: str = '0.5.0'#

forward_model_family: str = 'unknown'#

petitradtrans_version: str#

metadata: Mapping[str, Any]#

classmethod from_retrieval_config(retrieval_config: petitRADTRANS.retrieval.retrieval_config.RetrievalConfig, simulation_config: SimulationConfig | None = None, metadata: Mapping[str, Any] | None = None) → SBITask#

Create an SBI task from an existing retrieval configuration.

Parameters#

retrieval_config:: Retrieval configuration describing the free parameters, observation datasets, and runtime-native forward model family.
simulation_config:: Optional simulation policy overriding the default prior-predictive batch configuration.
metadata:: Optional task-level metadata. Reserved keys such as task_version, preprocessing_version, and forward_model_family override the default inferred values.

Returns#

SBITask: Immutable SBI task description with parameter layout, observation schema, and deterministic task fingerprint payload.

Notes#

The task inspects the configured retrieval datasets to infer modalities and a forward-model-family string that later participates in artifact compatibility checks.

property parameter_names: tuple[str, Ellipsis]#: Return the free parameter names in canonical task order.

property line_opacity_modes: Mapping[str, str]#: Return the line-opacity mode for each configured dataset. Can be c-k or lbl.

_forward_model_signature() → dict[str, Any]#

Signature of everything that determines the simulated spectra. This provides a sanity check to make sure that the forward model you’re using is consistent with the one used to generate the training data. If you change the forward model, you’ll also need to regenerate the training dataset.

Captures prior bounds, declared parameter units, and the source of each model-generating function.

property fingerprint_payload: Mapping[str, Any]#: Return the deterministic payload used to fingerprint the task family.

property task_fingerprint: petitRADTRANS.sbi.compatibility.TaskFingerprint#: Return the deterministic fingerprint of the task family. This just saves the metadata in a static class so that it can be used for compatibility checks later on.

validate_observation_schema(observation_schema: ObservationSchema) → petitRADTRANS.sbi.compatibility.CompatibilityReport#

Compare an external observation schema with the task expectation.

Parameters#

observation_schema:: Candidate schema to compare against the task’s required datasets and modalities.

Returns#

CompatibilityReport: Structured compatibility report containing mismatches, if any.

validate_artifact_metadata(artifact_metadata: Any) → petitRADTRANS.sbi.compatibility.CompatibilityReport#

Compare a task with persisted artifact metadata.

Parameters#

artifact_metadata:: Artifact-like object exposing task and preprocessing provenance fields. Any is accepted to avoid a hard dependency cycle.

Returns#

CompatibilityReport: Structured report describing whether the artifact is compatible with the current task definition.

The method accepts Any to avoid a hard import dependency back from the task module to the artifact module.

build_observation_states(model_contract: str = 'differentiable') → dict[str, petitRADTRANS.retrieval.runtime.ObservationState]#

Materialize immutable runtime observations for the task datasets.

Parameters#

model_contract:: Runtime contract requested from each retrieval dataset.

Returns#

dict[str, ObservationState]: Mapping from dataset name to immutable runtime observation state.

build_runtime() → petitRADTRANS.retrieval.runtime.RetrievalRuntime#

Construct a retrieval runtime matching the task definition. The retrieval runtime organises the forward model and parameters, in an immutable structure so that the retrieval can be vectorised and differentiable.

Returns#

RetrievalRuntime: Runtime object configured with the task’s parameter layout and retrieval datasets.

The default implementation mirrors the retrieval-runtime grouping logic used by the exact retrieval package and is sufficient for initial SBI simulation backends.

static _infer_modality(data: petitRADTRANS.retrieval.data.Data) → str#

static _infer_forward_model_family(retrieval_config: petitRADTRANS.retrieval.retrieval_config.RetrievalConfig) → str#