petitRADTRANS.sbi.observation

Observation containers and block builders for amortized inference.

Classes

ObservationModality	Supported observation block types for SBI conditioning.
ObservationBlock	Represent one modality-specific observation block.
EncodedObservation	Encoded representation consumed by a posterior estimator.
ObservationEncoder	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.

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

Represent one modality-specific observation block.

Attributes:

name:

Stable identifier of the block within a task.

modality:

Semantic block type used to dispatch encoding logic.

values:

Observed values after any task-level preprocessing.

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.ObservationEncoder

Bases: abc.ABC

Transform structured observation blocks into model-ready embeddings.

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.

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