petitRADTRANS.sbi.observation#

Observation containers and block builders for amortized inference.

This module implements classes to store real or simulated observations in structured, frozen data classes, enabling them to be passed through JAX transformations and JIT compilation. It also includes functions to build batches of observations, enabling parallelised processing of multiple samples, while also limiting peak memory usage.

Classes#

`ObservationModality`	Supported observation block types for SBI conditioning.
`ObservationBlock`	Representation of one (simulated or real) observation.
`EncodedObservation`	Encoded representation consumed by a posterior estimator.
`AbstractObservationEncoder`	Transform structured observation blocks into model-ready embeddings.
`PreStackedObservations`	Pre-stacked observation arrays ready to cross the JAX JIT boundary.

Functions#

`build_observation_block`(→ ObservationBlock)	Build one observation block with modality normalization.
`build_observation_blocks_from_sample`(...)	Build modality-aware observation blocks for one simulated sample.
`build_observation_block_batch`(...)	Build observation blocks for each sample in a batched payload.
`build_raw_observation_blocks`(→ list)	Build raw (un-normalized) spectrum ObservationBlocks for every dataset in

Module Contents#

class petitRADTRANS.sbi.observation.ObservationModality#

Bases: str, enum.Enum

Supported observation block types for SBI conditioning.

SPECTRUM = 'spectrum'#

PHOTOMETRY = 'photometry'#

TIME_SERIES = 'time_series'#

AUXILIARY = 'auxiliary'#

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

Encoded representation consumed by a posterior estimator.

embedding: Any#

mask: Any = None#

metadata: Mapping[str, Any]#

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

petitRADTRANS.sbi.observation.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.observation.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.observation.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.

class petitRADTRANS.sbi.observation.PreStackedObservations#

Pre-stacked observation arrays ready to cross the JAX JIT boundary.

All per-sample ObservationBlock data has been extracted into dense numpy arrays before this object is constructed, so JAX can trace through the encoder without re-compiling per batch. Registered as a pytree node to enable JIT compilation.

Attributes#

stacked_blocks:: One (values, uncertainties, coordinates, mask, log_median_flux, spectral_auxiliary_values, global_spectrum_scale) tuple per observation block. values, uncertainties, coordinates, mask, and spectral_auxiliary_values have shape (batch_size, n_wl). log_median_flux and global_spectrum_scale have shape (batch_size,).
modalities:: Modality value string for each block, e.g. 'spectrum'. Treated as static metadata by JAX (does not trigger retrace between batches since it is constant for a given SBI task).

stacked_blocks: tuple#

modalities: tuple#

petitRADTRANS.sbi.observation.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.