petitRADTRANS.sbi.flows.refinement

Test-time, per-observation posterior refinement.

Free functions that sharpen the amortized posterior for one specific target spectrum using the differentiable forward model: a flow-ELBO refinement (refine_posterior_on_observation()) and a MAP + Laplace / Gaussian-VI refinement (map_laplace_posterior_on_observation()). Each takes the trained ConditionalFlowPosterior as its first argument (posterior); the estimator exposes them as thin methods. Extracted from flows.posterior to keep the estimator focused on amortized inference.

Functions

_tile_prestacked_observation(→ Any)	Repeat a single (batch-1) prestacked observation n_tiles times.
refine_posterior_on_observation(...)	Semi-amortized test-time refinement of the posterior for one
importance_sampling_posterior_on_observation(...)	Likelihood importance sampling on top of the amortized posterior.
map_laplace_posterior_on_observation(...)	MAP + Laplace test-time inference for one observation (robust path).

Module Contents

petitRADTRANS.sbi.flows.refinement._tile_prestacked_observation(observation: Any, n_tiles: int) → Any

Repeat a single (batch-1) prestacked observation n_tiles times.

petitRADTRANS.sbi.flows.refinement.refine_posterior_on_observation(posterior, target_observation: Any, elbo_log_likelihood_fn: Callable[[jax.numpy.ndarray, Any], jax.numpy.ndarray], *, num_steps: int = 400, learning_rate: float = 0.0001, num_samples: int = 16, beta_start: float = 0.001, beta_anneal_fraction: float = 0.5, gradient_clip_norm: float = 1.0, gradient_clip_value: float = 1.0, seed: int = 0, verbose: bool = True) → petitRADTRANS.sbi.flows.posterior.ConditionalFlowPosterior

Semi-amortized test-time refinement of the posterior for one observation (flow-based ELBO variant).

Starting from the trained amortized model, optimize the single-observation tempered ELBO

ELBO_beta(x_obs) = E_{q(theta|c)}[ beta * log p(x_obs|theta)

• log q(theta|c) ]

adapting the flow while holding the encoder (hence the conditioning embedding c = encoder(x_obs)) fixed. beta is annealed soft -> 1 across the steps so the very small observational noise does not blow up the single-observation gradient early (which previously made this step diverge); a low learning rate and gradient clipping (both by global norm and by value) further stabilize it. The best beta=1 ELBO state seen during optimization is returned (not the last), so a late instability cannot degrade the result.

Returns a shallow copy of this posterior whose flow is specialized to target_observation; use it only for inference on that observation.

See also map_laplace_posterior_on_observation(), the more robust MAP + Laplace variant that sidesteps reverse-KL under-dispersion.

petitRADTRANS.sbi.flows.refinement.importance_sampling_posterior_on_observation(posterior, target_observation: Any, standardized_residual_fn: Callable[[jax.numpy.ndarray], jax.numpy.ndarray], *, n_proposal_samples: int = 16384, forward_batch_size: int = 256, n_posterior_samples: int = 4096, seed: int = 0, verbose: bool = True) → petitRADTRANS.sbi.flows.posterior.ConditionalFlowPosterior

Likelihood importance sampling on top of the amortized posterior.

Draws proposal samples from the amortized flow conditioned on the target observation, evaluates the exact Gaussian likelihood of each through the differentiable forward model, and reweights by log w = log L(theta) + log pi(theta) - log q(theta | x). This is the field-standard correction for amortized posteriors (Gebhard et al. 2025): unlike MAP+Laplace it cannot collapse to a degenerate covariance, and its sampling efficiency eps = ESS / N is a built-in trust diagnostic — eps near 1 means the amortized posterior already matches the true one; eps near 0 means the amortized posterior misses the true posterior mass and the importance estimate itself should not be trusted either.

Weights are computed in the flow’s unconstrained (logit-of-cube) space, where the uniform-cube prior is the standard logistic density and the proposal density is the flow’s own log_prob — both exact.

Parameters

target_observation:

Batch-1 prestacked target observation (conditions the proposal).

standardized_residual_fn:

Differentiable theta_cube[param_dim] -> (n_wavelengths,) returning noise-standardized residuals (model(theta) - x_obs)/sigma for the target; masked channels contribute 0.

n_proposal_samples:

Number of amortized proposal draws to weight.

forward_batch_size:

Forward-model chunk size (vmapped); bounds device memory exactly like the ELBO per-batch forward passes did.

n_posterior_samples:

Number of importance-resampled posterior draws stored on the returned estimator.

Returns

ConditionalFlowPosterior

A shallow copy of posterior whose _fixed_cube_samples hold the importance-resampled posterior, with the weighting diagnostics attached as _importance_sampling_diagnostics.

petitRADTRANS.sbi.flows.refinement.map_laplace_posterior_on_observation(posterior, target_observation: Any, standardized_residual_fn: Callable[[jax.numpy.ndarray], jax.numpy.ndarray], *, num_steps: int = 600, learning_rate: float = 0.05, beta_start: float = 0.001, beta_anneal_fraction: float = 0.5, n_posterior_samples: int = 4096, gradient_clip_value: float = 10.0, eigenvalue_floor: float = 1e-06, covariance_method: str = 'laplace', vi_num_steps: int = 300, vi_learning_rate: float = 0.01, vi_num_samples: int = 32, vi_freeze_mean: bool = True, vi_beta: float = 1.0, seed: int = 0, verbose: bool = True) → petitRADTRANS.sbi.flows.posterior.ConditionalFlowPosterior

MAP + Laplace test-time inference for one observation (robust path).

Optimizes the (beta-annealed -> 1) log-posterior in the unconstrained logit space u (theta_cube = sigmoid(u), uniform-cube prior -> logistic prior on u) to the MAP, initialized from the amortized posterior mean for the target, then forms a Gauss-Newton Laplace covariance from the forward-model Jacobian and the prior curvature:

precision Lambda = J^T J + diag(2 c (1 - c)), J = d r / d u

where r are the noise-standardized residuals. Posterior samples are drawn as u ~ N(u*, Lambda^-1) and mapped through the sigmoid. This gives the correct location (MAP) and correctly-scaled, per-direction local widths (tight where well-constrained, e.g. planet radius; moderate on degenerate directions), sidestepping the reverse-KL under-dispersion of the flow ELBO.

Returns a shallow copy of this posterior carrying the Laplace cube samples; its sample_posterior() returns those samples, so the existing inference/predictive pipeline can consume it unchanged. Use it only for inference on target_observation.

Parameters

target_observation:

Batch-1 prestacked target observation (used only for the amortized MAP initialization via the encoder).

standardized_residual_fn:

Differentiable (theta_cube[param_dim]) -> (n_wavelengths,) returning noise-standardized residuals (model(theta) - x_obs)/sigma for the target. Masked channels should contribute 0.