petitRADTRANS.sbi.flows#

Normalizing-flow transforms and flow-based posterior estimators.

This subpackage groups everything related to the conditional normalizing-flow machinery and the posterior estimators built on top of it:

transforms – low-level conditional flow transforms and flow stacks (affine / spline / autoregressive / NAF).
base – backend-agnostic posterior interfaces, shared dataclasses, and the persistence base class.
cube – parameter-space / unit-cube transforms.
regularization – spline regularization and checkpoint-selection stability diagnostics.
objectives – the training objective and the conditional encoder+flow model container.
refinement – test-time per-observation posterior refinement (MAP+Laplace / flow-ELBO).
posterior – ConditionalFlowPosterior.
matching – FlowMatchingPosterior.

The estimator-family registry and the generic load_posterior_estimator() loader live in base. Importing this package eagerly imports posterior and matching, which register their families, so a bare import petitRADTRANS.sbi is enough for load_posterior_estimator() to resolve either family.

Submodules#

Attributes#

POSTERIOR_METADATA_SCHEMA_VERSION

Classes#

`ConditionalAutoregressiveAffineTransform`	Context-conditioned affine autoregressive transform.
`ConditionalAutoregressiveFlow`	Stack of context-conditioned affine autoregressive transforms.
`ConditionalNeuralAutoregressiveFlow`	Stack of context-conditioned neural autoregressive transforms.
`ConditionalNeuralAutoregressiveTransform`	Context-conditioned neural autoregressive transform.
`ConditionalAffineCoupling`	Affine coupling layer conditioned on an observation embedding.
`ConditionalAffineFlow`	Stack of conditional affine couplings with alternating masks.
`ConditionalRationalQuadraticSplineCoupling`	Rational-quadratic spline coupling layer conditioned on an embedding.
`ConditionalScalarSplineTransform`	Context-conditioned scalar spline transform used for 1D posteriors.
`ConditionalSplineFlow`	Stack of conditional rational-quadratic spline transforms.
`CouplingConditioner`	Context-conditioned network producing coupling parameters.
`NeuralInverseDiagnostics`	Aggregate diagnostics for one neural autoregressive inverse pass.
`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.
`TrainingArtifacts`	Structured outputs of a posterior training run.
`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.

Functions#

`get_posterior_estimator_class`(→ type[PosteriorEstimator])	Return the registered estimator class for one estimator family.
`load_posterior_estimator`(→ PosteriorEstimator)	Load a saved posterior estimator without naming its concrete class.
`read_posterior_metadata`(→ dict[str, Any])	Read the persisted metadata envelope for one saved estimator.
`register_posterior_estimator`(→ None)	Register one estimator family for generic load dispatch.
`resolve_estimator_family_from_metadata`(→ str)	Infer the estimator family from current or legacy saved metadata.

Package Contents#

class petitRADTRANS.sbi.flows.ConditionalAutoregressiveAffineTransform(parameter_dim: int, context_dim: int, hidden_dim: int, key: jax.Array, depth: int = 5)#

Bases: equinox.Module

Context-conditioned affine autoregressive transform.

Notes#

The transform maps data to latent space in a masked-autoregressive manner, using previous data dimensions and the observation embedding to predict one affine transform per parameter dimension.

conditioners: tuple[_AutoregressiveDimensionConditioner, Ellipsis]#

parameter_dim: int#

forward(x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

inverse(z: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

class petitRADTRANS.sbi.flows.ConditionalAutoregressiveFlow(parameter_dim: int, context_dim: int, num_layers: int = 3, hidden_dim: int = 512, conditioner_depth: int = 5, key: jax.Array | None = None)#

Bases: _ConditionalFlowStack

Stack of context-conditioned affine autoregressive transforms.

Notes#

This flow is a lighter-weight autoregressive alternative to the spline stack. Each transform uses ELU MLPs to predict affine parameters for one dimension at a time, conditioned on previous dimensions and the observation embedding.

parameter_dim#

context_dim#

layers#

class petitRADTRANS.sbi.flows.ConditionalNeuralAutoregressiveFlow(parameter_dim: int, context_dim: int, num_layers: int = 3, hidden_dim: int = 512, conditioner_depth: int = 5, transform_units: int = 16, min_slope: float = 0.001, min_residual: float = 0.05, inverse_bisection_steps: int = 48, inverse_newton_steps: int = 3, key: jax.Array | None = None)#

Bases: _ConditionalFlowStack

Stack of context-conditioned neural autoregressive transforms.

Notes#

This backend uses a deep-sigmoidal-style monotonic scalar bijector per parameter dimension. It is more expressive than the affine autoregressive stack while remaining exactly invertible via scalar bisection.

transform_units: int#

min_slope: float#

min_residual: float#

inverse_bisection_steps: int#

inverse_newton_steps: int#

parameter_dim#

context_dim#

layers#

inverse_with_diagnostics(z: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, NeuralInverseDiagnostics]#

class petitRADTRANS.sbi.flows.ConditionalNeuralAutoregressiveTransform(parameter_dim: int, context_dim: int, hidden_dim: int, transform_units: int, key: jax.Array, depth: int = 5, min_slope: float = 0.001, min_residual: float = 0.05, inverse_base_bound: float = 4.0, inverse_bisection_steps: int = 48, inverse_expansion_steps: int = 16, inverse_newton_steps: int = 3)#

Bases: equinox.Module

Context-conditioned neural autoregressive transform.

Notes#

Each scalar transform uses a monotonic deep-sigmoidal-style network whose parameters are conditioned on previous dimensions and the observation embedding. Inversion is carried out with fixed-step bisection because the scalar map is strictly monotonic.

conditioners: tuple[_NeuralAutoregressiveDimensionConditioner, Ellipsis]#

parameter_dim: int#

transform_units: int#

min_slope: float#

min_residual: float#

inverse_base_bound: float#

inverse_bisection_steps: int#

inverse_expansion_steps: int#

inverse_newton_steps: int#

_parameters(conditioner: _NeuralAutoregressiveDimensionConditioner, prefix: jax.numpy.ndarray, context: jax.numpy.ndarray, dtype: jax.numpy.dtype) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

_evaluate_scalar(value: jax.numpy.ndarray, positive_slopes: jax.numpy.ndarray, offsets: jax.numpy.ndarray, weight_logits: jax.numpy.ndarray, residual_fraction: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

_inverse_scalar(target: jax.numpy.ndarray, positive_slopes: jax.numpy.ndarray, offsets: jax.numpy.ndarray, weight_logits: jax.numpy.ndarray, residual_fraction: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

_inverse_scalar_with_diagnostics(target: jax.numpy.ndarray, positive_slopes: jax.numpy.ndarray, offsets: jax.numpy.ndarray, weight_logits: jax.numpy.ndarray, residual_fraction: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

forward(x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

inverse(z: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

inverse_with_diagnostics(z: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, NeuralInverseDiagnostics]#

class petitRADTRANS.sbi.flows.ConditionalAffineCoupling#

Bases: equinox.Module

Affine coupling layer conditioned on an observation embedding.

Notes#

The mask selects which parameter dimensions remain fixed and which are transformed. The conditioner predicts shift and log-scale terms only from the masked input and context embedding. The mask is static (see _as_mask_tuple()) — it must never be trainable.

mask_values: tuple[float, Ellipsis]#

conditioner: CouplingConditioner#

property mask: jax.numpy.ndarray#

_parameters(masked_x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

forward(x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the forward affine transform.

Parameters#

x:: Parameter vector in data space.
context:: Observation embedding conditioning the transform.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Transformed vector and scalar log-determinant contribution.

inverse(y: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the inverse affine transform.

Parameters#

y:: Latent-space vector.
context:: Observation embedding conditioning the inverse transform.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Recovered data-space vector and inverse log-determinant.

class petitRADTRANS.sbi.flows.ConditionalAffineFlow(parameter_dim: int, context_dim: int, num_layers: int = 4, hidden_dim: int = 128, conditioner_depth: int = 2, key: jax.Array | None = None)#

Bases: _ConditionalFlowStack

Stack of conditional affine couplings with alternating masks.

Parameters#

parameter_dim:: Number of inferred parameters transformed by the flow.
context_dim:: Size of the observation embedding conditioning each layer.
num_layers:: Number of affine coupling layers.
hidden_dim:: Hidden width of each conditioner MLP.
key:: Optional JAX random key used for layer initialization.

Notes#

For one-dimensional problems the mask degenerates to a context-only affine transform. Higher-dimensional problems alternate masks across layers.

parameter_dim#

context_dim#

layers = ()#

class petitRADTRANS.sbi.flows.ConditionalRationalQuadraticSplineCoupling#

Bases: equinox.Module

Rational-quadratic spline coupling layer conditioned on an embedding.

Notes#

This is the main expressive transform used by spline-based SBI posteriors. The conditioner predicts per-dimension spline widths, heights, and derivatives for the unmasked subset of the parameter vector. The mask is static (see _as_mask_tuple()) — it must never be trainable.

mask_values: tuple[float, Ellipsis]#

conditioner: CouplingConditioner#

parameter_dim: int#

num_bins: int#

bound: float#

min_bin_width: float#

min_bin_height: float#

min_derivative: float#

max_derivative: float | None#

property mask: jax.numpy.ndarray#

_parameters(masked_x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

forward(x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the forward spline transform.

Parameters#

x:: Parameter vector in data space.
context:: Observation embedding conditioning the spline parameters.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Transformed vector and summed log-determinant contribution.

inverse(y: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the inverse spline transform.

Parameters#

y:: Latent-space vector.
context:: Observation embedding conditioning the inverse transform.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Recovered parameter vector and inverse log-determinant.

class petitRADTRANS.sbi.flows.ConditionalScalarSplineTransform#

Bases: equinox.Module

Context-conditioned scalar spline transform used for 1D posteriors.

Notes#

One-dimensional posterior models cannot rely on alternating coupling masks, so they use this dedicated context-only spline transform instead.

conditioner: _ScalarSplineConditioner#

num_bins: int#

bound: float#

min_bin_width: float#

min_bin_height: float#

min_derivative: float#

max_derivative: float | None#

_parameters(context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

static _as_scalar(value: jax.numpy.ndarray) → jax.numpy.ndarray#

forward(x: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the forward scalar spline transform.

Parameters#

x:: One-dimensional parameter vector.
context:: Observation embedding conditioning the spline parameters.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Transformed scalar vector and scalar log-determinant.

inverse(y: jax.numpy.ndarray, context: jax.numpy.ndarray) → tuple[jax.numpy.ndarray, jax.numpy.ndarray]#

Apply the inverse scalar spline transform.

Parameters#

y:: One-dimensional latent vector.
context:: Observation embedding conditioning the inverse transform.

Returns#

tuple[jnp.ndarray, jnp.ndarray]: Recovered scalar parameter vector and inverse log-determinant.

class petitRADTRANS.sbi.flows.ConditionalSplineFlow(parameter_dim: int, context_dim: int, num_layers: int = 6, hidden_dim: int = 128, conditioner_depth: int = 2, num_bins: int = 8, bound: float = 10.0, min_bin_width: float = 0.001, min_bin_height: float = 0.001, min_derivative: float = 0.001, max_derivative: float | None = None, base_distribution: str = 'gaussian', use_base_affine: bool = False, key: jax.Array | None = None)#

Bases: _ConditionalFlowStack

Stack of conditional rational-quadratic spline transforms.

Parameters#

parameter_dim:: Number of inferred parameters transformed by the flow.
context_dim:: Size of the conditioning observation embedding.
num_layers:: Number of spline transform layers.
hidden_dim:: Hidden width of the conditioner networks.
num_bins:: Number of rational-quadratic bins used per transformed dimension.
bound:: Finite support bound of the spline transform in latent coordinates.
min_bin_width, min_bin_height, min_derivative:: Numerical-stability floors applied to spline parameters.
key:: Optional JAX random key used for initialization.

Notes#

One-dimensional posteriors use a dedicated scalar spline stack, while higher-dimensional models use alternating masked coupling layers.

num_bins: int#

bound: float#

min_bin_width: float#

min_bin_height: float#

min_derivative: float#

max_derivative: float | None#

parameter_dim#

context_dim#

base_distribution = ''#

layers = ()#

class petitRADTRANS.sbi.flows.CouplingConditioner(input_dim: int, context_dim: int, output_dim: int, hidden_dim: int, key: jax.Array, depth: int = 2)#

Bases: equinox.Module

Context-conditioned network producing coupling parameters.

Parameters#

input_dim:: Size of the masked parameter vector presented to the conditioner.
context_dim:: Size of the observation embedding appended to the masked input.
output_dim:: Number of coupling parameters produced by the network.
hidden_dim:: Hidden width of the internal MLP.
key:: JAX random key used to initialize network weights.

Notes#

The conditioner concatenates masked parameter values with the observation context before passing them through a small MLP.

network: equinox.nn.MLP#

input_dim: int#

context_dim: int#

output_dim: int#

__call__(masked_x: jax.numpy.ndarray, context: jax.numpy.ndarray) → jax.numpy.ndarray#

class petitRADTRANS.sbi.flows.NeuralInverseDiagnostics#

Bases: NamedTuple

Aggregate diagnostics for one neural autoregressive inverse pass.

max_abs_residual: jax.numpy.ndarray#

mean_abs_residual: jax.numpy.ndarray#

max_bracket_width: jax.numpy.ndarray#

mean_bracket_width: jax.numpy.ndarray#

max_expansion_steps: jax.numpy.ndarray#

mean_expansion_steps: jax.numpy.ndarray#

converged_fraction: jax.numpy.ndarray#

petitRADTRANS.sbi.flows.POSTERIOR_METADATA_SCHEMA_VERSION = '0.2.0'#

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

Training batch passed to amortized posterior estimators.

parameters: Any#

observations: Any#

metadata: Mapping[str, Any]#

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

Posterior samples and optional per-sample diagnostics.

samples: Any#

log_probabilities: Any = None#

weights: Any = None#

metadata: Mapping[str, Any]#

class petitRADTRANS.sbi.flows.TrainingArtifacts#

Structured outputs of a posterior training run.

history: Mapping[str, Any]#

validation_metrics: Mapping[str, Any]#

metadata: Mapping[str, Any]#

petitRADTRANS.sbi.flows.get_posterior_estimator_class(estimator_family: str) → type[PosteriorEstimator]#: Return the registered estimator class for one estimator family.

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

petitRADTRANS.sbi.flows.read_posterior_metadata(input_directory: str) → dict[str, Any]#: Read the persisted metadata envelope for one saved estimator.

petitRADTRANS.sbi.flows.register_posterior_estimator(estimator_family: str, estimator_cls: type[PosteriorEstimator]) → None#: Register one estimator family for generic load dispatch.

petitRADTRANS.sbi.flows.resolve_estimator_family_from_metadata(metadata: dict[str, Any]) → str#: Infer the estimator family from current or legacy saved metadata.

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

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the autoregressive flow backend.

class petitRADTRANS.sbi.flows.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.flows.ConditionalNeuralAutoregressiveFlowPosterior(*args: Any, **kwargs: Any)#

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the neural autoregressive backend.

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