kinextract.bayesian.build_numpyro_model

kinextract.bayesian.build_numpyro_model(st: FitState, v_center: float)

Build the NumPyro model function for one LOSVD/template/continuum- offset posterior, at fixed regularization strength st.xlam.

b (the LOSVD histogram) and w (template weights) are both, physically, compositional/simplex-like quantities – a normalized histogram and a set of mixture weights, respectively (the forward model only ever uses w through the scale-invariant ratio (t @ w) / sum(w), so its absolute magnitude is unidentified by the likelihood, and b must satisfy sum(b) == 1). Each is sampled as an unconstrained latent Gaussian pushed through a softmax: z_b/z_w ~ Normal(0, 1), b = softmax(z_b), w = softmax(z_w). This exactly enforces non-negativity and sum() == 1 by construction, giving NUTS a naturally-scaled space to explore rather than a curved, box-constrained one – a bounded-Uniform parameterization with an ad hoc normalization penalty produces pathological NUTS trajectories (hundreds to thousands of leapfrog steps per transition, saturating the tree-depth ceiling regardless of target_accept_prob tuning), the classic signature of poorly-scaled, artificially box-constrained geometry. b/w are registered as numpyro.deterministic sites (not just local variables), so mcmc.get_samples() returns them directly under the "b"/"w" keys as their physical (non-negative, sum-to-1) values – every downstream consumer (_posterior_mean_vector(), kinextract.plotting.plot_losvd_posterior(), etc.) reads them exactly as it would any other named site.

The forward-model math (template mixture, LOSVD interpolation via the precomputed losvd_j0/j1/w tables, pixel-scatter convolution via ip_map/ip_mask, continuum multiplication, and the continuum_poly_mode additive/multiplicative correction) is reframed as a log-joint-density: the negative log-posterior (up to an additive constant) is 0.5*chi2 + 0.5*smooth – the 0.5 on both terms is what makes chi2 a genuine Gaussian log-likelihood and smooth (the already-quadratic wing-tapered second-difference roughness penalty, recentered per _wing_taper_lam_vec()) a genuine log-prior at the same relative strength that a given xlam means for the MAP-based xlam selection (kinextract.fitting._auto_select_xlam()) that picks it – omitting the 0.5 on smooth alone would leave the prior twice as strong, relative to the likelihood, than that selection intends. There is no separate normalization-penalty term: sum(b) == 1 already holds exactly by construction via the softmax parameterization above.

The continuum multiplier is a runtime argument of the returned model (model(cont)), not a value baked into the closure – passing it as an argument (with a fixed array shape) lets JAX/NumPyro reuse one compiled trace across calls instead of recompiling from scratch each time. In practice cont is always all-ones here (only pre-normalized-mode fits, cfg.fit_continuum=False, are supported by fit_state_bayesian()), but the model still accepts it as an explicit argument for the (unsupported) continuum-cofit case.

Parameters:

• st (FitState) – Fit state providing the data, template matrix, and precomputed LOSVD/pixel-shift interpolation tables.

• v_center (float) – Velocity to recenter the wing-taper regularization on (see estimate_velocity_xcorr()). Fixed for the whole fit.

Returns:

A NumPyro model function of one argument, model(cont), suitable for numpyro.infer.NUTS.

Return type:

callable