Skip to main content

gam_models/bms/
gradient_paths.rs

1use super::family::clamp_bernoulli_link_probability;
2use super::*;
3use gam_linalg::faer_ndarray::FaerEigh;
4use gam_linalg::matrix::{FiniteSignedWeightsView, LinearOperator};
5use gam_math::jet_scalar::SymmetricQuadraticCoefficients;
6use gam_math::jet_tower::Tower4;
7use gam_math::probability::normal_logcdf_derivatives;
8use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, escalate_ridge};
9
10pub(crate) fn standardize_latent_z_with_policy(
11    z: &Array1<f64>,
12    weights: &Array1<f64>,
13    context: &str,
14    policy: &LatentZPolicy,
15) -> Result<(Array1<f64>, LatentZNormalization), String> {
16    if z.len() != weights.len() {
17        return Err(format!(
18            "{context} latent-score normalization length mismatch: z={}, weights={}",
19            z.len(),
20            weights.len()
21        ));
22    }
23    let weight_sum = weights.iter().copied().sum::<f64>();
24    let weight_sq_sum = weights.iter().map(|&w| w * w).sum::<f64>();
25    if !(weight_sum.is_finite()
26        && weight_sum > 0.0
27        && weight_sq_sum.is_finite()
28        && weight_sq_sum > 0.0)
29    {
30        return Err(format!("{context} requires positive finite total weight"));
31    }
32    let effective_n = weight_sum * weight_sum / weight_sq_sum;
33    if !(effective_n.is_finite() && effective_n > 1.0) {
34        return Err(format!(
35            "{context} requires at least two effective observations for latent-score normalization"
36        ));
37    }
38    let mean = z
39        .iter()
40        .zip(weights.iter())
41        .map(|(&zi, &wi)| wi * zi)
42        .sum::<f64>()
43        / weight_sum;
44    let var = z
45        .iter()
46        .zip(weights.iter())
47        .map(|(&zi, &wi)| wi * (zi - mean) * (zi - mean))
48        .sum::<f64>()
49        / weight_sum;
50    let sd = var.sqrt();
51    if !(sd.is_finite() && sd > BMS_VARIANCE_FLOOR) {
52        return Err(format!(
53            "{context} requires z with positive finite weighted standard deviation"
54        ));
55    }
56    let target_norm = match policy.normalization {
57        LatentZNormalizationMode::None => LatentZNormalization { mean: 0.0, sd: 1.0 },
58        LatentZNormalizationMode::FitWeighted => LatentZNormalization { mean, sd },
59        LatentZNormalizationMode::Frozen {
60            mean: frozen_mean,
61            sd: frozen_sd,
62        } => LatentZNormalization {
63            mean: frozen_mean,
64            sd: frozen_sd,
65        },
66    };
67    let mean_tol = policy.mean_tol_multiplier / effective_n.sqrt();
68    let sd_tol = policy.sd_tol_multiplier / (2.0 * (effective_n - 1.0).max(1.0)).sqrt();
69    let check_msg = || {
70        format!(
71            "{context} requires z to already be approximately latent N(0,1) before identification normalization; got mean={mean:.6e}, sd={sd:.6e}, effective_n={effective_n:.1}, allowed_mean={mean_tol:.3e}, allowed_sd={sd_tol:.3e}"
72        )
73    };
74    if mean.abs() > mean_tol || (sd - 1.0).abs() > sd_tol {
75        match policy.check_mode {
76            LatentZCheckMode::Strict => return Err(check_msg()),
77            LatentZCheckMode::WarnOnly => log::warn!("{}", check_msg()),
78            LatentZCheckMode::Off => {}
79        }
80    }
81
82    let normalization = target_norm;
83    let z_std = normalization.apply(z, context)?;
84    // Standardized moments of z_std itself. `z_std` has weighted mean 0 and
85    // variance 1 only in `FitWeighted` mode; under `None`/`Frozen` its raw
86    // third/fourth moments are NOT the named statistics (a ×3-scaled Gaussian
87    // would read "excess_kurtosis≈240"), so center and scale by z_std's own
88    // weighted moments before labeling them skewness / excess kurtosis.
89    let std_mean = z_std
90        .iter()
91        .zip(weights.iter())
92        .map(|(&zi, &wi)| wi * zi)
93        .sum::<f64>()
94        / weight_sum;
95    let std_var = (z_std
96        .iter()
97        .zip(weights.iter())
98        .map(|(&zi, &wi)| wi * (zi - std_mean) * (zi - std_mean))
99        .sum::<f64>()
100        / weight_sum)
101        .max(f64::MIN_POSITIVE);
102    let skew = z_std
103        .iter()
104        .zip(weights.iter())
105        .map(|(&zi, &wi)| wi * (zi - std_mean).powi(3))
106        .sum::<f64>()
107        / weight_sum
108        / std_var.powf(1.5);
109    let kurt = z_std
110        .iter()
111        .zip(weights.iter())
112        .map(|(&zi, &wi)| wi * (zi - std_mean).powi(4))
113        .sum::<f64>()
114        / weight_sum
115        / (std_var * std_var)
116        - 3.0;
117    if skew.abs() > policy.max_abs_skew || kurt.abs() > policy.max_abs_excess_kurtosis {
118        let msg = format!(
119            "{context} requires z to be approximately Gaussian after identification normalization; got skewness={skew:.3}, excess_kurtosis={kurt:.3}"
120        );
121        match policy.check_mode {
122            LatentZCheckMode::Strict => return Err(msg),
123            LatentZCheckMode::WarnOnly => log::warn!("{}", msg),
124            LatentZCheckMode::Off => {}
125        }
126    }
127    if skew.abs() > 0.75 || kurt.abs() > 2.0 {
128        log::warn!(
129            "{context}: z has skewness={skew:.3} and excess kurtosis={kurt:.3}; latent-measure auto-selection will use empirical calibration unless stricter diagnostics pass"
130        );
131    }
132    Ok((z_std, normalization))
133}
134
135pub fn padded_deviation_seed(seed: &Array1<f64>, min_iqr: f64, pad_fraction: f64) -> Array1<f64> {
136    let mut sorted = seed.to_vec();
137    sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
138
139    if sorted.len() < 4 {
140        return seed.clone();
141    }
142
143    let n = sorted.len();
144    let q1 = sorted[n / 4];
145    let q3 = sorted[3 * n / 4];
146    let iqr = (q3 - q1).max(min_iqr);
147    let pad = pad_fraction * iqr;
148
149    let mut out = seed.to_vec();
150    out.push(sorted[0] - pad);
151    out.push(sorted[n - 1] + pad);
152    Array1::from_vec(out)
153}
154
155// ── Pooled 2-D probit pilot Newton solver tuning ─────────────────────────────
156//
157// `pooled_probit_baseline` solves a 2-parameter (intercept, slope) penalised
158// probit by damped Newton. The values below are the standard convergence /
159// safeguard knobs; they are deliberately conservative because the pilot is a
160// cheap warm-start for the full fit, not the production estimator.
161
162/// Maximum damped-Newton outer iterations for the pooled probit pilot. A 2-D
163/// strictly-convex probit converges in well under this; the cap only guards a
164/// pathological non-finite data configuration.
165const POOLED_PILOT_MAX_NEWTON_ITERS: usize = 50;
166/// Initial Levenberg ridge added to the 2×2 Hessian diagonal before the solve.
167pub(crate) const POOLED_PILOT_RIDGE_INIT: f64 = 1e-8;
168/// Below this absolute determinant the ridged 2×2 system is treated as
169/// singular and the ridge is escalated.
170pub(crate) const POOLED_PILOT_DET_FLOOR: f64 = 1e-18;
171/// Geometric factor by which the ridge grows when the system is singular.
172pub(crate) const POOLED_PILOT_RIDGE_GROWTH: f64 = 10.0;
173/// Ridge ceiling; exceeding it means the Hessian is unusable and the pilot
174/// fails rather than returning a meaningless step.
175pub(crate) const POOLED_PILOT_RIDGE_MAX: f64 = 1e6;
176/// Maximum backtracking-line-search halvings per Newton step.
177const POOLED_PILOT_MAX_BACKTRACKS: usize = 25;
178/// Backtracking step contraction factor.
179pub(crate) const POOLED_PILOT_BACKTRACK_SHRINK: f64 = 0.5;
180/// Objective-change tolerance below which a stalled (rejected) line search is
181/// accepted as converged instead of erroring.
182pub(crate) const POOLED_PILOT_STALL_TOL: f64 = 1e-10;
183/// Minimum-magnitude signed slope returned by the pilot, so the downstream
184/// `b/√(1+b²)` rigid seed never collapses to an exactly flat (zero-slope) link.
185pub(crate) const POOLED_PILOT_MIN_ABS_SLOPE: f64 = 1e-6;
186
187pub(super) fn pooled_probit_baseline(
188    y: &Array1<f64>,
189    z: &Array1<f64>,
190    weights: &Array1<f64>,
191) -> Result<(f64, f64), String> {
192    if y.len() != z.len() || y.len() != weights.len() {
193        return Err(format!(
194            "pooled bernoulli-marginal-slope pilot length mismatch: y={}, z={}, weights={}",
195            y.len(),
196            z.len(),
197            weights.len()
198        ));
199    }
200    let weight_sum = weights.iter().copied().sum::<f64>();
201    if !weight_sum.is_finite() || weight_sum <= 0.0 {
202        return Err(
203            "pooled bernoulli-marginal-slope pilot requires positive finite total weight"
204                .to_string(),
205        );
206    }
207    let prevalence = y
208        .iter()
209        .zip(weights.iter())
210        .map(|(&yi, &wi)| yi * wi)
211        .sum::<f64>()
212        / weight_sum;
213    let prevalence = prevalence.clamp(1e-6, 1.0 - 1e-6);
214    let z_mean = z
215        .iter()
216        .zip(weights.iter())
217        .map(|(&zi, &wi)| zi * wi)
218        .sum::<f64>()
219        / weight_sum;
220    let z_var = z
221        .iter()
222        .zip(weights.iter())
223        .map(|(&zi, &wi)| wi * (zi - z_mean) * (zi - z_mean))
224        .sum::<f64>()
225        / weight_sum;
226    let yz_cov = y
227        .iter()
228        .zip(z.iter())
229        .zip(weights.iter())
230        .map(|((&yi, &zi), &wi)| wi * (yi - prevalence) * (zi - z_mean))
231        .sum::<f64>()
232        / weight_sum;
233    let mut beta0 = standard_normal_quantile(prevalence).map_err(|e| {
234        format!("failed to initialize pooled bernoulli-marginal-slope pilot intercept: {e}")
235    })?;
236    let mut beta1 = if z_var > BMS_VARIANCE_FLOOR {
237        yz_cov / z_var
238    } else {
239        0.0
240    };
241
242    let objective_grad_hess =
243        |intercept: f64, slope: f64| -> Result<(f64, f64, f64, f64, f64, f64), String> {
244            let mut obj = 0.0;
245            let mut g0 = 0.0;
246            let mut g1 = 0.0;
247            let mut h00 = 0.0;
248            let mut h01 = 0.0;
249            let mut h11 = 0.0;
250            for ((&yi, &zi), &wi) in y.iter().zip(z.iter()).zip(weights.iter()) {
251                if wi == 0.0 {
252                    continue;
253                }
254                let eta = intercept + slope * zi;
255                let s = 2.0 * yi - 1.0;
256                let margin = s * eta;
257                let probit = normal_logcdf_derivatives(margin);
258                let logcdf = probit[0];
259                let lambda = probit[1];
260                let g_eta = -wi * s * lambda;
261                let h_eta = -wi * probit[2];
262                obj -= wi * logcdf;
263                g0 += g_eta;
264                g1 += g_eta * zi;
265                h00 += h_eta;
266                h01 += h_eta * zi;
267                h11 += h_eta * zi * zi;
268            }
269            Ok((obj, g0, g1, h00, h01, h11))
270        };
271
272    let mut obj_prev = f64::INFINITY;
273    for _ in 0..POOLED_PILOT_MAX_NEWTON_ITERS {
274        let (obj, g0, g1, h00, h01, h11) = objective_grad_hess(beta0, beta1)?;
275        if !obj.is_finite() || !g0.is_finite() || !g1.is_finite() {
276            return Err(
277                "pooled bernoulli-marginal-slope pilot produced non-finite objective or gradient"
278                    .to_string(),
279            );
280        }
281        let grad_max = g0.abs().max(g1.abs());
282        if grad_max < BMS_DERIV_TOL {
283            break;
284        }
285        // Ridge budget: the pre-migration loop grew δ from RIDGE_INIT by
286        // RIDGE_GROWTH until it exceeded RIDGE_MAX, so the trial count is the
287        // decade span of [RIDGE_INIT, RIDGE_MAX] inclusive.
288        let ridge_trials = (POOLED_PILOT_RIDGE_MAX / POOLED_PILOT_RIDGE_INIT)
289            .log10()
290            .ceil() as usize
291            + 1;
292        let (step0, step1) = escalate_ridge(
293            RidgeSchedule {
294                initial: POOLED_PILOT_RIDGE_INIT,
295                growth: POOLED_PILOT_RIDGE_GROWTH,
296                max_escalations: ridge_trials,
297            },
298            |ridge| {
299                let h00_r = h00 + ridge;
300                let h11_r = h11 + ridge;
301                let det = h00_r * h11_r - h01 * h01;
302                if !(det.is_finite() && det.abs() > POOLED_PILOT_DET_FLOOR) {
303                    return None;
304                }
305                let s0 = (h11_r * g0 - h01 * g1) / det;
306                let s1 = (-h01 * g0 + h00_r * g1) / det;
307                (s0.is_finite() && s1.is_finite()).then_some((s0, s1))
308            },
309        )
310        .map(|success| success.value)
311        .map_err(|_| "pooled bernoulli-marginal-slope pilot Hessian solve failed".to_string())?;
312        let accepted = backtracking_line_search::<_, String>(
313            BacktrackConfig {
314                contraction: POOLED_PILOT_BACKTRACK_SHRINK,
315                max_steps: POOLED_PILOT_MAX_BACKTRACKS,
316                ..BacktrackConfig::default()
317            },
318            |step_scale| {
319                let cand0 = beta0 - step_scale * step0;
320                let cand1 = beta1 - step_scale * step1;
321                let (cand_obj, _, _, _, _, _) = objective_grad_hess(cand0, cand1)?;
322                Ok(Some((cand_obj, (cand0, cand1))))
323            },
324            |_scale, cand_obj| cand_obj.is_finite() && cand_obj <= obj,
325        )?;
326        match accepted {
327            Some(step) => {
328                (beta0, beta1) = step.payload;
329                obj_prev = step.value;
330            }
331            None => {
332                if (obj_prev - obj).abs() < POOLED_PILOT_STALL_TOL {
333                    break;
334                }
335                return Err("pooled bernoulli-marginal-slope pilot line search failed".to_string());
336            }
337        }
338    }
339    let a = beta0;
340    // Signed slope: preserve direction from pilot probit.
341    let b = if beta1.abs() < POOLED_PILOT_MIN_ABS_SLOPE {
342        if beta1.is_sign_negative() {
343            -POOLED_PILOT_MIN_ABS_SLOPE
344        } else {
345            POOLED_PILOT_MIN_ABS_SLOPE
346        }
347    } else {
348        beta1
349    };
350    Ok((a / (1.0 + b * b).sqrt(), b))
351}
352
353// Compute a non-degenerate pilot η for the link-deviation cross-block
354// identifiability orthogonalisation.
355//
356// The rigid pooled probit pilot from `pooled_probit_baseline` is a scalar
357// pair `(a₀, b₀)`, so the rigid observed-scale linear predictor
358// `η_rigid[i] = a₀·√(1 + (s_f·b₀)²) + s_f·b₀·z[i]` is **exactly affine in z**
359// when the per-row offsets are zero. A degree-3 I-spline of an affine
360// function of `z` spans the same column space at training rows as a
361// degree-3 I-spline of `z` directly, so evaluating the link-deviation basis
362// at `η_rigid` and orthogonalising it against the score-warp basis (built
363// on `z`) produces a structurally singular cross-Gram — the candidate is
364// fully aliased even though at PIRLS time the link-deviation runtime is
365// re-evaluated at the current β-dependent η which carries genuine PC / age
366// structure that the score-warp cannot represent.
367//
368// One probit Gauss-Newton step from the rigid pilot, projected onto the
369// full marginal design at the W-IRLS working response, picks up that PC /
370// age structure cheaply (one `p_marg × p_marg` Cholesky plus a few matvecs
371// — `<<1 s` at large scale because `p_marg` is `O(10²)` whereas the
372// PIRLS dense Hessian build is `O(n·p²)` per cycle). The resulting
373// `η_pilot[i]` has the same row-by-row variation pattern PIRLS will see at
374// any non-degenerate β, so the orthogonalisation transform `T` drops only
375// the directions that are aliased *across all* β, not those that are
376// aliased only at the rigid (rank-1-in-z) pilot.
377/// IRLS Hessian row metric for the probit-style data Hessian at a fixed
378/// linear predictor `eta`: `w[i] = sample_weights[i] · φ(η_i)² / (μ_i·(1−μ_i))`.
379///
380/// This is the canonical row metric that the joint penalised Hessian sees
381/// during PIRLS for a probit GLM (and the dominant term for
382/// BernoulliMarginalSlope's data Hessian). Cross-block orthogonalisation
383/// against parametric anchors must use **this** metric — not a uniform
384/// W=spec.weights — for the joint Hessian to be block-orthogonal between
385/// parametric and flex spans. With a uniform W the orthogonalisation only
386/// kills the Euclidean alias; at PIRLS time `Aᵀ W_pirls C̃ ≠ 0` and the
387/// joint Hessian carries a near-null direction along the W-metric alias,
388/// which REML can drive to arbitrarily small eigenvalue by shrinking the
389/// flex block's smoothing parameter — β then runs away along the alias
390/// (the failure mode that manifests as `rho≈2.0`, constant `step_inf`,
391/// and `beta_inf` growing without bound during PIRLS).
392pub(super) fn pilot_irls_hessian_row_metric_at_eta(
393    eta_pilot: &Array1<f64>,
394    sample_weights: &Array1<f64>,
395) -> Array1<f64> {
396    let n = eta_pilot.len();
397    let mut w = Array1::<f64>::zeros(n);
398    for i in 0..n {
399        let eta = eta_pilot[i];
400        let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
401        let phi = normal_pdf(eta).max(1e-300);
402        let var = (mu * (1.0 - mu)).max(1e-300);
403        w[i] = sample_weights[i] * (phi * phi) / var;
404    }
405    w
406}
407
408/// Per-row rigid pooled-probit pilot η used to seed the IRLS Hessian
409/// metric for score-warp cross-block orthogonalisation. Score-warp's
410/// basis is evaluated at `z` (β-independent) so there is no GN-stepped
411/// pilot to share with the link-deviation path; the rigid pooled-probit
412/// pilot is a sensible β-independent reference at which to evaluate
413/// `W = p(1−p)·spec.weights` for the W-metric orthogonalisation.
414pub(super) fn rigid_pooled_probit_pilot_eta(
415    base_link: &InverseLink,
416    z: &Array1<f64>,
417    marginal_offset: &Array1<f64>,
418    logslope_offset: &Array1<f64>,
419    baseline_marginal: f64,
420    baseline_logslope: f64,
421    probit_scale: f64,
422) -> Result<Array1<f64>, String> {
423    let n = z.len();
424    let mut out = Array1::<f64>::zeros(n);
425    for i in 0..n {
426        let a_pre = baseline_marginal + marginal_offset[i];
427        let b_pre = baseline_logslope + logslope_offset[i];
428        let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
429            .map_err(|e| format!("rigid_pooled_probit_pilot_eta marginal link map: {e}"))?
430            .q;
431        out[i] = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
432    }
433    Ok(out)
434}
435
436/// Tikhonov ridge for the pilot IRLS marginal solve, as a fraction of the mean
437/// Hessian diagonal: `ridge = PILOT_RIDGE_DIAG_FRACTION * max(mean_diag, floor)`.
438/// Scaling by the diagonal keeps the ridge scale-invariant; the fraction is
439/// small enough to be numerically negligible against a well-conditioned design
440/// yet still regularise a near-singular pilot Gram.
441pub(crate) const PILOT_RIDGE_DIAG_FRACTION: f64 = 1e-6;
442/// Positivity floor on the mean Hessian diagonal used to scale the pilot ridge,
443/// so a degenerate (all-zero-diagonal) Gram still receives a tiny ridge.
444pub(crate) const PILOT_RIDGE_DIAG_FLOOR: f64 = 1e-12;
445
446pub(super) fn pilot_eta_for_link_dev_orthogonalisation(
447    base_link: &InverseLink,
448    y: &Array1<f64>,
449    z: &Array1<f64>,
450    weights: &Array1<f64>,
451    marginal_design: &DesignMatrix,
452    marginal_offset: &Array1<f64>,
453    logslope_offset: &Array1<f64>,
454    baseline_marginal: f64,
455    baseline_logslope: f64,
456    probit_scale: f64,
457) -> Result<Array1<f64>, String> {
458    use gam_linalg::faer_ndarray::FaerCholesky;
459
460    let n = y.len();
461    if marginal_design.nrows() != n {
462        return Err(format!(
463            "pilot_eta_for_link_dev_orthogonalisation: marginal design has {} rows, expected {}",
464            marginal_design.nrows(),
465            n,
466        ));
467    }
468    let mut working_eta = Array1::<f64>::zeros(n);
469    let mut w_irls = Array1::<f64>::zeros(n);
470    let mut residual = Array1::<f64>::zeros(n);
471    for i in 0..n {
472        let a_pre = baseline_marginal + marginal_offset[i];
473        let b_pre = baseline_logslope + logslope_offset[i];
474        let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
475            .map_err(|e| {
476                format!("pilot_eta_for_link_dev_orthogonalisation marginal link map: {e}")
477            })?
478            .q;
479        let eta = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
480        working_eta[i] = eta;
481        let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
482        let phi = normal_pdf(eta).max(1e-300);
483        let var = (mu * (1.0 - mu)).max(1e-300);
484        w_irls[i] = weights[i] * (phi * phi) / var;
485        residual[i] = (y[i] - mu) / phi;
486    }
487    let p_marg = marginal_design.ncols();
488    if p_marg == 0 {
489        return Ok(working_eta);
490    }
491    let xtwr = marginal_design.compute_xtwy(&w_irls, &residual)?;
492    let mut xtwx =
493        marginal_design.xt_diag_x_signed_op(FiniteSignedWeightsView::try_from_array(&w_irls)?)?;
494    let trace_diag: f64 = (0..p_marg).map(|i| xtwx[[i, i]]).sum();
495    let ridge =
496        (trace_diag / p_marg as f64).max(PILOT_RIDGE_DIAG_FLOOR) * PILOT_RIDGE_DIAG_FRACTION;
497    for i in 0..p_marg {
498        xtwx[[i, i]] += ridge;
499    }
500    let factor = xtwx
501        .cholesky(faer::Side::Lower)
502        .map_err(|e| format!("pilot_eta_for_link_dev_orthogonalisation Cholesky failed: {e}"))?;
503    let delta_beta_marg = factor.solvevec(&xtwr);
504    let marg_contrib = marginal_design.dot(&delta_beta_marg);
505    Ok(&working_eta + &marg_contrib)
506}
507
508pub(super) fn joint_setup(
509    data: ArrayView2<'_, f64>,
510    marginalspec: &TermCollectionSpec,
511    logslopespec: &TermCollectionSpec,
512    marginal_penalties: usize,
513    logslope_penalties: usize,
514    absorber_rho0: Option<f64>,
515    extra_rho0: &[f64],
516    kappa_options: &SpatialLengthScaleOptimizationOptions,
517) -> Result<ExactJointHyperSetup, gam_terms::basis::BasisError> {
518    let marginal_terms = spatial_length_scale_term_indices(marginalspec);
519    let logslope_terms = spatial_length_scale_term_indices(logslopespec);
520    let rho_dim = marginal_penalties + logslope_penalties + extra_rho0.len();
521    let mut rho0vec = Array1::<f64>::zeros(rho_dim);
522    // The #461 influence-absorber ridge is the TRAILING marginal coordinate
523    // (see `marginal_penalties_with_influence_ridge`); it is REML-learned like
524    // every other penalty but seeds at the ln(n) leakage scale instead of 0.
525    if let Some(seed) = absorber_rho0 {
526        assert!(
527            marginal_penalties > 0,
528            "an absorber rho0 seed requires at least one marginal penalty to land in"
529        );
530        rho0vec[marginal_penalties - 1] = seed;
531    }
532    for (idx, &value) in extra_rho0.iter().enumerate() {
533        rho0vec[marginal_penalties + logslope_penalties + idx] = value;
534    }
535    let rho_lower = Array1::<f64>::from_elem(rho_dim, -12.0);
536    let rho_upper = Array1::<f64>::from_elem(rho_dim, 12.0);
537    let marginal_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
538        marginalspec,
539        &marginal_terms,
540        kappa_options,
541    )
542    .reseed_from_data(data, marginalspec, &marginal_terms, kappa_options)?;
543    let logslope_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
544        logslopespec,
545        &logslope_terms,
546        kappa_options,
547    )
548    .reseed_from_data(data, logslopespec, &logslope_terms, kappa_options)?;
549    let mut values = marginal_kappa.as_array().to_vec();
550    values.extend(logslope_kappa.as_array().iter());
551    let marginal_dims = marginal_kappa.dims_per_term().to_vec();
552    let logslope_dims = logslope_kappa.dims_per_term().to_vec();
553    let mut dims = marginal_dims.clone();
554    dims.extend(logslope_dims.iter().copied());
555    let log_kappa0 = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(values), dims.clone());
556    // Bounds: concatenate per-block data-aware bounds in the same order.
557    let marginal_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
558        data,
559        marginalspec,
560        &marginal_terms,
561        &marginal_dims,
562        kappa_options,
563    )?;
564    let logslope_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
565        data,
566        logslopespec,
567        &logslope_terms,
568        &logslope_dims,
569        kappa_options,
570    )?;
571    let mut lower_vals = marginal_lower.as_array().to_vec();
572    lower_vals.extend(logslope_lower.as_array().iter());
573    let log_kappa_lower =
574        SpatialLogKappaCoords::new_with_dims(Array1::from_vec(lower_vals), dims.clone());
575    let marginal_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
576        data,
577        marginalspec,
578        &marginal_terms,
579        &marginal_dims,
580        kappa_options,
581    )?;
582    let logslope_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
583        data,
584        logslopespec,
585        &logslope_terms,
586        &logslope_dims,
587        kappa_options,
588    )?;
589    let mut upper_vals = marginal_upper.as_array().to_vec();
590    upper_vals.extend(logslope_upper.as_array().iter());
591    let log_kappa_upper = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(upper_vals), dims);
592    // Project seed onto bounds in case a user-provided spec.length_scale falls
593    // outside the data-derived ψ window; seed was a hint, not a hard constraint.
594    let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
595    Ok(ExactJointHyperSetup::new(
596        rho0vec,
597        rho_lower,
598        rho_upper,
599        log_kappa0,
600        log_kappa_lower,
601        log_kappa_upper,
602    ))
603}
604
605#[inline]
606pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth_numeric(
607    signed_margin: f64,
608    weight: f64,
609) -> (f64, f64, f64, f64) {
610    if weight == 0.0 || signed_margin == f64::INFINITY {
611        return (0.0, 0.0, 0.0, 0.0);
612    }
613    if signed_margin.is_nan() {
614        return (f64::NAN, f64::NAN, f64::NAN, f64::NAN);
615    }
616    let d = normal_logcdf_derivatives(signed_margin);
617    (
618        -weight * d[1],
619        -weight * d[2],
620        -weight * d[3],
621        -weight * d[4],
622    )
623}
624
625/// Exact probit derivative helper used by analytic jet code paths.
626///
627/// `+inf` is the saturated zero tail and is allowed. `-inf` and `NaN` are
628/// rejected instead of being silently collapsed, so exact callers fail fast
629/// rather than erasing curvature or domain errors. Numeric boundary behavior
630/// that needs to preserve `-inf` / `NaN` values lives in
631/// `signed_probit_neglog_derivatives_up_to_fourth_numeric`.
632pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth(
633    signed_margin: f64,
634    weight: f64,
635) -> Result<(f64, f64, f64, f64), String> {
636    if weight == 0.0 || signed_margin == f64::INFINITY {
637        return Ok((0.0, 0.0, 0.0, 0.0));
638    }
639    if !signed_margin.is_finite() {
640        return Err(format!(
641            "non-finite signed margin in exact probit derivative helper: {signed_margin}"
642        ));
643    }
644    Ok(signed_probit_neglog_derivatives_up_to_fourth_numeric(
645        signed_margin,
646        weight,
647    ))
648}
649
650/// Fused exact value+derivative stack for the signed-probit negative-log
651/// kernel: returns `[-w·logΦ(m), w·k1, w·k2, w·k3, w·k4]` in the `[f64; 5]`
652/// shape [`Tower4::compose_unary`] consumes.
653///
654/// This is the single-source replacement for the two-call pattern
655///
656/// ```ignore
657/// let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
658/// let (k1, k2, k3, k4) = signed_probit_neglog_derivatives_up_to_fourth(m, w)?;
659/// // → [-w*logcdf, k1, k2, k3, k4]
660/// ```
661///
662/// which evaluated the tail kernel twice on the same `m`. The centralized
663/// derivative stack evaluates it once and owns the cancellation-free left-tail
664/// continued fraction and right-tail log-magnitude recurrence shared by every
665/// probit consumer.
666///
667/// `+∞` is the saturated zero tail. At `−∞`, a positive weight has the exact
668/// limit `[+∞, −∞, +w, −0, −0]`; a negative signed contribution reverses the
669/// infinite and zero signs through ordinary scalar multiplication. `NaN`
670/// propagates unless the row is inactive (`weight == 0`).
671#[inline]
672pub(crate) fn signed_probit_neglog_unary_stack(signed_margin: f64, weight: f64) -> [f64; 5] {
673    if weight == 0.0 || signed_margin == f64::INFINITY {
674        return [0.0; 5];
675    }
676    if signed_margin.is_nan() {
677        return [f64::NAN; 5];
678    }
679    let d = normal_logcdf_derivatives(signed_margin);
680    [
681        -weight * d[0],
682        -weight * d[1],
683        -weight * d[2],
684        -weight * d[3],
685        -weight * d[4],
686    ]
687}
688
689#[inline]
690pub(super) fn rigid_observed_logslope(logslope: f64, probit_scale: f64) -> f64 {
691    probit_scale * logslope
692}
693
694#[inline]
695pub(super) fn rigid_observed_scale(logslope: f64, probit_scale: f64) -> f64 {
696    let observed_logslope = rigid_observed_logslope(logslope, probit_scale);
697    (1.0 + observed_logslope * observed_logslope).sqrt()
698}
699
700#[inline]
701pub(super) fn rigid_intercept_from_marginal(
702    marginal_eta: f64,
703    logslope: f64,
704    probit_scale: f64,
705) -> f64 {
706    marginal_eta * rigid_observed_scale(logslope, probit_scale)
707}
708
709#[inline]
710pub(super) fn rigid_prescale_intercept_from_marginal(
711    marginal_eta: f64,
712    logslope: f64,
713    probit_scale: f64,
714) -> f64 {
715    rigid_intercept_from_marginal(marginal_eta, logslope, probit_scale) / probit_scale
716}
717
718#[inline]
719pub(super) fn rigid_prescale_intercept_derivative_abs(
720    marginal_eta: f64,
721    logslope: f64,
722    probit_scale: f64,
723) -> f64 {
724    let c = rigid_observed_scale(logslope, probit_scale);
725    probit_scale * normal_pdf(marginal_eta) / c
726}
727
728#[inline]
729pub(super) fn rigid_observed_eta(
730    marginal_eta: f64,
731    logslope: f64,
732    z: f64,
733    probit_scale: f64,
734) -> f64 {
735    marginal_slope_standard_normal_scalar_eta(marginal_eta, logslope, z, probit_scale)
736}
737
738#[inline]
739pub(super) fn marginal_slope_standard_normal_scalar_eta(
740    q: f64,
741    slope: f64,
742    z: f64,
743    probit_scale: f64,
744) -> f64 {
745    let observed_slope = rigid_observed_logslope(slope, probit_scale);
746    q * (1.0 + observed_slope * observed_slope).sqrt() + observed_slope * z
747}
748
749pub(super) fn unary_derivatives_normal_cdf(x: f64) -> [f64; 5] {
750    let pdf = normal_pdf(x);
751    [
752        normal_cdf(x),
753        pdf,
754        -x * pdf,
755        (x * x - 1.0) * pdf,
756        (-x.powi(3) + 3.0 * x) * pdf,
757    ]
758}
759
760/// Streaming log-sum-exp update: accumulate `exp(log_term)` into a running
761/// `(log_max, sum)` pair representing `Σ exp(log_term_i) = exp(log_max) · sum`.
762///
763/// When `log_term` exceeds the running max, the partial sum is rescaled in
764/// place so the new max becomes the reference point. This keeps everything
765/// inside the dynamic range of f64 with no allocation.
766#[inline]
767pub(super) fn lse_accumulate(log_max: &mut f64, sum: &mut f64, log_term: f64) {
768    if !log_term.is_finite() {
769        return;
770    }
771    if log_term > *log_max {
772        if log_max.is_finite() {
773            *sum = *sum * (*log_max - log_term).exp() + 1.0;
774        } else {
775            *sum = 1.0;
776        }
777        *log_max = log_term;
778    } else {
779        *sum += (log_term - *log_max).exp();
780    }
781}
782
783#[derive(Clone, Copy, Debug, PartialEq, Eq)]
784pub enum MarginalSlopeCovarianceShape {
785    Diagonal,
786    Full,
787    LowRank,
788}
789
790#[derive(Clone, Debug, PartialEq)]
791enum MarginalSlopeCovarianceStorage {
792    Diagonal {
793        covariance: Array1<f64>,
794    },
795    Full {
796        covariance: Array2<f64>,
797        /// Row-oriented factor `B` with `Σ = BᵀB`.
798        square_root_factor: Array2<f64>,
799    },
800    /// Low-rank factor `L` with `Σ = LLᵀ`.
801    LowRank {
802        factor: Array2<f64>,
803    },
804}
805
806/// Immutable, validated covariance geometry for the physical log-slope vector.
807///
808/// Admission is the single validation boundary. Diagonal covariance entries
809/// remain the sole authority for their quadratic forms. Full covariances cache
810/// an eigensquare-root factor so subsequent quadratic forms are exact sums of
811/// squares; runtime code never repeats an eigendecomposition or applies a
812/// negative-value tolerance. The exact `1ᵀΣ1` shared-slope geometry is cached
813/// at the same boundary.
814#[derive(Clone, Debug)]
815pub struct MarginalSlopeCovariance {
816    storage: MarginalSlopeCovarianceStorage,
817    ones_quadratic_form: f64,
818}
819
820impl PartialEq for MarginalSlopeCovariance {
821    fn eq(&self, other: &Self) -> bool {
822        self.storage == other.storage
823    }
824}
825
826#[derive(Clone, Copy, Debug)]
827pub(crate) enum MarginalSlopeCovarianceRef<'a> {
828    Diagonal(&'a Array1<f64>),
829    Full(&'a Array2<f64>),
830    LowRank(&'a Array2<f64>),
831}
832
833impl MarginalSlopeCovariance {
834    pub fn diagonal(covariance: Array1<f64>) -> Result<Self, String> {
835        if covariance.is_empty() {
836            return Err("marginal-slope diagonal covariance is empty".to_string());
837        }
838        let mut ones_quadratic_form = 0.0;
839        for (axis, &value) in covariance.iter().enumerate() {
840            if !(value.is_finite() && value >= 0.0) {
841                return Err(format!(
842                    "marginal-slope diagonal covariance entry {axis} must be finite and non-negative, got {value}"
843                ));
844            }
845            ones_quadratic_form += value;
846        }
847        if !ones_quadratic_form.is_finite() {
848            return Err("marginal-slope diagonal covariance geometry overflowed".to_string());
849        }
850        Ok(Self {
851            storage: MarginalSlopeCovarianceStorage::Diagonal { covariance },
852            ones_quadratic_form,
853        })
854    }
855
856    pub fn full(covariance: Array2<f64>) -> Result<Self, String> {
857        if covariance.nrows() == 0 || covariance.nrows() != covariance.ncols() {
858            return Err(format!(
859                "marginal-slope full covariance must be non-empty and square, got {}x{}",
860                covariance.nrows(),
861                covariance.ncols(),
862            ));
863        }
864        for ((row, column), &value) in covariance.indexed_iter() {
865            if !value.is_finite() {
866                return Err(format!(
867                    "marginal-slope full covariance entry ({row},{column}) is non-finite"
868                ));
869            }
870        }
871        for row in 0..covariance.nrows() {
872            for column in (row + 1)..covariance.ncols() {
873                if covariance[[row, column]] != covariance[[column, row]] {
874                    return Err(format!(
875                        "marginal-slope full covariance must be exactly symmetric at ({row},{column}): upper={}, lower={}",
876                        covariance[[row, column]],
877                        covariance[[column, row]],
878                    ));
879                }
880            }
881        }
882        let (eigenvalues, eigenvectors) = covariance.eigh(faer::Side::Lower).map_err(|error| {
883            format!("marginal-slope covariance eigendecomposition failed: {error}")
884        })?;
885        let dimension = covariance.nrows();
886        let mut square_root_factor = Array2::<f64>::zeros((dimension, dimension));
887        for (eigen_axis, &eigenvalue) in eigenvalues.iter().enumerate() {
888            if !(eigenvalue.is_finite() && eigenvalue >= 0.0) {
889                return Err(format!(
890                    "marginal-slope full covariance must be positive semidefinite; eigenvalue {eigen_axis} is {eigenvalue}"
891                ));
892            }
893            let scale = eigenvalue.sqrt();
894            for axis in 0..dimension {
895                square_root_factor[[eigen_axis, axis]] = scale * eigenvectors[[axis, eigen_axis]];
896            }
897        }
898        let mut ones_quadratic_form = 0.0;
899        for factor_row in square_root_factor.rows() {
900            let projection = factor_row.sum();
901            ones_quadratic_form += projection * projection;
902        }
903        if !ones_quadratic_form.is_finite() {
904            return Err("marginal-slope full covariance geometry overflowed".to_string());
905        }
906        Ok(Self {
907            storage: MarginalSlopeCovarianceStorage::Full {
908                covariance,
909                square_root_factor,
910            },
911            ones_quadratic_form,
912        })
913    }
914
915    pub fn low_rank(factor: Array2<f64>) -> Result<Self, String> {
916        if factor.nrows() == 0 {
917            return Err("marginal-slope low-rank covariance factor has zero rows".to_string());
918        }
919        for ((row, column), &value) in factor.indexed_iter() {
920            if !value.is_finite() {
921                return Err(format!(
922                    "marginal-slope low-rank covariance factor entry ({row},{column}) is non-finite"
923                ));
924            }
925        }
926        let mut ones_quadratic_form = 0.0;
927        for factor_column in factor.columns() {
928            let projection = factor_column.sum();
929            ones_quadratic_form += projection * projection;
930        }
931        if !ones_quadratic_form.is_finite() {
932            return Err("marginal-slope low-rank covariance geometry overflowed".to_string());
933        }
934        Ok(Self {
935            storage: MarginalSlopeCovarianceStorage::LowRank { factor },
936            ones_quadratic_form,
937        })
938    }
939
940    pub fn to_dense(&self) -> Array2<f64> {
941        match &self.storage {
942            MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
943                Array2::from_diag(covariance)
944            }
945            MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.clone(),
946            MarginalSlopeCovarianceStorage::LowRank { factor } => factor.dot(&factor.t()),
947        }
948    }
949
950    pub fn shape(&self) -> MarginalSlopeCovarianceShape {
951        match &self.storage {
952            MarginalSlopeCovarianceStorage::Diagonal { .. } => {
953                MarginalSlopeCovarianceShape::Diagonal
954            }
955            MarginalSlopeCovarianceStorage::Full { .. } => MarginalSlopeCovarianceShape::Full,
956            MarginalSlopeCovarianceStorage::LowRank { .. } => MarginalSlopeCovarianceShape::LowRank,
957        }
958    }
959
960    pub fn dim(&self) -> usize {
961        match &self.storage {
962            MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => covariance.len(),
963            MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.nrows(),
964            MarginalSlopeCovarianceStorage::LowRank { factor } => factor.nrows(),
965        }
966    }
967
968    pub fn ones_quadratic_form(&self) -> f64 {
969        self.ones_quadratic_form
970    }
971
972    pub(crate) fn representation(&self) -> MarginalSlopeCovarianceRef<'_> {
973        match &self.storage {
974            MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
975                MarginalSlopeCovarianceRef::Diagonal(covariance)
976            }
977            MarginalSlopeCovarianceStorage::Full { covariance, .. } => {
978                MarginalSlopeCovarianceRef::Full(covariance)
979            }
980            MarginalSlopeCovarianceStorage::LowRank { factor } => {
981                MarginalSlopeCovarianceRef::LowRank(factor)
982            }
983        }
984    }
985
986    #[inline(always)]
987    pub(crate) fn quadratic_form_unchecked(&self, vector: &[f64]) -> f64 {
988        <Self as SymmetricQuadraticCoefficients>::quadratic_value(self, vector, |value| *value)
989    }
990
991    pub fn quadratic_form(&self, vector: &[f64]) -> Result<f64, String> {
992        if vector.len() != self.dim() {
993            return Err(format!(
994                "marginal-slope covariance dimension mismatch: vector={}, covariance={}",
995                vector.len(),
996                self.dim()
997            ));
998        }
999        if vector.iter().any(|value| !value.is_finite()) {
1000            return Err("marginal-slope covariance vector contains non-finite values".to_string());
1001        }
1002        let value = self.quadratic_form_unchecked(vector);
1003        if !value.is_finite() {
1004            return Err(format!(
1005                "marginal-slope covariance quadratic form is non-finite: {value}"
1006            ));
1007        }
1008        Ok(value)
1009    }
1010}
1011
1012enum VectorSupport {
1013    Zero,
1014    Singleton { axis: usize, value: f64 },
1015    Multiple,
1016}
1017
1018#[inline(always)]
1019fn vector_support(input: &[f64]) -> VectorSupport {
1020    let mut singleton = None;
1021    for (axis, &value) in input.iter().enumerate() {
1022        if value == 0.0 {
1023            continue;
1024        }
1025        if singleton.is_some() {
1026            return VectorSupport::Multiple;
1027        }
1028        singleton = Some((axis, value));
1029    }
1030    match singleton {
1031        None => VectorSupport::Zero,
1032        Some((axis, value)) => VectorSupport::Singleton { axis, value },
1033    }
1034}
1035
1036impl SymmetricQuadraticCoefficients for MarginalSlopeCovariance {
1037    fn dimension(&self) -> usize {
1038        self.dim()
1039    }
1040
1041    fn multiply(&self, input: &[f64], output: &mut [f64]) {
1042        assert_eq!(input.len(), self.dim());
1043        assert_eq!(output.len(), self.dim());
1044        match self.representation() {
1045            MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1046                for axis in 0..input.len() {
1047                    output[axis] = diagonal[axis] * input[axis];
1048                }
1049            }
1050            MarginalSlopeCovarianceRef::Full(matrix) => {
1051                match vector_support(input) {
1052                    VectorSupport::Zero => {
1053                        output.fill(0.0);
1054                        return;
1055                    }
1056                    VectorSupport::Singleton { axis, value } => {
1057                        for row in 0..input.len() {
1058                            output[row] = matrix[[row, axis]] * value;
1059                        }
1060                        return;
1061                    }
1062                    VectorSupport::Multiple => {}
1063                }
1064                for row in 0..input.len() {
1065                    let mut value = 0.0;
1066                    for column in 0..input.len() {
1067                        value += matrix[[row, column]] * input[column];
1068                    }
1069                    output[row] = value;
1070                }
1071            }
1072            MarginalSlopeCovarianceRef::LowRank(factor) => {
1073                output.fill(0.0);
1074                match vector_support(input) {
1075                    VectorSupport::Zero => return,
1076                    VectorSupport::Singleton { axis, value } => {
1077                        for rank in 0..factor.ncols() {
1078                            let projection = factor[[axis, rank]] * value;
1079                            for row in 0..input.len() {
1080                                output[row] += factor[[row, rank]] * projection;
1081                            }
1082                        }
1083                        return;
1084                    }
1085                    VectorSupport::Multiple => {}
1086                }
1087                for rank in 0..factor.ncols() {
1088                    let mut projection = 0.0;
1089                    for row in 0..input.len() {
1090                        projection += factor[[row, rank]] * input[row];
1091                    }
1092                    for row in 0..input.len() {
1093                        output[row] += factor[[row, rank]] * projection;
1094                    }
1095                }
1096            }
1097        }
1098    }
1099
1100    fn coefficient(&self, row: usize, column: usize) -> f64 {
1101        match self.representation() {
1102            MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1103                if row == column {
1104                    diagonal[row]
1105                } else {
1106                    0.0
1107                }
1108            }
1109            MarginalSlopeCovarianceRef::Full(matrix) => matrix[[row, column]],
1110            MarginalSlopeCovarianceRef::LowRank(factor) => {
1111                let mut value = 0.0;
1112                for rank in 0..factor.ncols() {
1113                    value += factor[[row, rank]] * factor[[column, rank]];
1114                }
1115                value
1116            }
1117        }
1118    }
1119
1120    fn visit_upper_triangle(
1121        &self,
1122        direction: &mut [f64],
1123        projected: &mut [f64],
1124        mut visit: impl FnMut(usize, usize, f64),
1125    ) {
1126        let dimension = self.dim();
1127        assert_eq!(direction.len(), dimension);
1128        assert_eq!(projected.len(), dimension);
1129        match self.representation() {
1130            MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1131                for column in 0..dimension {
1132                    for row in 0..=column {
1133                        visit(row, column, if row == column { diagonal[row] } else { 0.0 });
1134                    }
1135                }
1136            }
1137            MarginalSlopeCovarianceRef::Full(matrix) => {
1138                for column in 0..dimension {
1139                    for row in 0..=column {
1140                        visit(row, column, matrix[[row, column]]);
1141                    }
1142                }
1143            }
1144            MarginalSlopeCovarianceRef::LowRank(factor) => {
1145                for column in 0..dimension {
1146                    for row in 0..=column {
1147                        let mut value = 0.0;
1148                        for rank in 0..factor.ncols() {
1149                            value += factor[[row, rank]] * factor[[column, rank]];
1150                        }
1151                        visit(row, column, value);
1152                    }
1153                }
1154            }
1155        }
1156    }
1157
1158    fn quadratic_value<T, F>(&self, input: &[T], value: F) -> f64
1159    where
1160        F: Fn(&T) -> f64,
1161    {
1162        assert_eq!(input.len(), self.dim());
1163        match &self.storage {
1164            MarginalSlopeCovarianceStorage::Diagonal { covariance } => input
1165                .iter()
1166                .zip(covariance)
1167                .map(|(input, &covariance)| {
1168                    let input = value(input);
1169                    covariance * input * input
1170                })
1171                .sum(),
1172            MarginalSlopeCovarianceStorage::Full {
1173                square_root_factor, ..
1174            } => {
1175                let mut total = 0.0;
1176                for factor_row in square_root_factor.rows() {
1177                    let mut projection = 0.0;
1178                    for axis in 0..input.len() {
1179                        projection += factor_row[axis] * value(&input[axis]);
1180                    }
1181                    total += projection * projection;
1182                }
1183                total
1184            }
1185            MarginalSlopeCovarianceStorage::LowRank { factor } => {
1186                // Sigma = L L'. Evaluate x' Sigma x as ||L' x||^2 so the
1187                // primal runtime scalar remains matrix free and O(KR).
1188                let mut total = 0.0;
1189                for rank in 0..factor.ncols() {
1190                    let mut projection = 0.0;
1191                    for row in 0..input.len() {
1192                        projection += factor[[row, rank]] * value(&input[row]);
1193                    }
1194                    total += projection * projection;
1195                }
1196                total
1197            }
1198        }
1199    }
1200}
1201
1202// Marginal-slope probit identity.
1203//
1204// For a row with latent scores z | a ~ N(0, Sigma(a)) and probit index
1205//
1206//     eta = c(a) q(t, a) + r(a)' z,
1207//
1208// the preservation target is
1209//
1210//     E_z[Phi(-eta) | a] = Phi(-q(t, a)).
1211//
1212// If X = r' z is N(0, v) with v = r' Sigma r, then for independent
1213// E ~ N(0, 1),
1214//
1215//     E[Phi(-(c q + X))]
1216//       = P(E <= -c q - X)
1217//       = P(E + X <= -c q)
1218//       = Phi(-c q / sqrt(1 + v)).
1219//
1220// Thus the target holds for every q exactly when
1221//
1222//     c(a) = sqrt(1 + r(a)' Sigma(a) r(a)).
1223//
1224// `probit_scale` maps the raw log-slope surface to the observed probit
1225// gradient r(a). K=1 with diagonal variance 1 gives the original scalar
1226// formula sqrt(1 + r^2); full and low-rank covariances differ only in the
1227// shape-specific evaluation of the same quadratic form.
1228pub fn marginal_slope_covariance_from_scores(
1229    scores: ArrayView2<'_, f64>,
1230    weights: &Array1<f64>,
1231) -> Result<MarginalSlopeCovariance, String> {
1232    let (n, k) = scores.dim();
1233    if k == 0 {
1234        return Err("marginal-slope score matrix must have at least one column".to_string());
1235    }
1236    if weights.len() != n {
1237        return Err(format!(
1238            "marginal-slope covariance weight length mismatch: weights={}, rows={n}",
1239            weights.len()
1240        ));
1241    }
1242    let total_weight = weights.iter().copied().sum::<f64>();
1243    if !(total_weight.is_finite() && total_weight > 0.0) {
1244        return Err("marginal-slope covariance needs positive finite total weight".to_string());
1245    }
1246    let mut mean = Array1::<f64>::zeros(k);
1247    for i in 0..n {
1248        let weight = weights[i];
1249        if !(weight.is_finite() && weight >= 0.0) {
1250            return Err(format!(
1251                "marginal-slope covariance weight {i} must be finite and non-negative, got {weight}"
1252            ));
1253        }
1254        for j in 0..k {
1255            let score = scores[[i, j]];
1256            if !score.is_finite() {
1257                return Err(format!(
1258                    "marginal-slope covariance score ({i},{j}) is non-finite"
1259                ));
1260            }
1261            mean[j] += weight * score;
1262        }
1263    }
1264    mean.mapv_inplace(|value| value / total_weight);
1265
1266    let mut cov = Array2::<f64>::zeros((k, k));
1267    for i in 0..n {
1268        let weight = weights[i];
1269        for a in 0..k {
1270            let da = scores[[i, a]] - mean[a];
1271            for b in 0..=a {
1272                let value = weight * da * (scores[[i, b]] - mean[b]) / total_weight;
1273                cov[[a, b]] += value;
1274                if a != b {
1275                    cov[[b, a]] += value;
1276                }
1277            }
1278        }
1279    }
1280
1281    // Representation is geometry, not a statistical model-selection decision:
1282    // only an exactly diagonal matrix may discard its off-diagonal entries.
1283    // Every nonzero coupling is retained in the exact dense covariance.
1284    let is_diagonal = (0..k).all(|row| ((row + 1)..k).all(|column| cov[[row, column]] == 0.0));
1285    if is_diagonal {
1286        MarginalSlopeCovariance::diagonal(cov.diag().to_owned())
1287    } else {
1288        MarginalSlopeCovariance::full(cov)
1289    }
1290}
1291
1292pub fn marginal_slope_preserving_scale(
1293    slopes: &[f64],
1294    covariance: &MarginalSlopeCovariance,
1295    probit_scale: f64,
1296) -> Result<f64, String> {
1297    if !probit_scale.is_finite() {
1298        return Err(format!(
1299            "marginal-slope probit scale must be finite, got {probit_scale}"
1300        ));
1301    }
1302    let variance = probit_scale * probit_scale * covariance.quadratic_form(slopes)?;
1303    if !variance.is_finite() {
1304        return Err("marginal-slope preserving variance is non-finite".to_string());
1305    }
1306    Ok((1.0 + variance).sqrt())
1307}
1308
1309pub fn marginal_slope_probit_eta(
1310    q: f64,
1311    z: &[f64],
1312    slopes: &[f64],
1313    covariance: &MarginalSlopeCovariance,
1314    probit_scale: f64,
1315) -> Result<f64, String> {
1316    if z.len() != slopes.len() {
1317        return Err(format!(
1318            "marginal-slope score/slope dimension mismatch: z={}, slopes={}",
1319            z.len(),
1320            slopes.len()
1321        ));
1322    }
1323    if slopes.len() != covariance.dim() {
1324        return Err(format!(
1325            "marginal-slope covariance dimension mismatch: slopes={}, covariance={}",
1326            slopes.len(),
1327            covariance.dim()
1328        ));
1329    }
1330    if !q.is_finite() || z.iter().any(|value| !value.is_finite()) {
1331        return Err("marginal-slope probit eta inputs must be finite".to_string());
1332    }
1333    let scale = marginal_slope_preserving_scale(slopes, covariance, probit_scale)?;
1334    let linear = z
1335        .iter()
1336        .zip(slopes.iter())
1337        .map(|(&score, &slope)| probit_scale * slope * score)
1338        .sum::<f64>();
1339    Ok(q * scale + linear)
1340}
1341
1342/// Log-space residual evaluator for the empirical-frailty intercept calibration.
1343///
1344/// Solves, in log-space, the strictly-increasing equation
1345///
1346///   F(a) = log Σᵢ wᵢ Φ(a + b·zᵢ) − log μ★ = 0,
1347///
1348/// where `b = rigid_observed_logslope(slope, probit_scale)` and `(zᵢ, wᵢ)` are
1349/// the supplied quadrature nodes and (positive) weights.
1350///
1351/// Mathematical structure of `F`:
1352///   • `F ∈ C^∞(ℝ)`.
1353///   • `F` is strictly increasing: `F'(a) = (Σ wᵢ φᵢ) / (Σ wᵢ Φᵢ) > 0` everywhere.
1354///   • `F(a) → −∞` as `a → −∞`; `F(a) → log(Σ wᵢ) − log μ★ ≥ 0` as `a → +∞`.
1355///   • Unique root `a★ ∈ ℝ` exists for every `μ★ ∈ (0, 1)`.
1356///
1357/// Why log-space: the linear-space residual `Σ wᵢ Φᵢ − μ★` and its derivative
1358/// `Σ wᵢ φᵢ` are sums of strictly-positive `exp(−η²/2)`-scaled terms. When the
1359/// seed `a` puts every quadrature node `ηᵢ = a + b·zᵢ` into the deep tail
1360/// (|ηᵢ| ≳ 38), every term rounds to 0.0 in IEEE-754 and the derivative
1361/// underflows to exactly zero — destroying Newton's update direction.  The
1362/// log-space formulation evaluates `log φ(η) = −η²/2 − ½ log 2π` (always finite
1363/// for any finite η) and `log Φ(η)` via the `erfcx`-based `normal_logcdf`
1364/// (also always finite for any finite η).  All sums are accumulated by
1365/// streaming log-sum-exp, so `F`, `F'`, and `F''` are finite for every finite
1366/// `a` and the global Newton/Halley iteration converges from any seed.
1367///
1368/// Returns `(F, F', F'')`.  In the deep left tail Newton converges linearly
1369/// (Mills ratio: `F'(a) ≈ |a|`, step ≈ `|a|/2`); near the root convergence is
1370/// quadratic with Newton or cubic with Halley.
1371pub(super) fn empirical_rigid_calibration_eval(
1372    intercept: f64,
1373    log_target_mu: f64,
1374    slope: f64,
1375    probit_scale: f64,
1376    nodes: &[f64],
1377    weights: &[f64],
1378) -> Result<(f64, f64, f64), String> {
1379    if !intercept.is_finite() {
1380        return Err(format!(
1381            "empirical latent calibration: non-finite intercept {intercept}"
1382        ));
1383    }
1384    let observed_slope = rigid_observed_logslope(slope, probit_scale);
1385    const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_8; // 0.5 * ln(2π)
1386
1387    // Streaming LSE accumulators for log Σ wᵢ φᵢ and log Σ wᵢ Φᵢ.
1388    let mut log_max_phi = f64::NEG_INFINITY;
1389    let mut sum_phi = 0.0_f64;
1390    let mut log_max_cdf = f64::NEG_INFINITY;
1391    let mut sum_cdf = 0.0_f64;
1392
1393    // Streaming signed LSE for Σ wᵢ ηᵢ φᵢ, split into positive and negative
1394    // legs so the cancellation `pos − neg` happens once at the end on a
1395    // finite, well-scaled remainder.
1396    let mut log_max_pos = f64::NEG_INFINITY;
1397    let mut sum_pos = 0.0_f64;
1398    let mut log_max_neg = f64::NEG_INFINITY;
1399    let mut sum_neg = 0.0_f64;
1400
1401    for (&node, &weight) in nodes.iter().zip(weights.iter()) {
1402        if !(weight.is_finite() && weight > 0.0) {
1403            continue;
1404        }
1405        let eta = intercept + observed_slope * node;
1406        if !eta.is_finite() {
1407            return Err(format!(
1408                "empirical latent calibration: non-finite η at intercept={intercept}, slope={slope}, node={node}"
1409            ));
1410        }
1411        let log_w = weight.ln();
1412        let log_phi = -0.5 * eta * eta - HALF_LOG_2PI;
1413        let log_term_phi = log_w + log_phi;
1414        let log_term_cdf = log_w + normal_logcdf(eta);
1415
1416        lse_accumulate(&mut log_max_phi, &mut sum_phi, log_term_phi);
1417        lse_accumulate(&mut log_max_cdf, &mut sum_cdf, log_term_cdf);
1418
1419        if eta != 0.0 {
1420            let log_term_eta_phi = log_term_phi + eta.abs().ln();
1421            if eta > 0.0 {
1422                lse_accumulate(&mut log_max_pos, &mut sum_pos, log_term_eta_phi);
1423            } else {
1424                lse_accumulate(&mut log_max_neg, &mut sum_neg, log_term_eta_phi);
1425            }
1426        }
1427    }
1428
1429    if !(sum_phi.is_finite() && sum_cdf.is_finite() && sum_phi > 0.0 && sum_cdf > 0.0) {
1430        return Err(format!(
1431            "empirical latent calibration: log-space accumulation failed (sum_phi={sum_phi}, sum_cdf={sum_cdf}, intercept={intercept})"
1432        ));
1433    }
1434
1435    let log_s_phi = log_max_phi + sum_phi.ln();
1436    let log_s_cdf = log_max_cdf + sum_cdf.ln();
1437
1438    // F = log Σ wᵢ Φᵢ − log μ★
1439    let f = log_s_cdf - log_target_mu;
1440    // F' = exp(log Σ wᵢ φᵢ − log Σ wᵢ Φᵢ).
1441    //
1442    // F' is mathematically strictly positive everywhere — `Σ wᵢ φᵢ` and
1443    // `Σ wᵢ Φᵢ` are both sums of strictly-positive terms with positive weights.
1444    // In the far right tail, Mills ratio gives `φᵢ/Φᵢ → 0` exponentially, so
1445    // `log F' → −∞` and `(log F').exp()` IEEE-underflows to 0.0. Mathematically
1446    // it is a tiny positive number; floor it at `f64::MIN_POSITIVE` so the
1447    // monotone-root solver sees a strictly-positive derivative and routes
1448    // through its bracket-by-doubling phase (which only needs the *sign* of
1449    // `F'`, not its magnitude). Newton would propose `Δa = −F/F' = ±∞`, the
1450    // solver detects that and falls through to bracketing automatically.
1451    let log_f_prime = log_s_phi - log_s_cdf;
1452    let f_prime = if log_f_prime > -740.0 {
1453        log_f_prime.exp()
1454    } else {
1455        f64::MIN_POSITIVE
1456    };
1457
1458    // F'' = (d/da)(S_φ/S_Φ) = (S_φ' S_Φ − S_φ²)/S_Φ²
1459    //     = −(Σ wᵢ ηᵢ φᵢ)/S_Φ − (F')²
1460    // The η-weighted sum is cancellation-prone; combine its positive and
1461    // negative legs against the same `log_s_cdf` reference so the subtraction
1462    // happens on dimensionless quantities of bounded magnitude. When the ratio
1463    // also underflows (deep tail), the result is a clean numerical zero —
1464    // Halley reduces to Newton, which is what the solver does anyway.
1465    let exp_safe = |log_x: f64| -> f64 { if log_x > -740.0 { log_x.exp() } else { 0.0 } };
1466    let pos_over_cdf = if sum_pos > 0.0 {
1467        exp_safe(log_max_pos + sum_pos.ln() - log_s_cdf)
1468    } else {
1469        0.0
1470    };
1471    let neg_over_cdf = if sum_neg > 0.0 {
1472        exp_safe(log_max_neg + sum_neg.ln() - log_s_cdf)
1473    } else {
1474        0.0
1475    };
1476    let s_etaphi_over_s_cdf = pos_over_cdf - neg_over_cdf;
1477    let f_double_prime = -s_etaphi_over_s_cdf - f_prime * f_prime;
1478
1479    if !(f.is_finite() && f_prime.is_finite() && f_prime > 0.0 && f_double_prime.is_finite()) {
1480        return Err(format!(
1481            "empirical latent calibration: non-finite log-space state f={f}, f'={f_prime}, f''={f_double_prime} at intercept={intercept}"
1482        ));
1483    }
1484    Ok((f, f_prime, f_double_prime))
1485}
1486
1487pub(crate) fn empirical_intercept_from_marginal(
1488    target_mu: f64,
1489    target_q: f64,
1490    slope: f64,
1491    probit_scale: f64,
1492    nodes: &[f64],
1493    weights: &[f64],
1494    initial: Option<f64>,
1495) -> Result<f64, String> {
1496    if !(target_mu.is_finite() && target_mu > 0.0 && target_mu < 1.0) {
1497        return Err(format!(
1498            "empirical latent calibration requires target mu in (0,1), got {target_mu}"
1499        ));
1500    }
1501    let log_target_mu = target_mu.ln();
1502    let closed_form_seed = rigid_intercept_from_marginal(target_q, slope, probit_scale);
1503    let seed = initial.unwrap_or(closed_form_seed);
1504    let eval = |a: f64| {
1505        empirical_rigid_calibration_eval(a, log_target_mu, slope, probit_scale, nodes, weights)
1506    };
1507    // Convergence is on the log-space residual |F| = |log Σ wᵢ Φᵢ − log μ★|.
1508    // Near the root this is the relative error in the calibrated probability,
1509    // so 1e-13 in log-space corresponds to absolute residual μ★ · 1e-13 in
1510    // linear space — strictly tighter than the legacy 1e-13 absolute tolerance
1511    // for every μ★ ∈ (0, 1). The 4·ε floor keeps the contract meaningful when
1512    // μ★ approaches 1 (where log Σ Φᵢ approaches 0).
1513    let abs_tol = 1e-13_f64.max(4.0 * f64::EPSILON);
1514    let solve_from = |s: f64| {
1515        crate::monotone_root::solve_monotone_root(
1516            eval,
1517            s,
1518            "empirical latent intercept",
1519            abs_tol,
1520            64,
1521            48,
1522        )
1523        // Enclosing fn emits its own format!() rejection errors as String,
1524        // so the public return type stays Result<_, String>.
1525        .map_err(|e| e.to_string())
1526    };
1527    // A cached warm start can be poisoned across iterations: the per-row
1528    // `intercept_warm_starts` slot is shared by reference across line-search
1529    // trials and across outer-search seed validations, and is written after
1530    // every successful row-solve — including from rejected line-search trials
1531    // whose β/slope was wild. When that stale `a` is paired with the current
1532    // (much smaller) slope, the bracket-by-doubling phase can exhaust its
1533    // budget without crossing zero. Fall back to the deterministic
1534    // closed-form seed, which depends only on the current `(target_q, slope)`
1535    // and is bounded by the analytic rigid-probit geometry, so the cache
1536    // remains a pure speedup that cannot poison correctness.
1537    let (root, _, f_best) = match solve_from(seed) {
1538        Ok(v) => v,
1539        Err(first_err) => {
1540            if seed == closed_form_seed {
1541                return Err(first_err);
1542            }
1543            solve_from(closed_form_seed).map_err(|retry_err| {
1544                format!("{first_err}; closed-form retry from a={closed_form_seed:.6}: {retry_err}")
1545            })?
1546        }
1547    };
1548    if f_best.abs() > abs_tol {
1549        return Err(format!(
1550            "empirical latent intercept solve failed: log-residual={f_best:.3e} at a={root:.6}, target mu={target_mu:.6}"
1551        ));
1552    }
1553    Ok(root)
1554}
1555
1556#[inline]
1557pub(super) fn rigid_standard_normal_neglog_only(
1558    q: f64,
1559    g: f64,
1560    z: f64,
1561    y: f64,
1562    w: f64,
1563    probit_scale: f64,
1564) -> Result<f64, String> {
1565    let s = 2.0 * y - 1.0;
1566    let eta = marginal_slope_standard_normal_scalar_eta(q, g, z, probit_scale);
1567    let m = s * eta;
1568    let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
1569    if !logcdf.is_finite() {
1570        return Err(format!(
1571            "rigid probit neglog_only: non-finite log Φ at q={q}, g={g}, z={z}, y={y}"
1572        ));
1573    }
1574    Ok(-w * logcdf)
1575}
1576
1577/// The rigid standard-normal Bernoulli row negative log-likelihood, written
1578/// ONCE over the generic [`JetScalar`] interface (#932 scalar cutover).
1579///
1580/// Primaries `p = [q_eta = marginal η, g = slope]`. The body is exactly the
1581/// production likelihood — `ℓ = −w·logΦ((2y−1)·η)`, `η = q(η_marg)·√(1+(s·g)²)
1582/// + (s·g)·z` — composed with ONLY [`JetScalar`] ops, so it re-instantiates at
1583/// whatever order / representation a consumer needs:
1584///
1585/// * [`Order2`](super::super::jet_scalar::Order2) → `(v, g, H)`
1586///   ([`rigid_standard_normal_row_kernel`], the inner-Newton path);
1587/// * [`OneSeed`](super::super::jet_scalar::OneSeed) → contracted third
1588///   `Σ_c ℓ_{abc} dir_c` without materialising `t3` (the directional gate);
1589/// * [`TwoSeed`](super::super::jet_scalar::TwoSeed) → contracted fourth
1590///   `Σ_{cd} ℓ_{abcd} u_c v_d` without materialising `t4`;
1591/// * full [`Tower4`] → every uncontracted channel
1592///   ([`rigid_standard_normal_tower`], feeding the `third_full` / `fourth_full`
1593///   caches).
1594///
1595/// Every consumer derives from THIS one expression, so the value channel and
1596/// every derivative channel cannot desync (the #736 / #948 bug genus).
1597///
1598/// The marginal index `q(η_marg)` enters by composing the hand-certified link
1599/// derivative stack `[q, q1, q2, q3, q4]` onto the η primary (slot 0); the
1600/// margin transcendental enters by composing the certified
1601/// [`signed_probit_neglog_unary_stack`] onto the assembled signed margin — the
1602/// stability discipline of #932 (humans own primitive stability, the algebra
1603/// owns combinatorics). The caller MUST guard the signed-margin value against a
1604/// non-finite (non-`+∞`-excluded) NaN before calling; the seeded-evaluation
1605/// wrappers below do that.
1606#[inline]
1607pub(crate) fn rigid_standard_normal_row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1608    p: &[S; 2],
1609    marginal: BernoulliMarginalLinkMap,
1610    z: f64,
1611    y: f64,
1612    w: f64,
1613    probit_scale: f64,
1614) -> Result<S, String> {
1615    // The order-≤4 signed observed margin `m = (2y−1)·η`, written ONCE in
1616    // `rigid_standard_normal_signed_margin` over `S: JetScalar<2>` and shared
1617    // verbatim with the batched builder's Pass-A jet (#932 single source).
1618    let signed = rigid_standard_normal_signed_margin(p, marginal, z, y, probit_scale);
1619    // Preserve the production fail-fast: a NaN (non-`+∞`) signed margin is an
1620    // upstream domain failure, not a tail saturation.
1621    let m = signed.value();
1622    if !(m.is_finite() || m == f64::INFINITY) {
1623        return Err(format!(
1624            "non-finite signed margin in rigid probit row NLL: {m}"
1625        ));
1626    }
1627    // NLL = −w·logΦ(m) via the fused single-Mills-ratio probit neglog stack.
1628    Ok(signed.compose_unary(signed_probit_neglog_unary_stack(m, w)))
1629}
1630
1631/// The order-≤4 signed observed margin `m = (2y−1)·η` of one rigid
1632/// standard-normal Bernoulli row, written ONCE over `S: JetScalar<2>`:
1633/// `q(η_marg)` composed onto the η primary, observed slope `b = s·g`, scale
1634/// `c = √(1 + b²)`, `η = q·c + b·z`. This is the polynomial part shared by
1635/// every channel consumer — the per-row / contracted / full-tower generic NLL
1636/// ([`rigid_standard_normal_row_nll_generic`]) composes the probit-neglog
1637/// transcendental onto it, and the batched builder's Pass-A jet
1638/// ([`rigid_standard_normal_signed_jet`]) evaluates it at `Tower4<2>` — so the
1639/// signed margin has a single source (#932), with no second hand-packed jet.
1640#[inline]
1641pub(crate) fn rigid_standard_normal_signed_margin<S: gam_math::jet_scalar::JetScalar<2>>(
1642    p: &[S; 2],
1643    marginal: BernoulliMarginalLinkMap,
1644    z: f64,
1645    y: f64,
1646    probit_scale: f64,
1647) -> S {
1648    // q(η_marg): compose the link's q-as-function-of-η stack onto the η primary.
1649    let q = p[0].compose_unary([
1650        marginal.q,
1651        marginal.q1,
1652        marginal.q2,
1653        marginal.q3,
1654        marginal.q4,
1655    ]);
1656    let slope = p[1];
1657    // observed slope b = s·g, scale c = √(1 + b²).
1658    let observed_slope = slope.scale(probit_scale);
1659    let b2 = observed_slope.mul(&observed_slope);
1660    let c = b2.add(&S::constant(1.0)).sqrt();
1661    // η = q·c + (s·g)·z, signed margin m = (2y−1)·η.
1662    let eta = q.mul(&c).add(&observed_slope.scale(z));
1663    eta.scale(2.0 * y - 1.0)
1664}
1665
1666/// One row of rigid standard-normal Bernoulli data as a generic
1667/// [`RowProgram<2>`] (#932 production wiring).
1668///
1669/// This is the genuine production consumer of the generic program seam: the row
1670/// NLL is written ONCE in [`rigid_standard_normal_row_nll_generic`] over
1671/// `S: JetScalar<2>`, and this single-row program routes it through the
1672/// [`gam_math::jet_tower`] `generic_*` evaluators
1673/// ([`program_full_tower`](gam_math::jet_tower::program_full_tower) for
1674/// the uncontracted tensors, and the cheap order-2 / contracted scalars for the
1675/// value/grad/Hessian and directional channels). Primaries are
1676/// `[marginal η, slope g]`; the marginal link map and per-row data
1677/// `(z, y, w, probit_scale)` enter as constants on the body.
1678pub(crate) struct RigidStandardNormalRow {
1679    pub(crate) marginal: BernoulliMarginalLinkMap,
1680    pub(crate) g: f64,
1681    pub(crate) z: f64,
1682    pub(crate) y: f64,
1683    pub(crate) w: f64,
1684    pub(crate) probit_scale: f64,
1685}
1686
1687impl gam_math::jet_tower::RowProgram<2> for RigidStandardNormalRow {
1688    fn n_rows(&self) -> usize {
1689        1
1690    }
1691
1692    fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
1693        if row != 0 {
1694            return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1695        }
1696        Ok([self.marginal.eta_value(), self.g])
1697    }
1698
1699    fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
1700        &self,
1701        row: usize,
1702        p: &[S; 2],
1703    ) -> Result<S, String> {
1704        if row != 0 {
1705            return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1706        }
1707        rigid_standard_normal_row_nll_generic(
1708            p,
1709            self.marginal,
1710            self.z,
1711            self.y,
1712            self.w,
1713            self.probit_scale,
1714        )
1715    }
1716}
1717
1718#[inline]
1719pub(crate) fn rigid_standard_normal_tower(
1720    marginal: BernoulliMarginalLinkMap,
1721    g: f64,
1722    z: f64,
1723    y: f64,
1724    w: f64,
1725    probit_scale: f64,
1726) -> Result<Tower4<2>, String> {
1727    // #932 cutover: the full uncontracted tower comes from the SAME single
1728    // generic row-NLL expression every other channel consumer derives from,
1729    // routed through the generic program seam evaluated at the all-channels
1730    // `Tower4` scalar. `program_full_tower` seeds `[marginal η, g]` exactly as
1731    // the previous inline `Tower4::variable` form did, so this is bit-identical
1732    // while giving `RowProgram` a genuine production consumer.
1733    let program = RigidStandardNormalRow {
1734        marginal,
1735        g,
1736        z,
1737        y,
1738        w,
1739        probit_scale,
1740    };
1741    gam_math::jet_tower::program_full_tower(&program, 0).map(|tower| *tower)
1742}
1743
1744/// Branch-free `signed`-margin jet for the rigid standard-normal row kernel.
1745///
1746/// This is the order-≤4 polynomial part of [`rigid_standard_normal_tower`]
1747/// *before* the single transcendental compose: it builds the `Tower4<2>` of the
1748/// signed observed index `signed = (2y−1)·η`, `η = q·c(g) + g·(s·z)`,
1749/// `c(g) = √(1 + (s·g)²)`, with no `erfc`/`exp`/`ln` call. Splitting this off
1750/// lets the batched builder run all the cheap, branch-free jet products in one
1751/// SIMD-friendly pass and isolate the (branchy, transcendental) Mills-ratio
1752/// composition into its own tight pass. The returned jet is the SAME expression
1753/// (`rigid_standard_normal_signed_margin`) the per-row `rigid_standard_normal_tower`
1754/// signed margin evaluates, here at `Tower4<2>` — bit-identical by construction,
1755/// not a parallel hand-packed jet (#932 single source).
1756#[inline]
1757fn rigid_standard_normal_signed_jet(
1758    marginal: BernoulliMarginalLinkMap,
1759    g: f64,
1760    z: f64,
1761    y: f64,
1762    probit_scale: f64,
1763) -> Tower4<2> {
1764    // Seed `[marginal η, g]` exactly as the generic program's `primaries()`, then
1765    // evaluate the one shared signed-margin expression at the all-channels scalar.
1766    let p = [
1767        Tower4::<2>::variable(marginal.eta_value(), 0),
1768        Tower4::<2>::variable(g, 1),
1769    ];
1770    rigid_standard_normal_signed_margin(&p, marginal, z, y, probit_scale)
1771}
1772
1773/// Batched, two-pass builder of the rigid standard-normal row `Tower4<2>` jets
1774/// for a contiguous chunk of rows, written for the auto-vectorizer.
1775///
1776/// Per row the production path ([`rigid_standard_normal_tower`]) interleaves
1777/// (1) cheap branch-free jet products to form the `signed` margin jet, (2) ONE
1778/// branchy transcendental (`erfcx`/`exp`/`ln` via
1779/// [`signed_probit_neglog_unary_stack`]) that dominates the per-row scalar-ALU
1780/// budget across all `n ≈ 356k` rows, and (3) the branch-free Faà-di-Bruno
1781/// `compose_unary` tensor assembly. The interleaving keeps the compiler from
1782/// vectorizing the loop body because the transcendental's internal branches sit
1783/// between the two pure-FMA blocks.
1784///
1785/// This builder runs the same work as three *separate* loops over the chunk:
1786///
1787/// * Pass A — build every `signed` jet (branch-free, [`rigid_standard_normal_signed_jet`]),
1788///   spilling `signed.v` into a contiguous `margins` scratch buffer.
1789/// * Pass B — fill the per-row unary derivative stack `[d0..d4]` from
1790///   `margins`/`weights` (the transcendental, now back-to-back over a flat
1791///   `&[f64]` so branch prediction and the polynomial `k1..k4` portion stream).
1792/// * Pass C — `compose_unary` each `signed` jet against its stack (branch-free,
1793///   pure FMA over the dense tensors → the vectorizable hot block).
1794///
1795/// Every scalar operation, and its order, is identical to the per-row path, so
1796/// the produced jets are bit-for-bit equal; the win is making the n-row build
1797/// memory-bandwidth-bound rather than scalar-ALU/branch-bound. The `fill`
1798/// callback writes the consumer's per-row payload (e.g. `.t3` or `.t4`) from the
1799/// finished jet, so neither tensor is materialized into an intermediate `Vec`.
1800#[inline]
1801pub(super) fn rigid_standard_normal_towers_batch<T>(
1802    marginals: &[BernoulliMarginalLinkMap],
1803    slopes: &[f64],
1804    zs: &[f64],
1805    ys: &[f64],
1806    weights: &[f64],
1807    probit_scale: f64,
1808    out: &mut [T],
1809    mut fill: impl FnMut(&Tower4<2>) -> Result<T, String>,
1810) -> Result<(), String> {
1811    let chunk = marginals.len();
1812    if slopes.len() != chunk
1813        || zs.len() != chunk
1814        || ys.len() != chunk
1815        || weights.len() != chunk
1816        || out.len() != chunk
1817    {
1818        return Err(format!(
1819            "rigid_standard_normal_towers_batch length mismatch: marginals={chunk}, \
1820             slopes={}, zs={}, ys={}, weights={}, out={}",
1821            slopes.len(),
1822            zs.len(),
1823            ys.len(),
1824            weights.len(),
1825            out.len()
1826        ));
1827    }
1828
1829    // Pass A: branch-free signed-margin jets + flat margin scratch.
1830    let mut signed: Vec<Tower4<2>> = Vec::with_capacity(chunk);
1831    let mut margins: Vec<f64> = Vec::with_capacity(chunk);
1832    for i in 0..chunk {
1833        let jet =
1834            rigid_standard_normal_signed_jet(marginals[i], slopes[i], zs[i], ys[i], probit_scale);
1835        margins.push(jet.v);
1836        signed.push(jet);
1837    }
1838
1839    // Pass B: the transcendental, isolated over a flat margin slice. Each entry
1840    // is the exact `[d0..d4]` `compose_unary` consumes; the production path's
1841    // fail-fast on a non-finite (non-`+∞`) margin is preserved here.
1842    let mut stacks: Vec<[f64; 5]> = Vec::with_capacity(chunk);
1843    for i in 0..chunk {
1844        let m = margins[i];
1845        if !(m.is_finite() || m == f64::INFINITY) {
1846            return Err(format!(
1847                "non-finite signed margin in rigid probit tower batch: {m}"
1848            ));
1849        }
1850        stacks.push(signed_probit_neglog_unary_stack(m, weights[i]));
1851    }
1852
1853    // Pass C: branch-free dense compose + consumer fill.
1854    for i in 0..chunk {
1855        let tower = signed[i].compose_unary(stacks[i]);
1856        out[i] = fill(&tower)?;
1857    }
1858    Ok(())
1859}
1860
1861#[inline]
1862pub(super) fn rigid_standard_normal_row_kernel(
1863    marginal: BernoulliMarginalLinkMap,
1864    g: f64,
1865    z: f64,
1866    y: f64,
1867    w: f64,
1868    probit_scale: f64,
1869) -> Result<(f64, [f64; 2], [[f64; 2]; 2]), String> {
1870    // #932 cutover: value/gradient/Hessian derive from the SAME single generic
1871    // row-NLL expression (`rigid_standard_normal_row_nll_generic`) every other
1872    // channel consumer uses, routed through the `RowProgram` seam at the
1873    // packed `Order2<2>` scalar — there is no longer a hand-assembled `Tower2<2>`
1874    // here. Seeds `[marginal η, g]` exactly as the deleted inline form did, so it
1875    // is bit-identical (the `rigid_bernoulli_*_agrees_with_jet_tower_program_all_channels`
1876    // oracle pins v/g/H ≤ 1e-12), while sharing one definition with the third/
1877    // fourth/full-tower channels.
1878    let program = RigidStandardNormalRow {
1879        marginal,
1880        g,
1881        z,
1882        y,
1883        w,
1884        probit_scale,
1885    };
1886    gam_math::jet_tower::program_row_kernel(&program, 0)
1887}
1888
1889/// Mixed `(primary, z)` second derivative of the rigid standard-normal row
1890/// LOG-LIKELIHOOD score: the per-row 2-vector
1891/// `[∂²(log L)/∂q∂z, ∂²(log L)/∂g∂z]` in the primary coordinates `(q = marginal η,
1892/// g = slope)`, evaluated at this row's converged `(q, g)` and calibrated
1893/// latent score `z = ζ`.
1894///
1895/// SIGN CONVENTION (#1131). This returns the mixed partial of the
1896/// LOG-LIKELIHOOD score `score_β,i = ∂(log L_i)/∂β`, NOT of the negative
1897/// log-likelihood `ℓ = −log L`. Concretely the row jet evaluates the NLL
1898/// `ℓ = −w·log Φ(sign·η)` and we NEGATE its mixed `(primary, z)` Hessian entries,
1899/// so the returned 2-vector is `+∂²(log L_i)/∂(q,g)∂ζ_i = −∂²ℓ_i/∂(q,g)∂ζ_i`.
1900/// This is the convention under which the Murphy–Topel chain
1901/// `G = Σ_i s_i·(∂ζ_i/∂θ₁)` with `s_i = ∂score_β,i/∂ζ_i` and `Vb = H_β⁻¹`
1902/// (the NLL-Hessian inverse) gives the SIGNED sensitivity with the right sign:
1903/// the implicit-function theorem on the stationarity `∂(log L)/∂β = 0` yields
1904/// `∂β̂/∂θ₁ = −(∂²log L/∂β²)⁻¹·∂²(log L)/∂β∂θ₁ = +H_β⁻¹·G = +Vb·G`. (Had we
1905/// returned the NLL mixed partial instead, `Vb·G` would equal `−∂β̂/∂θ₁` — a
1906/// benign sign flip for the PSD quadratic SE `(Vb·G)V₁(Vb·G)ᵀ`, but wrong for
1907/// any signed consumer of the sensitivity.)
1908///
1909/// This is the #1028 Murphy–Topel generated-regressor channel: `score_β,i =
1910/// ∂(log L_i)/∂β = J_iᵀ·(∂(log L_i)/∂(q,g))`, so the per-row slope-score
1911/// sensitivity to the calibrated score is
1912/// `s_i = ∂score_β,i/∂ζ_i = J_iᵀ·(∂²(log L_i)/∂(q,g)∂ζ_i)`, and the primary
1913/// 2-vector returned here is exactly `∂²(log L_i)/∂(q,g)∂ζ_i`. The block-level
1914/// contraction `J_iᵀ` (marginal+logslope design rows) is applied by the caller.
1915///
1916/// It is computed by seeding `z` as a THIRD jet variable (index 2) in the SAME
1917/// order-≤2 jet algebra the value/gradient/Hessian path uses, carried by the
1918/// packed `Order2<3>`/`Tower2<3>` scalar rather than a dense `Tower4<3>`
1919/// (#932 row-jet machinery, packed-scalar perf cutover): the
1920/// rigid standard-normal observed index is `η = q·c(g) + g·(s·z)` with
1921/// `c(g) = √(1 + (s·g)²)`, `s = probit_scale`, and `ℓ = −w·log Φ(sign·η)`. The
1922/// converged-frame mixed partials of the NLL are the off-diagonal Hessian
1923/// entries `tower.h[q][z]` and `tower.h[g][z]`, read off in one composition and
1924/// NEGATED to the log-likelihood-score convention — the only extra cost over the
1925/// production `Tower4<2>` evaluation is the third jet axis.
1926#[inline]
1927pub(super) fn rigid_standard_normal_mixed_z_sensitivity(
1928    marginal: BernoulliMarginalLinkMap,
1929    g: f64,
1930    z: f64,
1931    y: f64,
1932    w: f64,
1933    probit_scale: f64,
1934) -> Result<[f64; 2], String> {
1935    // Three jet axes: q = marginal η (0), g = slope (1), z = latent score (2).
1936    //
1937    // #932 perf: this consumer reads ONLY the two mixed Hessian channels
1938    // `h[0][2]`/`h[1][2]`, so it needs only the value/gradient/Hessian stack —
1939    // the packed `Order2<3>` scalar (operating on its inner `Tower2<3>`), NOT a
1940    // dense `Tower4<3>` that would materialise the unused `K³`/`K⁴` `t3`/`t4`
1941    // tensors. The order-≤2 channels are bit-identical to the dense tower
1942    // (`Tower2::mul`/`compose_unary` match `Tower4` term-for-term), so the read
1943    // entries are unchanged; the `q3`/`q4` marginal-link channels are dropped
1944    // because no order-≤2 channel of the composed jet reads them.
1945    use gam_math::jet_tower::Tower2;
1946    let mut q = Tower2::<3>::constant(marginal.q);
1947    q.g[0] = marginal.q1;
1948    q.h[0][0] = marginal.q2;
1949    let slope = Tower2::<3>::variable(g, 1);
1950    let z_var = Tower2::<3>::variable(z, 2);
1951    let observed_logslope = slope * probit_scale;
1952    let c = (observed_logslope * observed_logslope + 1.0).sqrt();
1953    // η = q·c + g·(s·z): z enters linearly through the slope×z product, so the
1954    // mixed (q,z)/(g,z) curvature is carried entirely by the unary NLL chain and
1955    // the η-bilinear, exactly as in the Tower4<2> production path.
1956    let eta = q * c + slope * (z_var * probit_scale);
1957    let signed = eta * (2.0 * y - 1.0);
1958    // ONE transcendental per row (see `rigid_standard_normal_tower`).
1959    if !(signed.v.is_finite() || signed.v == f64::INFINITY) {
1960        return Err(format!(
1961            "rigid probit mixed-z sensitivity: non-finite signed margin {} at q={}, g={g}, z={z}, y={y}",
1962            signed.v, marginal.q
1963        ));
1964    }
1965    let stack = signed_probit_neglog_unary_stack(signed.v, w);
1966    if !stack[0].is_finite() {
1967        return Err(format!(
1968            "rigid probit mixed-z sensitivity: non-finite log Φ at q={}, g={g}, z={z}, y={y}",
1969            marginal.q
1970        ));
1971    }
1972    // Order-≤2 composition consumes only the leading `[f, f', f'']` of the
1973    // certified `[f64; 5]` derivative stack.
1974    let tower = signed.compose_unary([stack[0], stack[1], stack[2]]);
1975    // #1131: `tower` is the NLL `ℓ = −w·log Φ`, so `tower.h[·][z]` is the mixed
1976    // partial of the NLL. Negate to the LOG-LIKELIHOOD-score convention
1977    // `s = ∂²(log L)/∂(primary)∂z = −∂²ℓ/∂(primary)∂z`, under which the
1978    // downstream Murphy–Topel chain `Vb·G = +∂β̂/∂θ₁` carries the correct sign
1979    // (see the function doc). The SE is the PSD quadratic `(Vb·G)V₁(Vb·G)ᵀ` and
1980    // is invariant to this sign, so the reported standard errors are unchanged.
1981    let s_q = -tower.h[0][2];
1982    let s_g = -tower.h[1][2];
1983    if !(s_q.is_finite() && s_g.is_finite()) {
1984        return Err(format!(
1985            "rigid probit mixed-z sensitivity: non-finite ∂²(log L)/∂(q,g)∂z = [{s_q}, {s_g}] at q={}, g={g}, z={z}",
1986            marginal.q
1987        ));
1988    }
1989    Ok([s_q, s_g])
1990}
1991
1992/// Assemble the #1028 Murphy–Topel slope-score sensitivity matrix
1993/// `score_zeta_sensitivity` (`n × p_β`, row `i` = `s_i = ∂score_β,i/∂ζ_i`) for
1994/// the rigid standard-normal BMS kernel — the kernel the conditional
1995/// location-scale gate ALWAYS selects (`LatentMeasureKind::StandardNormal`).
1996///
1997/// where `s_i = ∂score_β,i/∂ζ_i` is the LOG-LIKELIHOOD-score sensitivity (see
1998/// the sign convention in [`rigid_standard_normal_mixed_z_sensitivity`], #1131).
1999/// For each row `i` the primary 2-vector `∂²(log L_i)/∂(q,g)∂ζ_i` is read off the
2000/// z-augmented row jet ([`rigid_standard_normal_mixed_z_sensitivity`]) at the
2001/// converged marginal index `q_i` (`marginal_eta[i]`) and slope `g_i`
2002/// (`slope_eta[i]`) and calibrated score `ζ_i` (`z[i]`), then contracted through
2003/// the block Jacobian `J_iᵀ` (the same marginal+logslope design-row scatter the
2004/// row kernel exposes via `jacobian_transpose_action`):
2005///
2006/// ```text
2007///   s_i[marginal_range]  = (∂²(log L_i)/∂q∂ζ_i) · marginal_design.row(i)
2008///   s_i[logslope_range]  = (∂²(log L_i)/∂g∂ζ_i) · logslope_design.row(i)
2009/// ```
2010///
2011/// `logslope_design` MUST be the reduced-basis design `G·T` actually fitted and
2012/// `p_beta` MUST equal `p_marginal + r`. Flex models use the separate exact
2013/// cubic-jet channel that includes score_warp/link_dev coordinates; this rigid
2014/// helper never accepts or zero-pads a wider covariance frame (#2303).
2015pub(super) fn rigid_standard_normal_score_zeta_sensitivity(
2016    base_link: &InverseLink,
2017    marginal_eta: &Array1<f64>,
2018    slope_eta: &Array1<f64>,
2019    z: &Array1<f64>,
2020    y: &Array1<f64>,
2021    weights: &Array1<f64>,
2022    probit_scale: f64,
2023    marginal_design: ArrayView2<'_, f64>,
2024    logslope_design: ArrayView2<'_, f64>,
2025    p_beta: usize,
2026) -> Result<Array2<f64>, String> {
2027    let n = marginal_eta.len();
2028    let p_m = marginal_design.ncols();
2029    let r = logslope_design.ncols();
2030    if slope_eta.len() != n
2031        || z.len() != n
2032        || y.len() != n
2033        || weights.len() != n
2034        || marginal_design.nrows() != n
2035        || logslope_design.nrows() != n
2036    {
2037        return Err(format!(
2038            "score_zeta_sensitivity row mismatch: marginal_eta={n}, slope_eta={}, z={}, y={}, \
2039             weights={}, marginal_design rows={}, logslope_design rows={}",
2040            slope_eta.len(),
2041            z.len(),
2042            y.len(),
2043            weights.len(),
2044            marginal_design.nrows(),
2045            logslope_design.nrows()
2046        ));
2047    }
2048    if p_m + r != p_beta {
2049        return Err(format!(
2050            "rigid score_zeta_sensitivity width mismatch: marginal({p_m}) + logslope({r}) != p_beta({p_beta})"
2051        ));
2052    }
2053    let mut s = Array2::<f64>::zeros((n, p_beta));
2054    for i in 0..n {
2055        let marginal = bernoulli_marginal_link_map(base_link, marginal_eta[i])?;
2056        let [s_q, s_g] = rigid_standard_normal_mixed_z_sensitivity(
2057            marginal,
2058            slope_eta[i],
2059            z[i],
2060            y[i],
2061            weights[i],
2062            probit_scale,
2063        )?;
2064        // J_iᵀ scatter into the reduced-frame coordinates: marginal block first,
2065        // then the reduced logslope block.
2066        if s_q != 0.0 {
2067            let m_row = marginal_design.row(i);
2068            for (j, &mij) in m_row.iter().enumerate() {
2069                s[[i, j]] = s_q * mij;
2070            }
2071        }
2072        if s_g != 0.0 {
2073            let g_row = logslope_design.row(i);
2074            for (j, &gij) in g_row.iter().enumerate() {
2075                s[[i, p_m + j]] = s_g * gij;
2076            }
2077        }
2078    }
2079    Ok(s)
2080}
2081
2082#[inline]
2083pub(super) fn rigid_standard_normal_third_full(
2084    marginal: BernoulliMarginalLinkMap,
2085    g: f64,
2086    z: f64,
2087    y: f64,
2088    w: f64,
2089    probit_scale: f64,
2090) -> Result<[[[f64; 2]; 2]; 2], String> {
2091    Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t3)
2092}
2093
2094/// Contract a symmetric 3-tensor on its third index with a primary-space
2095/// direction `d = (d_eta, d_g)`, producing the symmetric 2×2 contracted
2096/// matrix the outer-derivative pipeline consumes:
2097///   `M[a][b] = Σ_c T[a][b][c] · d[c]`.
2098#[inline]
2099pub(super) fn contract_third_full(t: &[[[f64; 2]; 2]; 2], d_eta: f64, d_g: f64) -> [[f64; 2]; 2] {
2100    [
2101        [
2102            t[0][0][0] * d_eta + t[0][0][1] * d_g,
2103            t[0][1][0] * d_eta + t[0][1][1] * d_g,
2104        ],
2105        [
2106            t[1][0][0] * d_eta + t[1][0][1] * d_g,
2107            t[1][1][0] * d_eta + t[1][1][1] * d_g,
2108        ],
2109    ]
2110}
2111
2112#[inline]
2113pub(super) fn rigid_standard_normal_fourth_full(
2114    marginal: BernoulliMarginalLinkMap,
2115    g: f64,
2116    z: f64,
2117    y: f64,
2118    w: f64,
2119    probit_scale: f64,
2120) -> Result<[[[[f64; 2]; 2]; 2]; 2], String> {
2121    // #932 single-sourcing: the full uncontracted fourth-order primary tensor is
2122    // the `.t4` channel of the SAME `Tower4<2>` row jet the value/gradient/Hessian
2123    // and the third-order tensor (`rigid_standard_normal_third_full` → `.t3`) are
2124    // read from. The marginal latent-coordinate chain `q(η)` is already seeded
2125    // into axis 0 of the tower (`q.g[0]=q1, q.h[0][0]=q2, q.t3[0][0][0]=q3,
2126    // q.t4[0][0][0][0]=q4` in `rigid_standard_normal_signed_jet`), so `.t4` is
2127    // delivered directly in the production `(η, g)` primary space — no separate
2128    // Faà-di-Bruno q-chain reassembly. This replaces the former hand-written
2129    // fourth-derivative chain rule with the mechanically-derived tower output,
2130    // exactly mirroring how `.t3` is consumed;
2131    // it is cross-checked term-for-term against the independent
2132    // `HandRigidProbitKernel` witness in
2133    // `rigid_standard_normal_tower_path_matches_hand_chain_witness`.
2134    Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t4)
2135}
2136
2137/// Combined uncontracted THIRD **and** FOURTH primary tensors for one rigid
2138/// standard-normal row, read off a SINGLE shared `Tower4<2>` jet.
2139///
2140/// `rigid_standard_normal_third_full` (→ `.t3`) and
2141/// `rigid_standard_normal_fourth_full` (→ `.t4`) each build a full
2142/// `rigid_standard_normal_tower` and discard the OTHER tensor — so a consumer
2143/// that needs both for the same `(row, β)` point (the outer Jeffreys/REML
2144/// derivative path warms both the `rigid_third_full` and `rigid_fourth_full`
2145/// caches in the same fit; see the paired `rigid_{third,fourth}_full_cached`
2146/// warm-up) pays the per-row Mills-ratio transcendental
2147/// (`signed_probit_neglog_unary_stack`, ~88% of the per-row scalar cost) TWICE
2148/// where ONCE suffices. The two tensors are the `.t3` / `.t4` channels of the
2149/// same tower, so this builder evaluates that tower ONCE and returns both.
2150///
2151/// Contract a symmetric 4-tensor on its last two indices with two
2152/// primary-space directions `u = (u_eta, u_g)` and `v = (v_eta, v_g)`,
2153/// producing the symmetric 2×2 matrix the outer-Hessian pipeline expects:
2154///   `M[a][b] = Σ_{c,d} T[a][b][c][d] · u[c] · v[d]`.
2155#[inline]
2156pub(super) fn contract_fourth_full(
2157    t: &[[[[f64; 2]; 2]; 2]; 2],
2158    u_eta: f64,
2159    u_g: f64,
2160    v_eta: f64,
2161    v_g: f64,
2162) -> [[f64; 2]; 2] {
2163    let mut out = [[0.0; 2]; 2];
2164    for a in 0..2 {
2165        for b in 0..2 {
2166            let mut sum = 0.0;
2167            sum += t[a][b][0][0] * u_eta * v_eta;
2168            sum += t[a][b][0][1] * u_eta * v_g;
2169            sum += t[a][b][1][0] * u_g * v_eta;
2170            sum += t[a][b][1][1] * u_g * v_g;
2171            out[a][b] = sum;
2172        }
2173    }
2174    out
2175}
2176
2177pub(super) fn ensure_finite_third_full_cache_row(
2178    t: &[[[f64; 2]; 2]; 2],
2179    context: &str,
2180) -> Result<(), String> {
2181    if t.iter().flatten().flatten().all(|value| value.is_finite()) {
2182        Ok(())
2183    } else {
2184        Err(format!(
2185            "{context}: warmed third-derivative cache row contains a non-finite value"
2186        ))
2187    }
2188}
2189
2190pub(super) fn ensure_finite_fourth_full_cache_row(
2191    t: &[[[[f64; 2]; 2]; 2]; 2],
2192    context: &str,
2193) -> Result<(), String> {
2194    if t.iter()
2195        .flatten()
2196        .flatten()
2197        .flatten()
2198        .all(|value| value.is_finite())
2199    {
2200        Ok(())
2201    } else {
2202        Err(format!(
2203            "{context}: warmed fourth-derivative cache row contains a non-finite value"
2204        ))
2205    }
2206}
2207
2208pub(crate) fn unary_derivatives_sqrt(x: f64) -> [f64; 5] {
2209    let s = x.max(1e-300).sqrt();
2210    let x1 = x.max(1e-300);
2211    let x2 = x1 * x1;
2212    let x3 = x2 * x1;
2213    [
2214        s,
2215        0.5 / s,
2216        -0.25 / (x1 * s),
2217        3.0 / (8.0 * x2 * s),
2218        -15.0 / (16.0 * x3 * s),
2219    ]
2220}
2221pub(crate) fn unary_derivatives_neglog_phi(x: f64, weight: f64) -> [f64; 5] {
2222    // Single source of truth for the signed-probit value+derivative stack:
2223    // one Mills-ratio transcendental feeds both logΦ and k1..k4 (the prior
2224    // body evaluated `signed_probit_logcdf_and_mills_ratio` twice). The
2225    // ±∞/NaN/zero-weight boundary limits are handled identically inside.
2226    signed_probit_neglog_unary_stack(x, weight)
2227}
2228
2229/// Derivatives of `log(x)` through 4th order.
2230///
2231/// # Contract
2232///
2233/// `x` must be strictly positive. `log` and its derivatives are undefined at
2234/// and below the boundary, so this function does NOT clamp: a previous version
2235/// silently replaced `x` by `x.max(1e-300)`, which fabricated enormous finite
2236/// derivatives (`1/1e-300` etc.) that are the derivatives of neither `log(x)`
2237/// nor `log(max(x, floor))`. Such a non-positive argument signals an upstream
2238/// domain failure (e.g. a monotonicity violation) that must surface, not be
2239/// masked. Every caller guarantees `x > 0` before invoking this:
2240/// the survival marginal-slope kernels evaluate `log` of the transformed time
2241/// derivative `q'(t)·√(1+b²)` only after passing `survival_derivative_guard`
2242/// (`q'(t) >= derivative_guard > 0`, `√(1+b²) > 0`). A non-positive `x`
2243/// therefore never reaches here on any supported path; were one to, the
2244/// function returns the honest IEEE result (`-inf`/`NaN`) — identical in debug
2245/// and release — rather than a finite fabrication.
2246pub(crate) fn unary_derivatives_log(x: f64) -> [f64; 5] {
2247    let x2 = x * x;
2248    let x3 = x2 * x;
2249    let x4 = x3 * x;
2250    [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
2251}
2252
2253/// Derivatives of log φ(x) = -½x² - ½ln(2π) through 4th order.
2254pub(crate) fn unary_derivatives_log_normal_pdf(x: f64) -> [f64; 5] {
2255    let c = 0.5 * (2.0 * std::f64::consts::PI).ln();
2256    [-0.5 * x * x - c, -x, -1.0, 0.0, 0.0]
2257}
2258
2259#[cfg(test)]
2260mod covariance_admission_tests {
2261    use super::*;
2262    use ndarray::array;
2263
2264    #[test]
2265    fn full_covariance_admission_rejects_one_ulp_asymmetry_932() {
2266        let upper = 0.25_f64;
2267        let lower = f64::from_bits(upper.to_bits() + 1);
2268        let error = MarginalSlopeCovariance::full(array![[1.0, upper], [lower, 1.0]])
2269            .expect_err("any asymmetric full operator must be rejected");
2270        assert!(error.contains("must be exactly symmetric"), "{error}");
2271    }
2272
2273    #[test]
2274    fn full_covariance_admission_rejects_indefinite_matrix_before_row_use_932() {
2275        let error = MarginalSlopeCovariance::full(array![[1.0, 2.0], [2.0, 1.0]])
2276            .expect_err("an indefinite full operator is not a covariance");
2277        assert!(error.contains("must be positive semidefinite"), "{error}");
2278    }
2279
2280    #[test]
2281    fn full_covariance_admission_accepts_exact_singular_psd_932() {
2282        MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, 0.0]])
2283            .expect("an exact singular PSD covariance is admissible");
2284    }
2285
2286    #[test]
2287    fn full_covariance_admission_accepts_coupled_singular_psd_932() {
2288        let covariance = MarginalSlopeCovariance::full(array![[1.0, 1.0], [1.0, 1.0]])
2289            .expect("a coupled singular PSD covariance is admissible");
2290        assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2291        assert_eq!(covariance.to_dense(), array![[1.0, 1.0], [1.0, 1.0]]);
2292    }
2293
2294    #[test]
2295    fn exact_nonzero_offdiagonal_classifier_retains_full_geometry_932() {
2296        let epsilon = 1.0e-14;
2297        let scores = array![[-1.0, -epsilon], [1.0, epsilon], [0.0, -1.0], [0.0, 1.0]];
2298        let covariance =
2299            marginal_slope_covariance_from_scores(scores.view(), &Array1::ones(4)).unwrap();
2300        let dense = covariance.to_dense();
2301        assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2302        assert_ne!(dense[[0, 1]], 0.0);
2303        assert_eq!(dense[[0, 1]], dense[[1, 0]]);
2304        let direction = [0.75, -1.25];
2305        let expected = direction[0] * (dense[[0, 0]] * direction[0] + dense[[0, 1]] * direction[1])
2306            + direction[1] * (dense[[1, 0]] * direction[0] + dense[[1, 1]] * direction[1]);
2307        let actual = covariance.quadratic_form(&direction).unwrap();
2308        assert!((actual - expected).abs() <= 2.0e-15);
2309    }
2310
2311    #[test]
2312    fn diagonal_covariance_entries_are_the_exact_geometry_authority_932() {
2313        let covariance = MarginalSlopeCovariance::diagonal(array![3.75]).unwrap();
2314        assert_eq!(covariance.ones_quadratic_form(), 3.75);
2315        assert_eq!(covariance.quadratic_form(&[1.0]).unwrap(), 3.75);
2316    }
2317
2318    #[test]
2319    fn equal_dense_covariance_quadratic_forms_match_all_representations_932() {
2320        let diagonal = MarginalSlopeCovariance::diagonal(array![1.2, 0.7]).unwrap();
2321        let full = MarginalSlopeCovariance::full(array![[1.2, 0.0], [0.0, 0.7]]).unwrap();
2322        let low_rank =
2323            MarginalSlopeCovariance::low_rank(array![[1.2_f64.sqrt(), 0.0], [0.0, 0.7_f64.sqrt()]])
2324                .unwrap();
2325        let direction = [0.35, -0.8];
2326        let expected = diagonal.quadratic_form(&direction).unwrap();
2327        for covariance in [&full, &low_rank] {
2328            let actual = covariance.quadratic_form(&direction).unwrap();
2329            assert!((actual - expected).abs() <= 2.0e-15);
2330            assert!(
2331                (covariance.ones_quadratic_form() - diagonal.ones_quadratic_form()).abs()
2332                    <= 2.0e-15
2333            );
2334        }
2335    }
2336}
2337
2338#[cfg(test)]
2339mod jet_tower_oracle_tests {
2340    //! #932 deployment step 2 for the BMS rigid Bernoulli `RowKernel<2>`.
2341    //!
2342    //! The production rigid standard-normal row kernel
2343    //! ([`rigid_standard_normal_row_kernel`] / `_third_full` / `_fourth_full`)
2344    //! reads value/grad/Hessian/third/fourth straight off ONE
2345    //! [`rigid_standard_normal_tower`] `Tower4<2>` — the strongest #932 form,
2346    //! where the production kernel literally *is* the single-expression jet.
2347    //! What was missing (unlike the two survival `RowKernel` families, which
2348    //! already carry `verify_kernel_channels` oracles) is an INDEPENDENT
2349    //! cross-check that this production tower is correct. This module adds it:
2350    //!
2351    //! * an independent [`RowProgram<2>`] that writes the row NLL
2352    //!   `ℓ = −w·logΦ((2y−1)·η)`, `η = q·√(1+(s·g)²) + s·g·z` ONCE over generic
2353    //!   generic jet arithmetic (a different composition order than the fused
2354    //!   production `signed` jet → exercises the Leibniz/Faà-di-Bruno layer
2355    //!   where the #736 cross-block sign-flip bug genus lives), and
2356    //! * a special-function-independent central-FD witness of the value channel
2357    //!   that re-derives `logΦ` from `libm::erfc`, pinning the probit derivative
2358    //!   stack itself (so the oracle does not merely re-use the production
2359    //!   transcendental).
2360
2361    use super::*;
2362
2363    #[test]
2364    fn signed_probit_stack_preserves_extreme_tail_derivatives_and_weight_sign() {
2365        let positive = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, 2.0);
2366        assert_eq!(
2367            positive,
2368            [f64::INFINITY, f64::NEG_INFINITY, 2.0, -0.0, -0.0]
2369        );
2370        let negative = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, -2.0);
2371        assert_eq!(negative, [f64::NEG_INFINITY, f64::INFINITY, -2.0, 0.0, 0.0]);
2372
2373        let right = signed_probit_neglog_unary_stack(38.6, 1.0);
2374        assert_eq!(right[1], -0.0);
2375        assert!(right[2] > 0.0 && right[2].is_subnormal());
2376        assert!(right[3] < 0.0 && right[3].is_subnormal());
2377        assert!(right[4] > 0.0 && right[4].is_subnormal());
2378
2379        let left = signed_probit_neglog_unary_stack(-1.0e100, 1.0);
2380        assert_eq!(left[1], -1.0e100);
2381        assert_eq!(left[2], 1.0);
2382        assert!(left[3] < 0.0 && left[3].is_finite());
2383        assert_eq!(left[4], -0.0);
2384    }
2385
2386    /// #932 combined third+fourth primary tensors read off ONE shared
2387    /// `rigid_standard_normal_tower` jet (the redundancy-free form of the
2388    /// separate `_third_full` / `_fourth_full` builds, bit-identical to them).
2389    /// Lives in this `#[cfg(test)]` module — its only consumers are the
2390    /// bit-identity checks below — so it is not a production `src` item with no
2391    /// production caller (production reads the separate builders) and is not dead
2392    /// code in the non-test lib build.
2393    fn rigid_standard_normal_third_and_fourth_full(
2394        marginal: BernoulliMarginalLinkMap,
2395        g: f64,
2396        z: f64,
2397        y: f64,
2398        w: f64,
2399        probit_scale: f64,
2400    ) -> Result<([[[f64; 2]; 2]; 2], [[[[f64; 2]; 2]; 2]; 2]), String> {
2401        let tower = rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?;
2402        Ok((tower.t3, tower.t4))
2403    }
2404    use gam_math::jet_tower::{
2405        KernelChannels, RowProgram, program_full_tower, verify_kernel_channels,
2406    };
2407
2408    /// Independent single-expression row NLL for the rigid standard-normal
2409    /// Bernoulli kernel, primaries `(q_eta = marginal η, g = slope)`.
2410    struct BernoulliRigidStandardNormalNllProgram {
2411        /// `(marginal η, slope g)` per row.
2412        primaries: Vec<[f64; 2]>,
2413        /// Per-row `(z latent score, y in {0,1}, w weight)`.
2414        z: Vec<f64>,
2415        y: Vec<f64>,
2416        w: Vec<f64>,
2417        probit_scale: f64,
2418    }
2419
2420    impl RowProgram<2> for BernoulliRigidStandardNormalNllProgram {
2421        fn n_rows(&self) -> usize {
2422            self.primaries.len()
2423        }
2424
2425        fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2426            self.primaries
2427                .get(row)
2428                .copied()
2429                .ok_or_else(|| format!("bernoulli rigid nll program: row {row} out of range"))
2430        }
2431
2432        fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
2433            &self,
2434            row: usize,
2435            p: &[S; 2],
2436        ) -> Result<S, String> {
2437            let z = self.z[row];
2438            let y = self.y[row];
2439            let w = self.w[row];
2440            let s = self.probit_scale;
2441            // q(η) via the family's own marginal link-map derivative stack,
2442            // composed through generic Leibniz on the η primary (independent of
2443            // the production signed-jet, which seeds the q tensor slots directly).
2444            let eta_marginal = p[0];
2445            let link = bernoulli_marginal_link_map(
2446                &InverseLink::Standard(gam_problem::StandardLink::Probit),
2447                eta_marginal.value(),
2448            )?;
2449            let q = eta_marginal.compose_unary([link.q, link.q1, link.q2, link.q3, link.q4]);
2450            let g = p[1];
2451            // observed slope b = s·g, scale c = √(1 + b²).
2452            let observed_slope = g.scale(s);
2453            let one_plus_slope_squared = observed_slope.mul(&observed_slope).add(&S::constant(1.0));
2454            let c = one_plus_slope_squared
2455                .compose_unary(unary_derivatives_sqrt(one_plus_slope_squared.value()));
2456            // η = q·c + b·z, signed margin m = (2y−1)·η.
2457            let eta = q.mul(&c).add(&observed_slope.scale(z));
2458            let signed = eta.scale(2.0 * y - 1.0);
2459            // NLL = −w·logΦ(m) via the documented probit neglog stack.
2460            Ok(signed.compose_unary(unary_derivatives_neglog_phi(signed.value(), w)))
2461        }
2462    }
2463
2464    /// Special-function-independent scalar row NLL `ℓ(q_eta, g)` using
2465    /// `libm::erfc`, for the central-FD value-channel witness.
2466    fn scalar_nll(eta_marginal: f64, g: f64, z: f64, y: f64, w: f64, s: f64) -> f64 {
2467        let link = bernoulli_marginal_link_map(
2468            &InverseLink::Standard(gam_problem::StandardLink::Probit),
2469            eta_marginal,
2470        )
2471        .unwrap();
2472        let observed_slope = g * s;
2473        let c = (observed_slope * observed_slope + 1.0).sqrt();
2474        let eta = link.q * c + observed_slope * z;
2475        let signed = (2.0 * y - 1.0) * eta;
2476        let cdf = 0.5 * libm::erfc(-signed / std::f64::consts::SQRT_2);
2477        -w * cdf.max(1e-300).ln()
2478    }
2479
2480    #[test]
2481    fn rigid_bernoulli_row_kernel_agrees_with_jet_tower_program_all_channels() {
2482        // Mixed responses, weights, latent scores, and slope regimes; the last
2483        // rows push the marginal index toward the normal tails while staying
2484        // finite. Probit marginal link, standard-normal latent measure.
2485        let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2486        let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2487        let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2488        let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2489        let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2490        let n = eta.len();
2491
2492        // Deterministic direction vectors (no RNG dependency).
2493        let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2494
2495        for &probit_scale in &[1.0_f64, 0.8] {
2496            let program = BernoulliRigidStandardNormalNllProgram {
2497                primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2498                z: z.to_vec(),
2499                y: y.to_vec(),
2500                w: w.to_vec(),
2501                probit_scale,
2502            };
2503
2504            for row in 0..n {
2505                let tower = program_full_tower(&program, row).expect("tower evaluation");
2506
2507                // Production scalar kernel channels (the hand path under audit).
2508                let marginal = bernoulli_marginal_link_map(
2509                    &InverseLink::Standard(gam_problem::StandardLink::Probit),
2510                    eta[row],
2511                )
2512                .expect("link map");
2513                let (value, gradient, hessian) = rigid_standard_normal_row_kernel(
2514                    marginal,
2515                    g[row],
2516                    z[row],
2517                    y[row],
2518                    w[row],
2519                    probit_scale,
2520                )
2521                .expect("production row kernel");
2522
2523                // One shared tower for BOTH the third and fourth tensors (the
2524                // #932 transcendental-de-dup builder): this is the redundancy-
2525                // free form of the former two separate
2526                // `rigid_standard_normal_{third,fourth}_full` calls, and it is
2527                // pinned bit-identically against them in
2528                // `rigid_third_and_fourth_full_shares_one_tower_bit_identical`.
2529                let (third_full, fourth_full) = rigid_standard_normal_third_and_fourth_full(
2530                    marginal,
2531                    g[row],
2532                    z[row],
2533                    y[row],
2534                    w[row],
2535                    probit_scale,
2536                )
2537                .expect("production third+fourth");
2538                let third: Vec<([f64; 2], [[f64; 2]; 2])> = dirs
2539                    .iter()
2540                    .map(|d| (*d, contract_third_full(&third_full, d[0], d[1])))
2541                    .collect();
2542
2543                let fourth: Vec<([f64; 2], [f64; 2], [[f64; 2]; 2])> = dirs
2544                    .iter()
2545                    .enumerate()
2546                    .map(|(i, u)| {
2547                        let v = dirs[(i + 1) % dirs.len()];
2548                        (
2549                            *u,
2550                            v,
2551                            contract_fourth_full(&fourth_full, u[0], u[1], v[0], v[1]),
2552                        )
2553                    })
2554                    .collect();
2555
2556                let claims = KernelChannels {
2557                    value,
2558                    gradient,
2559                    hessian,
2560                    third,
2561                    fourth,
2562                };
2563
2564                verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
2565                    panic!(
2566                        "probit_scale {probit_scale} row {row}: production rigid Bernoulli \
2567                         RowKernel disagrees with #932 jet-tower truth: {e}"
2568                    )
2569                });
2570
2571                // Special-function-independent FD witness of the value channel:
2572                // re-derives logΦ from `libm::erfc`, pinning the probit derivative
2573                // stack rather than re-using the production one.
2574                let h = 1e-3;
2575                let f = |de: f64, dg: f64| {
2576                    scalar_nll(
2577                        eta[row] + de,
2578                        g[row] + dg,
2579                        z[row],
2580                        y[row],
2581                        w[row],
2582                        probit_scale,
2583                    )
2584                };
2585                let f0 = f(0.0, 0.0);
2586                assert!(
2587                    (f0 - tower.v).abs() <= 1e-9 * f0.abs().max(1.0),
2588                    "row {row}: independent scalar NLL {f0:+.12e} != tower value {:+.12e}",
2589                    tower.v
2590                );
2591                // 5-point first-derivative stencils.
2592                let g_eta = (f(-2.0 * h, 0.0) - 8.0 * f(-h, 0.0) + 8.0 * f(h, 0.0)
2593                    - f(2.0 * h, 0.0))
2594                    / (12.0 * h);
2595                let g_g = (f(0.0, -2.0 * h) - 8.0 * f(0.0, -h) + 8.0 * f(0.0, h) - f(0.0, 2.0 * h))
2596                    / (12.0 * h);
2597                for (label, fd, ad) in [("∂η", g_eta, tower.g[0]), ("∂g", g_g, tower.g[1])] {
2598                    assert!(
2599                        (fd - ad).abs() <= 1e-5 * ad.abs().max(1.0),
2600                        "row {row} {label}: FD witness {fd:+.6e} != tower grad {ad:+.6e}"
2601                    );
2602                }
2603            }
2604        }
2605    }
2606
2607    /// #932 transcendental de-duplication: the combined
2608    /// [`rigid_standard_normal_third_and_fourth_full`] builder reads BOTH the
2609    /// third and fourth uncontracted tensors off ONE shared
2610    /// `rigid_standard_normal_tower` (one Mills-ratio transcendental per row),
2611    /// and must be BIT-IDENTICAL to the two separate single-tensor builders
2612    /// (`rigid_standard_normal_third_full` + `rigid_standard_normal_fourth_full`,
2613    /// two transcendentals). This pins the exactness of the redundancy
2614    /// elimination: `==`, max diff exactly 0.0 — same tower, no accuracy or
2615    /// generality change, only the redundant second transcendental removed.
2616    #[test]
2617    fn rigid_third_and_fourth_full_shares_one_tower_bit_identical() {
2618        let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2619        let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2620        let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2621        let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2622        let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2623        for &probit_scale in &[1.0_f64, 0.8] {
2624            for r in 0..eta.len() {
2625                let marginal = bernoulli_marginal_link_map(
2626                    &InverseLink::Standard(gam_problem::StandardLink::Probit),
2627                    eta[r],
2628                )
2629                .expect("link map");
2630                let t3_sep = rigid_standard_normal_third_full(
2631                    marginal,
2632                    g[r],
2633                    z[r],
2634                    y[r],
2635                    w[r],
2636                    probit_scale,
2637                )
2638                .expect("separate third");
2639                let t4_sep = rigid_standard_normal_fourth_full(
2640                    marginal,
2641                    g[r],
2642                    z[r],
2643                    y[r],
2644                    w[r],
2645                    probit_scale,
2646                )
2647                .expect("separate fourth");
2648                let (t3_comb, t4_comb) = rigid_standard_normal_third_and_fourth_full(
2649                    marginal,
2650                    g[r],
2651                    z[r],
2652                    y[r],
2653                    w[r],
2654                    probit_scale,
2655                )
2656                .expect("combined third+fourth");
2657                // Exact bitwise equality (same tower) — no tolerance.
2658                for a in 0..2 {
2659                    for b in 0..2 {
2660                        for c in 0..2 {
2661                            assert_eq!(
2662                                t3_comb[a][b][c], t3_sep[a][b][c],
2663                                "t3[{a}][{b}][{c}] row {r} scale {probit_scale} not bit-identical"
2664                            );
2665                            for d in 0..2 {
2666                                assert_eq!(
2667                                    t4_comb[a][b][c][d], t4_sep[a][b][c][d],
2668                                    "t4[{a}][{b}][{c}][{d}] row {r} scale {probit_scale} not bit-identical"
2669                                );
2670                            }
2671                        }
2672                    }
2673                }
2674            }
2675        }
2676    }
2677
2678    /// #932 production wiring: the rigid Bernoulli row, routed through the
2679    /// generic [`RowProgram<2>`] program seam and its cheap
2680    /// order-2 / contracted scalar evaluators (`program_row_kernel`,
2681    /// `program_third_contracted`, `program_fourth_contracted`,
2682    /// `program_full_tower`), must agree BIT-FOR-BIT with an independent generic
2683    /// [`RowProgram`] that uses a different composition order. Both write the
2684    /// same NLL — the contracted scalars fold the direction into differentiation,
2685    /// so this pins that the packed channels equal the corresponding contractions
2686    /// of the independent dense tower truth through a real production consumer.
2687    #[test]
2688    fn rigid_bernoulli_generic_program_matches_independent_program_all_channels() {
2689        use gam_math::jet_tower::{
2690            program_fourth_contracted, program_row_kernel, program_third_contracted,
2691        };
2692
2693        let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2694        let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2695        let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2696        let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2697        let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2698        let n = eta.len();
2699        let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2700
2701        let close = |a: f64, b: f64, label: &str| {
2702            let band = 1e-12 + 1e-12 * a.abs().max(b.abs());
2703            assert!(
2704                (a - b).abs() <= band,
2705                "{label}: generic {a:+.15e} vs Tower4-program {b:+.15e} (band {band:.3e})"
2706            );
2707        };
2708
2709        for &probit_scale in &[1.0_f64, 0.8] {
2710            // The independent generic program over all rows.
2711            let tower_program = BernoulliRigidStandardNormalNllProgram {
2712                primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2713                z: z.to_vec(),
2714                y: y.to_vec(),
2715                w: w.to_vec(),
2716                probit_scale,
2717            };
2718
2719            for row in 0..n {
2720                let truth = program_full_tower(&tower_program, row).expect("program tower");
2721
2722                let marginal = bernoulli_marginal_link_map(
2723                    &InverseLink::Standard(gam_problem::StandardLink::Probit),
2724                    eta[row],
2725                )
2726                .expect("link map");
2727                let program = RigidStandardNormalRow {
2728                    marginal,
2729                    g: g[row],
2730                    z: z[row],
2731                    y: y[row],
2732                    w: w[row],
2733                    probit_scale,
2734                };
2735
2736                // program_full_tower must reproduce the dense tower in EVERY
2737                // channel (v, g, H, t3, t4).
2738                let full = program_full_tower(&program, 0).expect("generic full tower");
2739                close(full.v, truth.v, "full value");
2740                for a in 0..2 {
2741                    close(full.g[a], truth.g[a], "full grad");
2742                    for b in 0..2 {
2743                        close(full.h[a][b], truth.h[a][b], "full hess");
2744                        for c in 0..2 {
2745                            close(full.t3[a][b][c], truth.t3[a][b][c], "full t3");
2746                            for d in 0..2 {
2747                                close(full.t4[a][b][c][d], truth.t4[a][b][c][d], "full t4");
2748                            }
2749                        }
2750                    }
2751                }
2752
2753                // program_row_kernel (Order2) must equal the tower's (v, g, H).
2754                let (val, grad, hess) =
2755                    program_row_kernel(&program, 0).expect("generic row kernel");
2756                close(val, truth.v, "order2 value");
2757                for a in 0..2 {
2758                    close(grad[a], truth.g[a], "order2 grad");
2759                    for b in 0..2 {
2760                        close(hess[a][b], truth.h[a][b], "order2 hess");
2761                    }
2762                }
2763
2764                // program_third_contracted (OneSeed) must equal the dense
2765                // tower's third contraction for each direction.
2766                for dir in &dirs {
2767                    let third = program_third_contracted(&program, 0, dir)
2768                        .expect("generic third contracted");
2769                    let truth3 = truth.third_contracted(dir);
2770                    for a in 0..2 {
2771                        for b in 0..2 {
2772                            close(third[a][b], truth3[a][b], "third contracted");
2773                        }
2774                    }
2775                }
2776
2777                // program_fourth_contracted (TwoSeed) must equal the dense
2778                // tower's fourth contraction for each direction pair.
2779                for (i, u) in dirs.iter().enumerate() {
2780                    let v = dirs[(i + 1) % dirs.len()];
2781                    let fourth = program_fourth_contracted(&program, 0, u, &v)
2782                        .expect("generic fourth contracted");
2783                    let truth4 = truth.fourth_contracted(u, &v);
2784                    for a in 0..2 {
2785                        for b in 0..2 {
2786                            close(fourth[a][b], truth4[a][b], "fourth contracted");
2787                        }
2788                    }
2789                }
2790            }
2791        }
2792    }
2793
2794    /// Original HAND value/gradient/Hessian path for the rigid standard-normal
2795    /// Bernoulli row, reconstructed verbatim from the pre-#932 production code
2796    /// (`RigidProbitKernel::new` + `rigid_transformed_gradient` +
2797    /// `rigid_transformed_hessian`, deleted in ee8a40b2a). This is the path the
2798    /// shipped jet kernel ([`rigid_standard_normal_row_kernel`]) replaced, kept
2799    /// here as an independent perf-and-correctness witness: the jet path must be
2800    /// numerically equal to it (≤1e-9 rel) and at least as fast (see
2801    /// `bench_rigid_vgh_jet_vs_hand`).
2802    fn hand_rigid_vgh(
2803        marginal: BernoulliMarginalLinkMap,
2804        g: f64,
2805        z: f64,
2806        y: f64,
2807        w: f64,
2808        probit_scale: f64,
2809    ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2810        let s = 2.0 * y - 1.0;
2811        let observed_logslope = probit_scale * g;
2812        let g2 = observed_logslope * observed_logslope;
2813        let c = (1.0 + g2).sqrt();
2814        let c1 = probit_scale * observed_logslope / c;
2815        let c_inv3 = 1.0 / (c * c * c);
2816        let c2 = probit_scale * probit_scale * c_inv3;
2817        let q = marginal.q;
2818        // η = q·c(g) + s_f·g·z, m = (2y−1)·η  (marginal_slope_standard_normal_scalar_eta).
2819        let eta = q * c + observed_logslope * z;
2820        let m = s * eta;
2821        let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
2822        // ONE transcendental via the Mills ratio (k1..k4 of the original 4th-order
2823        // kernel; only k1, k2 feed value/grad/Hessian, the rest is the waste the
2824        // jet Order2 path elides).
2825        let (k1, k2, _k3, _k4) =
2826            signed_probit_neglog_derivatives_up_to_fourth(m, w).expect("hand kernel");
2827        let u1 = s * k1;
2828        let u2 = k2;
2829        let eta_q = c;
2830        let eta_g = q * c1 + probit_scale * z;
2831        // value = −w·logΦ(m).
2832        let value = -w * logcdf;
2833        // rigid_transformed_gradient (in (η, g) primaries).
2834        let gradient = [u1 * eta_q * marginal.q1, u1 * eta_g];
2835        // primary_hessian in (q-index, g).
2836        let h00 = u2 * eta_q * eta_q;
2837        let h01 = u2 * eta_q * eta_g + u1 * c1;
2838        let h11 = u2 * eta_g * eta_g + u1 * q * c2;
2839        // rigid_transformed_hessian → (η, g).
2840        let grad_q = u1 * eta_q;
2841        let hessian = [
2842            [
2843                h00 * marginal.q1 * marginal.q1 + grad_q * marginal.q2,
2844                h01 * marginal.q1,
2845            ],
2846            [h01 * marginal.q1, h11],
2847        ];
2848        (value, gradient, hessian)
2849    }
2850
2851    /// The shipped jet value/grad/Hessian kernel must equal the original HAND
2852    /// path it replaced (≤1e-9 rel) on the standard fixture grid — a third,
2853    /// independent #932 single-source witness (the jet composes `q(η)` directly
2854    /// on the η primary; the hand path differentiates in the q-index then chains
2855    /// `q1/q2`, a different FP order, so this is a tolerance not a bit check).
2856    #[test]
2857    fn rigid_bernoulli_row_kernel_matches_hand_chain_witness() {
2858        let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2859        let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2860        let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2861        let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2862        let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2863        let close = |a: f64, b: f64, label: &str| {
2864            let band = 1e-12 + 1e-9 * a.abs().max(b.abs());
2865            assert!(
2866                (a - b).abs() <= band,
2867                "{label}: jet {a:+.15e} vs hand {b:+.15e} (band {band:.3e})"
2868            );
2869        };
2870        for &probit_scale in &[1.0_f64, 0.8] {
2871            for r in 0..eta.len() {
2872                let marginal = bernoulli_marginal_link_map(
2873                    &InverseLink::Standard(gam_problem::StandardLink::Probit),
2874                    eta[r],
2875                )
2876                .expect("link map");
2877                let (jv, jg, jh) = rigid_standard_normal_row_kernel(
2878                    marginal,
2879                    g[r],
2880                    z[r],
2881                    y[r],
2882                    w[r],
2883                    probit_scale,
2884                )
2885                .expect("jet kernel");
2886                let (hv, hg, hh) = hand_rigid_vgh(marginal, g[r], z[r], y[r], w[r], probit_scale);
2887                close(jv, hv, "value");
2888                for a in 0..2 {
2889                    close(jg[a], hg[a], "grad");
2890                    for b in 0..2 {
2891                        close(jh[a][b], hh[a][b], "hess");
2892                    }
2893                }
2894            }
2895        }
2896    }
2897
2898    /// #932 release speed gate for the rigid Bernoulli row: the shipped jet
2899    /// value/grad/Hessian kernel must beat the original hand chain it replaced
2900    /// (reconstructed verbatim above as [`hand_rigid_vgh`]). One measured cell
2901    /// per outcome branch; emits the harness-parsed `hand_over_production`
2902    /// token that the MSI release harness fails closed on whenever any cell is
2903    /// `<= 1`. (This supersedes the removed `#[ignore]`d microbench — release
2904    /// gates are plain non-ignored tests, same as the other `release_measure`
2905    /// cells.)
2906    #[test]
2907    fn release_measure_rigid_bernoulli_vgh_vs_hand_chain_932() {
2908        use std::time::Instant;
2909
2910        // (eta, g, z, y, w): one ordinary interior row per outcome branch —
2911        // y=1 and y=0 are distinct live sign branches of the Mills-ratio
2912        // kernel, so each is its own measured cell.
2913        let cases = [
2914            (0.3_f64, 0.2_f64, 0.4_f64, 1.0_f64, 1.0_f64),
2915            (-0.7, -0.5, -1.1, 0.0, 0.8),
2916        ];
2917        let probit_scale = 0.8;
2918
2919        // Feedback-coupled timing barrier (no `std::hint::black_box`): each
2920        // iteration nudges the log-slope primary by a negligible multiple of
2921        // the running checksum, and the checksum folds value, gradient, and
2922        // Hessian channels, so the pure row call can be neither hoisted nor
2923        // dropped while the measured regime stays bit-adjacent to the fixture.
2924        fn best_ns<F>(iterations: usize, base_g: f64, evaluate: F) -> f64
2925        where
2926            F: Fn(f64) -> (f64, [f64; 2], [[f64; 2]; 2]),
2927        {
2928            let mut best = f64::INFINITY;
2929            for _ in 0..5 {
2930                let mut checksum = 0.0_f64;
2931                let started = Instant::now();
2932                for _ in 0..iterations {
2933                    let (value, gradient, hessian) = evaluate(base_g + checksum * 1e-18);
2934                    checksum += value + gradient[0] + hessian[0][0];
2935                }
2936                assert!(
2937                    checksum.is_finite(),
2938                    "rigid Bernoulli release-measure checksum must stay finite"
2939                );
2940                best = best.min(started.elapsed().as_secs_f64());
2941            }
2942            best * 1e9 / iterations as f64
2943        }
2944
2945        let iterations = 2_000_000usize;
2946        for &(eta, g, z, y, w) in &cases {
2947            let marginal = bernoulli_marginal_link_map(
2948                &InverseLink::Standard(gam_problem::StandardLink::Probit),
2949                eta,
2950            )
2951            .expect("link map");
2952
2953            // Parity pin on the exact benchmarked inputs (the full-grid check
2954            // is `rigid_bernoulli_row_kernel_matches_hand_chain_witness`).
2955            let (jet_value, ..) =
2956                rigid_standard_normal_row_kernel(marginal, g, z, y, w, probit_scale)
2957                    .expect("jet kernel");
2958            let (hand_value, ..) = hand_rigid_vgh(marginal, g, z, y, w, probit_scale);
2959            let band = 1e-12 + 1e-9 * jet_value.abs().max(hand_value.abs());
2960            assert!(
2961                (jet_value - hand_value).abs() <= band,
2962                "y={y:.0} value: jet {jet_value:+.15e} vs hand {hand_value:+.15e}"
2963            );
2964
2965            let production_ns = best_ns(iterations, g, |perturbed_g| {
2966                rigid_standard_normal_row_kernel(marginal, perturbed_g, z, y, w, probit_scale)
2967                    .expect("jet kernel")
2968            });
2969            let hand_ns = best_ns(iterations, g, |perturbed_g| {
2970                hand_rigid_vgh(marginal, perturbed_g, z, y, w, probit_scale)
2971            });
2972            eprintln!(
2973                "RIGID-BERNOULLI-VGH-932 y={y:.0} production={production_ns:.2} ns/row \
2974                 hand={hand_ns:.2} ns/row hand_over_production={:.6}",
2975                hand_ns / production_ns,
2976            );
2977        }
2978    }
2979}
2980
2981#[cfg(test)]
2982mod flex_primary_hessian_oracle_tests {
2983    //! #932 correctness gate for the BMS-FLEX per-row primary Hessian assembled
2984    //! by hand product-rule in
2985    //! [`super::super::row_primary_hessian::BernoulliMarginalSlopeFamily::lower_bms_flex_row_order2_from_parts`]
2986    //! (`f_aa += w·φ·(η_aa − η·η_a·η_a)`, the `f_au`/`f_uv`/`a_uv` chain, and the
2987    //! final `d2_m·η_u·η_v + d1_m·s_y·η_uv` contraction).
2988    //!
2989    //! A prior audit found this hand Hessian had NO INDEPENDENT oracle: the only
2990    //! covering test (`families_bms_joint_hessian_hvp_correction_tests.rs`)
2991    //! asserts batched-vs-nonbatched self-consistency using the SAME hand code on
2992    //! both sides, so a dropped product-rule term would pass undetected. This
2993    //! module closes that gap with a finite-difference witness that NEVER runs the
2994    //! Hessian-assembly branch: it central-differences the flex GRADIENT — which
2995    //! is produced by an entirely separate code path (the `need_hessian = false`
2996    //! value/`eta_u`-scaling lines, none of which read the `f_aa`/`f_au`/`f_uv`
2997    //! product-rule accumulators) — and pins the analytic Hessian against it.
2998    //!
2999    //! The gradient itself is FD-validated transitively: it is the analytic
3000    //! gradient of the same per-row NLL, evaluated at the converged intercept,
3001    //! and the FD perturbation re-solves the intercept root per perturbed point
3002    //! (rebuilding the row context), so the difference quotient is the true
3003    //! mixed/second partial of the row negative log-likelihood — the independent
3004    //! truth the hand Hessian must reproduce.
3005
3006    use super::*;
3007    // `BernoulliMarginalSlopeFamily` (and the flex block-config helpers) live in
3008    // the sibling `super::family` module and are `pub(super)`; this oracle test
3009    // module's `use super::*` does not re-export them, so import the family
3010    // namespace explicitly. Mirrors `cell_moment_assembly.rs`'s
3011    // `use super::family::*`. Without this the flex oracle fixture fails to
3012    // resolve the family type (E0422/E0425/E0433) and blocks the whole lib build.
3013    use super::family::*;
3014    use gam_linalg::matrix::DenseDesignMatrix;
3015    use ndarray::Array1;
3016    use ndarray::Array2;
3017    use std::sync::Arc;
3018    use std::sync::Mutex;
3019
3020    /// Port of the integration-test flex fixture
3021    /// (`make_flex_hvp_cache_test_family`), kept in-crate so the oracle can run
3022    /// without the test crate (the family struct is `pub(super)`). Builds a small
3023    /// flex BMS family with both a score-warp and a link-deviation block so the
3024    /// flex Hessian assembly exercises every primary block (q, logslope, h, w).
3025    fn make_flex_oracle_family(
3026        n: usize,
3027    ) -> (BernoulliMarginalSlopeFamily, Vec<ParameterBlockState>) {
3028        let score_seed = Array1::linspace(-2.0, 2.0, n.max(6));
3029        let link_seed = Array1::linspace(-1.8, 1.8, n.max(6));
3030        let cfg = DeviationBlockConfig {
3031            num_internal_knots: 3,
3032            ..DeviationBlockConfig::default()
3033        };
3034        let score_prepared = build_score_warp_deviation_block_from_seed(&score_seed, &cfg)
3035            .expect("build score warp block");
3036        let link_prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
3037            &link_seed, &link_seed, &cfg,
3038        )
3039        .expect("build link deviation block");
3040
3041        let y: Array1<f64> =
3042            Array1::from_iter((0..n).map(|i| if (i * 17 + 3) % 7 >= 4 { 1.0 } else { 0.0 }));
3043        let weights: Array1<f64> =
3044            Array1::from_iter((0..n).map(|i| 0.75 + ((i * 11 + 5) % 5) as f64 * 0.05));
3045        let z: Array1<f64> =
3046            Array1::from_iter((0..n).map(|i| -1.7 + 3.4 * (i as f64 + 0.5) / n as f64));
3047        let marginal_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3048            if j == 0 {
3049                1.0
3050            } else {
3051                -0.4 + 0.8 * ((i * 19 + 7) % n) as f64 / n as f64
3052            }
3053        });
3054        let logslope_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3055            if j == 0 {
3056                1.0
3057            } else {
3058                0.3 - 0.6 * ((i * 23 + 11) % n) as f64 / n as f64
3059            }
3060        });
3061
3062        let family = BernoulliMarginalSlopeFamily {
3063            y: Arc::new(y),
3064            weights: Arc::new(weights),
3065            z: Arc::new(z.clone()),
3066            latent_measure: LatentMeasureKind::StandardNormal,
3067            gaussian_frailty_sd: Some(0.15),
3068            base_link: InverseLink::Standard(gam_problem::StandardLink::Probit),
3069            marginal_design: DesignMatrix::Dense(DenseDesignMatrix::from(marginal_x.clone())),
3070            logslope_design: DesignMatrix::Dense(DenseDesignMatrix::from(logslope_x.clone())),
3071            score_warp: Some(score_prepared.runtime.clone()),
3072            link_dev: Some(link_prepared.runtime.clone()),
3073            policy: gam_runtime::resource::ResourcePolicy::default_library(),
3074            cell_moment_lru: Arc::new(exact_kernel::CellMomentLruCache::new(1024)),
3075            cell_moment_cache_stats: Arc::new(exact_kernel::CellMomentCacheStats::default()),
3076            intercept_warm_starts: None,
3077            auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
3078            auto_subsample_last_rho: Arc::new(Mutex::new(None)),
3079        };
3080
3081        let beta_m = Array1::from_vec(vec![0.12, -0.04]);
3082        let beta_g = Array1::from_vec(vec![0.35, 0.03]);
3083        let beta_h = Array1::from_iter(
3084            (0..score_prepared.runtime.basis_dim()).map(|idx| 0.0015 * (idx as f64 + 1.0)),
3085        );
3086        let beta_w = Array1::from_iter(
3087            (0..link_prepared.runtime.basis_dim()).map(|idx| -0.001 * (idx as f64 + 1.0)),
3088        );
3089        let states = vec![
3090            ParameterBlockState {
3091                eta: marginal_x.dot(&beta_m),
3092                beta: beta_m,
3093            },
3094            ParameterBlockState {
3095                eta: logslope_x.dot(&beta_g),
3096                beta: beta_g,
3097            },
3098            ParameterBlockState {
3099                beta: beta_h,
3100                eta: Array1::zeros(z.len()),
3101            },
3102            ParameterBlockState {
3103                beta: beta_w,
3104                eta: Array1::zeros(z.len()),
3105            },
3106        ];
3107        (family, states)
3108    }
3109
3110    /// The flex primary gradient at a perturbed primary point. Perturbs primary
3111    /// coordinate `u` by `delta` (mutating the relevant block state — the
3112    /// marginal/logslope row η or a deviation β plus its design contribution
3113    /// where applicable), rebuilds the row context FRESH (re-solving the
3114    /// calibration intercept root at the perturbed point), and returns the
3115    /// analytic gradient. The Hessian-assembly branch is never run, so this is a
3116    /// genuinely independent witness for that branch.
3117    fn flex_gradient_at_perturbed(
3118        family: &BernoulliMarginalSlopeFamily,
3119        states: &[ParameterBlockState],
3120        primary: &super::super::hessian_paths::PrimarySlices,
3121        row: usize,
3122        u: usize,
3123        delta: f64,
3124    ) -> Array1<f64> {
3125        let mut states = states.to_vec();
3126        // Map the primary coordinate `u` onto the parameter that controls it.
3127        // q / logslope live in the per-row η of blocks 0 / 1; the deviation
3128        // bases live in the β of blocks 2 (score-warp) / 3 (link-wiggle), which
3129        // the row context reads via `score_beta` / `link_beta` (their η rows are
3130        // unused on the flex per-row path, so only β need move).
3131        if u == primary.q {
3132            states[0].eta[row] += delta;
3133        } else if u == primary.logslope {
3134            states[1].eta[row] += delta;
3135        } else if let Some(h_range) = primary.h.as_ref()
3136            && h_range.contains(&u)
3137        {
3138            states[2].beta[u - h_range.start] += delta;
3139        } else if let Some(w_range) = primary.w.as_ref()
3140            && w_range.contains(&u)
3141        {
3142            states[3].beta[u - w_range.start] += delta;
3143        } else {
3144            panic!("primary coordinate {u} out of range for flex oracle");
3145        }
3146        let row_ctx = family
3147            .build_row_exact_context_with_stats_and_cell_cache(row, &states, None, false)
3148            .expect("perturbed row context");
3149        let (_neglog, grad, _hess) = family
3150            .compute_row_primary_gradient_hessian(row, &states, primary, &row_ctx)
3151            .expect("perturbed gradient");
3152        grad
3153    }
3154
3155    /// NLL primary gradient after perturbing only the observed generated
3156    /// regressor for one row. The latent-measure calibration root is rebuilt,
3157    /// while every fitted coefficient stays fixed. Central-differencing this
3158    /// and negating gives the independent LOG-LIKELIHOOD score-z derivative
3159    /// used to verify the #2303 analytic channel.
3160    fn flex_nll_gradient_at_perturbed_z(
3161        family: &BernoulliMarginalSlopeFamily,
3162        states: &[ParameterBlockState],
3163        primary: &super::super::hessian_paths::PrimarySlices,
3164        row: usize,
3165        delta: f64,
3166    ) -> Array1<f64> {
3167        let mut perturbed = family.clone();
3168        let mut z = family.z.as_ref().clone();
3169        z[row] += delta;
3170        perturbed.z = Arc::new(z);
3171        let row_ctx = perturbed
3172            .build_row_exact_context_with_stats_and_cell_cache(row, states, None, false)
3173            .expect("z-perturbed row context");
3174        let mut scratch =
3175            super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3176        perturbed
3177            .lower_bms_flex_row_order2(row, states, primary, &row_ctx, None, false, &mut scratch)
3178            .expect("z-perturbed flex gradient");
3179        scratch.grad
3180    }
3181
3182    /// #2303: the Murphy–Topel observed-z channel must cover BOTH deviation
3183    /// blocks, not merely the rigid q/logslope coordinates. Compare every
3184    /// primary coordinate to an independent central difference of the score,
3185    /// then assert the coefficient-space scatter retains nonzero score-warp and
3186    /// link-deviation columns.
3187    #[test]
3188    fn flex_score_zeta_sensitivity_covers_all_active_deviation_blocks_2303() {
3189        let n = 12usize;
3190        let (family, states) = make_flex_oracle_family(n);
3191        let cache = family
3192            .build_exact_eval_cache(&states)
3193            .expect("flex exact eval cache");
3194        let primary = &cache.primary;
3195        let row = 5usize;
3196        let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3197        let mut scratch =
3198            super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3199        family
3200            .lower_bms_flex_row_order2(row, &states, primary, row_ctx, None, false, &mut scratch)
3201            .expect("analytic flex z-sensitivity");
3202        let analytic = scratch.score_zeta.clone();
3203
3204        let h = 1e-5_f64;
3205        let nll_plus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, h);
3206        let nll_minus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, -h);
3207        for u in 0..primary.total {
3208            // score = -NLL gradient.
3209            let finite_difference = -(nll_plus[u] - nll_minus[u]) / (2.0 * h);
3210            let scale = 1.0 + analytic[u].abs().max(finite_difference.abs());
3211            let relative_error = (analytic[u] - finite_difference).abs() / scale;
3212            assert!(
3213                relative_error <= 2e-6,
3214                "flex score-zeta primary {u}: analytic={} FD={} relative_error={relative_error}",
3215                analytic[u],
3216                finite_difference
3217            );
3218        }
3219        let h_range = primary.h.as_ref().expect("score-warp primary range");
3220        let w_range = primary.w.as_ref().expect("link-deviation primary range");
3221        assert!(
3222            analytic
3223                .slice(s![h_range.start..h_range.end])
3224                .iter()
3225                .any(|value| value.abs() > 1e-10),
3226            "score-warp z-sensitivity must not be zero-filled"
3227        );
3228        assert!(
3229            analytic
3230                .slice(s![w_range.start..w_range.end])
3231                .iter()
3232                .any(|value| value.abs() > 1e-10),
3233            "link-deviation z-sensitivity must not be zero-filled"
3234        );
3235
3236        let coefficient = family
3237            .flex_score_zeta_sensitivity(
3238                &states,
3239                &BlockwiseFitOptions::default(),
3240                cache.slices.total,
3241            )
3242            .expect("full coefficient score-zeta sensitivity");
3243        assert_eq!(coefficient.dim(), (n, cache.slices.total));
3244        for range in [
3245            cache.slices.h.as_ref().expect("score-warp beta range"),
3246            cache.slices.w.as_ref().expect("link-deviation beta range"),
3247        ] {
3248            assert!(
3249                coefficient
3250                    .slice(s![.., range.start..range.end])
3251                    .iter()
3252                    .any(|value| value.abs() > 1e-10),
3253                "active deviation coefficient range {range:?} must carry Murphy-Topel sensitivity"
3254            );
3255        }
3256
3257        let wrong_width = cache
3258            .slices
3259            .total
3260            .checked_sub(1)
3261            .expect("nonempty coefficient frame");
3262        let error = family
3263            .flex_score_zeta_sensitivity(&states, &BlockwiseFitOptions::default(), wrong_width)
3264            .expect_err("partial Murphy-Topel covariance frame must be rejected");
3265        assert!(
3266            error.contains("covariance/frame mismatch"),
3267            "unexpected partial-frame error: {error}"
3268        );
3269    }
3270
3271    /// The hand-assembled BMS-FLEX per-row primary Hessian must equal the
3272    /// central finite difference of the flex gradient at every fixture row.
3273    #[test]
3274    fn flex_primary_hessian_matches_central_fd_of_gradient() {
3275        let n = 12usize;
3276        let (family, states) = make_flex_oracle_family(n);
3277        let cache = family
3278            .build_exact_eval_cache(&states)
3279            .expect("flex exact eval cache");
3280        let primary = &cache.primary;
3281        let r = primary.total;
3282        assert!(
3283            r >= 4,
3284            "flex fixture must carry q + logslope + deviation blocks"
3285        );
3286
3287        // Central-difference step. The flex gradient is smooth in every primary
3288        // coordinate; 1e-4 balances truncation (O(h^2)) against the cancellation
3289        // floor of the per-perturbation intercept re-solve (~1e-12).
3290        let h = 1e-4;
3291        let mut max_rel = 0.0_f64;
3292
3293        // A handful of interior rows (avoid the strongest-tail endpoints where
3294        // the FD floor is loosest). Every primary coordinate is differenced.
3295        for &row in &[2usize, 5, 8] {
3296            let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3297            let (_neglog, _grad, analytic_hess) = family
3298                .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3299                .expect("analytic flex gradient + hessian");
3300
3301            for u in 0..r {
3302                let grad_plus = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3303                let grad_minus = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3304                for v in 0..r {
3305                    let fd = (grad_plus[v] - grad_minus[v]) / (2.0 * h);
3306                    let analytic = analytic_hess[[v, u]];
3307                    let denom = 1.0 + analytic.abs().max(fd.abs());
3308                    let rel = (analytic - fd).abs() / denom;
3309                    max_rel = max_rel.max(rel);
3310                    assert!(
3311                        rel <= 1e-6,
3312                        "flex hand Hessian H[{v}][{u}] = {analytic:.6e} disagrees with central \
3313                         FD of the gradient {fd:.6e} at row {row} (rel {rel:.3e}); a product-rule \
3314                         term is dropped or mis-signed"
3315                    );
3316                }
3317            }
3318        }
3319        // Surface the achieved tightness for the record.
3320        assert!(
3321            max_rel <= 1e-6,
3322            "flex Hessian FD oracle max rel {max_rel:.3e}"
3323        );
3324    }
3325
3326    /// ARBITER (diagnostic): is the H[0][0] flex-Hessian vs FD-of-gradient gap a
3327    /// REAL hand-derivation bug or just FD-truncation / intercept-re-solve noise
3328    /// in the witness? Sweep the central-difference step `h` on the worst entry
3329    /// (row 2, [q][q]); if the gap scales ~h^2 it is FD truncation (the analytic
3330    /// Hessian is right, the witness bound is just too tight); if it stays flat
3331    /// as h shrinks it is a genuine dropped/mis-signed term. Richardson-cancel
3332    /// the O(h^2) term and report the residual. Panics with the table so the
3333    /// harness surfaces the numbers (stdout is otherwise suppressed).
3334    #[test]
3335    fn arbiter_flex_hessian_h00_fd_step_scaling() {
3336        let n = 12usize;
3337        let (family, states) = make_flex_oracle_family(n);
3338        let cache = family
3339            .build_exact_eval_cache(&states)
3340            .expect("flex exact eval cache");
3341        let primary = &cache.primary;
3342        let row = 2usize;
3343        let u = primary.q; // intercept / q axis => H[0][0]
3344        let v = primary.q;
3345
3346        let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3347        let (_neglog, _grad, analytic_hess) = family
3348            .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3349            .expect("analytic flex gradient + hessian");
3350        let analytic = analytic_hess[[v, u]];
3351
3352        let fd_at = |h: f64| -> f64 {
3353            let gp = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3354            let gm = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3355            (gp[v] - gm[v]) / (2.0 * h)
3356        };
3357
3358        // Coarse and fine central-difference steps. If the analytic Hessian is
3359        // CORRECT and the witness gap is pure O(h^2) FD truncation, halving h
3360        // quarters the gap; the Richardson combination cancels that O(h^2) term
3361        // and lands on the analytic value to the intercept-re-solve floor
3362        // (~1e-9). If instead a hand product-rule term is dropped, the gap is
3363        // h-INDEPENDENT and the Richardson residual stays at the bug magnitude.
3364        let h = 1e-3_f64;
3365        let fd_h = fd_at(h);
3366        let fd_half = fd_at(h * 0.5);
3367        let fd_quarter = fd_at(h * 0.25);
3368        let gap_h = (analytic - fd_h).abs();
3369        let gap_half = (analytic - fd_half).abs();
3370        let gap_quarter = (analytic - fd_quarter).abs();
3371        let rich = (4.0 * fd_half - fd_h) / 3.0;
3372        let rich_gap = (analytic - rich).abs();
3373        let denom = analytic.abs().max(1.0);
3374
3375        // DIAGNOSTIC RECORD (shown on failure; this is the dispositive table):
3376        let record = format!(
3377            "FLEX H[0][0] ARBITER row 2: analytic={analytic:+.12e} \
3378             fd(h)={fd_h:+.12e} fd(h/2)={fd_half:+.12e} fd(h/4)={fd_quarter:+.12e} \
3379             gap(h)={gap_h:.3e} gap(h/2)={gap_half:.3e} gap(h/4)={gap_quarter:.3e} \
3380             ratio_h_over_half={:.3} ratio_half_over_quarter={:.3} \
3381             richardson={rich:+.12e} richardson_gap={rich_gap:.3e} (rich_rel={:.3e})",
3382            gap_h / gap_half.max(f64::MIN_POSITIVE),
3383            gap_half / gap_quarter.max(f64::MIN_POSITIVE),
3384            rich_gap / denom,
3385        );
3386
3387        // VERDICT: the analytic Hessian is correct iff the FD gap is O(h^2) — i.e.
3388        // the Richardson-extrapolated second derivative (truncation-cancelled)
3389        // matches it to the intercept-solve floor. A genuine dropped term leaves
3390        // a Richardson residual at the bug scale (~1e-5), failing this with the
3391        // record above so the harness surfaces the numbers.
3392        assert!(
3393            rich_gap / denom <= 1e-7,
3394            "{record}\nVERDICT: Richardson residual exceeds the FD-truncation floor — \
3395             the hand H[0][0] genuinely diverges (real dropped/mis-signed term), NOT FD noise"
3396        );
3397    }
3398}