petitRADTRANS.sbi.flows.objectives#

Conditional encoder+flow model container and the training objective.

Holds _PosteriorModel (the encoder + flow + auxiliary scale head bundled as one Equinox module) and the loss functions optimised by ConditionalFlowPosterior: the conditional negative-log-likelihood (NPE) objective with its optional spline and invertibility regularization, and the amortized-variational (ELBO) objective. Extracted from flows.posterior so the objective maths lives apart from the estimator’s plumbing.

Classes#

_PosteriorModel

Trainable encoder + conditional flow bundled as one Equinox module.

Functions#

_mean_spectral_log_scale(→ float | None)

Mean log-amplitude across the spectrum blocks of one observation.

_scale_targets_from_observations(...)

Per-sample log-amplitude targets for the auxiliary scale loss.

batch_embeddings(→ jax.numpy.ndarray)

Encode a batch of observations into conditioning vectors.

compute_loss(→ jax.numpy.ndarray)

Evaluation NPE loss for one batch (no regularization terms).

training_loss(→ jax.numpy.ndarray)

Training NPE loss for one batch (with spline / invertibility regularization).

loss_from_observations(→ jax.numpy.ndarray)

Conditional negative-log-likelihood (NPE) objective for one batch.

invertibility_penalty(→ jax.numpy.ndarray)

Differentiable inverse round-trip + log-det-closure penalty.

make_elbo_loss(→ Callable[[_PosteriorModel, ...)

Return an amortized-variational (ELBO) training loss.

compute_validation_diagnostics(→ dict[str, float])

Per-checkpoint flow invertibility / collapse diagnostics.

Module Contents#

class petitRADTRANS.sbi.flows.objectives._PosteriorModel#

Bases: equinox.Module

Trainable encoder + conditional flow bundled as one Equinox module.

Holds the observation encoder, the conditional flow, and an auxiliary scale_head (a small linear readout used by the optional amplitude-scale loss), together with the static regularization weights/targets read by the loss functions in this module. Bundling them keeps the whole trainable model a single pytree so Optax updates and Equinox (de)serialization act on it as a unit.

encoder: petitRADTRANS.sbi.encoders.ObservationEncoder#
flow: Any#
scale_head: equinox.nn.Linear#
aux_scale_loss_weight: float#
spline_small_bin_regularization_weight: float#
spline_derivative_regularization_weight: float#
spline_min_bin_ratio_target: float#
spline_min_derivative_target: float#
spline_max_derivative_target: float#
invertibility_loss_weight: float#
invertibility_loss_n_samples: int#
bnpe_balance_loss_weight: float#
petitRADTRANS.sbi.flows.objectives._mean_spectral_log_scale(blocks: list[petitRADTRANS.sbi.observation.ObservationBlock]) → float | None#

Mean log-amplitude across the spectrum blocks of one observation.

Reads each spectrum block’s log_spectrum_scale (or log_median_flux) metadata and averages the finite values, giving a single scalar that summarises the observation’s overall flux scale. Used as the regression target of the optional auxiliary scale loss. Returns None when no spectrum block carries a finite scale.

petitRADTRANS.sbi.flows.objectives._scale_targets_from_observations(observations: Any) → jax.numpy.ndarray | None#

Per-sample log-amplitude targets for the auxiliary scale loss.

Extracts one scalar log-scale per observation in the batch – from the pre-stacked log_scale channel (field index 4) of each spectrum block, or via _mean_spectral_log_scale() for block-list observations. The result is the regression target the scale_head is trained against when aux_scale_loss_weight > 0, which keeps an explicit amplitude signal in the embedding. Returns None for observation representations that carry no spectral scale.

petitRADTRANS.sbi.flows.objectives.batch_embeddings(model: _PosteriorModel, observations: Any) → jax.numpy.ndarray#

Encode a batch of observations into conditioning vectors.

Dispatches over the supported observation representations and returns a (batch, embedding_dim) float32 array: already-encoded arrays pass through; PreStackedObservations use the JIT-friendly encode_from_prestacked path; a list of equally-structured block lists uses the vmapped encode_stacked_batch; otherwise each observation is encoded individually. This is the single entry point the loss and refinement paths use to turn observations into the flow’s conditioning context.

petitRADTRANS.sbi.flows.objectives.compute_loss(model: _PosteriorModel, batch: petitRADTRANS.sbi.flows.base.PosteriorBatch) → jax.numpy.ndarray#

Evaluation NPE loss for one batch (no regularization terms).

Plain conditional negative log-likelihood, used for validation / model selection where the optimisation-only regularization terms should not perturb the reported value. See loss_from_observations().

petitRADTRANS.sbi.flows.objectives.training_loss(model: _PosteriorModel, batch: petitRADTRANS.sbi.flows.base.PosteriorBatch) → jax.numpy.ndarray#

Training NPE loss for one batch (with spline / invertibility regularization).

Identical to compute_loss() but enables the regularization terms that only shape the optimisation (spline bin/derivative penalties and the invertibility penalty). This is the function minimised by the trainer.

petitRADTRANS.sbi.flows.objectives.loss_from_observations(model: _PosteriorModel, parameters: jax.numpy.ndarray, observations: Any, parameter_space: str, *, include_spline_regularization: bool) → jax.numpy.ndarray#

Conditional negative-log-likelihood (NPE) objective for one batch.

Encodes observations into conditioning vectors, evaluates the flow’s log-density of parameters in the requested parameter_space (applying the unit-cube change of variables when needed), and returns the mean negative log-likelihood over the finite entries. Non-finite log-probabilities are masked out rather than allowed to poison the average, and the log-density is clipped to a wide finite range for numerical stability.

When include_spline_regularization is set, two optimisation-only terms are added (and only if their weights are positive): the spline bin-ratio / derivative penalties from _spline_regularization_terms(), which keep a rational-quadratic spline well-conditioned, and the invertibility penalty (invertibility_penalty()), which keeps an iterative-inverse flow invertible as it sharpens. When model.aux_scale_loss_weight > 0 an auxiliary mean-squared error between the scale_head prediction and the per-observation log-amplitude target (_scale_targets_from_observations()) is also added, preserving an explicit amplitude signal in the embedding.

Parameters#

model:

The encoder + flow bundle (_PosteriorModel).

parameters:

Parameter batch of shape (batch, n_parameters), in parameter_space coordinates.

observations:

Batch of observations in any representation accepted by batch_embeddings().

parameter_space:

'physical', 'cube', or 'unconstrained' – selects whether the cube change-of-variables correction is applied.

include_spline_regularization:

Whether to add the optimisation-only regularization terms (used for training, disabled for evaluation).

Returns#

jnp.ndarray

Scalar loss (mean negative log-likelihood plus any enabled penalties).

petitRADTRANS.sbi.flows.objectives.invertibility_penalty(model: _PosteriorModel, embeddings: jax.numpy.ndarray, n_samples: int) → jax.numpy.ndarray#

Differentiable inverse round-trip + log-det-closure penalty.

Draws base samples z ~ N(0, I), maps them to flow space with the (iterative) inverse, then back to latent space with the analytic forward. The inverse output is detached with stop_gradient so no gradient flows through the bisection solver; the gradient instead trains the forward map to be consistent with its own inverse. This penalises exactly the inverse_forward round-trip error and the inverse/forward log-det closure that the validation diagnostics report – the quantities that diverge as a neural-autoregressive flow sharpens (val_loss keeps falling while the sampler that draws posterior samples becomes non-invertible, so the deployed posterior predicts nonsense).

The penalty is bounded in scale: the reconstruction term is a mean squared error in latent space and the log-det term is normalised by the parameter dimension, so a single invertibility_loss_weight controls both. For flows with an exact analytic inverse (affine/spline) the penalty is ~0 and merely costs a little compute.

petitRADTRANS.sbi.flows.objectives.make_elbo_loss(log_likelihood_fn: Callable[[jax.numpy.ndarray, Any], jax.numpy.ndarray], parameter_space: str, *, num_samples: int = 1) → Callable[[_PosteriorModel, petitRADTRANS.sbi.flows.base.PosteriorBatch, Any], jax.numpy.ndarray]#

Return an amortized-variational (ELBO) training loss.

Maximizes, for observations x drawn from the (prior-predictive) dataset, the (tempered) ELBO

ELBO_beta(x) = E_{q(theta|x)}[ beta * log p(x|theta) + log p(theta)

  • log q(theta|x) ]

with a single- (or few-) sample reparameterized estimator. q(theta|x) is the conditional flow: latents z are drawn from the flow base and pushed through flow.inverse to reparameterized samples, so the gradient flows through both the flow and the encoder.

The returned loss takes a third argument beta (a scalar in (0, 1]): the inverse likelihood temperature. beta scales the log-likelihood, which for a Gaussian is exactly equivalent to inflating the noise to sigma_eff = sigma / sqrt(beta) (a power posterior p_beta whose target is the true posterior at beta = 1). The entropy -log q is kept at full weight. Annealing beta from small -> 1 tames the gradient stiffness from the very small observational noise (which would otherwise swamp the entropy term ~10^6 : 1 and drive an over-confident, poorly-located posterior), lets the entropy shape the width during annealing, and recovers the untempered posterior once beta = 1.

The -log q (entropy) term diverges to -inf as q collapses to a point mass, so a delta posterior is structurally impossible. log p(x|theta) is supplied by the injected differentiable forward-model likelihood evaluated at the physical parameters reconstructed from the sampled cube coordinates. With parameter_space='cube' the prior is uniform on the unit hypercube, so log p(theta_cube) = 0 and is dropped.

Parameters#

log_likelihood_fn:

(theta_cube, observations) -> (batch,) returning the Gaussian observational log-likelihood log p(x | theta) for each batch element. Receives unit-cube coordinates (shape (batch, dim)) and the batch’s prestacked observations; it owns the cube->physical transform, the differentiable forward model, and the noise model. Must be JAX-differentiable w.r.t. theta_cube.

parameter_space:

Must be "cube".

num_samples:

Number of reparameterized posterior draws per observation used to estimate the ELBO expectation. 1 is standard for amortized VI; larger values reduce gradient variance at a proportional forward-model cost.

petitRADTRANS.sbi.flows.objectives.compute_validation_diagnostics(model: _PosteriorModel, batch: petitRADTRANS.sbi.flows.base.PosteriorBatch, *, verbose_diagnostics: bool = False) → dict[str, float]#

Per-checkpoint flow invertibility / collapse diagnostics.

Probes the trained flow on a small slice of a validation batch and returns a dictionary of scalar metrics that quantify how trustworthy the flow’s sampler and density are – the inputs to the checkpoint-selection gate (_flow_validation_stability_metrics() / _flow_validation_selection_metric()). The metrics fall into a few groups:

  • Round-trip errors – forward_inverse_* and inverse_forward_* max/mean absolute errors and the log-det closure of the forward/inverse pair. A flow whose analytic forward and (possibly iterative) inverse disagree produces samples that do not match its own density; these errors grow as a neural-autoregressive flow sharpens, which is the dominant way a low-noise posterior silently degrades.

  • Iterative-solver health – residuals, bracket widths, expansion steps, and converged/failure fractions for flows inverted by bisection.

  • Support occupancy – cube_edge_hit_rate: the fraction of sampled coordinates pinned at the unit-cube edges, a signature of density collapse.

  • Spline conditioning – the statistics from _spline_parameter_statistics() (bin widths/heights, derivative extremes, floor-hit fractions).

Parameters#

model:

The encoder + flow bundle to probe.

batch:

Validation batch; only the leading _VALIDATION_DIAGNOSTIC_SAMPLES rows are used.

verbose_diagnostics:

When True, include the extra per-dimension cube-edge breakdown used by the verbose training diagnostics.

Returns#

dict[str, float]

Flat mapping of metric name to scalar value (nan for metrics that do not apply to the current flow family or an empty batch).