petitRADTRANS.sbi.task

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.

Classes#

`NoiseModelConfig`	Describe how observational noise is injected during simulation.
`ObservationValueConstraint`	Admissible range for simulated observation values.
`ObservationSchema`	Capture the supported observation family for an amortized task.
`SimulationConfig`	Control how a task generates prior-predictive simulations.
`SBITask`	Bundle the immutable ingredients required for an SBI problem.

Module Contents#

class petitRADTRANS.sbi.task.NoiseModelConfig#

Describe how observational noise is injected during simulation.

Attributes:

mode:: Short identifier of the noise model implementation.
parameters:: Backend-specific parameters used to instantiate the noise model.
seed:: Optional deterministic seed for repeatable simulation pipelines.

mode: str = 'observational'#

parameters: Mapping[str, Any]#

seed: int | None = None#

class petitRADTRANS.sbi.task.ObservationValueConstraint#

Admissible range for simulated observation values.

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#

Capture the supported observation family for an amortized task.

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 a task generates prior-predictive simulations.

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]#

class petitRADTRANS.sbi.task.SBITask#

Bundle the immutable ingredients required for an SBI problem.

The task owns the retrieval configuration, canonical parameter layout, observation schema, and simulation policy needed to generate training data and evaluate amortized posteriors.

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.

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.

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.

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#