petitRADTRANS.sbi#

Simulation-based inference interfaces for petitRADTRANS.

The modules in petitRADTRANS.sbi provide a production-oriented architecture for amortized inference workflows built on top of the existing retrieval runtime. The package is intentionally lightweight at this stage: it defines the core task, simulation, dataset, and posterior interfaces without committing to one training backend or storage engine.

Submodules#

Exceptions#

TaskCompatibilityError

Raised when an observation or artifact is incompatible with a task.

Classes#

`LocalSensitivityPointReport`	Local linear-identifiability summary around one representative point.
`LocalSensitivityReport`	Aggregate local information-content diagnostics for one observation.
`PosteriorInformativenessReport`	Joint contraction + calibration verdict per parameter.
`PosteriorPredictiveReport`	Aggregate posterior-predictive summaries for held-out observations.
`SimulationBasedCalibrationReport`	Rank-based SBC summary over a held-out set of observations.
`HDF5SimulationDatasetStore`	HDF5-backed store for simulation corpora.
`SimulationDatasetStore`	Backend-independent interface for reading and writing simulation data.
`DatasetSplit`	Named dataset partitions used during training and evaluation.
`NormalizedObservationDatasetReader`	Lightweight reader yielding normalized ObservationBlock batches for training.
`ObservationEncoder`	Learned hierarchical encoder dispatching over modalities.
`PhotometryPointEncoder`	Learned photometry encoder using per-point MLP features and pooling.
`EncoderConfig`	Observation-embedding encoder.
`FlowConfig`	Normalizing-flow family and geometry.
`RefinementConfig`	Test-time, per-observation posterior refinement (any objective).
`SBIConfig`	End-to-end configuration for `petitRADTRANS.sbi.run_sbi()`.
`StabilityConfig`	Checkpoint-selection stability thresholds (flow invertibility / collapse).
`TrainingConfig`	Optimisation schedule and training objective.
`AmortizedRetrieval`	Serve trained SBI models through a retrieval-like interface.
`AmortizedRetrievalResult`	Return type for amortized inference queries.
`OODDiagnostic`	Describe whether an observation is inside the training support.
`ObservationBlock`	Representation of one (simulated or real) observation.
`AbstractObservationEncoder`	Transform structured observation blocks into model-ready embeddings.
`ObservationModality`	Supported observation block types for SBI conditioning.
`ConditionalAutoregressiveFlowPosterior`	Posterior estimator specialized to the autoregressive flow backend.
`ConditionalFlowPosterior`	Concrete amortized posterior using a conditional flow backend.
`ConditionalNeuralAutoregressiveFlowPosterior`	Posterior estimator specialized to the neural autoregressive backend.
`ConditionalSplineFlowPosterior`	Posterior estimator specialized to the spline flow backend.
`FlowMatchingPosterior`	Conditional flow-matching posterior skeleton.
`PersistentPosteriorEstimator`	Shared persistence helper for estimator backends with on-disk artifacts.
`PosteriorBatch`	Training batch passed to amortized posterior estimators.
`PosteriorEstimator`	Backend-agnostic interface for amortized posterior models.
`PosteriorSamples`	Posterior samples and optional per-sample diagnostics.
`TaskPreprocessingMetadata`	Serializable preprocessing metadata for an SBI task family.
`ProposalSampler`	Interface for simulation proposals beyond the prior distribution.
`RuntimeSimulator`	Concrete simulator backed by the retrieval runtime.
`SimulationBatch`	Container for one batch of simulations.
`Simulator`	Base simulator for SBI dataset generation and validation.
`NoiseModelConfig`	Describe how observational noise is injected during simulation.
`ObservationSchema`	Describes the type of observation being used. For SBI
`ObservationValueConstraint`	Admissible range for simulated observation values.
`SBITask`	An SBITask is the core unit of the SBI workflow. It is an immutable datastructure
`SimulationConfig`	Control how an SBITask generates prior-predictive forward models.
`EarlyStoppingConfig`	Early-stopping policy for SBI training.
`SBITrainer`	Reusable optimisation loop for amortized SBI posteriors.
`ObservationScaleReport`	Result of `check_observation_in_training_range()`.

Functions#

`classify_posterior_informativeness`(...)	Classify each parameter by contraction and SBC calibration.
`generate_local_sensitivity_report`(, ...)	Diagnose local physical identifiability around representative posterior points.
`generate_posterior_predictive_report`(...)	Generate posterior-predictive summaries for a held-out dataset split.
`generate_sbc_report`(, max_cases, seed, data_parallel)	Run SBC over a dataset reader using normalized observation batches.
`local_sensitivity_report_to_payload`(→ dict[str, Any])	Convert a local sensitivity report into a JSON-serializable payload.
`serialize_ood_diagnostic`(→ dict[str, Any] \| None)	Serialize an OOD diagnostic to a JSON-safe dict (or None).
`summarize_posterior_contraction`(→ dict[str, dict[str, ...)	Per-parameter posterior contraction relative to the prior.
`summarize_posterior_samples`(→ dict[str, dict[str, ...)	Per-parameter posterior p16/median/p84, with optional truth values.
`summarize_predictive_report`(→ dict[str, Any])	Serialize a posterior-predictive report's scalar metrics to JSON-safe dicts.
`summarize_training_diagnostics`(→ dict[str, Any])	Summarize per-epoch stability/inverse-error diagnostics of a trained flow.
`save_predictive_spectrum`(→ pathlib.Path)	Write a posterior-predictive spectrum (wavelength, observed, uncertainty,
`generate_simulation_dataset`(→ GeneratedSimulationDataset)	Generate and persist a simulation corpus in one call.
`estimate_preloaded_split_memory_gib`(→ float)	Estimate host RAM (GiB) to eagerly preload one split as float32 arrays.
`infer_split_counts`(→ dict[str, int])	Split a simulation budget into train/validation/test counts.
`load_posterior_estimator`(→ PosteriorEstimator)	Load a saved posterior estimator without naming its concrete class.
`build_conditional_flow_posterior`(...)	Build a `ConditionalFlowPosterior` from grouped config dataclasses.
`run_sbi`(→ pathlib.Path)	Run the full amortized-SBI validation workflow and return the run directory.
`build_observation_block`(→ ObservationBlock)	Build one observation block with modality normalization.
`build_observation_block_batch`(...)	Build observation blocks for each sample in a batched payload.
`build_observation_blocks_from_sample`(...)	Build modality-aware observation blocks for one simulated sample.
`build_raw_observation_blocks`(→ list)	Build raw (un-normalized) spectrum ObservationBlocks for every dataset in
`plot_local_sensitivity_fisher_correlations`(...)	Plot Fisher-correlation heatmaps for each representative point.
`plot_local_sensitivity_jacobians`(→ tuple[Any, ...)	Plot whitened Jacobian heatmaps for each representative point.
`plot_local_sensitivity_singular_values`(→ tuple[Any, ...)	Plot singular spectra of the whitened Jacobian for each point.
`plot_posterior_corner`(→ tuple[Any, numpy.ndarray])	Plot a lower-triangular corner view of posterior structure.
`plot_posterior_marginals`(→ tuple[Any, numpy.ndarray])	Plot one histogram per posterior dimension.
`plot_posterior_predictive_report`(→ tuple[Any, ...)	Plot observed values against posterior-predictive means and intervals.
`plot_sbc_rank_histograms`(→ tuple[Any, numpy.ndarray])	Plot one SBC rank histogram per inferred parameter.
`save_validation_plots`(→ None)	Render and save the standard validation plot set (marginals, corner, SBC
`fit_task_preprocessing`(→ TaskPreprocessingMetadata)	Fit preprocessing statistics from training observation blocks.
`normalize_observation_block`(...)	Normalize one observation block with fitted preprocessing statistics.
`normalize_observation_blocks`(...)	Normalize a list of observation blocks.
`check_observation_in_training_range`(...)	Check the observation's flux scale lies within the training distribution.

Package Contents#

class petitRADTRANS.sbi.LocalSensitivityPointReport#

Local linear-identifiability summary around one representative point.

Attributes#

label:: Short human-readable label for the representative point.
parameters:: Physical parameter vector at which the Jacobian was evaluated.
finite_difference_steps:: Per-parameter finite-difference step sizes used during Jacobian construction.
finite_difference_schemes:: Per-parameter scheme labels such as 'central' or 'forward'.
whitened_jacobian:: Observation Jacobian divided by observational uncertainty, with shape (n_observation_values, n_parameters).
singular_values:: Singular values of the whitened Jacobian.
effective_rank:: Number of singular values larger than the configured relative cutoff.
condition_number:: Condition number inferred from the singular spectrum.
fisher_matrix:: Approximate Fisher information matrix J^T J.
fisher_covariance:: Damped pseudo-inverse of the Fisher matrix.
fisher_correlation:: Correlation matrix derived from the approximate Fisher covariance.
parameter_sensitivity_norm:: Per-parameter square-root Fisher diagonal.
local_sigma:: Approximate local standard deviation from the Fisher covariance.
metadata:: Auxiliary diagnostics such as the ridge term and failed columns.

label: str#

parameters: numpy.ndarray#

finite_difference_steps: numpy.ndarray#

finite_difference_schemes: tuple[str, Ellipsis]#

whitened_jacobian: numpy.ndarray#

singular_values: numpy.ndarray#

effective_rank: int#

condition_number: float#

fisher_matrix: numpy.ndarray#

fisher_covariance: numpy.ndarray#

fisher_correlation: numpy.ndarray#

parameter_sensitivity_norm: numpy.ndarray#

local_sigma: numpy.ndarray#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.LocalSensitivityReport#

Aggregate local information-content diagnostics for one observation.

Attributes#

parameter_names:: Parameter ordering used throughout the report.
posterior_mean:: Posterior mean in physical parameter space.
posterior_std:: Posterior standard deviation in physical parameter space.
posterior_median:: Posterior median in physical parameter space.
posterior_iqr:: Posterior interquartile range in physical parameter space.
representative_points:: Local sensitivity summaries evaluated at representative posterior points such as the posterior mean and highest-density sample.
aggregate_local_sigma:: Median local Fisher sigma across representative points.
aggregate_parameter_sensitivity_norm:: Median parameter sensitivity norm across representative points.
posterior_to_local_sigma_ratio:: Ratio between posterior standard deviation and local Fisher sigma.
parameter_diagnostics:: Per-parameter heuristic summary separating weak data constraints from broader-than-local posterior structure.
metadata:: Auxiliary metadata such as quantile levels and observation slices.

parameter_names: tuple[str, Ellipsis]#

posterior_mean: numpy.ndarray#

posterior_std: numpy.ndarray#

posterior_median: numpy.ndarray#

posterior_iqr: numpy.ndarray#

representative_points: tuple[LocalSensitivityPointReport, Ellipsis]#

aggregate_local_sigma: numpy.ndarray#

aggregate_parameter_sensitivity_norm: numpy.ndarray#

posterior_to_local_sigma_ratio: numpy.ndarray#

parameter_diagnostics: Mapping[str, Mapping[str, Any]]#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.PosteriorInformativenessReport#

Joint contraction + calibration verdict per parameter.

contraction is the ratio of the posterior 68% width to the prior 68% width. On its own it is misleading: a posterior that has collapsed to a narrow but wrong region has a tiny contraction yet is useless. Combining it with the SBC normalized mean-rank error separates the three regimes.

parameter_names: list[str]#

contraction: numpy.ndarray#

sbc_normalized_rank_error: numpy.ndarray | None#

verdict: list[str]#

class petitRADTRANS.sbi.PosteriorPredictiveReport#

Aggregate posterior-predictive summaries for held-out observations.

Attributes#

observed_values:: Observed values kept on the original observation scale.
predictive_mean:: Posterior-predictive mean curves or vectors.
predictive_std:: Posterior-predictive standard deviations.
interval_lower, interval_upper:: Central predictive interval bounds for the requested level.
interval_coverage:: Per-dataset fraction of observed points covered by the predictive interval.
mean_absolute_error:: Mean absolute deviation between predictive mean and observed values.
metadata:: Auxiliary metadata such as split name, number of cases, and parameter space used during predictive generation.

observed_values: Mapping[str, numpy.ndarray]#

predictive_mean: Mapping[str, numpy.ndarray]#

predictive_std: Mapping[str, numpy.ndarray]#

interval_lower: Mapping[str, numpy.ndarray]#

interval_upper: Mapping[str, numpy.ndarray]#

interval_coverage: Mapping[str, float]#

mean_absolute_error: Mapping[str, float]#

mean_absolute_error_sigma: Mapping[str, float]#

median_interval_width_over_uncertainty: Mapping[str, float]#

crps: Mapping[str, float]#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.SimulationBasedCalibrationReport#

Rank-based SBC summary over a held-out set of observations. The rank is essentially the number of posterior samples that are less than the true α, and should be uniformly distributed if the posterior is calibrated.

Attributes#

ranks:: Integer rank of the ground-truth parameter within posterior samples for each held-out case and parameter dimension.
rank_histogram_counts:: Per-parameter SBC histogram counts.
posterior_means:: Posterior mean for each held-out case.
truths:: Ground-truth parameters paired with each posterior sample set.
coverages:: Coverage summaries for the requested nominal interval levels.
mean_rank:: Mean empirical rank for each parameter dimension.
normalized_mean_rank_error:: Absolute mean-rank error normalized by the expected average rank.
metadata:: Auxiliary run metadata such as split name and number of posterior draws.

ranks: numpy.ndarray#

rank_histogram_counts: numpy.ndarray#

posterior_means: numpy.ndarray#

truths: numpy.ndarray#

coverages: tuple[CoverageLevelReport, Ellipsis]#

mean_rank: numpy.ndarray#

normalized_mean_rank_error: numpy.ndarray#

metadata: Mapping[str, Any]#

petitRADTRANS.sbi.classify_posterior_informativeness(contraction: Any, sbc_normalized_rank_error: Any | None = None, parameter_names: Sequence[str] | None = None, *, contraction_informative: float = 0.9, contraction_overconfident: float = 0.1, sbc_miscalibrated: float = 0.3) → PosteriorInformativenessReport#

Classify each parameter by contraction and SBC calibration.

Verdicts:

'uninformative' – contraction >= contraction_informative: the marginal is effectively the prior.
'overconfident' – contraction <= contraction_overconfident and the SBC normalized mean-rank error >= sbc_miscalibrated: the posterior collapsed to a narrow, miscalibrated region (confidently wrong). This is exactly the failure a bare contraction metric mislabels as “constrained”.
'constrained' – informative and (no SBC supplied, or SBC acceptable).
'unknown' – contraction not finite.

sbc_normalized_rank_error is optional; without it the function can only distinguish informative from uninformative (an informative marginal is reported as 'constrained' but cannot be confirmed calibrated).

petitRADTRANS.sbi.generate_local_sensitivity_report(task: petitRADTRANS.sbi.task.SBITask, posterior_samples: Any, observation_blocks: Sequence[Any], posterior_log_probabilities: Any = None, parameter_space: str = 'physical', simulator: petitRADTRANS.sbi.simulator.RuntimeSimulator | None = None, quantile_levels: Sequence[float] = (0.1, 0.5, 0.9), finite_difference_relative_step: float = 0.001, finite_difference_std_fraction: float = 0.1, finite_difference_absolute_floor: float = 1e-05, max_step_reduction_attempts: int = 6, svd_relative_tolerance: float = 0.001, posterior_underexploited_ratio_threshold: float = 1.5, weak_sensitivity_fraction_threshold: float = 0.15, seed: int | None = None) → LocalSensitivityReport#

Diagnose local physical identifiability around representative posterior points.

The report evaluates a deterministic simulator Jacobian at representative posterior points, whitens it by observational uncertainty, and derives a Fisher-style local covariance approximation for each point.

petitRADTRANS.sbi.generate_posterior_predictive_report(task: petitRADTRANS.sbi.task.SBITask, posterior: Any, dataset_reader: petitRADTRANS.sbi.dataset.NormalizedObservationDatasetReader, split: petitRADTRANS.sbi.dataset.DatasetSplit = DatasetSplit.TEST, n_posterior_samples: int = 256, interval_level: float = 0.9, max_cases: int | None = None, seed: int | None = None, simulator: petitRADTRANS.sbi.simulator.RuntimeSimulator | None = None, n_predictive_forward_model_samples: int | None = None, checkpoint_directory: str | pathlib.Path | None = None, data_parallel: bool | None = None) → PosteriorPredictiveReport#

Generate posterior-predictive summaries for a held-out dataset split.

Parameters#

task:: SBI task defining parameter transforms and the simulator configuration.
posterior:: Trained posterior estimator used to sample held-out predictive draws.
dataset_reader:: Reader providing normalized observations and preprocessing metadata.
split:: Held-out split used for the predictive report.
n_posterior_samples:: Number of posterior draws generated per held-out observation.
interval_level:: Central predictive interval level reported for each dataset.
max_cases:: Optional cap on the number of held-out observations evaluated.
seed:: Optional base seed used to make predictive sampling reproducible.
simulator:: Optional runtime simulator override.
n_predictive_forward_model_samples:: Number of posterior draws passed through the forward model per held-out case. When None all n_posterior_samples draws are forwarded. Subsampling here is the primary lever for keeping the total number of petitRADTRANS calls to a manageable level when evaluating many held-out cases.
checkpoint_directory:: Optional directory for per-case checkpoints. When provided, each completed case is written to disk as a compressed .npz file and skipped on resume. This makes the expensive forward-model loop restartable after interruption.

Returns#

PosteriorPredictiveReport: Aggregate posterior-predictive summary over the requested split.

Notes#

Observations are normalized internally for posterior encoding but compared on the original observation scale in the returned report.

petitRADTRANS.sbi.generate_sbc_report(posterior: Any, dataset_reader: petitRADTRANS.sbi.dataset.NormalizedObservationDatasetReader, split: petitRADTRANS.sbi.dataset.DatasetSplit = DatasetSplit.TEST, n_posterior_samples: int = 256, batch_size: int = 32, parameter_space: str | None = None, levels: Sequence[float] = (0.5, 0.8, 0.95), max_cases: int | None = None, seed: int | None = None, data_parallel: bool | None = None) → SimulationBasedCalibrationReport#

Run SBC over a dataset reader using normalized observation batches.

Parameters#

posterior:: Trained posterior estimator exposing encode_observation and sample_posterior.
dataset_reader:: Reader yielding normalized held-out observations and matched parameter values.
split:: Dataset split used for the SBC evaluation.
n_posterior_samples:: Number of posterior draws generated per held-out case.
batch_size:: Reader batch size used during report generation.
parameter_space:: Optional parameter space override. Defaults to the posterior’s own configured parameter space.
levels:: Coverage levels summarized in the returned report.
max_cases:: Optional cap on the number of held-out observations evaluated.
seed:: Optional base seed used to generate reproducible posterior draws.

Returns#

SimulationBasedCalibrationReport: SBC summary computed from the requested dataset split.

petitRADTRANS.sbi.local_sensitivity_report_to_payload(report: LocalSensitivityReport) → dict[str, Any]#: Convert a local sensitivity report into a JSON-serializable payload.

petitRADTRANS.sbi.serialize_ood_diagnostic(ood_diagnostic: Any) → dict[str, Any] | None#: Serialize an OOD diagnostic to a JSON-safe dict (or None).

petitRADTRANS.sbi.summarize_posterior_contraction(posterior_samples: numpy.ndarray, parameter_names: list[str], retrieval_config: Any) → dict[str, dict[str, float | None]]#

Per-parameter posterior contraction relative to the prior.

contraction is the ratio of the posterior central 68% width to the prior central 68% width. Values near 1.0 mean the marginal is effectively the prior (uninformative); values well below 1.0 mean the observation constrains that parameter. This is the diagnostic that separates an informative posterior from a calibrated-but-uninformative one – something SBC rank histograms alone cannot reveal, since a posterior that returns the prior is trivially calibrated.

petitRADTRANS.sbi.summarize_posterior_samples(posterior_samples: numpy.ndarray, parameter_names: list[str], truth_parameter_values: Mapping[str, float] | None = None) → dict[str, dict[str, float | None]]#: Per-parameter posterior p16/median/p84, with optional truth values.

petitRADTRANS.sbi.summarize_predictive_report(predictive_report: PosteriorPredictiveReport) → dict[str, Any]#: Serialize a posterior-predictive report’s scalar metrics to JSON-safe dicts.

petitRADTRANS.sbi.summarize_training_diagnostics(posterior: Any) → dict[str, Any]#: Summarize per-epoch stability/inverse-error diagnostics of a trained flow.

petitRADTRANS.sbi.save_predictive_spectrum(output_path: pathlib.Path, retrieval_config: Any, predictive_report: Any) → pathlib.Path#

Write a posterior-predictive spectrum (wavelength, observed, uncertainty, predictive mean and interval) as a plain-text table.

The dataset is taken from the predictive report itself, so the saver is generic across observations.

exception petitRADTRANS.sbi.TaskCompatibilityError#

Bases: ValueError

Raised when an observation or artifact is incompatible with a task.

class petitRADTRANS.sbi.HDF5SimulationDatasetStore(chunk_size: int = 256)#

Bases: SimulationDatasetStore

HDF5-backed store for simulation corpora.

Stores all simulation data for a corpus in a single .h5 file, keeping file counts at 1 regardless of the number of simulations or splits. Requires h5py (already present in the jaxprt environment).

chunk_size = 256#

static _require_h5py() → None#

create_writer(manifest: SimulationDatasetManifest, mode: str = 'w') → HDF5SimulationDatasetWriter#: Create a writer for a new simulation dataset.

open(manifest_or_uri: SimulationDatasetManifest | str) → HDF5StoredSimulationDataset#: Open a stored dataset for training or evaluation.

class petitRADTRANS.sbi.SimulationDatasetStore#

Bases: abc.ABC

Backend-independent interface for reading and writing simulation data.

abstractmethod create_writer(manifest: SimulationDatasetManifest) → SimulationDatasetWriter#: Create a writer for a new simulation dataset.

abstractmethod open(manifest_or_uri: SimulationDatasetManifest | str) → Any#: Open a stored dataset for training or evaluation.

petitRADTRANS.sbi.generate_simulation_dataset(task: petitRADTRANS.sbi.task.SBITask, storage_uri: str, n_simulations: int, chunk_size: int = 256, split_counts: Mapping[DatasetSplit | str, int] | None = None, split_fractions: Mapping[DatasetSplit | str, float] | None = None, split_policy: SplitSamplingPolicy | str = SplitSamplingPolicy.SEQUENTIAL, split_seed: int | None = None, simulator: Any = None, include_preprocessing_metadata: bool = True, dataset_version: str = '0.1.0', resume: bool = False, data_parallel: bool | None = None, store_covariance: bool = False) → GeneratedSimulationDataset#

Generate and persist a simulation corpus in one call.

All simulation data is written to a single .h5 file via HDF5SimulationDatasetStore, keeping the file count at 1 regardless of the number of simulations.

Parameters#

data_parallel:: When True (or None with multiple JAX devices), the vmapped RT kernel is distributed across devices using jax.pmap. The effective per-iteration batch size is automatically scaled by the device count so each device processes the configured simulation_config.batch_size samples.
store_covariance:: When True the full covariance matrix for each simulated spectrum is written to disk. When False (default) only the covariance diagonal is stored under the covariance field to reduce storage pressure.

class petitRADTRANS.sbi.DatasetSplit#

Bases: str, enum.Enum

Named dataset partitions used during training and evaluation.

TRAIN = 'train'#

VALIDATION = 'validation'#

TEST = 'test'#

BENCHMARK = 'benchmark'#

class petitRADTRANS.sbi.NormalizedObservationDatasetReader#

Lightweight reader yielding normalized ObservationBlock batches for training.

dataset: HDF5StoredSimulationDataset#

preprocessing_metadata: petitRADTRANS.sbi.preprocessing.TaskPreprocessingMetadata#

target_observations: Mapping[str, Any] | None = None#

scale_check: str = 'raise'#

scale_check_max_orders_outside: float = 2.0#

scale_check_split: DatasetSplit#

noise_resampling: bool = False#

_split_cache: dict#

__post_init__() → None#

_validate_observation_scale() → None#

Run the pre-training observation/training flux-scale check.

Invoked automatically from __post_init__ whenever target_observations is supplied. Uses a (cheap) split’s raw spectra as the training reference and raises/warns per scale_check if the observation lies many orders of magnitude outside the simulated band.

iter_batches(split: DatasetSplit = DatasetSplit.TRAIN, batch_size: int = 32, shuffle: bool = False, seed: int | None = None, parameter_space: str = 'physical', encoder: Any = None) → Iterator[petitRADTRANS.sbi.flows.base.PosteriorBatch]#

Yield mini-batches of normalized observations and matched parameters.

Parameters#

split:: Dataset split to iterate over.
batch_size:: Number of samples yielded in each batch.
shuffle:: Whether to shuffle sample order within the requested split.
seed:: Optional random seed used when shuffle=True.
parameter_space:: Parameter representation returned in each batch. Supported values are 'physical', 'cube', and 'unconstrained'.
encoder:: Optional encoder used to convert block lists into dense embedding arrays before batches are yielded.

Returns#

Iterator[PosteriorBatch]: Iterator over batches containing parameters, observations, and small metadata dictionaries describing the batch provenance.

Notes#

Splits with fewer than _STREAMING_THRESHOLD rows are cached in RAM after the first read so that repeated epoch calls (e.g. validation) pay only slicing cost. Larger splits are streamed directly from disk one batch at a time to avoid loading the full dataset into memory.

_iter_batches_cached(split: DatasetSplit, batch_size: int, shuffle: bool, seed: int | None, parameter_space: str, encoder: Any) → Iterator[petitRADTRANS.sbi.flows.base.PosteriorBatch]#: Iterate using the full-split RAM cache (small splits).

_iter_batches_streaming(split: DatasetSplit, batch_size: int, shuffle: bool, seed: int | None, parameter_space: str, encoder: Any) → Iterator[petitRADTRANS.sbi.flows.base.PosteriorBatch]#

Iterate by streaming rows directly from disk (large splits).

When shuffle=True a global index permutation is computed in RAM (one integer per simulation row — negligible memory) and used to fetch HDF5 rows in sorted sub-windows of batch_size, satisfying h5py’s monotonic-index requirement while still presenting shuffled order to the training loop.

static _select_parameters(split_data: dict, parameter_space: str) → Any#: Return the parameter array for the requested coordinate space.

read_raw_spectral_split(split: DatasetSplit, parameter_space: str = 'physical', chunk_size: int = 8192) → dict[str, Any] | None#

Load raw (noise-free) spectra for training-time noise resampling.

Returns the un-normalized deterministic_values and uncertainties of every block, so the trainer can draw a fresh noise realization per batch (values = deterministic + eps * uncertainties) and normalize on the fly with normalize_spectral_block_matrix. Coordinates are pre-normalized here (they never change) and masks mirror the per-sample block semantics used by the frozen-noise path.

Returns None (with a printed warning) when the resampling path is unsupported: a non-spectral block, missing preprocessing statistics, or a dataset generated before deterministic_values were stored — the caller then falls back to the frozen-noise path.

petitRADTRANS.sbi.estimate_preloaded_split_memory_gib(n_samples: int, n_wavelength_points: int, n_parameters: int) → float#: Estimate host RAM (GiB) to eagerly preload one split as float32 arrays.

petitRADTRANS.sbi.infer_split_counts(total_simulations: int) → dict[str, int]#: Split a simulation budget into train/validation/test counts.

class petitRADTRANS.sbi.ObservationEncoder(embedding_dim: int = 128, spectrum_embedding_dim: int = 64, photometry_embedding_dim: int = 64, hidden_dim: int = 128, spectrum_encoder_type: str = 'convolution', n_wavelengths: int = 233, key: jax.Array | None = None)#

Bases: equinox.Module, petitRADTRANS.sbi.observation.AbstractObservationEncoder

Learned hierarchical encoder dispatching over modalities.

Parameters#

embedding_dim:: Size of the final joint observation embedding.
spectrum_embedding_dim:: Intermediate embedding size used by the spectral sub-encoder.
photometry_embedding_dim:: Intermediate embedding size used by the photometry sub-encoder.
hidden_dim:: Hidden width shared across the component encoders and aggregator.
key:: Optional JAX random key used to initialize all submodules.

Notes#

Spectrum and photometry blocks are handled by dedicated sub-encoders and then merged with a permutation-invariant aggregator. Unsupported modalities currently fall back to resized raw-value vectors.

spectrum_encoder: equinox.Module#

photometry_encoder: PhotometryPointEncoder#

aggregator: DatasetSetAggregator#

embedding_dim: int#

_encode_block(block: petitRADTRANS.sbi.observation.ObservationBlock) → jax.numpy.ndarray#

Encode one observation block with modality-aware dispatch.

Parameters#

block:: Observation block to encode.

Returns#

jnp.ndarray: Fixed-width embedding for the supplied block.

encode(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → petitRADTRANS.sbi.observation.EncodedObservation#

Encode one structured observation made of multiple blocks.

Parameters#

blocks:: Observation blocks associated with one target system.

Returns#

EncodedObservation: Aggregated embedding and light metadata describing the block set.

encode_stacked_batch(blocks_batch: list[list[petitRADTRANS.sbi.observation.ObservationBlock]]) → jax.numpy.ndarray#

Encode a batch of identically-structured observations using vmap.

All observations must share the same block structure (same number of blocks, same array shapes per block). This is always the case for SBI tasks where every observation is produced by the same forward model.

Parameters#

blocks_batch:: Outer list indexes samples; inner list holds the per-block observation data for one sample.

Returns#

jnp.ndarray: Float32 array of shape (n_samples, embedding_dim).

encode_from_prestacked(obs: Any) → jax.numpy.ndarray#

Encode a batch of observations from pre-stacked arrays.

Accepts a PreStackedObservations instance whose array fields have already been extracted from ObservationBlock objects outside the JAX JIT boundary. Only the vmapped XLA computation runs here, enabling the training step to be compiled once and reused across all batches.

Parameters#

obs:: Pre-stacked observation container with stacked_blocks (one (values, uncertainties, coordinates, mask, log_scale, absolute_values) tuple per block, each of shape (batch_size, n_wl) except the scalar log_scale arrays) and modalities (static tuple of modality value strings).

Returns#

jnp.ndarray: Float32 array of shape (batch_size, embedding_dim).

class petitRADTRANS.sbi.PhotometryPointEncoder(embedding_dim: int = 64, hidden_dim: int = 96, key: jax.Array | None = None)#

Bases: equinox.Module

Learned photometry encoder using per-point MLP features and pooling.

Parameters#

embedding_dim:: Size of the returned photometric embedding.
hidden_dim:: Hidden width of the per-point MLP.
key:: Optional JAX random key used for initialization.

Notes#

Each photometric point is represented by value, uncertainty, coordinate, and an inferred width feature before permutation-invariant pooling.

point_mlp: equinox.nn.MLP#

output_projection: equinox.nn.Linear#

embedding_dim: int#

encode_block(block: petitRADTRANS.sbi.observation.ObservationBlock) → jax.numpy.ndarray#

Encode one photometric observation block.

Parameters#

block:: Photometry-like observation block to encode.

Returns#

jnp.ndarray: Dense embedding for the supplied photometric block.

_encode_block_raw(values: jax.numpy.ndarray, uncertainties: jax.numpy.ndarray, coordinates: jax.numpy.ndarray, mask: jax.numpy.ndarray) → jax.numpy.ndarray#

Encode pre-processed block arrays (vmappable — no ObservationBlock input).

Parameters#

values, uncertainties, coordinates:: Float32 1-D arrays already produced by _as_vector. uncertainties and coordinates may have length 0.
mask:: Boolean 1-D array already produced by _safe_mask.

Returns#

jnp.ndarray: Dense photometric embedding of shape (embedding_dim,).

petitRADTRANS.sbi.load_posterior_estimator(input_directory: str) → PosteriorEstimator#: Load a saved posterior estimator without naming its concrete class.

class petitRADTRANS.sbi.EncoderConfig#

Observation-embedding encoder.

Configures the network that turns observation blocks into the fixed-length embedding the flow conditions on. Defaults to the convolution encoder that explicitly factors a spectrum’s amplitude from its shape (see petitRADTRANS.sbi.encoders.SpectralConvolutionEncoder), so the posterior cannot collapse to the prior on shape-determined parameters.

Attributes#

spectrum_encoder_type:: Spectral encoder architecture: 'convolution' (dilated-conv + attention with amplitude/shape factorization) or 'resmlp' (plain residual MLP over the whole normalized spectrum, the published-recipe embedding of Vasist et al. 2023; hidden_dim sets its first block width, 512 reproducing their 2x512/3x256/5x128 ladder).
embedding_dim:: Size of the final joint observation embedding consumed by the flow as its conditioning context.
spectrum_embedding_dim:: Intermediate embedding size of the spectral sub-encoder.
photometry_embedding_dim:: Intermediate embedding size of the photometric sub-encoder.
hidden_dim:: Hidden width of the encoder networks.

spectrum_encoder_type: str = 'convolution'#

embedding_dim: int = 80#

spectrum_embedding_dim: int = 64#

photometry_embedding_dim: int = 32#

hidden_dim: int = 128#

class petitRADTRANS.sbi.FlowConfig#

Normalizing-flow family and geometry.

Selects the conditional density-estimator hypothesis class and its shape. See petitRADTRANS.sbi.flows.transforms for the trade-offs between families.

Attributes#

family:: Flow family: 'neural_autoregressive' (default; most expressive, iterative inverse), 'spline' (rational-quadratic splines, analytic inverse), 'affine' (cheapest), or 'autoregressive' (affine autoregressive).
num_transforms:: Number of flow transform layers. More layers add capacity at higher cost.
hidden_dim:: Hidden width of the per-layer conditioner MLPs.
conditioner_depth:: Number of hidden layers in each conditioner MLP.
num_spline_bins:: Number of rational-quadratic bins per transformed dimension (family='spline'); more bins give finer per-dimension shape.
spline_bound:: Half-width of the finite spline support [-bound, bound] in latent coordinates (identity outside). Automatically raised to cover the logit-of-cube range for a uniform-prior cube flow.
base_distribution:: Flow base density, 'gaussian' (default) or 'logistic'. A logistic base matches logit-of-uniform coordinates so an untrained cube flow already equals the prior.
use_base_affine:: If True (spline family), prepend a diagonal-affine layer so the flow can set each parameter’s location/scale directly – helps directions constrained far more tightly than the prior.
autoregressive_transform_units:: Number of monotonic mixture units per scalar map (family='neural_autoregressive'); higher = more flexible scalar transforms.
neural_autoregressive_min_slope:: Positive floor on the neural-autoregressive monotonic slopes (keeps the map strictly increasing and well-conditioned).
neural_autoregressive_min_residual:: Identity-residual fraction in (0, 1) blended into each neural-autoregressive scalar map; keeps it invertible as the posterior sharpens.
neural_autoregressive_inverse_bisection_steps:: Bisection steps used to numerically invert each neural-autoregressive scalar map when sampling.
spline_small_bin_regularization_weight:: Weight of the penalty discouraging collapsed spline bins (spline flows only). None defaults to 1e-2 for family='spline' (falling back to spline_entropy_regularization_weight when that legacy weight is set), and 0.0 otherwise. Sharpening the conditional density is otherwise unpenalized by the NPE objective, so spline flows keep a mild anti-collapse penalty on by default.
spline_derivative_regularization_weight:: Weight of the penalty keeping spline slopes within the target derivative range. None defaults to 1e-2 for family='spline' and 0.0 otherwise.
spline_entropy_regularization_weight:: Legacy alias supplying the default for spline_small_bin_regularization_weight.
spline_entropy_floor:: Target floor (in [0, 1]) for the normalised spline-bin entropy.
spline_min_bin_ratio_target:: Minimum bin-size ratio (relative to a uniform bin) below which the small-bin penalty activates.
spline_min_derivative_target, spline_max_derivative_target:: Target range for spline boundary slopes. spline_max_derivative_target is also the hard slope cap that prevents unbounded-density collapse; it must exceed spline_min_derivative_target.

family: str = 'neural_autoregressive'#

num_transforms: int = 4#

hidden_dim: int = 256#

conditioner_depth: int = 3#

num_spline_bins: int = 8#

spline_bound: float = 10.0#

base_distribution: str = 'gaussian'#

use_base_affine: bool = False#

autoregressive_transform_units: int = 16#

neural_autoregressive_min_slope: float = 0.005#

neural_autoregressive_min_residual: float = 0.1#

neural_autoregressive_inverse_bisection_steps: int = 64#

spline_small_bin_regularization_weight: float | None = None#

spline_derivative_regularization_weight: float | None = None#

spline_entropy_regularization_weight: float = 0.0#

spline_entropy_floor: float = 0.6#

spline_min_bin_ratio_target: float = 0.2#

spline_min_derivative_target: float = 0.25#

spline_max_derivative_target: float = 5.0#

__post_init__() → None#

class petitRADTRANS.sbi.RefinementConfig#

Test-time, per-observation posterior refinement (any objective).

method="none" (default) uses the amortized posterior directly for the target inference – the behaviour of a plain NPE run. "map_laplace", "importance_sampling", or "elbo" refine the posterior for the specific target spectrum with the differentiable forward model, starting from the amortized fit; only the target inference uses the refined posterior (validation/SBC stay amortized).

Attributes#

method:: Refinement strategy: "none" (default; use the amortized posterior directly), "map_laplace" (MAP point estimate + Laplace/Gaussian-VI covariance), "importance_sampling" (exact-likelihood reweighting of amortized proposal draws – cannot collapse, and its sampling efficiency is a built-in trust diagnostic), or "elbo" (flow-ELBO refinement).
is_num_proposal_samples:: Amortized proposal draws weighted by the exact likelihood (method="importance_sampling").
is_forward_batch_size:: Forward-model chunk size for the importance weights (bounds device memory like the ELBO per-batch forward passes).
map_num_steps:: Number of optimization steps for the MAP point estimate (method="map_laplace").
map_learning_rate:: Learning rate for the MAP optimization.
map_beta_start:: Initial likelihood temperature for the MAP anneal.
map_beta_anneal_fraction:: Fraction of MAP steps over which the temperature anneals to 1.
map_covariance_method:: How the posterior covariance is estimated around the MAP point: "laplace" (Hessian/Gauss-Newton) or "gaussian_vi" (Gaussian variational fit).
map_vi_num_steps, map_vi_learning_rate, map_vi_num_samples:: Controls for the "gaussian_vi" covariance estimate (steps, learning rate, reparameterized samples per step).
test_time_validation:: If True, also run the actual test-time method (MAP+Laplace) on held-out simulations to measure its accuracy and calibration.
test_time_validation_cases:: Number of held-out cases used for that validation.
test_time_validation_samples:: Posterior samples drawn per validation case.
elbo_num_steps:: Number of optimization steps for flow-ELBO refinement (method="elbo").
elbo_learning_rate:: Learning rate for the flow-ELBO refinement.
elbo_num_samples:: Reparameterized samples per refinement step.
elbo_beta_anneal_fraction:: Fraction of refinement steps over which the likelihood temperature anneals (beta_start is taken from TrainingConfig.elbo_beta_start).

method: str = 'none'#

is_num_proposal_samples: int = 16384#

is_forward_batch_size: int = 256#

map_num_steps: int = 600#

map_learning_rate: float = 0.05#

map_beta_start: float = 0.001#

map_beta_anneal_fraction: float = 0.5#

map_covariance_method: str = 'laplace'#

map_vi_num_steps: int = 300#

map_vi_learning_rate: float = 0.01#

map_vi_num_samples: int = 32#

test_time_validation: bool = False#

test_time_validation_cases: int = 16#

test_time_validation_samples: int = 1024#

elbo_num_steps: int = 400#

elbo_learning_rate: float = 0.0001#

elbo_num_samples: int = 32#

elbo_beta_anneal_fraction: float = 0.5#

class petitRADTRANS.sbi.SBIConfig#

End-to-end configuration for petitRADTRANS.sbi.run_sbi().

Bundles the experiment definition (a built RetrievalConfig carrying the priors and model-generating function) with the grouped estimator configs and the workflow knobs, so a benchmark script declares only what differs from the library defaults and calls one function – no module-global / monkeypatch plumbing.

Attributes#

retrieval_config:: Built RetrievalConfig defining the free parameters/priors, observation datasets, and the differentiable model-generating function. The single largest scientific input.
output_root:: Root directory under which the shared datasets/ cache and per-run runs/<run_label>/ directories are created.
flow, encoder, training, stability, refinement:: The grouped estimator configs (FlowConfig, EncoderConfig, TrainingConfig, StabilityConfig, RefinementConfig).
n_simulations:: Total number of prior-predictive simulations to generate (split into train/validation/test).
simulation_batch_size:: Per-device simulator batch size during dataset generation.
seed:: Master random seed; per-stage seeds are derived from it for reproducibility.
n_posterior_samples:: Number of posterior draws for the delivered target inference.
n_predictive_forward_model_samples:: Number of posterior draws pushed through the forward model for posterior-predictive summaries (bounds the expensive pRT calls).
validation_max_cases:: Maximum held-out test cases used for the SBC and posterior-predictive validation reports.
run_label:: Explicit run directory name; when None a label is derived from the run config hash.
load_completed_model:: If True, skip training and load a previously trained, fingerprint- compatible posterior.
resume_predictive_checkpoints:: If True (with load_completed_model), reuse the source run directory and resume per-case predictive checkpoints.
generate_plots:: Whether to render the validation plot set.
plot_dpi:: Resolution of saved figures.
save_checkpoint_guardrails:: Whether to run inference for each saved checkpoint kind and persist a guardrail spectrum/plot/summary (makes checkpoint-selection regressions visible without retraining).
guardrail_posterior_samples, guardrail_predictive_forward_model_samples:: Smaller sample budgets used for the (cheaper) guardrail inference.
startup_scale_check:: Behaviour of the pre-training observation/training flux-scale guard: "raise", "warn", or "off".
verbose_training_diagnostics:: Whether the trainer emits per-epoch diagnostic artifacts/plots.
training_diagnostics_plot_interval:: Epoch interval between verbose training-diagnostic plots.
preprocessing_version:: Version tag stamped onto the fitted preprocessing metadata (participates in cache invalidation).
observation_value_constraint:: Per-dataset prior-predictive value constraint for rejection sampling. When None the workflow applies “value > 0” (positive emission flux / transit depth).
spectrum_x_label:: X-axis label for predictive-spectrum plots.
spectrum_y_label:: Y-axis label for predictive-spectrum plots; when None it is auto-derived (emission flux when scattering_in_emission is set, otherwise transit depth).
truth_parameter_values:: Optional known ground-truth parameter values, recorded in the summary and overlaid on plots when available.

retrieval_config: Any#

output_root: pathlib.Path#

flow: FlowConfig#

encoder: EncoderConfig#

training: TrainingConfig#

stability: StabilityConfig#

refinement: RefinementConfig#

n_simulations: int = 1000000#

simulation_batch_size: int = 256#

seed: int = 12345#

n_posterior_samples: int = 4096#

n_predictive_forward_model_samples: int = 256#

validation_max_cases: int = 128#

run_label: str | None = None#

load_completed_model: bool = False#

resume_predictive_checkpoints: bool = False#

generate_plots: bool = True#

plot_dpi: int = 160#

save_checkpoint_guardrails: bool = True#

guardrail_posterior_samples: int = 512#

guardrail_predictive_forward_model_samples: int = 96#

startup_scale_check: str = 'raise'#

verbose_training_diagnostics: bool = True#

training_diagnostics_plot_interval: int = 10#

preprocessing_version: str = '0.6.0'#

observation_value_constraint: Any = None#

noise_resample_during_training: bool = True#

spectrum_x_label: str = 'Wavelength [$\\mu$m]'#

spectrum_y_label: str | None = None#

truth_parameter_values: Mapping[str, float]#

to_run_config() → dict[str, Any]#: Flat, JSON-serializable provenance dict (also drives the run hash).

class petitRADTRANS.sbi.StabilityConfig#

Checkpoint-selection stability thresholds (flow invertibility / collapse).

During training, candidate checkpoints are scored against these thresholds; a checkpoint that violates them is treated as unstable/collapsed and penalised (or disqualified) when selecting the best model. Setting a field to None disables that check.

Attributes#

inverse_forward_max_abs_error:: Maximum tolerated inverse->forward round-trip error. A flow whose inverse disagrees with its forward map produces samples that do not match its own density.
inverse_forward_logdet_closure_max_abs_error:: Maximum tolerated inverse/forward log-det closure error.
cube_edge_hit_rate:: Maximum tolerated fraction of sampled coordinates pinned at the unit-cube edges (a signature of density collapse).
calibration_coverage_error:: Maximum tolerated deviation of the per-epoch empirical calibration probe (central-80%-interval coverage over held-out validation simulations) from its nominal 0.8. Selecting checkpoints on validation NLL alone deploys the sharpest epoch even when it has become overconfident; this gate makes “best” mean calibrated and low-loss.

inverse_forward_max_abs_error: float | None = 0.03#

inverse_forward_logdet_closure_max_abs_error: float | None = 0.03#

cube_edge_hit_rate: float | None = 0.02#

calibration_coverage_error: float | None = 0.25#

class petitRADTRANS.sbi.TrainingConfig#

Optimisation schedule and training objective.

Defines what is optimised and how the optimizer is scheduled. Note this is the public, grouped training config; the trainer’s internal TrainingSchemeConfig is separate.

Attributes#

num_epochs:: Maximum number of passes over the training split.
batch_size:: Number of simulations per optimization step.
learning_rate:: Peak AdamW learning rate.
min_learning_rate:: Floor of the learning-rate schedule.
weight_decay:: AdamW weight-decay coefficient.
gradient_clip_norm:: Global gradient-norm clip value; None disables clipping.
early_stopping_patience:: Non-improving epochs tolerated before stopping; None disables early stopping.
use_cosine_schedule:: Enable cosine learning-rate decay (otherwise constant after warmup).
warmup_fraction:: Fraction of the schedule spent linearly warming up the learning rate (used when warmup_epochs is not set).
warmup_epochs:: Absolute number of warmup epochs; overrides warmup_fraction when set.
lr_schedule_total_epochs:: Horizon (in epochs) over which the schedule decays; defaults to num_epochs when None.
training_objective:: 'npe' (default; conditional negative log-likelihood, robust) or 'elbo' (amortized variational inference, which requires parameter_space='cube' and a differentiable forward-model likelihood at fit time).
aux_scale_loss_weight:: Weight of the auxiliary loss training the scale head to predict each observation’s log-amplitude; 0 disables it.
invertibility_loss_weight:: Weight of the differentiable inverse round-trip / log-det-closure penalty added to the NPE loss; keeps an iterative-inverse flow (NAF) invertible as it sharpens. 0 (default) disables it.
invertibility_loss_n_samples:: Number of base samples used to estimate the invertibility penalty.
bnpe_balance_loss_weight:: Weight of the optional BNPE balance regularizer that pushes the amortized posterior toward conservative (non-overconfident) coverage. 0 (default) disables it; only active for parameter_space='cube'.
parameter_space:: Coordinates the flow is trained in: 'cube' (default; uniform-prior unit hypercube) or 'unconstrained'. ELBO requires 'cube'.
elbo_num_samples:: Reparameterized posterior draws per observation for the ELBO estimator (training_objective='elbo').
elbo_beta_start, elbo_beta_end:: Start/end of the likelihood-temperature (beta) anneal, each in (0, 1]. Ending below 1 keeps the amortized fit out of the noise-dominated, over-confident regime; test-time refinement then supplies the sharp final fit.
elbo_anneal_fraction:: Fraction of epochs over which beta anneals log-linearly from elbo_beta_start to elbo_beta_end before holding.

num_epochs: int = 96#

batch_size: int = 1024#

learning_rate: float = 0.0002#

min_learning_rate: float = 2e-06#

weight_decay: float = 0.001#

gradient_clip_norm: float | None = 2.0#

early_stopping_patience: int | None = 20#

use_cosine_schedule: bool = True#

warmup_fraction: float = 0.02#

warmup_epochs: float | None = 1.0#

lr_schedule_total_epochs: float | None = None#

training_objective: str = 'npe'#

aux_scale_loss_weight: float = 0.0#

invertibility_loss_weight: float = 0.0#

invertibility_loss_n_samples: int = 64#

bnpe_balance_loss_weight: float = 0.0#

parameter_space: str = 'cube'#

elbo_num_samples: int = 1#

elbo_beta_start: float = 0.001#

elbo_anneal_fraction: float = 0.6#

elbo_beta_end: float = 1.0#

petitRADTRANS.sbi.build_conditional_flow_posterior(*, parameter_dim: int, n_wavelengths: int, flow: FlowConfig | None = None, encoder: EncoderConfig | None = None, training: TrainingConfig | None = None, stability: StabilityConfig | None = None, seed: int = 0, checkpoint_directory: str | None = None, checkpoint_backend: str = 'equinox', resume_from_checkpoint: bool = False, verbose_diagnostics: bool = False, diagnostics_output_directory: str | None = None, diagnostics_plot_interval: int = 1, task_metadata: Mapping[str, Any] | None = None) → petitRADTRANS.sbi.flows.ConditionalFlowPosterior#

Build a ConditionalFlowPosterior from grouped config dataclasses.

Expands the themed config groups back into the estimator’s ~50 flat constructor arguments, so the common path is a couple of lines instead of a long keyword ceremony. Any config left as None uses that group’s defaults.

Parameters#

parameter_dim:: Number of inferred free parameters (e.g. len(task.parameter_names)).
n_wavelengths:: Number of spectral points per observation block (sizes the encoder); must match the task’s observation schema.
flow, encoder, training, stability:: Grouped configs; None substitutes the corresponding default dataclass.
seed:: Base random seed for flow/encoder initialization and sampling.
checkpoint_directory:: Directory for resumable trainer checkpoints; None disables checkpointing.
checkpoint_backend:: Checkpoint backend ('equinox' here by default, vs the estimator’s own 'auto' default).
resume_from_checkpoint:: Whether fit should resume from the latest checkpoint.
verbose_diagnostics:: Emit per-epoch diagnostic artifacts/plots during training.
diagnostics_output_directory:: Directory for the verbose diagnostic artifacts.
diagnostics_plot_interval:: Epoch interval between verbose diagnostic plots.
task_metadata:: Optional metadata persisted alongside the trained model (e.g. the task name).

Returns#

ConditionalFlowPosterior: An untrained estimator ready for ConditionalFlowPosterior.fit().

petitRADTRANS.sbi.run_sbi(config: petitRADTRANS.sbi.config.SBIConfig) → pathlib.Path#

Run the full amortized-SBI validation workflow and return the run directory.

This is the single end-to-end entry point: given one SBIConfig (the built RetrievalConfig plus the grouped estimator/workflow knobs), it builds the task, obtains a simulation dataset, trains or loads the conditional flow posterior, validates it, runs inference on the real target observation, and writes every artifact to a self-contained run directory. It is designed to be deterministic (seeded) and resumable, and to fail loudly on stale or mismatched inputs via the task fingerprint.

Workflow stages#

Task construction. Build an SBITask from config.retrieval_config, attaching the noise model and the per-dataset prior-predictive value constraint (default: strictly positive flux / transit depth). The task’s fingerprint is derived from the priors, units, and forward-model source code.
Run/dataset resolution. Resolve the output layout under config.output_root: a fingerprint-keyed dataset path under datasets/ and a per-run directory under runs/<run_label>. When config.load_completed_model is set, locate the source run (explicit run_label or the latest completed run) and, if config.resume_predictive_checkpoints is set, reuse its directory.
Dataset generation or reuse. Reuse an existing .h5 dataset when its fingerprint matches (rebuilding the preprocessing sidecar if missing); otherwise generate config.n_simulations simulations split into train/validation/test, fitting preprocessing on the train split. Dataset resume is disabled when value-constraint rejection sampling is active (rejection changes RNG consumption).
Reader construction. Wrap the dataset in a NormalizedObservationDatasetReader, which on construction runs the observation/training flux-scale guard and aborts before training on a units/scale mismatch or a stale/corrupt dataset.
Train or load the posterior. Either load a previously trained, fingerprint-compatible posterior, or build and fit a ConditionalFlowPosterior (the npe objective, or elbo with a lazily-built differentiable forward-model likelihood), then save it.
Checkpoint guardrails. For each saved checkpoint kind (best-selection / best-loss / best-stable), run inference on the target and persist a predictive spectrum, plot, and summary so a checkpoint-selection regression is visible without retraining.
Held-out validation. Compute a simulation-based calibration (SBC) report and a posterior-predictive report on the held-out test split.
Optional test-time refinement (any objective). When config.refinement.method is set, sharpen the posterior for the specific target spectrum with the differentiable forward model – map_laplace (MAP + Laplace / Gaussian-VI, optionally validated on held-out cases) or elbo (flow-ELBO). Validation/SBC stay amortized; only the delivered target inference uses the refined posterior.
Target inference. Run AmortizedRetrieval on the real observation to produce posterior samples (mapped to physical parameter space), a posterior-predictive summary, and an OOD diagnostic.
Reporting & persistence. Classify per-parameter informativeness (contraction cross-checked against SBC calibration), optionally render the validation plot set, and write the samples, predictive spectrum, summary JSON, and run manifest; finally mark the run completed and update the latest pointer.

Parameters#

config:: Fully-populated SBIConfig. It carries the experiment definition (the built RetrievalConfig with priors and model-generating function), the grouped estimator configs (flow, encoder, training, stability, refinement), and the workflow knobs (simulation count, sample counts, plotting, guardrails, resume/load flags, labels, and optional known truths).

Returns#

pathlib.Path: The run directory (output_root/runs/<run_label>) containing all artifacts produced by the run.

Side effects#

Writes to disk under config.output_root:

datasets/<task>_<n>_<fingerprint>.h5 and a .preprocessing.json sidecar (shared across runs of the same task family).
runs/<run_label>/ containing posterior_model/ (weights + metadata.json), checkpoints/ and predictive_checkpoints/, the target *.npy samples/log-probabilities, the target predictive spectrum, amortized_summary.json, the checkpoint-guardrail artifacts, the optional test_time_validation_report.json and validation plots, and a run manifest. The latest run pointer under output_root is updated on completion.

Raises#

FileExistsError: If the target run directory already exists and is not being reused for a predictive-checkpoint resume.
RuntimeError: If a cached dataset or a loaded posterior fingerprint does not match the current task, or if inference unexpectedly returns non-physical samples.
ValueError: If the ELBO objective is requested for a task whose runtime has more than one model group or observation (not yet supported by the ELBO likelihood).

Notes#

Progress is reported to stdout throughout (device list, split counts, memory estimates, per-stage status, validation metrics, and the per-parameter informativeness verdict). The function is the orchestration layer only: the science lives in the simulator, encoder, flow, objective, and calibration modules it calls.

class petitRADTRANS.sbi.AmortizedRetrieval(task: petitRADTRANS.sbi.task.SBITask, posterior_estimator: petitRADTRANS.sbi.flows.PosteriorEstimator, simulator: petitRADTRANS.sbi.simulator.RuntimeSimulator | None = None, preprocessing_metadata: petitRADTRANS.sbi.preprocessing.TaskPreprocessingMetadata | Mapping[str, Any] | None = None)#

Serve trained SBI models through a retrieval-like interface.

Parameters#

task:: SBI task compatible with the trained posterior.
posterior_estimator:: Trained posterior estimator used for encoding and sampling.
simulator:: Optional simulator override used for posterior-predictive generation.
preprocessing_metadata:: Optional preprocessing metadata or payload. When omitted the inference service attempts to recover the preprocessing payload saved inside the posterior artifact.

Notes#

Raw user-facing observation blocks are normalized internally before encoding when preprocessing metadata is available, but posterior-predictive comparisons remain on the original observation scale.

task#

posterior_estimator#

simulator = None#

preprocessing_metadata#

classmethod from_retrieval_config(retrieval_config: Any, *, n_wavelengths: int, flow: petitRADTRANS.sbi.config.FlowConfig | None = None, encoder: petitRADTRANS.sbi.config.EncoderConfig | None = None, training: petitRADTRANS.sbi.config.TrainingConfig | None = None, stability: petitRADTRANS.sbi.config.StabilityConfig | None = None, simulation_config: petitRADTRANS.sbi.task.SimulationConfig | None = None, seed: int = 0, checkpoint_directory: str | None = None, verbose_diagnostics: bool = False, diagnostics_output_directory: str | None = None, task_metadata: Mapping[str, Any] | None = None) → AmortizedRetrieval#

Assemble an amortized retrieval from a RetrievalConfig + configs.

This collapses the former build_run_config / build_posterior_estimator ceremony into one call: it builds the SBITask and an untrained ConditionalFlowPosterior from the grouped config dataclasses. Call fit() with a normalized simulation dataset reader to train, then infer() on an observation.

n_wavelengths is the number of spectral points per observation block (e.g. int(numpy.size(data.wavelengths))); it sizes the encoder.

fit(dataset: Any, **fit_kwargs: Any) → AmortizedRetrieval#

Train the posterior estimator on a normalized simulation dataset reader.

Thin pass-through to posterior_estimator.fit so the common path reads AmortizedRetrieval.from_retrieval_config(...).fit(reader).infer(obs). Returns self for chaining. (Dataset generation + preprocessing remain explicit, deliberately separate steps.)

_resolve_preprocessing_metadata(preprocessing_metadata: petitRADTRANS.sbi.preprocessing.TaskPreprocessingMetadata | Mapping[str, Any] | None) → petitRADTRANS.sbi.preprocessing.TaskPreprocessingMetadata | None#

prepare_observation_blocks(observation_blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → list[petitRADTRANS.sbi.observation.ObservationBlock]#

Normalize user-facing observation blocks when preprocessing metadata is available.

Parameters#

observation_blocks:: Raw observation blocks provided by the caller.

Returns#

list[ObservationBlock]: Either the original blocks or normalized copies, depending on whether preprocessing metadata is available.

infer(observation_blocks: list[petitRADTRANS.sbi.observation.ObservationBlock], n_posterior_samples: int = 1000, include_posterior_predictive: bool = False, posterior_predictive_interval_level: float = 0.9, n_predictive_forward_model_samples: int | None = None, seed: int | None = None) → AmortizedRetrievalResult#

Infer a posterior for one structured observation.

Parameters#

observation_blocks:: Raw observation blocks for one target system. These contain the observed fluxes, wavelengths, and uncertainties, and are normalized internally when preprocessing metadata is available.
n_posterior_samples:: Number of posterior draws to generate.
include_posterior_predictive:: Whether to also generate a posterior-predictive summary.
posterior_predictive_interval_level:: Predictive interval level used when generating the optional posterior-predictive report.
n_predictive_forward_model_samples:: Number of posterior draws passed through the forward model when generating the posterior-predictive summary. When None all n_posterior_samples draws are forwarded. Set this to a small value (e.g. 50–200) to avoid a multi-hour hang when n_posterior_samples is large.
seed:: Optional seed for posterior and posterior-predictive sampling.

Returns#

AmortizedRetrievalResult: Posterior samples and any optional predictive metadata.

posterior_predictive(posterior: petitRADTRANS.sbi.flows.PosteriorSamples, observation_blocks: list[petitRADTRANS.sbi.observation.ObservationBlock], interval_level: float = 0.9, n_predictive_forward_model_samples: int | None = None, seed: int | None = None) → petitRADTRANS.sbi.calibration.PosteriorPredictiveReport#

Generate posterior-predictive draws for a fitted observation.

Parameters#

posterior:: Posterior samples associated with one observed system.
observation_blocks:: Original observation blocks used for user-facing comparison.
interval_level:: Central predictive interval level to report.
n_predictive_forward_model_samples:: Number of posterior draws passed through the forward model. When None all samples in posterior are used. Set this to a small value (e.g. 50–200) to keep the number of expensive petitRADTRANS calls manageable.
seed:: Optional seed for predictive simulation.

Returns#

PosteriorPredictiveReport: Predictive summary for the supplied posterior samples.

diagnose_domain(observation_blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → OODDiagnostic | None#

Estimate whether the observation is in-distribution for the model.

Parameters#

observation_blocks:: Observation blocks to diagnose.

Returns#

OODDiagnostic | None: Robust support-distance diagnostic derived from preprocessing statistics, or None when preprocessing metadata is unavailable.

class petitRADTRANS.sbi.AmortizedRetrievalResult#

Return type for amortized inference queries.

Attributes#

posterior:: Posterior samples and any attached diagnostics.
posterior_predictive:: Optional posterior-predictive report for the same observation.
ood_diagnostic:: Optional in/out-of-distribution assessment.
metadata:: Additional metadata such as preprocessing usage during inference.

posterior: petitRADTRANS.sbi.flows.PosteriorSamples#

posterior_predictive: Any = None#

ood_diagnostic: OODDiagnostic | None = None#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.OODDiagnostic#

Describe whether an observation is inside the training support.

Attributes#

score:: Scalar out-of-distribution score.
threshold:: Optional threshold used to convert the score into a pass/fail decision.
passed:: Optional boolean indicating whether the observation passed the OOD test.
metadata:: Auxiliary diagnostic metadata.

score: float#

threshold: float | None = None#

passed: bool | None = None#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.ObservationBlock#

Representation of one (simulated or real) observation.

The modality can be ‘spectrum’, ‘photometry’, ‘time_series’, or ‘auxiliary’. However, for SBI usage we’re generally only concerned with spectrum.

Attributes:

name:: Stable identifier of the block within a task.
modality:: Description of the type of observation, can be one of ‘spectrum’, ‘photometry’, ‘time_series’, or ‘auxiliary’.
values:: Observed values after any task-level preprocessing. This stores either transit depths or fluxes, modulo some processing.
uncertainties:: Optional per-element uncertainty representation.
coordinates:: Optional coordinate arrays such as wavelengths or timestamps.
mask:: Optional mask applied to the values.
metadata:: Additional instrument and preprocessing metadata.

name: str#

modality: ObservationModality#

values: Any#

uncertainties: Any = None#

coordinates: Any = None#

mask: Any = None#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.AbstractObservationEncoder#

Bases: abc.ABC

Transform structured observation blocks into model-ready embeddings. This allows an observation to be represented in a shared latent space, with the encoders implemented in encorders.py.

abstractmethod encode(blocks: list[ObservationBlock]) → EncodedObservation#: Encode a list of observation blocks into a shared representation.

batch_encode(observations: list[list[ObservationBlock]]) → list[EncodedObservation]#: Encode multiple observations using repeated single-item encoding.

class petitRADTRANS.sbi.ObservationModality#

Bases: str, enum.Enum

Supported observation block types for SBI conditioning.

SPECTRUM = 'spectrum'#

PHOTOMETRY = 'photometry'#

TIME_SERIES = 'time_series'#

AUXILIARY = 'auxiliary'#

petitRADTRANS.sbi.build_observation_block(name: str, modality: ObservationModality | str, values: Any, uncertainties: Any = None, coordinates: Any = None, mask: Any = None, metadata: Mapping[str, Any] | None = None) → ObservationBlock#: Build one observation block with modality normalization.

petitRADTRANS.sbi.build_observation_block_batch(observation_payloads: Mapping[str, Mapping[str, Any]], modalities: Mapping[str, str]) → list[list[ObservationBlock]]#: Build observation blocks for each sample in a batched payload.

petitRADTRANS.sbi.build_observation_blocks_from_sample(observation_payloads: Mapping[str, Mapping[str, Any]], modalities: Mapping[str, str], sample_index: int) → list[ObservationBlock]#: Build modality-aware observation blocks for one simulated sample.

petitRADTRANS.sbi.build_raw_observation_blocks(retrieval_config: Any) → list#: Build raw (un-normalized) spectrum ObservationBlocks for every dataset in a retrieval configuration – the target observation(s) for amortized inference and the startup scale guard.

class petitRADTRANS.sbi.ConditionalAutoregressiveFlowPosterior(*args: Any, **kwargs: Any)#

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the autoregressive flow backend.

class petitRADTRANS.sbi.ConditionalFlowPosterior(parameter_dim: int, embedding_dim: int = 128, num_coupling_layers: int = 4, hidden_dim: int = 128, conditioner_depth: int = 2, autoregressive_transform_units: int = 16, neural_autoregressive_min_slope: float = 0.001, neural_autoregressive_min_residual: float = 0.05, neural_autoregressive_inverse_bisection_steps: int = 48, learning_rate: float = 0.001, batch_size: int = 32, num_epochs: int = 5, parameter_space: str = 'unconstrained', flow_family: str = 'spline', num_spline_bins: int = 8, spline_bound: float = 10.0, base_distribution: str = 'gaussian', use_base_affine: bool = False, training_objective: str = 'npe', elbo_num_samples: int = 1, elbo_beta_start: float = 0.001, elbo_beta_end: float = 1.0, elbo_anneal_fraction: float = 0.6, early_stopping_patience: int | None = None, early_stopping_min_delta: float = 0.0, checkpoint_directory: str | None = None, checkpoint_backend: str = 'auto', resume_from_checkpoint: bool = False, gradient_clip_norm: float | None = 1.0, aux_scale_loss_weight: float = 0.0, spline_small_bin_regularization_weight: float | None = None, spline_min_bin_ratio_target: float = 0.2, spline_entropy_regularization_weight: float = 0.0, spline_derivative_regularization_weight: float = 0.0, spline_entropy_floor: float = 0.6, spline_min_derivative_target: float = 0.25, spline_max_derivative_target: float = 5.0, invertibility_loss_weight: float = 0.0, invertibility_loss_n_samples: int = 64, bnpe_balance_loss_weight: float = 0.0, weight_decay: float = 0.0, use_cosine_schedule: bool = False, warmup_fraction: float = 0.02, warmup_epochs: float | None = None, min_learning_rate: float = 1e-06, lr_schedule_total_epochs: float | None = None, stable_inverse_forward_max_abs_error_threshold: float | None = 0.0001, stable_inverse_forward_logdet_closure_max_abs_error_threshold: float | None = 0.0001, stable_cube_edge_hit_rate_threshold: float | None = 0.0, stable_calibration_coverage_error_threshold: float | None = 0.25, spectrum_encoder_type: str = 'convolution', n_wavelengths: int = 233, spectrum_embedding_dim: int = 64, photometry_embedding_dim: int = 64, encoder_hidden_dim: int | None = None, seed: int = 0, task_metadata: Mapping[str, Any] | None = None, verbose_diagnostics: bool = False, diagnostics_output_directory: str | None = None, diagnostics_plot_interval: int = 1)#

Bases: petitRADTRANS.sbi.flows.base.PersistentPosteriorEstimator

Concrete amortized posterior using a conditional flow backend.

Parameters are grouped below by concern. Most callers should construct this object through petitRADTRANS.sbi.config.build_conditional_flow_posterior() (which expands the grouped config dataclasses), but every knob is also a direct constructor argument.

Parameters#

parameter_dim:: Number of inferred free parameters represented by the posterior. The only required argument.
flow_family:: Conditional density-transform family (see petitRADTRANS.sbi.flows.transforms for the trade-offs). One of 'spline' (default; rational-quadratic splines, analytic inverse), 'affine' (cheapest, smooth posteriors), 'autoregressive' (affine autoregressive), or 'neural_autoregressive' (most expressive, iterative inverse).
num_coupling_layers:: Number of flow transform layers. More layers increase capacity at higher cost.
hidden_dim:: Hidden width of the flow conditioner MLPs (and the default for the encoder’s hidden width).
conditioner_depth:: Number of hidden layers in each flow conditioner MLP.
num_spline_bins:: Number of rational-quadratic bins per transformed dimension (flow_family='spline'). More bins = finer per-dimension shape.
spline_bound:: Half-width of the finite spline support [-bound, bound] in latent coordinates; outside it the map is the identity. For a uniform-prior cube flow this is automatically raised to cover the logit-of-cube range.
base_distribution:: Flow base density, 'gaussian' (default) or 'logistic'. 'logistic' matches logit-of-uniform coordinates so an untrained cube flow already equals the prior.
use_base_affine:: If True (spline family), prepend a ConditionalDiagonalAffine layer so the flow can set each parameter’s location/scale directly – helps directions constrained far more tightly than the prior.
autoregressive_transform_units:: Number of monotonic mixture units per scalar map (flow_family='neural_autoregressive'); higher = more flexible scalar transforms.
neural_autoregressive_min_slope:: Positive floor on the neural-autoregressive monotonic slopes (keeps the map strictly increasing and well-conditioned).
neural_autoregressive_min_residual:: Identity-residual fraction in (0, 1) blended into each neural autoregressive scalar map; keeps it invertible as the posterior sharpens.
neural_autoregressive_inverse_bisection_steps:: Bisection steps used to numerically invert each neural-autoregressive scalar map when sampling.
parameter_space:: Coordinates the flow is trained in: 'physical', 'cube' (unit hypercube / uniform-prior coordinates, with the change-of-variables correction applied), or 'unconstrained' (default). 'elbo' training requires 'cube'.
training_objective:: 'npe' (default) trains the conditional negative log-likelihood; 'elbo' trains an amortized variational objective and requires parameter_space='cube' plus an elbo_log_likelihood_fn at fit().
learning_rate:: Peak optimizer (AdamW) learning rate.
batch_size:: Number of simulations per optimization step.
num_epochs:: Maximum number of passes over the training split.
weight_decay:: AdamW weight-decay coefficient.
gradient_clip_norm:: Global gradient-norm clip value; None disables clipping.
use_cosine_schedule:: Enable cosine learning-rate decay (otherwise the rate is constant after warmup).
warmup_fraction:: Fraction of the schedule spent linearly warming up the learning rate (used when warmup_epochs is not given).
warmup_epochs:: Absolute number of warmup epochs; overrides warmup_fraction when set.
min_learning_rate:: Floor of the learning-rate schedule.
lr_schedule_total_epochs:: Horizon (in epochs) over which the schedule decays; defaults to num_epochs.
early_stopping_patience:: Number of non-improving epochs tolerated before stopping; None disables early stopping.
early_stopping_min_delta:: Minimum validation improvement that counts as progress for early stopping.
elbo_num_samples:: Reparameterized posterior draws per observation used to estimate the ELBO (training_objective='elbo').
elbo_beta_start, elbo_beta_end:: Start/end of the likelihood-temperature (beta) anneal, each in (0, 1]. Ending below 1 keeps the amortized fit out of the noise-dominated, over-confident regime.
elbo_anneal_fraction:: Fraction of epochs over which beta anneals log-linearly from elbo_beta_start to elbo_beta_end before holding.
aux_scale_loss_weight:: Weight of the auxiliary loss training the scale_head to predict each observation’s log-amplitude; 0 disables it.
spline_small_bin_regularization_weight:: Weight of the penalty discouraging collapsed spline bins. Defaults to spline_entropy_regularization_weight when None.
spline_min_bin_ratio_target:: Minimum bin-size ratio (relative to a uniform bin) below which the small-bin penalty activates.
spline_entropy_regularization_weight:: Legacy alias; supplies the default for spline_small_bin_regularization_weight.
spline_derivative_regularization_weight:: Weight of the penalty keeping spline slopes within the target derivative range.
spline_entropy_floor:: Target floor (in [0, 1]) for the normalised spline-bin entropy.
spline_min_derivative_target, spline_max_derivative_target:: Target range for spline boundary slopes. spline_max_derivative_target is also used as the hard slope cap that prevents unbounded-density collapse; max must exceed min.
invertibility_loss_weight:: Weight of the differentiable inverse round-trip / log-det-closure penalty that keeps an iterative-inverse flow invertible as it sharpens; 0 disables it.
invertibility_loss_n_samples:: Number of base samples used to estimate the invertibility penalty.
bnpe_balance_loss_weight:: Weight of the optional BNPE balance regularizer (Delaunoy et al.) that pushes the posterior toward conservative (non-overconfident) coverage; 0 (default) disables it. Only active for parameter_space='cube', where the prior log-density is exactly zero.
stable_inverse_forward_max_abs_error_threshold:: Checkpoint-selection gate: maximum tolerated inverse->forward round-trip error; None disables this check.
stable_inverse_forward_logdet_closure_max_abs_error_threshold:: Checkpoint-selection gate on the inverse/forward log-det closure error; None disables it.
stable_cube_edge_hit_rate_threshold:: Checkpoint-selection gate on the fraction of samples pinned at the cube edges (collapse indicator); None disables it.
stable_calibration_coverage_error_threshold:: Checkpoint-selection gate on the per-epoch empirical-calibration probe: maximum tolerated deviation of the central-80%-interval coverage from 0.8 over held-out validation simulations. Guards against selecting a sharper-but-overconfident epoch on validation NLL alone; None disables it.
spectrum_encoder_type:: Spectral encoder architecture: 'convolution' (amplitude/shape factorizing).
n_wavelengths:: Number of spectral points per observation block; must match the task and sizes the encoder.
spectrum_embedding_dim, photometry_embedding_dim:: Intermediate embedding sizes of the spectral and photometric sub-encoders.
embedding_dim:: Size of the final joint observation embedding consumed by the flow as its conditioning context.
encoder_hidden_dim:: Hidden width of the encoder networks; defaults to hidden_dim when None.
checkpoint_directory:: Directory for resumable trainer checkpoints; None disables checkpointing.
checkpoint_backend:: Checkpoint backend. 'auto' prefers Orbax and falls back to Equinox serialization.
resume_from_checkpoint:: Whether fit() should resume from the latest checkpoint in checkpoint_directory.
seed:: Base random seed for flow/encoder initialization and posterior sampling.
task_metadata:: Optional user metadata persisted alongside the trained model.
verbose_diagnostics:: Emit per-epoch diagnostic artifacts and plots during training.
diagnostics_output_directory:: Directory for the verbose diagnostic artifacts; None uses the trainer default.
diagnostics_plot_interval:: Epoch interval between verbose diagnostic plots.

Notes#

The posterior stores task fingerprinting, observation schema, and preprocessing payload information when those are available from the training dataset reader. That metadata is later reused by inference and artifact registration paths.

estimator_family = 'conditional_flow'#

embedding_dim = 128#

num_coupling_layers = 4#

hidden_dim = 128#

conditioner_depth = 2#

autoregressive_transform_units = 16#

neural_autoregressive_min_slope#

neural_autoregressive_min_residual#

neural_autoregressive_inverse_bisection_steps = 48#

learning_rate#

batch_size = 32#

num_epochs = 5#

flow_family = ''#

effective_flow_family = ''#

base_distribution = ''#

use_base_affine = False#

training_objective = ''#

elbo_num_samples#

elbo_beta_start#

elbo_beta_end#

elbo_anneal_fraction#

num_spline_bins = 8#

early_stopping_patience = None#

early_stopping_min_delta#

checkpoint_directory = None#

checkpoint_backend = 'auto'#

resume_from_checkpoint = False#

gradient_clip_norm = 1.0#

aux_scale_loss_weight#

spline_entropy_regularization_weight#

spline_small_bin_regularization_weight#

spline_derivative_regularization_weight#

spline_min_bin_ratio_target#

spline_entropy_floor#

spline_min_derivative_target#

spline_max_derivative_target#

invertibility_loss_weight#

invertibility_loss_n_samples = 64#

bnpe_balance_loss_weight#

weight_decay#

use_cosine_schedule = False#

warmup_fraction#

warmup_epochs = None#

min_learning_rate#

lr_schedule_total_epochs = None#

stable_inverse_forward_max_abs_error_threshold = None#

stable_inverse_forward_logdet_closure_max_abs_error_threshold = None#

stable_cube_edge_hit_rate_threshold = None#

stable_calibration_coverage_error_threshold = None#

spectrum_encoder_type = ''#

n_wavelengths = 233#

spectrum_embedding_dim = 64#

photometry_embedding_dim = 64#

encoder_hidden_dim#

verbose_diagnostics = False#

diagnostics_output_directory = None#

diagnostics_plot_interval = 1#

model#

_build_flow(key: jax.Array) → Any#

_validation_diagnostics(model: petitRADTRANS.sbi.flows.objectives._PosteriorModel, batch: petitRADTRANS.sbi.flows.base.PosteriorBatch) → dict[str, float]#

Per-checkpoint flow invertibility / collapse diagnostics.

Thin wrapper around petitRADTRANS.sbi.flows.objectives.compute_validation_diagnostics(), forwarding this estimator’s verbose_diagnostics flag.

refine_posterior_on_observation(*args: Any, **kwargs: Any) → petitRADTRANS.sbi.flows.base.PosteriorSamples#

Test-time flow-ELBO posterior refinement for one observation.

Thin wrapper around petitRADTRANS.sbi.flows.refinement.refine_posterior_on_observation(), which holds the implementation.

map_laplace_posterior_on_observation(*args: Any, **kwargs: Any) → petitRADTRANS.sbi.flows.base.PosteriorSamples#

Test-time MAP + Laplace / Gaussian-VI posterior refinement.

Thin wrapper around petitRADTRANS.sbi.flows.refinement.map_laplace_posterior_on_observation(), which holds the implementation.

importance_sampling_posterior_on_observation(*args: Any, **kwargs: Any) → petitRADTRANS.sbi.flows.base.PosteriorSamples#

Test-time likelihood importance sampling on the amortized posterior.

Thin wrapper around petitRADTRANS.sbi.flows.refinement.importance_sampling_posterior_on_observation(), which holds the implementation.

fit(dataset: Any, *, elbo_log_likelihood_fn: Callable[[jax.numpy.ndarray, Any], jax.numpy.ndarray] | None = None, elbo_num_samples: int | None = None) → petitRADTRANS.sbi.flows.base.TrainingArtifacts#

Train the posterior on one normalized simulation dataset reader.

Parameters#

dataset:: Reader-like object that yields PosteriorBatch instances via iter_batches and exposes dataset manifest metadata when available.
elbo_log_likelihood_fn:: Required when training_objective='elbo'. A differentiable (theta_cube, observations) -> (batch,) callable returning the Gaussian observational log-likelihood through the forward model. See _make_elbo_loss(). Ignored for the default NPE objective.
elbo_num_samples:: Number of reparameterized posterior draws per observation for the ELBO estimator (default 1).

Returns#

TrainingArtifacts: Training history, validation metrics, and trainer metadata for the completed optimization run.

Notes#

If the reader exposes preprocessing metadata or a manifest fingerprint, those are cached on the posterior so they can be saved and reused by the inference and artifact layers.

_build_estimator_config() → dict[str, Any]#: Return backend-specific configuration for metadata persistence.

_build_serialized_metadata(artifact_metadata) → dict[str, Any]#

static _resolve_estimator_config(metadata: Mapping[str, Any]) → dict[str, Any]#

hydrate_loaded_metadata(metadata: Mapping[str, Any]) → None#

classmethod from_serialized_metadata(metadata: Mapping[str, Any]) → ConditionalFlowPosterior#: Rebuild an estimator instance from persisted metadata only.

save_backend_state(output_path: pathlib.Path) → None#: Persist backend-specific model state into the output directory.

load_backend_state(input_path: pathlib.Path) → None#: Restore backend-specific model state from the input directory.

encode_observation(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → petitRADTRANS.sbi.observation.EncodedObservation#

Encode one structured observation into the posterior context space.

Parameters#

blocks:: Observation blocks describing one spectral/photometric observation.

Returns#

EncodedObservation: Aggregated embedding and lightweight metadata describing the input observation family.

batch_encode_observation(blocks_list: list[list[petitRADTRANS.sbi.observation.ObservationBlock]]) → list[petitRADTRANS.sbi.observation.EncodedObservation]#

Encode a batch of observations using a single vmapped forward pass.

Parameters#

blocks_list:: List of per-sample observation block lists.

Returns#

list[EncodedObservation]: One encoded observation per input sample.

sample_posterior(observation: petitRADTRANS.sbi.observation.EncodedObservation, n_samples: int, seed: int | None = None) → petitRADTRANS.sbi.flows.base.PosteriorSamples#

Draw posterior samples conditioned on one encoded observation.

Parameters#

observation:: Encoded observation produced by encode_observation().
n_samples:: Number of posterior draws to generate.
seed:: Optional random seed overriding the model-level default seed.

Returns#

PosteriorSamples: Samples in the posterior’s configured parameter space. Non-finite outputs are clamped to large finite values for downstream stability.

batch_sample_posterior(embeddings: Any, n_samples: int, base_seed: int = 0) → numpy.ndarray#

Draw posterior samples for a batch of encoded observations.

Runs a single JIT-compiled vmapped call over all contexts rather than looping sample_posterior once per observation, eliminating the Python overhead of per-observation dispatch.

Parameters#

embeddings:: Float32 array of shape (batch_size, embedding_dim) produced by stacking EncodedObservation.embedding vectors.
n_samples:: Number of posterior draws per observation.
base_seed:: Base random seed. Each observation in the batch receives a unique sub-key derived from this seed.

Returns#

np.ndarray: Array of shape (batch_size, n_samples, parameter_dim) with non-finite values clamped to large finite values.

log_prob(observation: petitRADTRANS.sbi.observation.EncodedObservation, parameters: Any) → Any#

Evaluate posterior log-density for one or many parameter vectors.

Parameters#

observation:: Encoded observation that defines the posterior context.
parameters:: One parameter vector or a batch of parameter vectors in the posterior’s configured parameter space.

Returns#

Any: Scalar log-density or a vector of log-densities matching the input batch structure.

class petitRADTRANS.sbi.ConditionalNeuralAutoregressiveFlowPosterior(*args: Any, **kwargs: Any)#

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the neural autoregressive backend.

class petitRADTRANS.sbi.ConditionalSplineFlowPosterior(*args: Any, **kwargs: Any)#

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the spline flow backend.

Notes#

This convenience subclass forces flow_family='spline' while keeping the rest of the ConditionalFlowPosterior API unchanged.

class petitRADTRANS.sbi.FlowMatchingPosterior(parameter_dim: int, embedding_dim: int = 128, hidden_dim: int = 128, num_velocity_layers: int = 3, learning_rate: float = 0.001, batch_size: int = 32, num_epochs: int = 5, parameter_space: str = 'unconstrained', integration_steps: int = 32, early_stopping_patience: int | None = None, early_stopping_min_delta: float = 0.0, checkpoint_directory: str | None = None, checkpoint_backend: str = 'auto', resume_from_checkpoint: bool = False, seed: int = 0, task_metadata: Mapping[str, Any] | None = None)#

Bases: petitRADTRANS.sbi.flows.base.PersistentPosteriorEstimator

Conditional flow-matching posterior skeleton.

Warning

This estimator is experimental. It does not expose log_prob and the ODE integration scheme is a simple midpoint rule. Expect the API and numerical behaviour to change in future releases.

This estimator family trains a conditional vector field on straight-line interpolation paths between Gaussian noise and target parameters, then generates posterior samples by integrating the learned field from noise to the terminal parameter state.

estimator_family = 'flow_matching'#

embedding_dim = 128#

hidden_dim = 128#

num_velocity_layers = 3#

learning_rate#

batch_size = 32#

num_epochs = 5#

integration_steps = 32#

early_stopping_patience = None#

early_stopping_min_delta#

checkpoint_directory = None#

checkpoint_backend = 'auto'#

resume_from_checkpoint = False#

model#

static _batch_embeddings(model: _FlowMatchingModel, observations: Any) → jax.numpy.ndarray#

static _loss(model: _FlowMatchingModel, batch: petitRADTRANS.sbi.flows.base.PosteriorBatch) → jax.numpy.ndarray#

fit(dataset: Any) → petitRADTRANS.sbi.flows.base.TrainingArtifacts#: Train the posterior estimator on a simulation dataset.

_build_estimator_config() → dict[str, Any]#: Return backend-specific configuration for metadata persistence.

static _resolve_estimator_config(metadata: Mapping[str, Any]) → dict[str, Any]#

_build_serialized_metadata(artifact_metadata) → dict[str, Any]#

classmethod from_serialized_metadata(metadata: Mapping[str, Any]) → FlowMatchingPosterior#: Rebuild an estimator instance from persisted metadata only.

save_backend_state(output_path: pathlib.Path) → None#: Persist backend-specific model state into the output directory.

load_backend_state(input_path: pathlib.Path) → None#: Restore backend-specific model state from the input directory.

encode_observation(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → petitRADTRANS.sbi.observation.EncodedObservation#: Encode a structured observation into the estimator input space.

batch_encode_observation(blocks_list: list[list[petitRADTRANS.sbi.observation.ObservationBlock]]) → list[petitRADTRANS.sbi.observation.EncodedObservation]#: Encode a batch of observations using a single vmapped forward pass.

sample_posterior(observation: petitRADTRANS.sbi.observation.EncodedObservation, n_samples: int, seed: int | None = None) → petitRADTRANS.sbi.flows.base.PosteriorSamples#: Sample the amortized posterior for one encoded observation.

log_prob(observation: petitRADTRANS.sbi.observation.EncodedObservation, parameters: Any) → Any#: Evaluate posterior log-density when supported by the backend.

class petitRADTRANS.sbi.PersistentPosteriorEstimator(parameter_dim: int, parameter_space: str = 'unconstrained', seed: int = 0, task_metadata: Mapping[str, Any] | None = None)#

Bases: PosteriorEstimator

Shared persistence helper for estimator backends with on-disk artifacts.

estimator_family = 'persistent_estimator'#

metadata_schema_version = '0.2.0'#

parameter_dim#

parameter_space = 'unconstrained'#

seed = 0#

task_metadata#

training_artifacts: TrainingArtifacts | None = None#

task_name: str | None#

task_version: str | None = None#

task_fingerprint: str | None = None#

observation_schema: Mapping[str, Any]#

preprocessing_metadata_payload: Mapping[str, Any]#

artifact_metadata: petitRADTRANS.sbi.artifacts.ArtifactMetadata | None = None#

abstractmethod _build_estimator_config() → dict[str, Any]#: Return backend-specific configuration for metadata persistence.

classmethod from_serialized_metadata(metadata: Mapping[str, Any]) → PersistentPosteriorEstimator#

Abstractmethod:

Rebuild an estimator instance from persisted metadata only.

abstractmethod save_backend_state(output_path: pathlib.Path) → None#: Persist backend-specific model state into the output directory.

abstractmethod load_backend_state(input_path: pathlib.Path) → None#: Restore backend-specific model state from the input directory.

static _load_training_artifacts(metadata: Mapping[str, Any]) → TrainingArtifacts | None#

hydrate_loaded_metadata(metadata: Mapping[str, Any]) → None#

_build_artifact_metadata_payload() → dict[str, Any]#

build_artifact_metadata(version: str) → petitRADTRANS.sbi.artifacts.ArtifactMetadata#: Assemble registry metadata for the currently trained estimator.

_build_serialized_metadata(artifact_metadata: petitRADTRANS.sbi.artifacts.ArtifactMetadata) → dict[str, Any]#

save(output_directory: str, artifact_registry: petitRADTRANS.sbi.artifacts.ArtifactRegistry | None = None, artifact_version: str = '0.1.0') → None#: Persist model weights, metadata, and optional artifact registration.

classmethod load(input_directory: str) → PersistentPosteriorEstimator#: Restore a saved persistent estimator from disk.

class petitRADTRANS.sbi.PosteriorBatch#

Training batch passed to amortized posterior estimators.

parameters: Any#

observations: Any#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.PosteriorEstimator#

Bases: abc.ABC

Backend-agnostic interface for amortized posterior models.

abstractmethod fit(dataset: Any) → TrainingArtifacts#: Train the posterior estimator on a simulation dataset.

abstractmethod encode_observation(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → petitRADTRANS.sbi.observation.EncodedObservation#: Encode a structured observation into the estimator input space.

abstractmethod sample_posterior(observation: petitRADTRANS.sbi.observation.EncodedObservation, n_samples: int, seed: int | None = None) → PosteriorSamples#: Sample the amortized posterior for one encoded observation.

abstractmethod log_prob(observation: petitRADTRANS.sbi.observation.EncodedObservation, parameters: Any) → Any#: Evaluate posterior log-density when supported by the backend.

abstractmethod save(output_directory: str) → None#: Persist trained model weights and metadata.

classmethod load(input_directory: str) → PosteriorEstimator#

Abstractmethod:

Restore a saved estimator from disk.

class petitRADTRANS.sbi.PosteriorSamples#

Posterior samples and optional per-sample diagnostics.

samples: Any#

log_probabilities: Any = None#

weights: Any = None#

metadata: Mapping[str, Any]#

petitRADTRANS.sbi.plot_local_sensitivity_fisher_correlations(report: petitRADTRANS.sbi.calibration.LocalSensitivityReport, figsize: tuple[float, float] | None = None) → tuple[Any, numpy.ndarray]#: Plot Fisher-correlation heatmaps for each representative point.

petitRADTRANS.sbi.plot_local_sensitivity_jacobians(report: petitRADTRANS.sbi.calibration.LocalSensitivityReport, figsize: tuple[float, float] | None = None) → tuple[Any, numpy.ndarray]#: Plot whitened Jacobian heatmaps for each representative point.

petitRADTRANS.sbi.plot_local_sensitivity_singular_values(report: petitRADTRANS.sbi.calibration.LocalSensitivityReport, figsize: tuple[float, float] | None = None) → tuple[Any, numpy.ndarray]#: Plot singular spectra of the whitened Jacobian for each point.

petitRADTRANS.sbi.plot_posterior_corner(samples: Any, parameter_names: Sequence[str] | None = None, bins: int = 40, figsize: tuple[float, float] | None = None, max_points: int = 8192) → tuple[Any, numpy.ndarray]#

Plot a lower-triangular corner view of posterior structure.

Parameters#

samples:: Posterior samples with shape (n_samples, n_dim) or a scalar-vector equivalent.
parameter_names:: Optional display names for each posterior dimension.
bins:: Number of bins used for one- and two-dimensional histograms.
figsize:: Optional Matplotlib figure size.
max_points:: Maximum number of posterior draws plotted. Larger sample sets are evenly subsampled to keep rendering costs bounded.

Returns#

tuple[Any, np.ndarray]: Matplotlib figure and axes array for further customization or saving.

petitRADTRANS.sbi.plot_posterior_marginals(samples: Any, parameter_names: Sequence[str] | None = None, bins: int = 40, figsize: tuple[float, float] | None = None) → tuple[Any, numpy.ndarray]#

Plot one histogram per posterior dimension.

Parameters#

samples:: Posterior samples with shape (n_samples, n_dim) or a scalar-vector equivalent.
parameter_names:: Optional display names for each posterior dimension.
bins:: Number of histogram bins per dimension.
figsize:: Optional Matplotlib figure size.

Returns#

tuple[Any, np.ndarray]: Matplotlib figure and axes array for further customization or saving.

petitRADTRANS.sbi.plot_posterior_predictive_report(report: petitRADTRANS.sbi.calibration.PosteriorPredictiveReport, dataset_names: Sequence[str] | None = None, figsize: tuple[float, float] | None = None, x_coordinates: Mapping[str, Any] | Any | None = None, x_label: str = 'Wavelength [$\\mu$m]', y_label: str = 'Observation scale') → tuple[Any, numpy.ndarray]#

Plot observed values against posterior-predictive means and intervals.

Parameters#

report:: Posterior-predictive report to visualize.
dataset_names:: Optional subset of dataset names to plot. Defaults to all datasets in the report.
figsize:: Optional Matplotlib figure size.
x_coordinates:: Physical x-axis sample positions (e.g. wavelength in micron). May be a single array applied to every dataset, or a mapping from dataset name to array. When None – or when a supplied array does not match the number of plotted samples – the helper falls back to the integer sample index. Supplying wavelengths replaces the bare coordinate index on the x-axis.
x_label:: Label for the x-axis (defaults to a wavelength label). Used only when an x_coordinates array is actually applied.
y_label:: Label for the y-axis. Pass the observable-specific unit, e.g. "Flux [W m^-2 um^-1]" for emission or "Transit depth" for transmission.

Returns#

tuple[Any, np.ndarray]: Matplotlib figure and axes array.

Notes#

Spectrum plots adopt the full petitRADTRANS.plotlib.style style (fonts, ticks and colour cycle) for the duration of the call via a temporary rc context, so the global Matplotlib state is left untouched on return.

For batched reports the helper currently visualizes only the first case for each dataset, which keeps the plot compact for quick inspection.

petitRADTRANS.sbi.plot_sbc_rank_histograms(report: petitRADTRANS.sbi.calibration.SimulationBasedCalibrationReport, parameter_names: Sequence[str] | None = None, figsize: tuple[float, float] | None = None) → tuple[Any, numpy.ndarray]#

Plot one SBC rank histogram per inferred parameter.

Parameters#

report:: SBC report containing per-parameter rank histogram counts.
parameter_names:: Optional display names for each inferred parameter.
figsize:: Optional Matplotlib figure size.

Returns#

tuple[Any, np.ndarray]: Matplotlib figure and axes array.

petitRADTRANS.sbi.save_validation_plots(run_root: Any, posterior_samples: numpy.ndarray, parameter_names: Sequence[str], validation_sbc_report: petitRADTRANS.sbi.calibration.SimulationBasedCalibrationReport, validation_predictive_report: petitRADTRANS.sbi.calibration.PosteriorPredictiveReport, target_predictive_report: petitRADTRANS.sbi.calibration.PosteriorPredictiveReport, dataset_names: Sequence[str], spectrum_wavelengths: numpy.ndarray | None = None, x_label: str = 'Wavelength [$\\mu$m]', y_label: str = 'Observation scale', dpi: int = 160) → None#: Render and save the standard validation plot set (marginals, corner, SBC rank histograms, and validation/target posterior-predictive spectra) into run_root. A no-op if matplotlib is unavailable.

class petitRADTRANS.sbi.TaskPreprocessingMetadata#

Serializable preprocessing metadata for an SBI task family.

version: str#

blocks: Mapping[str, BlockNormalizationStats]#

metadata: Mapping[str, Any]#

to_payload() → dict[str, Any]#

classmethod from_payload(payload: Mapping[str, Any]) → TaskPreprocessingMetadata#

petitRADTRANS.sbi.fit_task_preprocessing(training_observations: list[list[petitRADTRANS.sbi.observation.ObservationBlock]], version: str = '0.5.0', metadata: Mapping[str, Any] | None = None) → TaskPreprocessingMetadata#: Fit preprocessing statistics from training observation blocks.

petitRADTRANS.sbi.normalize_observation_block(block: petitRADTRANS.sbi.observation.ObservationBlock, preprocessing_metadata: TaskPreprocessingMetadata) → petitRADTRANS.sbi.observation.ObservationBlock#

Normalize one observation block with fitted preprocessing statistics.

Uses robust (median/IQR) normalization for values and uncertainties when available, falling back to mean/std for backward compatibility with older preprocessing metadata. Coordinates always use mean/std normalization.

petitRADTRANS.sbi.normalize_observation_blocks(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock], preprocessing_metadata: TaskPreprocessingMetadata) → list[petitRADTRANS.sbi.observation.ObservationBlock]#: Normalize a list of observation blocks.

class petitRADTRANS.sbi.ProposalSampler#

Bases: abc.ABC

Interface for simulation proposals beyond the prior distribution.

abstractmethod sample(n_samples: int, task: petitRADTRANS.sbi.task.SBITask) → Any#: Draw free-parameter vectors for the given task.

class petitRADTRANS.sbi.RuntimeSimulator(task: petitRADTRANS.sbi.task.SBITask, runtime: Any | None = None, seed: int | None = None, data_parallel: bool | None = None)#

Bases: BatchedSimulator

Concrete simulator backed by the retrieval runtime.

The simulator samples from the retrieval prior using JAX random utilities and projects deterministic forward-model outputs into the observation space using RetrievalRuntime.

_rng_key#

_validate_task_support() → None#

_next_key() → jax.Array#

static _row_invalid_value_mask(values: Any, mask: Any, constraint: Any) → numpy.ndarray#: Sanity checking simulated observations. Uses the ObservationValueConstraints to check if the values are within the specified bounds. Failing simulations will be rejected and resampled. This function returns a boolean mask indicating which rows are invalid.

_invalid_sample_mask(observations: Mapping[str, Mapping[str, Any]], n_samples: int) → tuple[numpy.ndarray, dict[str, int]]#

static _slice_batch_rows(batch: SimulationBatch, row_indices: numpy.ndarray) → SimulationBatch#

static _concatenate_simulation_batches(batches: list[SimulationBatch], diagnostics: Mapping[str, Any] | None = None) → SimulationBatch#

_sample_prior_parameter_batch(n_samples: int) → dict[str, jax.numpy.ndarray]#

sample_parameters(n_samples: int, proposal: ProposalSampler | None = None) → Any#: Sample free-parameter vectors from the task prior or a proposal.

_advance_rng_for_batch(n_samples: int) → None#

Advance the PRNG state as if simulate_batch(n_samples) were called.

This performs only cheap key-splitting – no forward-model evaluation – making it suitable for fast dataset-resume skipping.

The number of _next_key calls matches the pattern inside simulate_batch → _sample_prior_parameter_batch (1 key) and _simulate_parameter_matrix (1 key per observation that adds noise).

_apply_noise(values: Any, uncertainties: Any, covariance: Any) → tuple[Any, Any, Any]#

_apply_noise_batched(values_batch: Any, uncertainties_batch: Any, covariance_batch: Any) → tuple[Any, Any, Any]#

Apply noise to a batch of simulation outputs.

Args:: values_batch: Shape (n_samples, n_wavelengths). uncertainties_batch: Shape (n_samples, n_wavelengths) or None. covariance_batch: Shape (n_samples, n_wl, n_wl) or None.
Returns:: Noisy values, uncertainties, and covariance with the same leading batch dimension as values_batch.

static _stack_observation_payloads(payloads: list[dict[str, Any]]) → dict[str, Any]#

_simulate_parameter_matrix(parameter_matrix: Any, cube_parameters: Any = None, unconstrained_parameters: Any = None, data_parallel: bool = False) → SimulationBatch#

simulate_from_parameters(parameters: Any) → SimulationBatch#: Run the forward model and noise pipeline for pre-specified parameters.

simulate_batch(n_samples: int, proposal: ProposalSampler | None = None) → SimulationBatch#: Sample parameters and preserve prior-space coordinates when available.

class petitRADTRANS.sbi.SimulationBatch#

Container for one batch of simulations.

Attributes:

parameters:: Array-like free-parameter matrix with leading dimension equal to the number of simulated samples.
observations:: Task-conditioned simulated observations.
log_likelihood:: Optional scalar likelihood values associated with each sample.
diagnostics:: Additional runtime diagnostics such as clipping flags or NaN counts.

parameters: Any#

observations: Any#

cube_parameters: Any = None#

unconstrained_parameters: Any = None#

log_likelihood: Any = None#

diagnostics: Mapping[str, Any]#

property n_samples: int#: Return the number of simulated samples represented by the batch.

class petitRADTRANS.sbi.Simulator(task: petitRADTRANS.sbi.task.SBITask)#

Bases: abc.ABC

Base simulator for SBI dataset generation and validation.

task#

abstractmethod sample_parameters(n_samples: int, proposal: ProposalSampler | None = None) → Any#: Sample free-parameter vectors from the task prior or a proposal.

abstractmethod simulate_from_parameters(parameters: Any) → SimulationBatch#: Run the forward model and noise pipeline for pre-specified parameters.

simulate_batch(n_samples: int, proposal: ProposalSampler | None = None) → SimulationBatch#: Sample parameters and simulate one batch in a single call.

class petitRADTRANS.sbi.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.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.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.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#

class petitRADTRANS.sbi.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]#

class petitRADTRANS.sbi.EarlyStoppingConfig#

Early-stopping policy for SBI training.

Attributes#

patience:: Number of consecutive non-improving epochs tolerated before stopping.
min_delta:: Minimum decrease in the monitored metric that counts as an improvement.

patience: int#

min_delta: float = 0.0#

class petitRADTRANS.sbi.SBITrainer(config: TrainingSchemeConfig)#

Reusable optimisation loop for amortized SBI posteriors.

Configured once with a TrainingSchemeConfig and reused via fit(), which is given the model, a dataset reader, and the estimator’s loss/diagnostic/selection callbacks. The trainer owns the AdamW optimiser and schedule, the JIT-compiled train/eval steps, optional multi-device sharding, the per-epoch validation + diagnostics pass, the five checkpoint kinds (see the module docstring), and early stopping. Checkpointing is enabled only when the config provides a checkpoint_directory.

config#

checkpoint_directory = None#

checkpoint_backend#

_checkpoint_backend_fallback_reason: str | None = None#

property _latest_checkpoint_directory: pathlib.Path | None#: Path of the latest (resume) checkpoint, or None if disabled.

property _best_checkpoint_directory: pathlib.Path | None#: Path of the best (deployed-model) checkpoint, or None if disabled.

_checkpoint_directory_for_kind(checkpoint_kind: str) → pathlib.Path | None#: Resolve the subdirectory for one checkpoint kind, or None when no checkpoint directory is configured.

_save_checkpoint(checkpoint_kind: str, state: dict[str, Any], metadata: dict[str, Any]) → None#

Persist one checkpoint kind, transparently downgrading Orbax→Equinox.

No-op when checkpointing is disabled. If the Orbax backend rejects a leaf type mid-run (_is_orbax_unsupported_type_error()), the trainer permanently switches to the Equinox backend, records the fallback reason in the metadata, and retries the save.

_restore_checkpoint(template_state: dict[str, Any], checkpoint_kind: str) → tuple[dict[str, Any], dict[str, Any]] | None#

Restore one checkpoint kind, or None if it is missing/disabled.

Honours the backend name recorded in the checkpoint’s metadata (so a checkpoint written by Equinox after an Orbax fallback restores correctly), and returns the (state, metadata) pair from that backend.

fit(model: Any, dataset: Any, loss_fn: Callable[[Any, Any], Any], eval_loss_fn: Callable[[Any, Any], Any] | None = None, eval_diagnostic_fn: Callable[[Any, Any], Mapping[str, float]] | None = None, selection_metric_fn: Callable[[float, Mapping[str, float] | None], float] | None = None, selection_metric_name: str | None = None, stability_metric_fn: Callable[[float, Mapping[str, float] | None], Mapping[str, float]] | None = None, stability_flag_key: str = 'checkpoint_is_stable', loss_conditioning_schedule: Callable[[int, int], float] | None = None) → tuple[Any, dict[str, list[float]], dict[str, list[float]], dict[str, Any]]#

Optimise one posterior model against a dataset reader.

Runs the full epoch loop: (optionally resume from latest) → per-epoch train pass → validation loss + diagnostics → selection/stability scoring → checkpoint writes → early-stopping check, returning the selected (“best”) model and the histories/metadata.

Parameters#

model:: Trainable Equinox-style model to optimise (e.g. the encoder+flow container).
dataset:: Reader exposing iter_batches(split=..., batch_size=..., ...) and a dataset.manifest.splits mapping with 'train' / 'validation' counts.
loss_fn:: Training loss. Called as (model, batch) by default, or (model, batch, cond) when loss_conditioning_schedule is given.
eval_loss_fn:: Validation loss; defaults to loss_fn. Always evaluated at conditioning 1.0 so checkpoint selection tracks the untempered objective.
eval_diagnostic_fn:: Optional (model, batch) -> {metric: value} evaluated on a few validation batches per epoch; results are aggregated (_aggregate_diagnostic_metrics()) into the epoch metrics (e.g. the flow invertibility / collapse diagnostics).
selection_metric_fn:: Optional (validation_loss, epoch_metrics) -> float producing the scalar minimised to pick best_selection. When None the raw validation loss is used.
selection_metric_name:: Name under which the selection metric is logged/stored; defaults to 'validation_selection_metric' (or 'train_loss' with no validation split).
stability_metric_fn:: Optional (validation_loss, epoch_metrics) -> {metric: value} adding stability flags (notably checkpoint_is_stable) to the epoch metrics; drives the best_stable checkpoint and the per-epoch stable flag.
stability_flag_key:: Key in the epoch metrics whose value (>= 0.5) marks an epoch stable.
loss_conditioning_schedule:: Optional (epoch, num_epochs) -> float feeding a per-epoch scalar to the training loss (e.g. likelihood-temperature beta for VI annealing); evaluation always uses 1.0.

Returns#

tuple[Any, dict[str, list[float]], dict[str, list[float]], dict[str, Any]]

(best_model, train_history, validation_metrics_history, metadata):

best_model – the deployed (best) model.
train_history – {'train_loss': [...]} per epoch.
validation_metrics_history – per-key lists for every validation metric (loss, selection metric, diagnostics, stability flags).
metadata – the best/selection/loss/stable epochs and metrics, the selected best_checkpoint_kind, early-stopping/resume flags, step counts, and the checkpoint backend used.

Notes#

With no validation split the trainer selects on training loss. The five checkpoint kinds and the selection-metric semantics are described in the module docstring. Training/validation data are pre-loaded and pre-stacked into RAM where possible (falling back to streaming from disk otherwise), and batches are sharded across devices when data_parallel is set.

property _min_delta: float#: Minimum-improvement threshold for best-checkpoint comparisons (0 when early stopping is not configured).

_should_stop(patience_counter: int) → bool#: Whether training should stop: True once the non-improving-epoch count reaches the configured early-stopping patience.

class petitRADTRANS.sbi.ObservationScaleReport#

Result of check_observation_in_training_range().

log10_observation_scale: float#

log10_training_p1: float#

log10_training_p50: float#

log10_training_p99: float#

orders_outside: float#

max_orders_outside: float#

passed: bool#

petitRADTRANS.sbi.check_observation_in_training_range(training_values: Any, observation_values: Any, *, max_orders_outside: float = 2.0, on_fail: str = 'raise', max_training_rows: int = 8192) → ObservationScaleReport#

Check the observation’s flux scale lies within the training distribution.

For every simulated spectrum a robust amplitude (median absolute flux) is computed, and the [p1, p99] range of those amplitudes in log10 space defines the simulated scale band. The observation’s amplitude must lie within max_orders_outside decades of that band. A units/scale bug – e.g. training spectra at ~1e-20 while the observation is ~1e-15 – puts the observation ~3+ decades outside and is rejected here, before training, instead of only via the post-training OOD diagnostic.

Parameters#

training_values:: Simulated spectra, shape (n_samples, n_wavelengths) (raw flux scale, matching observation_values).
observation_values:: The observed spectrum, shape (n_wavelengths,).
max_orders_outside:: How many decades the observation amplitude may fall outside the simulated [p1, p99] band before the check fails.
on_fail:: "raise" (default) aborts with a descriptive error, "warn" prints a warning and continues, "off" disables the check (still returns a report).
max_training_rows:: Training rows are evenly subsampled to at most this many before computing statistics, keeping the check cheap on a multi-hundred-thousand-row split.

Returns#

ObservationScaleReport: The computed log10 scale statistics and pass/fail verdict.