Skip to main content

gam_models/bms/
gradient_paths.rs

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