petitRADTRANS.sbi.flows.regularization#

Spline regularization and checkpoint-selection stability diagnostics.

Training-time helpers for the conditional spline flow: the rational-quadratic spline regularization terms (bin-ratio / derivative penalties) added to the NPE loss, the per-checkpoint spline parameter statistics, and the flow-invertibility / collapse metrics used to gate and score checkpoint selection. Also holds the small batch-slicing helpers shared by the validation diagnostics. Extracted from flows.posterior to keep that module focused on the estimator itself.

Attributes#

`_VALIDATION_DIAGNOSTIC_SAMPLES`
`_DIAGNOSTIC_EDGE_THRESHOLD`
`_VALIDATION_CALIBRATION_DRAWS`
`_REGULARIZATION_TAIL_FRACTION`
`_DIAGNOSTIC_FLOOR_RELATIVE_TOLERANCE`
`_ELBO_STABILITY_INVERSE_FORWARD_MAX_ABS_ERROR`
`_ELBO_STABILITY_LOGDET_CLOSURE_MAX_ABS_ERROR`
`_ELBO_STABILITY_CUBE_EDGE_HIT_RATE`
`_PATHOLOGICAL_INVERSE_FORWARD_MAX_ABS_ERROR`
`_SELECTION_PENALTY_WEIGHT`
`_SELECTION_PENALTY_CAP`
`_CALIBRATION_SELECTION_PENALTY_WEIGHT`
`_PATHOLOGICAL_CALIBRATION_COVERAGE_ERROR`

Functions#

`_slice_observations_for_diagnostics`(→ Any)	Return the first `n_samples` observations, preserving the container type.
`_slice_batch_for_diagnostics`(...)	Sub-sample a batch to at most `max_samples` rows for cheap diagnostics.
`_normalized_categorical_entropy`(→ jax.numpy.ndarray)	Entropy of `softmax(logits)` normalised to `[0, 1]` by `log(n_categories)`.
`_weighted_average`(→ jax.numpy.ndarray)	Weighted mean of `values` (plain mean when `weights` is `None`).
`_upper_tail_average`(→ jax.numpy.ndarray)	Mean of the largest `tail_fraction` of `values` (optionally weighted).
`_log_barrier_below_target`(→ jax.numpy.ndarray)	Squared log-barrier penalising `values` that fall below `target`.
`_minimum_bin_ratio_statistics`(...)	Per-sample and global minimum bin-size ratios relative to a uniform bin.
`_spline_layer_regularization_terms`(...)	Regularization penalties and conditioning statistics for one spline layer.
`_spline_regularization_terms`(...)	Batch spline regularization terms added to the NPE training loss.
`_spline_parameter_statistics`(→ dict[str, float])	Diagnostic statistics of a spline flow's bin and derivative parameters.
`_flow_validation_selection_metric`(→ float)	Checkpoint-selection metric: validation loss plus a bounded instability penalty.
`_flow_validation_stability_metrics`(→ dict[str, float])	Classify a checkpoint as stable and/or pathological from its diagnostics.

Module Contents#

petitRADTRANS.sbi.flows.regularization._VALIDATION_DIAGNOSTIC_SAMPLES = 64#

petitRADTRANS.sbi.flows.regularization._DIAGNOSTIC_EDGE_THRESHOLD = 1e-06#

petitRADTRANS.sbi.flows.regularization._VALIDATION_CALIBRATION_DRAWS = 64#

petitRADTRANS.sbi.flows.regularization._REGULARIZATION_TAIL_FRACTION = 0.25#

petitRADTRANS.sbi.flows.regularization._DIAGNOSTIC_FLOOR_RELATIVE_TOLERANCE = 1.05#

petitRADTRANS.sbi.flows.regularization._ELBO_STABILITY_INVERSE_FORWARD_MAX_ABS_ERROR = 5.0#

petitRADTRANS.sbi.flows.regularization._ELBO_STABILITY_LOGDET_CLOSURE_MAX_ABS_ERROR = 10.0#

petitRADTRANS.sbi.flows.regularization._ELBO_STABILITY_CUBE_EDGE_HIT_RATE = 0.5#

petitRADTRANS.sbi.flows.regularization._PATHOLOGICAL_INVERSE_FORWARD_MAX_ABS_ERROR = 2.0#

petitRADTRANS.sbi.flows.regularization._SELECTION_PENALTY_WEIGHT = 1.0#

petitRADTRANS.sbi.flows.regularization._SELECTION_PENALTY_CAP = 12.0#

petitRADTRANS.sbi.flows.regularization._CALIBRATION_SELECTION_PENALTY_WEIGHT = 4.0#

petitRADTRANS.sbi.flows.regularization._PATHOLOGICAL_CALIBRATION_COVERAGE_ERROR = 0.6#

petitRADTRANS.sbi.flows.regularization._slice_observations_for_diagnostics(observations: Any, n_samples: int) → Any#: Return the first n_samples observations, preserving the container type.

petitRADTRANS.sbi.flows.regularization._slice_batch_for_diagnostics(batch: petitRADTRANS.sbi.flows.base.PosteriorBatch, max_samples: int = _VALIDATION_DIAGNOSTIC_SAMPLES) → petitRADTRANS.sbi.flows.base.PosteriorBatch#: Sub-sample a batch to at most max_samples rows for cheap diagnostics.

petitRADTRANS.sbi.flows.regularization._normalized_categorical_entropy(logits: jax.numpy.ndarray) → jax.numpy.ndarray#

Entropy of softmax(logits) normalised to [0, 1] by log(n_categories).

A spline-bin uniformity measure: 1 when the predicted bin widths/heights are uniform, decaying toward 0 as probability mass collapses onto a few bins.

petitRADTRANS.sbi.flows.regularization._weighted_average(values: jax.numpy.ndarray, weights: jax.numpy.ndarray | None = None) → jax.numpy.ndarray#

Weighted mean of values (plain mean when weights is None).

Weights are broadcast to the value shape and the result is normalised by their sum (floored at 1 to avoid division by zero).

petitRADTRANS.sbi.flows.regularization._upper_tail_average(values: jax.numpy.ndarray, weights: jax.numpy.ndarray | None = None, *, tail_fraction: float = _REGULARIZATION_TAIL_FRACTION) → jax.numpy.ndarray#

Mean of the largest tail_fraction of values (optionally weighted).

Concentrates a penalty on the worst-offending elements – e.g. the few spline bins closest to collapse – instead of letting a global average dilute them. With weights only positive-weight elements are eligible; at least one element is always included in the tail.

petitRADTRANS.sbi.flows.regularization._log_barrier_below_target(values: jax.numpy.ndarray, target: float) → jax.numpy.ndarray#

Squared log-barrier penalising values that fall below target.

Returns relu(log(target) - log(value))**2: zero once a value reaches the target, growing quadratically in log-space as it drops below it.

petitRADTRANS.sbi.flows.regularization._minimum_bin_ratio_statistics(bin_sizes: jax.numpy.ndarray, uniform_bin_size: Any, active_weights: jax.numpy.ndarray | None = None) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

Per-sample and global minimum bin-size ratios relative to a uniform bin.

Each ratio is a spline bin size divided by the uniform bin size; small values flag bins collapsing toward their floor. Returns (minimum_ratios, weighted_mean_minimum_ratio, global_minimum_ratio); active_weights restricts the mean and global minimum to the unmasked dimensions.

petitRADTRANS.sbi.flows.regularization._spline_layer_regularization_terms(widths: jax.numpy.ndarray, heights: jax.numpy.ndarray, unnormalized_derivatives: jax.numpy.ndarray, *, min_derivative: float, active_weights: jax.numpy.ndarray | None, min_bin_ratio_target: float, min_derivative_target: float, max_derivative_target: float) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

Regularization penalties and conditioning statistics for one spline layer.

From a layer’s normalised bin widths / heights and raw unnormalized_derivatives, builds two penalties: a small-bin penalty (a log-barrier on the minimum width/height ratios that fall below min_bin_ratio_target, averaged over the worst tail) and a derivative penalty pushing the softplus-mapped slopes into [min_derivative_target, max_derivative_target] in log-space. Together they stop the spline from developing vanishing bins or extreme slopes as the posterior sharpens. active_weights masks the transformed (non-identity) dimensions of a coupling layer.

Returns#

tuple: (small_bin_penalty, width_penalty, height_penalty, derivative_penalty, mean_min_width_ratio, mean_min_height_ratio, global_min_width_ratio, global_min_height_ratio) – the first four are penalties, the last four are diagnostic ratios.

petitRADTRANS.sbi.flows.regularization._spline_regularization_terms(flow: Any, inputs: jax.numpy.ndarray, embeddings: jax.numpy.ndarray, *, min_bin_ratio_target: float, min_derivative_target: float, max_derivative_target: float) → tuple[jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray, jax.numpy.ndarray]#

Batch spline regularization terms added to the NPE training loss.

Walks the flow’s layers and, for every rational-quadratic spline layer (coupling or scalar), accumulates the per-layer penalties from _spline_layer_regularization_terms(), evaluated at each layer’s own inputs as the batch is pushed forward through the flow. The penalties keep the spline well-conditioned – no vanishing bins, slopes in a sane range – so a sharpening posterior does not develop pathological transforms. Returns zeros (and nan diagnostics) for non-spline flows.

Parameters#

flow:: Conditional flow; penalties apply only to its spline layers.
inputs:: Flow-space parameter batch of shape (batch, n_parameters).
embeddings:: Conditioning vectors of shape (batch, embedding_dim).
min_bin_ratio_target, min_derivative_target, max_derivative_target:: Targets defining where each penalty begins to bite.

Returns#

tuple: (small_bin_penalty, width_penalty, height_penalty, derivative_penalty, mean_min_width_ratio, mean_min_height_ratio, global_min_width_ratio, global_min_height_ratio) aggregated over the batch.

petitRADTRANS.sbi.flows.regularization._spline_parameter_statistics(flow: Any, inputs: jax.numpy.ndarray, embeddings: jax.numpy.ndarray) → dict[str, float]#

Diagnostic statistics of a spline flow’s bin and derivative parameters.

Pushes the batch through the flow and, for each spline layer, records the extremes and floor-hit fractions of the normalised bin widths/heights and the boundary derivatives, plus the mean normalised width/height entropy and mean absolute log-derivative. These feed the checkpoint diagnostics – e.g. spline_max_derivative participates in the pathology gate. Returns a dict of nan for non-spline flows.

petitRADTRANS.sbi.flows.regularization._flow_validation_selection_metric(validation_loss: float, diagnostics: Mapping[str, Any] | None) → float#

Checkpoint-selection metric: validation loss plus a bounded instability penalty.

Lower is better. Returns the raw validation loss for a stable checkpoint; a large constant (1e6) for a pathological/collapsed one so it can never be selected; and, for an unstable-but-usable checkpoint, the loss plus a smooth, capped penalty that grows with how far each round-trip / log-det-closure / cube-edge-hit diagnostic exceeds its threshold. The cap keeps a much-better-loss checkpoint selectable rather than letting a marginally-over-threshold diagnostic veto it.

petitRADTRANS.sbi.flows.regularization._flow_validation_stability_metrics(validation_loss: float, diagnostics: Mapping[str, Any] | None, *, inverse_forward_max_abs_error_threshold: float | None, inverse_forward_logdet_closure_max_abs_error_threshold: float | None, cube_edge_hit_rate_threshold: float | None, calibration_coverage_error_threshold: float | None = None) → dict[str, float]#

Classify a checkpoint as stable and/or pathological from its diagnostics.

Compares the per-checkpoint diagnostics against the supplied thresholds and returns a flat metric dict containing, per checked quantity, a pass flag and signed margin, plus two summary gates: checkpoint_is_stable (no soft threshold violated) and checkpoint_is_pathological (density collapse – a non-finite loss, most samples pinned at the cube edges, exploded spline derivatives, or an inverse round-trip error so large the map is effectively non-invertible). The selection metric uses the soft gate; the hard pathology gate additionally disqualifies a checkpoint and bars it from resetting early-stopping patience, so a collapsing run is stopped rather than allowed to masquerade its diverging loss as improvement.

Parameters#

validation_loss:: The checkpoint’s validation loss (a non-finite value is itself pathological).
diagnostics:: Per-checkpoint diagnostics (see compute_validation_diagnostics()).

inverse_forward_max_abs_error_threshold: inverse_forward_logdet_closure_max_abs_error_threshold: cube_edge_hit_rate_threshold:

Soft-gate thresholds; None disables the corresponding check.

calibration_coverage_error_threshold:: Maximum tolerated deviation of the empirical central-80%-interval coverage from 0.8 on held-out validation simulations (the per-epoch calibration probe). None disables the check.