petitRADTRANS.sbi.flows.posterior

Conditional normalizing-flow posterior estimator.

Classes

ConditionalFlowPosterior	Concrete amortized posterior using a conditional flow backend.
ConditionalSplineFlowPosterior	Posterior estimator specialized to the spline flow backend.
ConditionalAutoregressiveFlowPosterior	Posterior estimator specialized to the autoregressive flow backend.
ConditionalNeuralAutoregressiveFlowPosterior	Posterior estimator specialized to the neural autoregressive backend.

Module Contents

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

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the autoregressive flow backend.

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

Bases: ConditionalFlowPosterior

Posterior estimator specialized to the neural autoregressive backend.