Skip to main content

gam_models/bms/
block_specs.rs

1use super::family::*;
2use super::gradient_paths::*;
3use super::hessian_paths::{new_cell_moment_cache_stats, new_cell_moment_lru_cache};
4use super::install_flex::validate_spec;
5use super::*;
6use crate::marginal_slope_orthogonal::influence_absorber_log_lambda;
7use faer::Side;
8use gam_linalg::faer_ndarray::{FaerEigh, fast_ab, fast_atb, fast_xt_diag_x};
9
10/// Sup-norm of the FITTED marginal linear predictor `η = X·β` at which the
11/// probit model becomes numerically degenerate. This is the decisive
12/// separation quantity — raw coefficient magnitude is NOT, because the
13/// thin-plate / Duchon marginal bases are non-orthonormal and ill-conditioned,
14/// so a smooth, bounded fitted surface can carry large-and-cancelling
15/// coefficients (e.g. `β=[60,-60]` on two collinear columns yields a bounded
16/// `Xβ`). The probit information `φ(η)²/[Φ(η)(1−Φ(η))]` only collapses once
17/// `|η|` is enormous: `Φ(−35) ≈ 1e−268` is still representable, but by `|η|≈38`
18/// the tail probability underflows to exactly 0 in `f64` and the per-row
19/// Fisher weight vanishes. We trip the guard at 35 — comfortably inside the
20/// representable range yet far beyond any legitimately fitted predictor (a
21/// converged penalized probit surface keeps `|η|` at single/low-double digits),
22/// so a true separating direction (whose `|η|→∞`) still trips it while a
23/// well-fitted ill-conditioned surface does not.
24pub(crate) const BMS_PROBIT_SEPARATION_ETA_INF: f64 = 35.0;
25
26// ── Canonical-gauge priority ladder (issue #322) ─────────────────────────────
27//
28// The priority-ordered RRQR in `canonicalize_for_identifiability` presents
29// higher-priority blocks first and routes any shared cross-block alias drop
30// into the lowest-priority block that still spans the aliased direction. The
31// values below form a single ordered ladder so the relationships that the
32// architecture depends on (anchors > parametric surfaces > flex deviations,
33// marginal > logslope, score_warp > link_dev) are expressed once here rather
34// than re-derived from comments at each `gauge_priority:` site. The ladder
35// mirrors the survival marginal-slope entry (time=200 / marginal=150 /
36// logslope=120 / score_warp=80 / link_dev=60).
37
38/// Audit-only anchor blocks sit at the top of the ladder so the candidate
39/// flex block always yields to them in the cross-block identifiability audit.
40pub(super) const GAUGE_PRIORITY_ANCHOR: u8 = 200;
41/// Marginal surface: strictly above the logslope surface so a shared affine
42/// direction is demoted out of logslope, never out of marginal.
43pub(super) const GAUGE_PRIORITY_MARGINAL: u8 = 150;
44/// Logslope surface: one rung below the marginal surface.
45pub(super) const GAUGE_PRIORITY_LOGSLOPE: u8 = 120;
46/// Candidate flex block under audit: below every parametric anchor so the
47/// audit demotes the candidate when it aliases an anchor.
48pub(super) const GAUGE_PRIORITY_CANDIDATE_FLEX: u8 = 100;
49/// `score_warp_dev`: above `link_dev` because in mixed-flex configurations
50/// link_dev is the residualised block and should yield first.
51pub(super) const GAUGE_PRIORITY_SCORE_WARP_DEV: u8 = 80;
52/// Default for any deviation auxiliary block not otherwise named; below the
53/// parametric default so shared affine directions never demote a parametric
54/// block.
55pub(super) const GAUGE_PRIORITY_DEVIATION_DEFAULT: u8 = 70;
56/// `link_dev`: lowest rung, yields first among the flex deviation blocks.
57pub(super) const GAUGE_PRIORITY_LINK_DEV: u8 = 60;
58
59/// Floor on the relative outer tolerance used by the exact-joint spatial
60/// length-scale optimiser. The user's `rel_tol` drives most of the fit, but
61/// the exact spatial outer loop is a coarse 1-D search over a smooth profiled
62/// objective; tightening it below this floor only burns cycles on noise in the
63/// inner-solve-reported objective without moving the selected length scale.
64pub(crate) const EXACT_SPATIAL_OUTER_TOL_FLOOR: f64 = 1e-6;
65
66// ── BlockEffectiveJacobian impls for BMS ─────────────────────────────────────
67//
68// BMS has a single Bernoulli output per row (n_outputs = 1). The observed η is
69//
70//   η_i = q_i · c_i + s·g_i · z_i
71//
72// where
73//   q_i   = marginal_design[i,:] · β_m + offset_m[i]      (marginal η)
74//   g_i   = logslope_design[i,:] · β_s + offset_s[i]      (log-slope η)
75//   s     = probit_frailty_scale(gaussian_frailty_sd)
76//   c_i   = sqrt(1 + (s·g_i)²)
77//
78// Per-block Jacobians ∂η_i / ∂β_block:
79//
80//   Marginal block  → ∂η_i/∂β_m = c_i · M_i
81//     (M_i = marginal_design row i; c_i is β-dependent but does not involve β_m)
82//
83//   Logslope block  → ∂η_i/∂β_s = (q_i · s²·g_i / c_i + s·z_i) · G_i
84//     (G_i = logslope_design row i)
85//
86// score_warp_dev and link_dev blocks use IFT-corrected η, but their
87// contribution to the identifiability audit is captured by the raw design
88// columns (the IFT correction adds a direction already in the anchor span at
89// compile time). These blocks leave jacobian_callback = None and rely on
90// effective_design (= raw design) for the flat audit.
91
92/// β-dependent Jacobian for the BMS marginal block.
93///
94/// ∂η_i/∂β_m = c_i · M[i,:]
95/// where c_i = sqrt(1 + (s · g_i)²),
96///       g_i = G[i,:] · β_s + offset_s[i],
97///       s   = state.probit_frailty_scale.
98///
99/// `probit_frailty_scale` is read from the evaluation state at call time (not
100/// captured at construction) so the callback remains correct across outer-loop
101/// σ updates without rebuilding the block spec.
102///
103/// Designs are pre-densified at construction to avoid repeated materialisation.
104pub struct BmsMarginalJacobian {
105    /// Dense marginal design: n × p_m.
106    pub marginal_dense: Arc<Array2<f64>>,
107    /// Dense logslope design: n × p_s.
108    pub logslope_dense: Arc<Array2<f64>>,
109    pub offset_m: Array1<f64>,
110    pub offset_s: Array1<f64>,
111    /// Number of marginal columns (= size of β_m slice in the full β vector).
112    pub p_marginal: usize,
113}
114
115impl BmsMarginalJacobian {
116    pub fn new(
117        marginal_dense: Arc<Array2<f64>>,
118        logslope_dense: Arc<Array2<f64>>,
119        offset_m: Array1<f64>,
120        offset_s: Array1<f64>,
121        p_marginal: usize,
122    ) -> Self {
123        Self {
124            marginal_dense,
125            logslope_dense,
126            offset_m,
127            offset_s,
128            p_marginal,
129        }
130    }
131}
132
133impl BlockEffectiveJacobian for BmsMarginalJacobian {
134    fn effective_jacobian_rows(
135        &self,
136        state: &FamilyLinearizationState<'_>,
137        rows: std::ops::Range<usize>,
138    ) -> Result<Array2<f64>, String> {
139        let beta = state.beta;
140        let s = state.probit_frailty_scale;
141        let p_m = self.p_marginal;
142        let p_s_block = self.logslope_dense.ncols();
143        let beta_s_raw = if beta.len() > p_m {
144            &beta[p_m..]
145        } else {
146            &[][..]
147        };
148        let p_s_use = p_s_block.min(beta_s_raw.len());
149        let beta_s = &beta_s_raw[..p_s_use];
150        let n = self.marginal_dense.nrows();
151        let rows = rows.start.min(n)..rows.end.min(n);
152        let p_block = self.marginal_dense.ncols();
153
154        // ∂η_i/∂β_m = c_i · M[i,:], with c_i = sqrt(1 + (s·g_i)²) and
155        //   g_i = G[i, :p_s_use] · β_s + offset_s[i].
156        //
157        // This block owns the logslope design and offset, so g_i — and hence
158        // c_i — is fully self-computable at the current β for every row,
159        // including β = 0 (where g_i = offset_s[i], which carries the fitted
160        // logslope baseline and is generically nonzero).  There is no external
161        // scalar this block cannot reconstruct, so the Jacobian is evaluated
162        // directly from owned data with no caller-supplied contract.
163        let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_block));
164        for i in rows.clone() {
165            let g_i = self.offset_s[i]
166                + self
167                    .logslope_dense
168                    .row(i)
169                    .slice(ndarray::s![..p_s_use])
170                    .dot(&ArrayView1::from(beta_s));
171            let sg = s * g_i;
172            let c_i = (1.0 + sg * sg).sqrt();
173            // J[i,:] = c_i · M[i,:]
174            let m_row = self.marginal_dense.row(i);
175            out.row_mut(i - rows.start).assign(&m_row.mapv(|x| c_i * x));
176        }
177        Ok(out)
178    }
179
180    fn n_outputs(&self) -> usize {
181        1
182    }
183
184    fn locks_raw_width_reduction(&self) -> bool {
185        // The BMS family reads its raw-width internal `marginal_design` and
186        // `validate_exact_block_state_shapes` asserts `beta.len() ==
187        // marginal_design.ncols()`. The block's spec carries the real (non
188        // zero-placeholder) marginal design, so the `#933` gauge-composed
189        // reduction cannot silently absorb a dropped column into the callback —
190        // a reduced β would desynchronise the raw-width layout the family reads.
191        // Keep it at raw width and let the robust Jeffreys curvature regularise
192        // any weak cross-block direction (mirrors the survival marginal-slope
193        // time-wiggle block).
194        true
195    }
196}
197
198/// β-dependent Jacobian for the BMS logslope block.
199///
200/// ∂η_i/∂β_s = (q_i · s²·g_i / c_i + s·z_i) · G[i,:]
201/// where q_i = M[i,:] · β_m + offset_m[i],
202///       g_i = G[i,:] · β_s + offset_s[i],
203///       c_i = sqrt(1 + (s·g_i)²),
204///       s   = state.probit_frailty_scale.
205///
206/// `probit_frailty_scale` is read from the evaluation state at call time.
207///
208/// Designs are pre-densified at construction to avoid repeated materialisation.
209pub struct BmsLogslopeJacobian {
210    /// Dense marginal design: n × p_m.
211    pub marginal_dense: Arc<Array2<f64>>,
212    /// Dense logslope design: n × p_s.
213    pub logslope_dense: Arc<Array2<f64>>,
214    pub offset_m: Array1<f64>,
215    pub offset_s: Array1<f64>,
216    pub z: Arc<Array1<f64>>,
217    /// Number of marginal columns (= start of β_s in the full β vector).
218    pub p_marginal: usize,
219}
220
221impl BmsLogslopeJacobian {
222    pub fn new(
223        marginal_dense: Arc<Array2<f64>>,
224        logslope_dense: Arc<Array2<f64>>,
225        offset_m: Array1<f64>,
226        offset_s: Array1<f64>,
227        z: Arc<Array1<f64>>,
228        p_marginal: usize,
229    ) -> Self {
230        Self {
231            marginal_dense,
232            logslope_dense,
233            offset_m,
234            offset_s,
235            z,
236            p_marginal,
237        }
238    }
239}
240
241impl BlockEffectiveJacobian for BmsLogslopeJacobian {
242    fn effective_jacobian_rows(
243        &self,
244        state: &FamilyLinearizationState<'_>,
245        rows: std::ops::Range<usize>,
246    ) -> Result<Array2<f64>, String> {
247        let beta = state.beta;
248        let s = state.probit_frailty_scale;
249        let p_m = self.p_marginal;
250        let p_m_use = p_m.min(beta.len());
251        let beta_m = &beta[..p_m_use];
252        let beta_s_raw = if beta.len() > p_m {
253            &beta[p_m..]
254        } else {
255            &[][..]
256        };
257        let p_s_block = self.logslope_dense.ncols();
258        let p_s_use = p_s_block.min(beta_s_raw.len());
259        let beta_s = &beta_s_raw[..p_s_use];
260        let n = self.logslope_dense.nrows();
261        let rows = rows.start.min(n)..rows.end.min(n);
262
263        // ∂η_i/∂β_s = (q_i · s²·g_i / c_i + s·z_i) · G[i,:] where
264        //   q_i = M[i,:] · β_m + offset_m[i],
265        //   g_i = G[i,:] · β_s + offset_s[i],
266        //   c_i = sqrt(1 + (s·g_i)²),
267        //   z_i = self.z[i].
268        //
269        // This block owns the marginal design, logslope design, both offsets,
270        // and z, so q_i, g_i, c_i, z_i are all closed-form functions of the
271        // current β and owned data — for every row at every β, including β = 0
272        // (where g_i = offset_s[i] carries the nonzero fitted logslope
273        // baseline).  The Jacobian is therefore evaluated directly from owned
274        // data with no caller-supplied scalar contract.
275        let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_s_block));
276        for i in rows.clone() {
277            let q_i = self.offset_m[i]
278                + self
279                    .marginal_dense
280                    .row(i)
281                    .slice(ndarray::s![..p_m_use])
282                    .dot(&ArrayView1::from(beta_m));
283            let g_i = self.offset_s[i]
284                + self
285                    .logslope_dense
286                    .row(i)
287                    .slice(ndarray::s![..p_s_use])
288                    .dot(&ArrayView1::from(beta_s));
289            let sg = s * g_i;
290            let c_i = (1.0 + sg * sg).sqrt();
291            let z_i = self.z[i];
292            // per-row scalar factor: q_i · s²·g_i / c_i + s·z_i
293            let factor = q_i * s * s * g_i / c_i + s * z_i;
294            // J[i,:] = factor · G[i,:]
295            let g_row = self.logslope_dense.row(i);
296            out.row_mut(i - rows.start)
297                .assign(&g_row.mapv(|x| factor * x));
298        }
299        Ok(out)
300    }
301
302    fn n_outputs(&self) -> usize {
303        1
304    }
305
306    fn locks_raw_width_reduction(&self) -> bool {
307        // The BMS family reads its raw-width internal `logslope_design` and
308        // `validate_exact_block_state_shapes` asserts `beta.len() ==
309        // logslope_design.ncols()`. An intercept-only logslope (`logslope_formula
310        // = "1"`) is a constant column perfectly aliased with the marginal
311        // intercept; the effective reduced-reparam leaves a width-1 block at raw
312        // width (a zero-width reduction would delete the score-effect surface), so
313        // the canonicaliser would otherwise column-reduce this block to width 0
314        // and hand the family a length-0 β against a width-1 design. Lock it to
315        // raw width — the robust Jeffreys curvature regularises the weak aliased
316        // direction (mirrors the survival marginal-slope time-wiggle block).
317        true
318    }
319}
320
321/// Horizontally stack the absorbed influence columns `Z̃_infl` onto the raw
322/// marginal design `M`, yielding the widened additive marginal-index design
323/// `[M | Z̃_infl]` (#461). When `influence_columns` is `None` the original
324/// dense design is returned unchanged. The influence columns shift the marginal
325/// index `α(x)` additively, so the de-nested probit kernel — which reads the
326/// marginal index from `block_states[0].eta` and reconstructs `∂q/∂β_m` from
327/// `self.marginal_design` (a matched (design, β) pair) — picks them up with no
328/// kernel-site change; the widened `p_marginal` keeps every per-row Jacobian /
329/// gradient / Hessian projection consistent.
330pub(crate) fn widen_marginal_dense_with_influence(
331    marginal_dense: &Arc<Array2<f64>>,
332    influence_columns: Option<&Array2<f64>>,
333) -> Result<Arc<Array2<f64>>, String> {
334    let Some(z_infl) = influence_columns else {
335        return Ok(Arc::clone(marginal_dense));
336    };
337    let n = marginal_dense.nrows();
338    if z_infl.nrows() != n {
339        return Err(format!(
340            "influence block: residualised columns have {} rows, marginal design has {n}",
341            z_infl.nrows()
342        ));
343    }
344    let p_m = marginal_dense.ncols();
345    let p1 = z_infl.ncols();
346    let mut widened = Array2::<f64>::zeros((n, p_m + p1));
347    widened
348        .slice_mut(s![.., ..p_m])
349        .assign(marginal_dense.as_ref());
350    widened.slice_mut(s![.., p_m..]).assign(z_infl);
351    Ok(Arc::new(widened))
352}
353
354/// Tolerance (relative to the dominant retained eigenvalue) below which a
355/// reduced-basis direction of the W-orthogonalised effective logslope Gram is
356/// treated as a confounded null direction and dropped. Directions whose
357/// effective weighted image is (near-)explained by the marginal span collapse to
358/// ~0 eigenvalue in `Gtt` (see [`build_reduced_logslope_reparam`]); this keeps
359/// the cut well above floating-point noise but well below any genuine surviving
360/// logslope curvature.
361pub(crate) const LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL: f64 = 1.0e-6;
362
363/// An exact reduced-basis reparameterization of the BMS logslope design through
364/// the family's OWN internal `logslope_design` geometry, expressed as a single
365/// linear map `T` (`p_logslope × r`, `r ≤ p_logslope`).
366///
367/// # Why a reduced basis (not a dense design swap)
368///
369/// The structural confound is that the score-weighted logslope channel
370/// `diag(factor)·G·β_s` overlaps the effective marginal channel `diag(c)·M·β_m`
371/// in the PIRLS row metric `W`, leaving the joint penalised Hessian rank-soft
372/// along the shared direction. The shared-solver primitive
373/// [`OrthogonalReparam`](gam_solve::orthogonal_reparam::OrthogonalReparam)
374/// forms `C̃ = C − M·B`, exactly W-orthogonal to `span(M)` — but `C̃` is a dense
375/// design the BMS family's row kernel does NOT consume: the family reads
376/// `η_logslope = G·β_s` from its own `logslope_design` and reconstructs the
377/// per-row Jacobian `factor_i · G_i` from that same matrix. A block-level design
378/// swap is therefore ignored by the family, and feeding a rank-deficient `C̃` at
379/// full width desynchronises the inner identifiable-subspace reduction from the
380/// stored design width.
381///
382/// This builds instead a TRUE reparameterization the family consumes: a
383/// full-rank reduced logslope design `G_reduced = G·T` (width `r`) plus the
384/// penalty projection `S_reduced = Tᵀ S T`. The map is constructed so that the
385/// directions of raw logslope coefficient space whose effective weighted image
386/// is W-explained by the marginal span are removed (they carry ~zero curvature
387/// in the W-orthogonalised effective Gram), and the surviving `r` directions are
388/// full-rank.
389///
390/// # The math
391///
392/// At the rigid pilot, the effective Jacobians are
393///
394/// ```text
395///     M_eff = diag(c) · M        (n × p_m),   c_i = sqrt(1 + (s·g_i)²)
396///     G_eff = diag(f) · G        (n × p_g),   f_i = q_i·s²·g_i/c_i + s·z_i
397/// ```
398///
399/// In the row metric `W` the component of the effective logslope design that is
400/// W-orthogonal to `span(M_eff)` has the raw-coordinate Gram
401///
402/// ```text
403///     Gtt = G_effᵀ W G_eff − (G_effᵀ W M_eff)(M_effᵀ W M_eff + εI)⁻¹(M_effᵀ W G_eff)
404/// ```
405///
406/// (a `p_g × p_g` PSD matrix in the raw logslope coefficient coordinates). Its
407/// range = the logslope directions that survive the confound removal; its null
408/// space = the confounded directions absorbed by the marginal span. The reduced
409/// transform `T` is the orthonormal eigenbasis of `Gtt` for eigenvalues above a
410/// relative tolerance; `r = rank(Gtt)`. The new design `G_reduced = G·T`, the
411/// reparameterized penalty `S_reduced = Tᵀ S T`, and the round-trip
412/// `β_logslope = T·β'` make the family's geometry consistent at width `r` and
413/// recover the original-basis logslope coefficients for prediction/reporting.
414#[derive(Debug, Clone)]
415pub(super) struct ReducedLogslopeReparam {
416    /// Reduced transform `T` (`p_logslope × r`). `G_reduced = G·T`,
417    /// `β_logslope = T·β'`, `S_reduced = Tᵀ S T`.
418    transform: Array2<f64>,
419}
420
421impl ReducedLogslopeReparam {
422    /// Original (full) logslope width `p_logslope`.
423    #[inline]
424    pub(super) fn original_cols(&self) -> usize {
425        self.transform.nrows()
426    }
427
428    /// Reduced width `r`.
429    #[inline]
430    pub(super) fn reduced_cols(&self) -> usize {
431        self.transform.ncols()
432    }
433
434    /// Map a reduced-basis logslope coefficient `β'` (length `r`) back to the
435    /// original logslope basis `β_logslope = T·β'` (length `p_logslope`), so
436    /// prediction/reporting are unchanged-in-meaning.
437    pub(super) fn recover_original_logslope_beta(
438        &self,
439        beta_reduced: &Array1<f64>,
440    ) -> Result<Array1<f64>, String> {
441        if beta_reduced.len() != self.reduced_cols() {
442            return Err(format!(
443                "reduced logslope reparam: β' length ({}) != reduced width ({})",
444                beta_reduced.len(),
445                self.reduced_cols()
446            ));
447        }
448        Ok(self.transform.dot(beta_reduced))
449    }
450}
451
452/// Build the reduced-basis logslope reparameterization (see
453/// [`ReducedLogslopeReparam`]) from the rigid-pilot EFFECTIVE Jacobian geometry,
454/// in the PIRLS row metric `W`. Extracts the dense pilot designs and delegates
455/// the geometry to [`reduced_logslope_transform_effective`]. Returns `Ok(None)`
456/// when there is no logslope/marginal span or no effective-confounded direction
457/// to remove (`r == p_g`); the caller then keeps the raw design. A fully
458/// confounded block (`r == 0`: the entire effective logslope image is
459/// W-explained by the effective marginal span) is an error: the data identify
460/// only the sum of the marginal and score-slope surfaces, so any fitted
461/// decomposition would be an arbitrary penalty artifact.
462fn build_reduced_logslope_reparam(
463    marginal_design: &TermCollectionDesign,
464    logslope_design: &TermCollectionDesign,
465    z: &Array1<f64>,
466    row_metric: &Array1<f64>,
467    marginal_offset: &Array1<f64>,
468    logslope_offset: &Array1<f64>,
469    marginal_baseline: f64,
470    logslope_baseline: f64,
471    probit_scale: f64,
472) -> Result<Option<ReducedLogslopeReparam>, String> {
473    let marginal = marginal_design
474        .design
475        .try_to_dense_arc("build_reduced_logslope_reparam::marginal")?;
476    let logslope = logslope_design
477        .design
478        .try_to_dense_arc("build_reduced_logslope_reparam::logslope")?;
479    let n = marginal.nrows();
480    if logslope.nrows() != n
481        || z.len() != n
482        || row_metric.len() != n
483        || marginal_offset.len() != n
484        || logslope_offset.len() != n
485    {
486        return Err(format!(
487            "reduced logslope reparam row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
488            marginal.nrows(),
489            logslope.nrows(),
490            z.len(),
491            row_metric.len(),
492            marginal_offset.len(),
493            logslope_offset.len(),
494        ));
495    }
496    let p_m = marginal.ncols();
497    let p_g = logslope.ncols();
498    if p_m == 0 || p_g == 0 {
499        return Ok(None);
500    }
501    if !marginal_baseline.is_finite()
502        || !logslope_baseline.is_finite()
503        || !probit_scale.is_finite()
504        || probit_scale <= 0.0
505        || z.iter().any(|v| !v.is_finite())
506        || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
507        || marginal_offset.iter().any(|v| !v.is_finite())
508        || logslope_offset.iter().any(|v| !v.is_finite())
509    {
510        return Err(
511            "reduced logslope reparam requires finite pilot geometry and finite non-negative row metric"
512                .to_string(),
513        );
514    }
515
516    // The joint Hessian the inner solve factorises is built from the EFFECTIVE
517    // BMS Jacobians, not the raw design. At the rigid pilot,
518    //   ∂η_i/∂β_m = c_i · M_i,   c_i = sqrt(1 + (s·g_i)²)
519    //   ∂η_i/∂β_s = f_i · G_i,   f_i = q_i·s²·g_i/c_i + s·z_i
520    // so a raw logslope direction `v` is rank-soft in the joint Hessian iff its
521    // EFFECTIVE image `diag(f)·G·v` is W-explained by `span(diag(c)·M)` — NOT iff
522    // raw `G·v` is W-explained by `span(M)`. Auditing the raw design removes the
523    // wrong directions; the reduced basis is built from the effective Schur Gram.
524    // The pure-array geometry lives in `reduced_logslope_transform_effective` so
525    // it can be unit-tested directly against the raw-vs-effective counterexample.
526    match reduced_logslope_transform_effective(
527        marginal.view(),
528        logslope.view(),
529        z,
530        row_metric,
531        marginal_offset,
532        logslope_offset,
533        marginal_baseline,
534        logslope_baseline,
535        probit_scale,
536    )? {
537        ReducedLogslopeOutcome::Reduced(transform) => {
538            Ok(Some(ReducedLogslopeReparam { transform }))
539        }
540        ReducedLogslopeOutcome::FullRank => Ok(None),
541        ReducedLogslopeOutcome::FullyConfounded => Err(
542            "BMS score-slope block is fully confounded with the marginal index: every \
543             effective logslope direction diag(f)·G·v is W-explained by the effective \
544             marginal span at the rigid pilot, so the data identify only the sum of the \
545             marginal and score-slope surfaces and the smoothing penalty would select an \
546             arbitrary decomposition between them. Refusing to fit; remove the score-slope \
547             terms or supply covariates that separate them from the marginal index."
548                .to_string(),
549        ),
550    }
551}
552
553/// Distinct outcomes of the effective logslope confound audit. `FullRank`
554/// (nothing to reduce, `r == p_g`) and `FullyConfounded` (`r == 0`) must not
555/// share a signal: keeping the raw design is correct for the former, while for
556/// the latter the data identify only a combination of the marginal and
557/// score-slope surfaces, and silently retaining the raw columns would let the
558/// penalty pick an arbitrary decomposition and report it as an estimate.
559#[derive(Debug)]
560pub(crate) enum ReducedLogslopeOutcome {
561    /// Every effective logslope direction carries its own curvature; keep the
562    /// raw design.
563    FullRank,
564    /// `0 < r < p_g`: the reduced basis `T` (`p_g × r`) spanning the
565    /// identifiable directions.
566    Reduced(Array2<f64>),
567    /// `r == 0`: the entire effective logslope image is W-explained by the
568    /// effective marginal span — the block is unidentified.
569    FullyConfounded,
570}
571
572/// Build the reduced logslope basis `T` (p_g × r) from the EFFECTIVE BMS pilot
573/// geometry, in the PIRLS row metric `W`. `T`'s columns span the raw logslope
574/// coefficient directions whose effective image `diag(f)·G·v` is NOT W-explained
575/// by `span(diag(c)·M)` — i.e. the directions the joint Hessian retains real
576/// curvature along. Returns [`ReducedLogslopeOutcome::FullRank`] when there is
577/// nothing to reduce (`r == p_g`, raw design kept) and
578/// [`ReducedLogslopeOutcome::FullyConfounded`] when the entire effective
579/// logslope image collapses into the effective marginal span (`r == 0`), so the
580/// caller can refuse the unidentified block instead of conflating the two
581/// cases.
582///
583/// At the rigid pilot the effective Jacobians are
584///     M_eff = diag(c) · M,   c_i = sqrt(1 + (s·g_i)²)
585///     G_eff = diag(f) · G,   f_i = q_i·s²·g_i/c_i + s·z_i
586/// and the raw-coordinate Gram of the logslope component W-orthogonal to
587/// `span(M_eff)` is the Schur complement
588///     Gtt = G_effᵀ W G_eff − (G_effᵀ W M_eff)(M_effᵀ W M_eff + εI)⁻¹(M_effᵀ W G_eff).
589/// `T` is the orthonormal eigenbasis of `Gtt` for eigenvalues above a tolerance
590/// relative to the effective logslope energy scale.
591pub(crate) fn reduced_logslope_transform_effective(
592    marginal: ArrayView2<'_, f64>,
593    logslope: ArrayView2<'_, f64>,
594    z: &Array1<f64>,
595    row_metric: &Array1<f64>,
596    marginal_offset: &Array1<f64>,
597    logslope_offset: &Array1<f64>,
598    marginal_baseline: f64,
599    logslope_baseline: f64,
600    probit_scale: f64,
601) -> Result<ReducedLogslopeOutcome, String> {
602    let n = marginal.nrows();
603    let p_m = marginal.ncols();
604    let p_g = logslope.ncols();
605    if p_m == 0 || p_g == 0 {
606        return Ok(ReducedLogslopeOutcome::FullRank);
607    }
608
609    // Effective pilot Jacobians M_eff = diag(c)·M and G_eff = diag(f)·G.
610    let mut m_eff = Array2::<f64>::zeros((n, p_m));
611    let mut g_eff = Array2::<f64>::zeros((n, p_g));
612    for i in 0..n {
613        let q_i = marginal_offset[i] + marginal_baseline;
614        let g_i = logslope_offset[i] + logslope_baseline;
615        let sg = probit_scale * g_i;
616        let c_i = (1.0 + sg * sg).sqrt();
617        let f_i = q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
618        for j in 0..p_m {
619            m_eff[[i, j]] = c_i * marginal[[i, j]];
620        }
621        for j in 0..p_g {
622            g_eff[[i, j]] = f_i * logslope[[i, j]];
623        }
624    }
625
626    // C = G_effᵀ W G_eff (raw-coordinate effective logslope Gram); its diagonal
627    // sets the energy scale for the relative kept-direction tolerance.
628    let c_gram = fast_xt_diag_x(&g_eff, row_metric);
629    let energy_scale = (0..p_g).map(|i| c_gram[[i, i]]).fold(0.0_f64, f64::max);
630    if !energy_scale.is_finite() {
631        return Err(
632            "reduced logslope reparam: effective logslope Gram produced non-finite energy"
633                .to_string(),
634        );
635    }
636    if energy_scale <= 0.0 {
637        // A zero effective logslope image (f_i·G_i ≡ 0 in W) carries no joint-
638        // Hessian curvature at all — trivially W-explained by any span.
639        return Ok(ReducedLogslopeOutcome::FullyConfounded);
640    }
641
642    // A = M_effᵀ W M_eff + εI (ridge relative to the marginal effective energy so
643    // the Schur solve is well-posed even when the marginal pilot Gram is
644    // rank-soft; the ridge only under-removes, i.e. is conservative).
645    let mut a_gram = fast_xt_diag_x(&m_eff, row_metric);
646    let a_scale = (0..p_m).map(|i| a_gram[[i, i]]).fold(0.0_f64, f64::max);
647    let a_ridge = (a_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL).max(f64::EPSILON);
648    for i in 0..p_m {
649        a_gram[[i, i]] += a_ridge;
650    }
651
652    // B = M_effᵀ W G_eff (p_m × p_g);  Gtt = C − Bᵀ A⁻¹ B (p_g × p_g, PSD).
653    let b_cross = gam_linalg::faer_ndarray::fast_xt_diag_y(&m_eff, row_metric, &g_eff);
654    let a_view = gam_linalg::faer_ndarray::FaerArrayView::new(&a_gram);
655    let a_factor =
656        gam_linalg::faer_ndarray::factorize_symmetricwith_fallback(a_view.as_ref(), Side::Lower)
657            .map_err(|e| {
658                format!(
659                    "reduced logslope reparam: effective marginal Gram factorization failed: {e}"
660                )
661            })?;
662    let b_view = gam_linalg::faer_ndarray::FaerArrayView::new(&b_cross);
663    let solved = a_factor.solve(b_view.as_ref()); // A⁻¹ B  (p_m × p_g)
664    let a_inv_b = Array2::from_shape_fn((p_m, p_g), |(i, j)| solved[(i, j)]);
665    let schur = fast_atb(&b_cross, &a_inv_b); // Bᵀ A⁻¹ B  (p_g × p_g)
666    let mut stt = &c_gram - &schur;
667    stt = (&stt + &stt.t()) * 0.5;
668    if stt.iter().any(|v| !v.is_finite()) {
669        return Err(
670            "reduced logslope reparam: effective Schur Gram produced non-finite entries"
671                .to_string(),
672        );
673    }
674
675    let (evals, evecs) = stt
676        .eigh(Side::Lower)
677        .map_err(|e| format!("reduced logslope reparam: eigendecomposition failed: {e:?}"))?;
678    // A `Gtt` eigenvalue far below the effective logslope energy scale means that
679    // direction's effective logslope column is W-explained by the effective
680    // marginal span — exactly the joint-Hessian rank-soft confounded direction.
681    let tol = energy_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL;
682    let mut kept: Vec<usize> = (0..evals.len()).filter(|&i| evals[i] > tol).collect();
683    kept.sort_by(|&a, &b| {
684        evals[b]
685            .partial_cmp(&evals[a])
686            .unwrap_or(std::cmp::Ordering::Equal)
687    });
688    let r = kept.len();
689    // r == p_g: no effective-confounded direction to remove — keep the raw
690    // design. r == 0: the whole effective logslope image is in the effective
691    // marginal span — a distinct outcome the caller must refuse, never a
692    // keep-the-raw-design signal (the raw columns would be an arbitrary
693    // penalty-selected decomposition of an unidentified sum).
694    if r == p_g {
695        return Ok(ReducedLogslopeOutcome::FullRank);
696    }
697    if r == 0 {
698        return Ok(ReducedLogslopeOutcome::FullyConfounded);
699    }
700    let mut transform = Array2::<f64>::zeros((p_g, r));
701    for (out_col, &src) in kept.iter().enumerate() {
702        transform.column_mut(out_col).assign(&evecs.column(src));
703    }
704    if transform.iter().any(|v| !v.is_finite()) {
705        return Err(
706            "reduced logslope reparam: reduced transform produced non-finite entries".to_string(),
707        );
708    }
709    Ok(ReducedLogslopeOutcome::Reduced(transform))
710}
711
712/// Apply a [`ReducedLogslopeReparam`] to a logslope `TermCollectionDesign`,
713/// producing a new design at the reduced width `r`: the design becomes
714/// `G_reduced = G·T`, and every blockwise penalty `S` is reparameterized to
715/// `S_reduced = Tᵀ S T` over the full reduced column range `0..r`. The reduced
716/// penalty's null space is recomputed from its numerical rank so the REML
717/// log-determinant accounting stays consistent at the reduced width.
718fn reparameterize_logslope_design_reduced(
719    logslope_design: &TermCollectionDesign,
720    reparam: &ReducedLogslopeReparam,
721) -> Result<TermCollectionDesign, String> {
722    let g = logslope_design
723        .design
724        .try_to_dense_arc("reparameterize_logslope_design_reduced::logslope")?;
725    let p_g = g.ncols();
726    if p_g != reparam.original_cols() {
727        return Err(format!(
728            "reduced logslope reparam width mismatch: design has {p_g} cols, transform expects {}",
729            reparam.original_cols()
730        ));
731    }
732    let t = &reparam.transform;
733    let r = reparam.reduced_cols();
734    // G_reduced = G·T   (n × r).
735    let g_reduced = fast_ab(&g, t);
736
737    // Reparameterize each penalty: embed its local block at full width p_g, then
738    // form S_reduced = Tᵀ S T (r × r) over the whole reduced column range.
739    let mut new_penalties: Vec<gam_terms::smooth::BlockwisePenalty> =
740        Vec::with_capacity(logslope_design.penalties.len());
741    let mut new_nullspace_dims: Vec<usize> = Vec::with_capacity(logslope_design.penalties.len());
742    for bp in &logslope_design.penalties {
743        let mut full = Array2::<f64>::zeros((p_g, p_g));
744        full.slice_mut(s![bp.col_range.clone(), bp.col_range.clone()])
745            .assign(&bp.local);
746        // S_reduced = Tᵀ (S) T.
747        let st = fast_ab(&full, t); // p_g × r
748        let mut s_reduced = fast_atb(t, &st); // r × r
749        s_reduced = (&s_reduced + &s_reduced.t()) * 0.5;
750        // Null-space dimension of the reduced penalty = r − rank(S_reduced).
751        let (evals, _) = s_reduced
752            .eigh(Side::Lower)
753            .map_err(|e| format!("reduced logslope penalty eigendecomposition failed: {e:?}"))?;
754        let max_eval = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
755        let pen_tol = (max_eval * 1.0e-12).max(f64::EPSILON);
756        let rank = evals.iter().filter(|&&v| v.abs() > pen_tol).count();
757        let nullspace_dim = r.saturating_sub(rank);
758        new_penalties.push(gam_terms::smooth::BlockwisePenalty::new(0..r, s_reduced));
759        new_nullspace_dims.push(nullspace_dim);
760    }
761
762    let new_design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(g_reduced));
763    // The reduced logslope block is a single dense smooth-like surface over the
764    // reparameterized coordinates; it carries no parametric/random-effect/
765    // intercept structure of its own (those live in the marginal block), so the
766    // structural ranges collapse to empty and the smooth metadata is cleared.
767    // The penalties + nullspace_dims above are what the joint REML consumes.
768    Ok(TermCollectionDesign {
769        design: new_design,
770        // Reparameterization changes only the coefficient chart G -> G T.
771        // The known affine row function is coefficient-independent and must
772        // therefore pass through unchanged.
773        affine_offset: logslope_design.affine_offset.clone(),
774        penalties: new_penalties,
775        nullspace_dims: new_nullspace_dims,
776        penaltyinfo: Vec::new(),
777        dropped_penaltyinfo: Vec::new(),
778        coefficient_lower_bounds: None,
779        linear_constraints: None,
780        intercept_range: 0..0,
781        linear_ranges: Vec::new(),
782        linear_function_masses: Vec::new(),
783        random_effect_ranges: Vec::new(),
784        random_effect_levels: Vec::new(),
785        smooth: gam_terms::smooth::SmoothDesign {
786            term_designs: Vec::new(),
787            penalties: Vec::new(),
788            nullspace_dims: Vec::new(),
789            penaltyinfo: Vec::new(),
790            dropped_penaltyinfo: Vec::new(),
791            terms: Vec::new(),
792            coefficient_lower_bounds: None,
793            linear_constraints: None,
794        },
795    })
796}
797
798/// Re-embed the term-collection marginal penalties at the (possibly widened)
799/// block dimension `p_m [+ p₁]`, then append the #461 fixed-ridge absorber:
800///
801///  (#461, only with influence columns) the REML-learned absorber identity on
802///  the influence columns `p_m..p_m+p₁`.
803///
804/// The two former gam#754 pinned ridges — the marginal nullspace-shrinkage ridge
805/// and the marginal↔logslope overlap ridge — are DELETED: robustness is now
806/// unconditional, so the full-identifiable-span Jeffreys term (`Z_J = I`, see
807/// `jeffreys_subspace_from_penalty`) supplies automatic O(n)-scaled curvature on
808/// every under-identified direction (subsuming the nullspace ridge), and the
809/// exact orthogonal reparameterization of the logslope design (now unconditional,
810/// see `build_reduced_logslope_reparam`) resolves the marginal↔logslope confound
811/// by construction (subsuming the overlap ridge).
812///
813/// The genuine marginal smooth penalties keep their `col_range` (marginal
814/// columns stay in `0..p_m`). Returns `(penalties, nullspace_dims,
815/// initial_log_lambdas)` to install on the marginal block. With influence
816/// columns present, `rho_marginal` carries one TRAILING coordinate for the
817/// absorber ridge: its precision is REML-learned like every other penalty
818/// (SPEC: shrinkage is explicit or REML-selected, never a pinned magic
819/// constant), seeded at the ln(n) leakage scale by `joint_setup`.
820pub(crate) fn marginal_penalties_with_influence_ridge(
821    design: &TermCollectionDesign,
822    rho_marginal: &Array1<f64>,
823    influence_columns: Option<&Array2<f64>>,
824) -> Result<(Vec<PenaltyMatrix>, Vec<usize>, Array1<f64>), String> {
825    let p_m = design.design.ncols();
826    let p1 = influence_columns.map(|z| z.ncols()).unwrap_or(0);
827    let total_dim = p_m + p1;
828    let expected_rho = design.penalties.len() + usize::from(p1 > 0);
829    if rho_marginal.len() != expected_rho {
830        return Err(format!(
831            "marginal rho width {} != smooth penalties {} + absorber slot {}",
832            rho_marginal.len(),
833            design.penalties.len(),
834            usize::from(p1 > 0),
835        ));
836    }
837    // Re-embed each marginal penalty at the (widened) total dimension (col_range
838    // unchanged: marginal columns remain 0..p_m).
839    let mut penalties: Vec<PenaltyMatrix> = design
840        .penalties
841        .iter()
842        .map(|bp| bp.to_penalty_matrix(total_dim))
843        .collect();
844    let mut nullspace_dims = design.nullspace_dims.clone();
845    let log_lambdas = rho_marginal.to_vec();
846
847    // (#461) absorber ridge: identity on the influence columns only. Full rank
848    // (nullspace 0); its log λ is the trailing `rho_marginal` coordinate, so
849    // the outer REML selects the absorber precision like any other
850    // random-effect variance (large λ recovers the null correction; the
851    // residualized columns carry no marginal-span signal by construction).
852    if p1 > 0 {
853        penalties.push(PenaltyMatrix::Blockwise {
854            local: Array2::<f64>::eye(p1),
855            col_range: p_m..total_dim,
856            total_dim,
857        });
858        nullspace_dims.push(0);
859    }
860
861    Ok((penalties, nullspace_dims, Array1::from_vec(log_lambdas)))
862}
863
864/// Widen an optional β warm-start hint to the influence-widened marginal
865/// dimension, zero-filling the absorber coefficients `γ` (#461).
866pub(crate) fn widen_marginal_beta_hint(
867    beta_hint: Option<Array1<f64>>,
868    p_marginal_widened: usize,
869) -> Option<Array1<f64>> {
870    beta_hint.map(|hint| {
871        if hint.len() == p_marginal_widened {
872            hint
873        } else {
874            let mut widened = Array1::<f64>::zeros(p_marginal_widened);
875            let copy = hint.len().min(p_marginal_widened);
876            widened
877                .slice_mut(s![..copy])
878                .assign(&hint.slice(s![..copy]));
879            widened
880        }
881    })
882}
883
884/// Sup-norm of the fitted marginal linear predictor `η = X·β` restricted to a
885/// subset of the marginal design's columns. The mask selects which columns of
886/// `design.design` (length `design.design.ncols()`) contribute; coefficients
887/// beyond `ncols` (the fixed-ridge influence absorber) never enter the marginal
888/// predictor and are excluded by construction. Returns `0.0` for an empty
889/// design. This is the decisive separation quantity: the probit Fisher weight
890/// collapses with `|η|`, not with `‖β‖` (an ill-conditioned non-orthonormal
891/// Duchon/thin-plate basis carries large cancelling coefficients on a smooth,
892/// bounded surface).
893fn marginal_fitted_eta_sup_norm(design: &TermCollectionDesign, masked_beta: &Array1<f64>) -> f64 {
894    let x = &design.design;
895    let n = x.nrows();
896    if n == 0 || x.ncols() == 0 {
897        return 0.0;
898    }
899    let mut sup = 0.0_f64;
900    for row in 0..n {
901        let eta = x.dot_row_view(row, masked_beta.view());
902        if eta.is_finite() {
903            sup = sup.max(eta.abs());
904        }
905    }
906    sup
907}
908
909/// Build a copy of the marginal block β truncated to `design.design.ncols()`
910/// (drops any fixed-ridge absorber tail) so it can drive `X·β`.
911fn marginal_design_beta(
912    design: &TermCollectionDesign,
913    block_beta: ArrayView1<'_, f64>,
914) -> Array1<f64> {
915    let ncols = design.design.ncols();
916    let mut masked = Array1::<f64>::zeros(ncols);
917    let copy = ncols.min(block_beta.len());
918    masked
919        .slice_mut(s![..copy])
920        .assign(&block_beta.slice(s![..copy]));
921    masked
922}
923
924/// Zero every entry of `beta` outside the parametric (penalty-nullspace)
925/// marginal columns — the intercept and the single-penalty linear terms. These
926/// are the directions an unpenalized fit can genuinely separate along (no
927/// smoothness penalty bounds them), so their fitted contribution is tested on
928/// the same η scale.
929fn mask_parametric_columns(
930    design: &TermCollectionDesign,
931    spec: &TermCollectionSpec,
932    full: &Array1<f64>,
933) -> Array1<f64> {
934    let ncols = design.design.ncols();
935    let mut masked = Array1::<f64>::zeros(ncols);
936    if design.intercept_range.len() == 1 {
937        let idx = design.intercept_range.start;
938        if idx < ncols {
939            masked[idx] = full[idx];
940        }
941    }
942    for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
943        if linear.double_penalty {
944            continue;
945        }
946        for col in range.clone() {
947            if col < ncols {
948                masked[col] = full[col];
949            }
950        }
951    }
952    masked
953}
954
955/// Decide whether the converged marginal fit has genuinely separated, using the
956/// FITTED predictor sup-norm `|η|∞` (not raw `|β|∞`). Two arms share the
957/// numerical-degeneracy threshold [`BMS_PROBIT_SEPARATION_ETA_INF`]:
958///   - parametric arm: the penalty-nullspace columns' fitted contribution
959///     (an unpenalized direction can run to infinity);
960///   - full arm: the whole marginal surface's fitted predictor.
961/// The raw `|β|∞` (and its term label) is reported only as diagnostic context;
962/// it never gates the abort. When `|η|∞` is below threshold the converged
963/// penalized fit is numerically trustworthy and this returns `None` — no error,
964/// even if individual coefficients are large.
965pub(crate) fn bernoulli_marginal_slope_runaway_error_from_beta(
966    block_beta: ArrayView1<'_, f64>,
967    design: &TermCollectionDesign,
968    spec: &TermCollectionSpec,
969    inner_converged: bool,
970    eval_label: &str,
971) -> Option<String> {
972    let full_beta = marginal_design_beta(design, block_beta);
973    let parametric_beta = mask_parametric_columns(design, spec, &full_beta);
974
975    let eta_parametric = marginal_fitted_eta_sup_norm(design, &parametric_beta);
976    let eta_full = marginal_fitted_eta_sup_norm(design, &full_beta);
977
978    let (eta_inf, explanation) = if eta_parametric >= BMS_PROBIT_SEPARATION_ETA_INF {
979        (
980            eta_parametric,
981            "an unpenalized parametric marginal direction has no stable finite probit optimum and its fitted predictor has run to the probit underflow scale",
982        )
983    } else if eta_full >= BMS_PROBIT_SEPARATION_ETA_INF {
984        (
985            eta_full,
986            "a marginal direction is trading off against the logslope surface; this is the under-constrained marginal/logslope coupling that appears when the score is correlated with the shared surface covariates",
987        )
988    } else {
989        // |η|∞ is bounded: even if raw coefficients are large (ill-conditioned
990        // non-orthonormal basis with cancellation), the converged penalized
991        // probit fit is numerically trustworthy. Do NOT abort.
992        return None;
993    };
994
995    let inner_status = if inner_converged {
996        "the inner solve reached a KKT certificate at this separation-scale predictor"
997    } else {
998        "the inner solve failed while already carrying a separation-scale predictor"
999    };
1000    // Raw |β|∞ context (decisive quantity is |η|∞ above).
1001    let beta_abs = full_beta
1002        .iter()
1003        .copied()
1004        .filter(|v| v.is_finite())
1005        .fold(0.0_f64, |acc, v| acc.max(v.abs()));
1006
1007    Some(format!(
1008        "bernoulli marginal-slope probit marginal/logslope runaway detected in block \
1009         'marginal_surface' during {eval_label}: the fitted marginal predictor has \
1010         |η|∞={eta_inf:.3e} (numerical-degeneracy threshold \
1011         {BMS_PROBIT_SEPARATION_ETA_INF:.1}; raw |β|∞={beta_abs:.3e} is reported for \
1012         context only and does not gate this diagnostic). The joint design is \
1013         identifiable; {explanation}. {inner_status}. The robust Jeffreys curvature \
1014         path is already installed for this fit, so this diagnostic means the current \
1015         coupled surface still drives the linear predictor to the probit underflow \
1016         scale rather than a request for an external bias-reduction prior. Reduce or \
1017         reparameterize the coupled marginal/logslope surface, or use a \
1018         lower-dimensional logslope interaction. This is not a \
1019         Matérn/Duchon polynomial-nullspace or cross-block gauge-priority \
1020         failure."
1021    ))
1022}
1023
1024pub(crate) fn bernoulli_marginal_slope_runaway_error(
1025    warm_start: &CustomFamilyWarmStart,
1026    design: &TermCollectionDesign,
1027    spec: &TermCollectionSpec,
1028    inner_converged: bool,
1029    eval_label: &str,
1030) -> Option<String> {
1031    let block_beta = warm_start.block_beta_view(0)?;
1032    bernoulli_marginal_slope_runaway_error_from_beta(
1033        block_beta,
1034        design,
1035        spec,
1036        inner_converged,
1037        eval_label,
1038    )
1039}
1040
1041#[cfg(test)]
1042mod runaway_tests {
1043    use super::*;
1044    use gam_linalg::faer_ndarray::{
1045        FaerArrayView, factorize_symmetricwith_fallback, fast_xt_diag_y,
1046    };
1047    use gam_terms::smooth::{LinearCoefficientGeometry, LinearTermSpec};
1048
1049    // The marginal↔logslope overlap penalty is no longer installed as a pinned
1050    // ridge (subsumed by the now-unconditional exact logslope orthogonalisation in
1051    // `build_reduced_logslope_reparam`). The geometry helper is retained here under
1052    // the test module because the basis-independence/weight-orthogonality unit tests
1053    // below exercise it directly as the canonical overlap-direction reference.
1054    pub(crate) fn marginal_logslope_overlap_penalty(
1055        marginal_design: &DesignMatrix,
1056        logslope_design: &DesignMatrix,
1057        z: &Array1<f64>,
1058        row_metric: &Array1<f64>,
1059        marginal_offset: &Array1<f64>,
1060        logslope_offset: &Array1<f64>,
1061        marginal_baseline: f64,
1062        logslope_baseline: f64,
1063        probit_scale: f64,
1064    ) -> Result<Option<Array2<f64>>, String> {
1065        let marginal =
1066            marginal_design.try_to_dense_arc("marginal_logslope_overlap_penalty::marginal")?;
1067        let logslope =
1068            logslope_design.try_to_dense_arc("marginal_logslope_overlap_penalty::logslope")?;
1069        let n = marginal.nrows();
1070        if logslope.nrows() != n
1071            || z.len() != n
1072            || row_metric.len() != n
1073            || marginal_offset.len() != n
1074            || logslope_offset.len() != n
1075        {
1076            return Err(format!(
1077                "marginal/logslope overlap penalty row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
1078                marginal.nrows(),
1079                logslope.nrows(),
1080                z.len(),
1081                row_metric.len(),
1082                marginal_offset.len(),
1083                logslope_offset.len(),
1084            ));
1085        }
1086        let p_m = marginal.ncols();
1087        let p_g = logslope.ncols();
1088        if p_m == 0 || p_g == 0 {
1089            return Ok(None);
1090        }
1091        if !marginal_baseline.is_finite()
1092            || !logslope_baseline.is_finite()
1093            || !probit_scale.is_finite()
1094            || probit_scale <= 0.0
1095            || z.iter().any(|v| !v.is_finite())
1096            || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
1097            || marginal_offset.iter().any(|v| !v.is_finite())
1098            || logslope_offset.iter().any(|v| !v.is_finite())
1099        {
1100            return Err(
1101                "marginal/logslope overlap penalty requires finite pilot geometry and finite non-negative row metric"
1102                    .to_string(),
1103            );
1104        }
1105
1106        let mut marginal_effective = Array2::<f64>::zeros((n, p_m));
1107        let mut effective_logslope = Array2::<f64>::zeros((n, p_g));
1108        for i in 0..n {
1109            let q_i = marginal_offset[i] + marginal_baseline;
1110            let g_i = logslope_offset[i] + logslope_baseline;
1111            let sg = probit_scale * g_i;
1112            let c_i = (1.0 + sg * sg).sqrt();
1113            let logslope_factor =
1114                q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
1115            for j in 0..p_m {
1116                marginal_effective[[i, j]] = c_i * marginal[[i, j]];
1117            }
1118            for j in 0..p_g {
1119                effective_logslope[[i, j]] = logslope_factor * logslope[[i, j]];
1120            }
1121        }
1122        if effective_logslope.iter().all(|v| v.abs() <= f64::EPSILON) {
1123            return Ok(None);
1124        }
1125
1126        let mut gram = fast_xt_diag_x(&effective_logslope, row_metric);
1127        let gram_scale = gram.diag().iter().copied().fold(0.0_f64, f64::max);
1128        if !gram_scale.is_finite() || gram_scale <= 0.0 {
1129            return Ok(None);
1130        }
1131        let projection_ridge = (gram_scale * 1.0e-10).max(f64::EPSILON);
1132        for i in 0..p_g {
1133            gram[[i, i]] += projection_ridge;
1134        }
1135        let cross = fast_xt_diag_y(&effective_logslope, row_metric, &marginal_effective);
1136        let gram_view = FaerArrayView::new(&gram);
1137        let factor = factorize_symmetricwith_fallback(gram_view.as_ref(), Side::Lower)
1138            .map_err(|e| format!("marginal/logslope overlap Gram factorization failed: {e}"))?;
1139        let rhsview = FaerArrayView::new(&cross);
1140        let coeffs_mat = factor.solve(rhsview.as_ref());
1141        let coeffs = Array2::from_shape_fn((p_g, p_m), |(i, j)| coeffs_mat[(i, j)]);
1142        let projected_marginal = fast_ab(&effective_logslope, &coeffs);
1143        let mut penalty = fast_xt_diag_y(&marginal_effective, row_metric, &projected_marginal);
1144        penalty = (&penalty + &penalty.t()) * 0.5;
1145        let max_abs = penalty.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1146        if !max_abs.is_finite() || max_abs <= 1.0e-12 {
1147            return Ok(None);
1148        }
1149        Ok(Some(penalty))
1150    }
1151
1152    // The raw-vs-effective counterexample. With q=g=0, s=1: c_i=1 (so M_eff=M)
1153    // and f_i=z_i (so G_eff=diag(z)·G). Pick M=[1,1,1]ᵀ and a two-column logslope
1154    // G whose RAW columns are both linearly independent of M (raw orthogonalising
1155    // [M|G] would keep BOTH — the old code returned None, no reduction), but whose
1156    // first EFFECTIVE column diag(z)·G[:,0] equals M_eff exactly. The effective
1157    // audit must therefore drop exactly one direction (r=1), proving it removes
1158    // the joint-Hessian rank-soft direction the raw audit could not see.
1159    #[test]
1160    pub(crate) fn effective_reduction_drops_score_weighted_confound_raw_audit_misses() {
1161        // G col0 = [1,2,3], col1 = [1,2,9]  (row-major rows: [1,1],[2,2],[3,9]).
1162        let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1163        let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 1.0, 2.0, 2.0, 3.0, 9.0]).unwrap();
1164        let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1165        let w = Array1::<f64>::ones(3);
1166        let zero = Array1::<f64>::zeros(3);
1167
1168        // diag(z)·G[:,0] = [1·1, 0.5·2, (1/3)·3] = [1,1,1] = M_eff (fully aliased);
1169        // diag(z)·G[:,1] = [1·1, 0.5·2, (1/3)·9] = [1,1,3] (NOT in span([1,1,1])).
1170        let reparam = match reduced_logslope_transform_effective(
1171            m.view(),
1172            g.view(),
1173            &z,
1174            &w,
1175            &zero,
1176            &zero,
1177            0.0,
1178            0.0,
1179            1.0,
1180        )
1181        .expect("effective reduction must succeed")
1182        {
1183            ReducedLogslopeOutcome::Reduced(t) => t,
1184            other => panic!(
1185                "effective audit must reduce the score-weighted confound (raw audit would not), got {}",
1186                match other {
1187                    ReducedLogslopeOutcome::FullRank => "FullRank",
1188                    ReducedLogslopeOutcome::FullyConfounded => "FullyConfounded",
1189                    ReducedLogslopeOutcome::Reduced(_) => unreachable!(),
1190                }
1191            ),
1192        };
1193        assert_eq!(
1194            reparam.ncols(),
1195            1,
1196            "exactly one effective-identifiable logslope direction should survive"
1197        );
1198
1199        // The surviving raw direction's EFFECTIVE image diag(z)·G·t must carry the
1200        // non-constant ([1,1,3]) content — i.e. it is the identifiable direction,
1201        // not the [1,1,1] confound. Its row variance must be clearly positive.
1202        let g_eff = {
1203            let mut e = Array2::<f64>::zeros((3, 2));
1204            for i in 0..3 {
1205                for j in 0..2 {
1206                    e[[i, j]] = z[i] * g[[i, j]];
1207                }
1208            }
1209            e
1210        };
1211        let img = g_eff.dot(&reparam.column(0));
1212        let mean = img.iter().sum::<f64>() / 3.0;
1213        let var = img.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / 3.0;
1214        assert!(
1215            var > 1.0e-6,
1216            "kept direction must be the identifiable (non-constant) effective column, var={var}"
1217        );
1218    }
1219
1220    // The single-column fully-confounded case: G=[1,2,3]ᵀ, z=[1,1/2,1/3] gives
1221    // G_eff=[1,1,1]=M_eff, so the entire effective logslope image is in the
1222    // effective marginal span (r==0). The helper must report the distinct
1223    // FullyConfounded outcome — NOT the FullRank keep-the-raw-design signal —
1224    // because the data identify only the sum of the two surfaces and the
1225    // caller must refuse the block.
1226    #[test]
1227    pub(crate) fn effective_reduction_fully_confounded_single_column_is_distinct_outcome() {
1228        let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1229        let g = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
1230        let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1231        let w = Array1::<f64>::ones(3);
1232        let zero = Array1::<f64>::zeros(3);
1233        let outcome = reduced_logslope_transform_effective(
1234            m.view(),
1235            g.view(),
1236            &z,
1237            &w,
1238            &zero,
1239            &zero,
1240            0.0,
1241            0.0,
1242            1.0,
1243        )
1244        .expect("effective reduction must succeed");
1245        assert!(
1246            matches!(outcome, ReducedLogslopeOutcome::FullyConfounded),
1247            "fully effective-confounded logslope must surface the distinct FullyConfounded outcome"
1248        );
1249    }
1250
1251    // No effective confound: both effective logslope columns stay independent of
1252    // M_eff, so nothing is reduced (r==p_g ⇒ None) and healthy fits are untouched.
1253    #[test]
1254    pub(crate) fn effective_reduction_no_confound_returns_none() {
1255        let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1256        // diag(z)·col gives non-constant images for both columns under z below.
1257        let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]).unwrap();
1258        let z = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1259        let w = Array1::<f64>::ones(3);
1260        let zero = Array1::<f64>::zeros(3);
1261        let outcome = reduced_logslope_transform_effective(
1262            m.view(),
1263            g.view(),
1264            &z,
1265            &w,
1266            &zero,
1267            &zero,
1268            0.0,
1269            0.0,
1270            1.0,
1271        )
1272        .expect("effective reduction must succeed");
1273        assert!(
1274            matches!(outcome, ReducedLogslopeOutcome::FullRank),
1275            "no effective confound ⇒ FullRank (raw design kept unchanged)"
1276        );
1277    }
1278
1279    #[test]
1280    pub(crate) fn spatial_joint_setup_counts_only_learned_penalties_in_rho() {
1281        let data = Array2::<f64>::zeros((3, 1));
1282        let empty_terms = TermCollectionSpec {
1283            linear_terms: Vec::new(),
1284            random_effect_terms: Vec::new(),
1285            smooth_terms: Vec::new(),
1286        };
1287        let setup = joint_setup(
1288            data.view(),
1289            &empty_terms,
1290            &empty_terms,
1291            2,
1292            3,
1293            Some(2.5),
1294            &[0.4],
1295            &SpatialLengthScaleOptimizationOptions::default(),
1296        )
1297        .expect("empty spatial geometry is valid");
1298
1299        assert_eq!(
1300            setup.rho_dim(),
1301            6,
1302            "BMS spatial setup rho holds every learned marginal/logslope/auxiliary penalty; the #461 absorber ridge occupies the trailing marginal slot"
1303        );
1304        assert_eq!(
1305            setup.theta0()[1],
1306            2.5,
1307            "absorber ridge seeds the trailing marginal rho coordinate at the ln(n) leakage scale"
1308        );
1309    }
1310
1311    #[test]
1312    pub(crate) fn overlap_penalty_targets_score_weighted_logslope_span() {
1313        let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1314            Array2::from_shape_vec((4, 1), vec![0.0, 1.0, 2.0, 3.0]).unwrap(),
1315        ));
1316        let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1317            Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1318        ));
1319        let z = Array1::from_vec(vec![0.0, 1.0, 2.0, 3.0]);
1320        let row_metric = Array1::ones(4);
1321        let offsets = Array1::zeros(4);
1322
1323        let penalty = marginal_logslope_overlap_penalty(
1324            &marginal,
1325            &logslope,
1326            &z,
1327            &row_metric,
1328            &offsets,
1329            &offsets,
1330            0.0,
1331            0.0,
1332            1.0,
1333        )
1334        .expect("overlap penalty should build")
1335        .expect("marginal signal lies in the pilot logslope Jacobian span");
1336
1337        assert_eq!(penalty.dim(), (1, 1));
1338        assert!((penalty[[0, 0]] - 14.0).abs() < 1.0e-6);
1339    }
1340
1341    #[test]
1342    pub(crate) fn overlap_penalty_skips_weight_orthogonal_channels() {
1343        let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1344            Array2::from_shape_vec((4, 1), vec![-1.0, 1.0, -1.0, 1.0]).unwrap(),
1345        ));
1346        let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1347            Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1348        ));
1349        let z = Array1::ones(4);
1350        let row_metric = Array1::ones(4);
1351        let offsets = Array1::zeros(4);
1352
1353        let penalty = marginal_logslope_overlap_penalty(
1354            &marginal,
1355            &logslope,
1356            &z,
1357            &row_metric,
1358            &offsets,
1359            &offsets,
1360            0.0,
1361            0.0,
1362            1.0,
1363        )
1364        .expect("overlap penalty should build");
1365
1366        assert!(penalty.is_none());
1367    }
1368
1369    // ── Fitted-η separation guard fixtures ───────────────────────────────
1370    //
1371    // The runaway guard tests the FITTED marginal predictor sup-norm
1372    // `|η|∞ = max_i |X[i,:]·β|`, not raw `|β|∞`. These helpers build minimal
1373    // `TermCollectionDesign` / `TermCollectionSpec` pairs from a dense design so
1374    // the criterion is exercised deterministically with no data files.
1375
1376    fn dense_marginal_design(
1377        x: Array2<f64>,
1378        intercept_range: std::ops::Range<usize>,
1379        linear_ranges: Vec<(String, std::ops::Range<usize>)>,
1380    ) -> TermCollectionDesign {
1381        let nrows = x.nrows();
1382        TermCollectionDesign {
1383            design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
1384            affine_offset: Array1::zeros(nrows),
1385            penalties: Vec::new(),
1386            nullspace_dims: Vec::new(),
1387            penaltyinfo: Vec::new(),
1388            dropped_penaltyinfo: Vec::new(),
1389            coefficient_lower_bounds: None,
1390            linear_constraints: None,
1391            intercept_range,
1392            linear_ranges,
1393            linear_function_masses: Vec::new(),
1394            random_effect_ranges: Vec::new(),
1395            random_effect_levels: Vec::new(),
1396            smooth: gam_terms::smooth::SmoothDesign {
1397                term_designs: Vec::new(),
1398                penalties: Vec::new(),
1399                nullspace_dims: Vec::new(),
1400                penaltyinfo: Vec::new(),
1401                dropped_penaltyinfo: Vec::new(),
1402                terms: Vec::new(),
1403                coefficient_lower_bounds: None,
1404                linear_constraints: None,
1405            },
1406        }
1407    }
1408
1409    fn linear_term(name: &str, feature_col: usize) -> LinearTermSpec {
1410        LinearTermSpec {
1411            name: name.to_string(),
1412            feature_col,
1413            feature_cols: vec![feature_col],
1414            categorical_levels: vec![],
1415            double_penalty: false,
1416            coefficient_geometry: LinearCoefficientGeometry::default(),
1417            coefficient_min: None,
1418            coefficient_max: None,
1419            frozen_function_mass: None,
1420        }
1421    }
1422
1423    fn empty_spec() -> TermCollectionSpec {
1424        TermCollectionSpec {
1425            linear_terms: Vec::new(),
1426            random_effect_terms: Vec::new(),
1427            smooth_terms: Vec::new(),
1428        }
1429    }
1430
1431    /// Regression lock for the false-positive fix: an ill-conditioned
1432    /// non-orthonormal basis (two identical/collinear columns) with a large
1433    /// cancelling coefficient `β=[60,-60]` yields `Xβ ≡ 0` — a perfectly
1434    /// bounded fitted predictor. The guard MUST NOT fire even though raw
1435    /// `|β|∞=60` is far above the old `40.0` coefficient threshold. This is the
1436    /// exact pathology that aborted valid biobank fits.
1437    #[test]
1438    pub(crate) fn runaway_guard_silent_when_huge_beta_cancels_to_bounded_eta() {
1439        // Two identical columns ⇒ Xβ = (β0+β1)·col; β=[60,-60] ⇒ Xβ ≡ 0.
1440        let x = Array2::<f64>::from_shape_vec((4, 2), vec![1.0; 8]).unwrap();
1441        let design = dense_marginal_design(x, 0..0, Vec::new());
1442        let beta = Array1::from_vec(vec![60.0, -60.0]);
1443
1444        let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1445            beta.view(),
1446            &design,
1447            &empty_spec(),
1448            true,
1449            "regression-fixture",
1450        );
1451        assert!(
1452            msg.is_none(),
1453            "huge cancelling β with bounded fitted η must NOT trip the runaway guard; got {msg:?}"
1454        );
1455    }
1456
1457    /// A genuinely separating full marginal surface: a single column with a
1458    /// large coefficient drives `|η|∞ = 40 ≥ 35`, so the guard fires and names
1459    /// the marginal/logslope coupling explanation.
1460    #[test]
1461    pub(crate) fn runaway_guard_fires_when_fitted_eta_exceeds_threshold() {
1462        let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1463        let design = dense_marginal_design(x, 0..0, Vec::new());
1464        let beta = Array1::from_vec(vec![40.0]);
1465
1466        let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1467            beta.view(),
1468            &design,
1469            &empty_spec(),
1470            true,
1471            "separation-fixture",
1472        )
1473        .expect("fitted |η|∞=40 ≥ 35 must trip the runaway guard");
1474
1475        assert!(msg.contains("marginal/logslope runaway"));
1476        assert!(msg.contains("|η|∞"));
1477        assert!(msg.contains("4.000e1"));
1478        assert!(msg.contains("score is correlated with the shared surface covariates"));
1479        assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1480        assert!(msg.contains("KKT certificate"));
1481    }
1482
1483    /// An unpenalized parametric direction can genuinely separate (no smoothness
1484    /// penalty bounding it). When its fitted contribution reaches the η scale
1485    /// the parametric arm fires first and names the parametric explanation.
1486    #[test]
1487    pub(crate) fn runaway_guard_names_unpenalized_parametric_direction_via_fitted_eta() {
1488        let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1489        let design = dense_marginal_design(x, 0..0, vec![("sex".to_string(), 0..1)]);
1490        let mut spec = empty_spec();
1491        spec.linear_terms.push(linear_term("sex", 0));
1492        let beta = Array1::from_vec(vec![41.0]);
1493
1494        let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1495            beta.view(),
1496            &design,
1497            &spec,
1498            true,
1499            "parametric-fixture",
1500        )
1501        .expect("parametric fitted |η|∞=41 ≥ 35 must trip the runaway guard");
1502
1503        assert!(msg.contains("unpenalized parametric marginal direction"));
1504        assert!(msg.contains("|η|∞"));
1505        assert!(msg.contains("robust Jeffreys curvature path is already installed"));
1506        assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1507    }
1508
1509    /// A non-converged inner solve with a BOUNDED fitted predictor must NOT
1510    /// surface the separation error — the non-convergence is reported through
1511    /// the existing downstream path, not as a runaway.
1512    #[test]
1513    pub(crate) fn runaway_guard_silent_for_nonconverged_but_bounded_eta() {
1514        let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1515        let design = dense_marginal_design(x, 0..0, Vec::new());
1516        let beta = Array1::from_vec(vec![5.0]);
1517
1518        let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1519            beta.view(),
1520            &design,
1521            &empty_spec(),
1522            false,
1523            "nonconverged-fixture",
1524        );
1525        assert!(
1526            msg.is_none(),
1527            "bounded fitted η must not raise the separation error even when the inner solve did not converge; got {msg:?}"
1528        );
1529    }
1530
1531    /// A genuinely separating fit that ALSO failed to converge still surfaces
1532    /// the runaway error, and reports the non-converged inner status.
1533    #[test]
1534    pub(crate) fn runaway_guard_fires_for_nonconverged_separating_eta() {
1535        let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1536        let design = dense_marginal_design(x, 0..0, Vec::new());
1537        let beta = Array1::from_vec(vec![50.0]);
1538
1539        let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1540            beta.view(),
1541            &design,
1542            &empty_spec(),
1543            false,
1544            "nonconverged-separating-fixture",
1545        )
1546        .expect("separating |η|∞ at non-convergence must still trip the guard");
1547
1548        assert!(msg.contains(
1549            "the inner solve failed while already carrying a separation-scale predictor"
1550        ));
1551    }
1552
1553    /// Regression lock for gam#370: the pre-fit identifiability audit evaluates
1554    /// every block's effective Jacobian at the empty/zero β with
1555    /// `family_scalars: None`. The BMS marginal and logslope blocks own the
1556    /// logslope design + offset, so `g_i = offset_s[i]` (the fitted logslope
1557    /// baseline, generically nonzero) at β = 0 — and the old code hard-errored
1558    /// ("requires BmsFamilyScalars when beta != 0 … got family_scalars: None")
1559    /// because it read `any_nonzero_g` and demanded a caller-supplied scalar it
1560    /// could in fact reconstruct itself. Both blocks must now self-compute the
1561    /// closed-form `c_i = sqrt(1 + (s·g_i)²)` (and the logslope factor) from
1562    /// owned data and return a finite Jacobian, NOT an error. This makes every
1563    /// `logslope_formula` / `linkwiggle(...)` BMS fit reachable through the
1564    /// Python `gamfit.fit` API (the audit no longer aborts before the fit).
1565    #[test]
1566    pub(crate) fn bms_block_jacobians_self_compute_at_audit_empty_beta_nonzero_logslope_baseline() {
1567        use std::sync::Arc;
1568        let n = 4usize;
1569        let marginal =
1570            Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1571        let logslope =
1572            Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1573        let offset_m = Array1::<f64>::zeros(n);
1574        // Nonzero logslope baseline absorbed into offset_s — the exact pooled
1575        // probit pilot value that is "essentially never exactly 0".
1576        let g_baseline = 0.3_f64;
1577        let offset_s = Array1::<f64>::from_elem(n, g_baseline);
1578        let z = Arc::new(Array1::from_vec(vec![-0.7, 0.2, 0.9, 1.4]));
1579        let s = 1.0_f64;
1580
1581        // The pre-fit audit linearization state: EMPTY β, no family scalars.
1582        let beta: Vec<f64> = Vec::new();
1583        let state = FamilyLinearizationState {
1584            beta: &beta,
1585            family_scalars: None,
1586            channel_hessian: None,
1587            probit_frailty_scale: s,
1588        };
1589
1590        let marginal_jac = BmsMarginalJacobian::new(
1591            Arc::clone(&marginal),
1592            Arc::clone(&logslope),
1593            offset_m.clone(),
1594            offset_s.clone(),
1595            1,
1596        );
1597        let j_m = marginal_jac
1598            .effective_jacobian_rows(&state, 0..n)
1599            .expect("BMS marginal Jacobian must self-compute at audit empty β (gam#370)");
1600        // ∂η_i/∂β_m = c_i · M[i,:], c_i = sqrt(1 + (s·g_baseline)²), M[i,0]=1.
1601        let c_expected = (1.0 + (s * g_baseline).powi(2)).sqrt();
1602        assert_eq!(j_m.dim(), (n, 1));
1603        for i in 0..n {
1604            assert!(
1605                (j_m[[i, 0]] - c_expected).abs() < 1e-12,
1606                "marginal J[{i}] = {} != closed-form c_i = {c_expected}",
1607                j_m[[i, 0]]
1608            );
1609        }
1610
1611        let logslope_jac = BmsLogslopeJacobian::new(
1612            Arc::clone(&marginal),
1613            Arc::clone(&logslope),
1614            offset_m,
1615            offset_s,
1616            Arc::clone(&z),
1617            1,
1618        );
1619        let j_s = logslope_jac
1620            .effective_jacobian_rows(&state, 0..n)
1621            .expect("BMS logslope Jacobian must self-compute at audit empty β (gam#370)");
1622        // ∂η_i/∂β_s = (q_i·s²·g_i/c_i + s·z_i)·G[i,:]; at empty β, q_i = 0 and
1623        // g_i = g_baseline, so the factor is s²·0·… + s·z_i = s·z_i (G[i,0]=1).
1624        assert_eq!(j_s.dim(), (n, 1));
1625        for i in 0..n {
1626            let expected = s * z[i];
1627            assert!(
1628                (j_s[[i, 0]] - expected).abs() < 1e-12,
1629                "logslope J[{i}] = {} != closed-form factor {expected}",
1630                j_s[[i, 0]]
1631            );
1632            assert!(j_s[[i, 0]].is_finite());
1633        }
1634    }
1635}
1636
1637pub(crate) fn build_marginal_blockspec_bms(
1638    design: &TermCollectionDesign,
1639    baseline: f64,
1640    offset: &Array1<f64>,
1641    rho: Array1<f64>,
1642    beta_hint: Option<Array1<f64>>,
1643    logslope_design: &TermCollectionDesign,
1644    logslope_offset: &Array1<f64>,
1645    logslope_baseline: f64,
1646    p_marginal: usize,
1647    influence_columns: Option<&Array2<f64>>,
1648) -> Result<ParameterBlockSpec, String> {
1649    let offset_m = offset + baseline;
1650    let offset_s = logslope_offset + logslope_baseline;
1651    let raw_marginal_dense = design
1652        .design
1653        .try_to_dense_arc("build_marginal_blockspec_bms::marginal")?;
1654    let marginal_dense =
1655        widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1656    let logslope_dense = logslope_design
1657        .design
1658        .try_to_dense_arc("build_marginal_blockspec_bms::logslope")?;
1659    let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsMarginalJacobian {
1660        marginal_dense: Arc::clone(&marginal_dense),
1661        logslope_dense,
1662        offset_m: offset_m.clone(),
1663        offset_s,
1664        p_marginal,
1665    });
1666    let (penalties, nullspace_dims, initial_log_lambdas) =
1667        marginal_penalties_with_influence_ridge(design, &rho, influence_columns)?;
1668    Ok(ParameterBlockSpec {
1669        name: "marginal_surface".to_string(),
1670        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1671            (*marginal_dense).clone(),
1672        )),
1673        offset: offset_m,
1674        penalties,
1675        nullspace_dims,
1676        initial_log_lambdas,
1677        initial_beta: widen_marginal_beta_hint(beta_hint, p_marginal),
1678        // Canonical-gauge architecture (issue #322): give marginal_surface
1679        // strictly higher priority than logslope_surface so the priority-
1680        // ordered RRQR in `canonicalize_for_identifiability` presents
1681        // marginal columns first and routes any cross-block alias drop into
1682        // logslope.  Equal priorities (the previous default of 100/100)
1683        // produced a same-priority `hard_alias_pair` whenever a
1684        // high-dimensional smooth — e.g. `s(x, type=duchon, centers>=6)`
1685        // in the location block — accidentally spanned the logslope basis
1686        // direction, leaving the joint Hessian with a structural null and
1687        // the spectral Newton solve refusing to step.  The values mirror
1688        // the canonical-gauge entry for survival marginal-slope
1689        // (marginal=150, logslope=120).
1690        gauge_priority: GAUGE_PRIORITY_MARGINAL,
1691        jacobian_callback: Some(callback),
1692        stacked_design: None,
1693        stacked_offset: None,
1694    })
1695}
1696
1697pub(crate) fn build_logslope_blockspec_bms(
1698    design: &TermCollectionDesign,
1699    baseline: f64,
1700    offset: &Array1<f64>,
1701    rho: Array1<f64>,
1702    beta_hint: Option<Array1<f64>>,
1703    marginal_design: &TermCollectionDesign,
1704    marginal_offset: &Array1<f64>,
1705    marginal_baseline: f64,
1706    z: Arc<Array1<f64>>,
1707    p_marginal: usize,
1708    influence_columns: Option<&Array2<f64>>,
1709) -> Result<ParameterBlockSpec, String> {
1710    let offset_s = offset + baseline;
1711    let offset_m = marginal_offset + marginal_baseline;
1712    let raw_marginal_dense = marginal_design
1713        .design
1714        .try_to_dense_arc("build_logslope_blockspec_bms::marginal")?;
1715    // The logslope Jacobian reconstructs q_i = M·β_m + offset_m; with the
1716    // absorbed influence columns folded into the marginal index, the marginal
1717    // reference design and p_marginal MUST be the widened [M | Z̃] / (p_m+p₁)
1718    // so β_m slices to the absorber and q_i carries the Z̃·γ shift (#461).
1719    let marginal_dense =
1720        widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1721    let logslope_dense = design
1722        .design
1723        .try_to_dense_arc("build_logslope_blockspec_bms::logslope")?;
1724    let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsLogslopeJacobian {
1725        marginal_dense,
1726        logslope_dense: Arc::clone(&logslope_dense),
1727        offset_m,
1728        offset_s: offset_s.clone(),
1729        z,
1730        p_marginal,
1731    });
1732    Ok(ParameterBlockSpec {
1733        name: "logslope_surface".to_string(),
1734        design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1735            (*logslope_dense).clone(),
1736        )),
1737        offset: offset_s,
1738        penalties: design.penalties_as_penalty_matrix(),
1739        nullspace_dims: design.nullspace_dims.clone(),
1740        initial_log_lambdas: rho,
1741        initial_beta: beta_hint,
1742        // Canonical-gauge architecture (issue #322): logslope is strictly
1743        // lower priority than marginal so the priority-ordered RRQR in
1744        // `canonicalize_for_identifiability` demotes a shared cross-block
1745        // direction here, not in marginal.  Mirrors the survival-mgs
1746        // value (marginal=150, logslope=120).  See the matching comment
1747        // on `build_marginal_blockspec_bms` for the failure mode this
1748        // resolves.
1749        gauge_priority: GAUGE_PRIORITY_LOGSLOPE,
1750        jacobian_callback: Some(callback),
1751        stacked_design: None,
1752        stacked_offset: None,
1753    })
1754}
1755
1756pub(crate) fn build_deviation_aux_blockspec(
1757    name: &str,
1758    prepared: &DeviationPrepared,
1759    rho: Array1<f64>,
1760    beta_hint: Option<Array1<f64>>,
1761) -> Result<ParameterBlockSpec, String> {
1762    let mut block = prepared.block.clone();
1763    block.initial_log_lambdas = Some(rho);
1764    let candidate_beta = beta_hint.or_else(|| Some(Array1::<f64>::zeros(block.design.ncols())));
1765    block.initial_beta = candidate_beta
1766        .map(|beta| {
1767            let zero = Array1::<f64>::zeros(beta.len());
1768            project_monotone_feasible_beta(&prepared.runtime, &zero, &beta, name)
1769        })
1770        .transpose()?;
1771    let mut spec = block.intospec(name)?;
1772    // Deviation auxiliary blocks (score_warp_dev, link_dev, and any
1773    // future flex block routed through this builder) model pure
1774    // shape modifications on top of parametric anchors. They must
1775    // never own a shared affine direction with the parametric
1776    // (time / marginal / logslope) blocks. The canonical-gauge
1777    // selector drops shared directions from blocks with lower
1778    // gauge_priority first; assigning a value below the parametric
1779    // default (GAUGE_PRIORITY_CANDIDATE_FLEX) realises that contract
1780    // automatically.
1781    spec.gauge_priority = match name {
1782        "link_dev" => GAUGE_PRIORITY_LINK_DEV,
1783        // score_warp_dev gets a slightly higher priority than link_dev
1784        // because in mixed-flex configurations (both blocks present)
1785        // link_dev is the residualised one (orthogonalised against the
1786        // parametric anchors PLUS the already-prepared score_warp
1787        // basis at construction time); link_dev should therefore yield
1788        // first when an alias still survives into the joint design.
1789        "score_warp_dev" => GAUGE_PRIORITY_SCORE_WARP_DEV,
1790        _ => GAUGE_PRIORITY_DEVIATION_DEFAULT,
1791    };
1792    Ok(spec)
1793}
1794
1795pub(crate) fn push_deviation_aux_blockspecs(
1796    blocks: &mut Vec<ParameterBlockSpec>,
1797    rho: &Array1<f64>,
1798    cursor: &mut usize,
1799    score_warp_prepared: Option<&DeviationPrepared>,
1800    link_dev_prepared: Option<&DeviationPrepared>,
1801    score_warp_beta_hint: Option<Array1<f64>>,
1802    link_dev_beta_hint: Option<Array1<f64>>,
1803) -> Result<(), String> {
1804    fn take_rho_slice(
1805        rho: &Array1<f64>,
1806        cursor: &mut usize,
1807        count: usize,
1808        block_name: &str,
1809    ) -> Result<Array1<f64>, String> {
1810        let start = *cursor;
1811        let end = start.checked_add(count).ok_or_else(|| {
1812            format!("{block_name} penalty-rho range overflow: start={start}, count={count}")
1813        })?;
1814        if end > rho.len() {
1815            return Err(format!(
1816                "{block_name} penalty-rho range {start}..{end} exceeds rho length {}",
1817                rho.len()
1818            ));
1819        }
1820        let slice = rho.slice(s![start..end]).to_owned();
1821        *cursor = end;
1822        Ok(slice)
1823    }
1824
1825    if let Some(prepared) = score_warp_prepared {
1826        let rho_h = take_rho_slice(
1827            rho,
1828            cursor,
1829            prepared.block.penalties.len(),
1830            "score_warp_dev",
1831        )?;
1832        blocks.push(build_deviation_aux_blockspec(
1833            "score_warp_dev",
1834            prepared,
1835            rho_h,
1836            score_warp_beta_hint,
1837        )?);
1838    }
1839    if let Some(prepared) = link_dev_prepared {
1840        let rho_w = take_rho_slice(rho, cursor, prepared.block.penalties.len(), "link_dev")?;
1841        blocks.push(build_deviation_aux_blockspec(
1842            "link_dev",
1843            prepared,
1844            rho_w,
1845            link_dev_beta_hint,
1846        )?);
1847    }
1848    Ok(())
1849}
1850
1851#[cfg(test)]
1852mod deviation_penalty_layout_tests {
1853    use super::*;
1854
1855    fn prepared_with_penalty_orders(orders: Vec<usize>) -> DeviationPrepared {
1856        let seed = Array1::linspace(-1.0, 1.0, 48);
1857        let config = DeviationBlockConfig {
1858            degree: 3,
1859            num_internal_knots: 6,
1860            penalty_order: *orders.first().expect("test requires a penalty order"),
1861            penalty_orders: orders,
1862            double_penalty: false,
1863            monotonicity_eps: 0.0,
1864        };
1865        build_score_warp_deviation_block_from_seed(&seed, &config)
1866            .expect("test deviation block must build")
1867    }
1868
1869    #[test]
1870    fn composed_score_link_influence_rho_layout_advances_by_emitted_counts_2315() {
1871        // Deliberately use unequal component counts and disjoint sentinels. A
1872        // length-only assertion cannot detect the old bug: link_dev received the
1873        // right-looking slice, but failing to advance the cursor made the following
1874        // influence absorber silently reuse link_dev's first rho coordinate.
1875        let score_warp = prepared_with_penalty_orders(vec![1, 2]);
1876        let link_dev = prepared_with_penalty_orders(vec![1, 2, 3]);
1877        assert_eq!(score_warp.block.penalties.len(), 2);
1878        assert_eq!(link_dev.block.penalties.len(), 3);
1879
1880        // Coordinate zero stands for an already-consumed core penalty. The final
1881        // coordinate is the influence absorber's trailing ridge.
1882        let rho = Array1::from_vec(vec![-101.0, 11.0, 12.0, 21.0, 22.0, 23.0, 31.0]);
1883        let mut cursor = 1usize;
1884        let mut blocks = Vec::new();
1885        push_deviation_aux_blockspecs(
1886            &mut blocks,
1887            &rho,
1888            &mut cursor,
1889            Some(&score_warp),
1890            Some(&link_dev),
1891            None,
1892            None,
1893        )
1894        .expect("composed deviation layout must be realized");
1895
1896        assert_eq!(blocks.len(), 2);
1897        assert_eq!(blocks[0].name, "score_warp_dev");
1898        assert_eq!(
1899            blocks[0].initial_log_lambdas.as_slice(),
1900            Some(&[11.0, 12.0][..])
1901        );
1902        assert_eq!(blocks[1].name, "link_dev");
1903        assert_eq!(
1904            blocks[1].initial_log_lambdas.as_slice(),
1905            Some(&[21.0, 22.0, 23.0][..])
1906        );
1907        assert_eq!(
1908            cursor, 6,
1909            "the next consumer must start after every emitted deviation penalty"
1910        );
1911
1912        let influence_rho = rho.slice(s![cursor..cursor + 1]).to_owned();
1913        assert_eq!(influence_rho.as_slice(), Some(&[31.0][..]));
1914        assert_ne!(
1915            influence_rho[0], blocks[1].initial_log_lambdas[0],
1916            "the influence absorber must not reuse link_dev's first rho coordinate"
1917        );
1918    }
1919}
1920
1921fn inner_fit(
1922    family: &BernoulliMarginalSlopeFamily,
1923    blocks: &[ParameterBlockSpec],
1924    options: &BlockwiseFitOptions,
1925) -> Result<UnifiedFitResult, String> {
1926    let mut options = options.clone();
1927    // BMS carries fixed physical ridge penalties that regularize coefficient
1928    // geometry but are not REML coordinates. The exact hyper-Hessian route can
1929    // stall after that projection; the family has a dedicated exact-gradient
1930    // path with full-data polish, so make it the primary nested smoother.
1931    options.use_outer_hessian = false;
1932    options.outer_tol = options.outer_tol.max(2.0e-5);
1933    fit_custom_family(family, blocks, &options).map_err(|e| e.to_string())
1934}
1935
1936fn inner_fit_from_certified_outer(
1937    family: &BernoulliMarginalSlopeFamily,
1938    blocks: &[ParameterBlockSpec],
1939    options: &BlockwiseFitOptions,
1940    mode: CustomFamilyJointHyperModeSelection,
1941    theta: &Array1<f64>,
1942    outer: &gam_solve::rho_optimizer::CertifiedOuterResult,
1943) -> Result<UnifiedFitResult, String> {
1944    let mut options = crate::outer_subsample::exact_outer_options_for_row_set(
1945        options,
1946        &crate::row_kernel::RowSet::All,
1947    );
1948    options.use_outer_hessian = false;
1949    options.outer_tol = options.outer_tol.max(2.0e-5);
1950    fit_custom_family_fixed_log_lambdas_from_mode_selection(
1951        family, blocks, &options, mode, theta, outer,
1952    )
1953    .map_err(|error| error.to_string())
1954}
1955
1956pub fn fit_bernoulli_marginal_slope_terms(
1957    data: ArrayView2<'_, f64>,
1958    spec: BernoulliMarginalSlopeTermSpec,
1959    options: &BlockwiseFitOptions,
1960    kappa_options: &SpatialLengthScaleOptimizationOptions,
1961    policy: &gam_runtime::resource::ResourcePolicy,
1962) -> Result<BernoulliMarginalSlopeFitResult, String> {
1963    let mut spec = spec;
1964    let data_view = data;
1965    validate_spec(data_view, &spec)?;
1966    // Freeze the measure-jet representer length-scale dial on the coupled
1967    // marginal + log-slope surfaces (#1116). A shared mjs basis feeds BOTH
1968    // blocks; a design-moving ℓ on those shared covariates lets the outer
1969    // search reach a sharp ℓ where a marginal smooth direction trades off
1970    // against the log-slope into a separation-scale runaway (|β|→1e3). The auto
1971    // (frozen) ℓ is well-conditioned here — the pre-#1116 behavior this
1972    // restores. ℓ-learning stays on for single-surface (e.g. Gaussian) fits,
1973    // where there is no marginal/log-slope coupling to destabilize.
1974    let mjs_frozen_marginal =
1975        gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.marginalspec);
1976    let mjs_frozen_logslope =
1977        gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.logslopespec);
1978    if mjs_frozen_marginal + mjs_frozen_logslope > 0 {
1979        log::info!(
1980            "[BMS spatial] froze measure-jet length-scale learning on {} marginal + {} log-slope \
1981             term(s): the coupled surface keeps ℓ at its conditioned auto value (#1116)",
1982            mjs_frozen_marginal,
1983            mjs_frozen_logslope
1984        );
1985    }
1986    let mut effective_kappa_options = kappa_options.clone();
1987    // Honor explicit `length_scale=X` in the user's formula: when every
1988    // spatial term in BOTH the marginal mean and log-slope blocks carries
1989    // a user-supplied scalar length scale and no per-axis anisotropy is
1990    // requested, there is nothing for the joint-spatial outer optimizer
1991    // to do. Routing through it anyway spends ~80 outer ARC iters stalled
1992    // at the user's chosen ρ (the n-block ARC's first proposed step lands
1993    // at the box corner and never recovers), then falls through to the
1994    // ρ-only "custom family" path which is what we wanted all along.
1995    // Short-circuit straight to the ρ-only path.
1996    let kappa_locked_marginal =
1997        gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.marginalspec);
1998    let kappa_locked_logslope =
1999        gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.logslopespec);
2000    if effective_kappa_options.enabled && kappa_locked_marginal && kappa_locked_logslope {
2001        log::info!(
2002            "[BMS spatial] disabling κ/ψ optimization: every spatial term has an \
2003             explicit length_scale and no anisotropy; user-supplied kernel scale is fixed"
2004        );
2005        effective_kappa_options.enabled = false;
2006    }
2007    let flex_spatial_pilot_path = (spec.score_warp.is_some() || spec.link_dev.is_some())
2008        && spec.y.len() >= BMS_FLEX_SPATIAL_OUTER_PILOT_ROW_THRESHOLD
2009        && effective_kappa_options.enabled;
2010    if flex_spatial_pilot_path {
2011        let marginal_terms = spatial_length_scale_term_indices(&spec.marginalspec);
2012        let logslope_terms = spatial_length_scale_term_indices(&spec.logslopespec);
2013        let marginal_updates = apply_spatial_anisotropy_pilot_initializer(
2014            data_view,
2015            &mut spec.marginalspec,
2016            &marginal_terms,
2017            effective_kappa_options.pilot_subsample_threshold,
2018            &effective_kappa_options,
2019        )
2020        .map_err(|error| error.to_string())?;
2021        let logslope_updates = apply_spatial_anisotropy_pilot_initializer(
2022            data_view,
2023            &mut spec.logslopespec,
2024            &logslope_terms,
2025            effective_kappa_options.pilot_subsample_threshold,
2026            &effective_kappa_options,
2027        )
2028        .map_err(|error| error.to_string())?;
2029        effective_kappa_options.enabled = false;
2030        log::info!(
2031            "[BMS spatial] n={} flex=true pilot_geometry_updates={} iterative_spatial_outer=false reason=large-flex-spatial-pilot",
2032            spec.y.len(),
2033            marginal_updates + logslope_updates,
2034        );
2035    }
2036    let (z_standardized, z_normalization) = standardize_latent_z_with_policy(
2037        &spec.z,
2038        &spec.weights,
2039        "bernoulli-marginal-slope",
2040        &spec.latent_z_policy,
2041    )?;
2042    spec.z = z_standardized;
2043    let sigma_learnable = matches!(
2044        &spec.frailty,
2045        FrailtySpec::GaussianShift {
2046            scale: FrailtyScale::Learned { .. }
2047        }
2048    );
2049    let initial_sigma = match &spec.frailty {
2050        FrailtySpec::GaussianShift {
2051            scale: FrailtyScale::Fixed { sigma },
2052        } => Some(*sigma),
2053        FrailtySpec::GaussianShift {
2054            scale: FrailtyScale::Learned { initial_sigma },
2055        } => Some(*initial_sigma),
2056        FrailtySpec::None => None,
2057        FrailtySpec::HazardMultiplier { .. } => {
2058            return Err(
2059                "internal: validate_spec should have rejected unsupported marginal-slope frailty"
2060                    .to_string(),
2061            );
2062        }
2063    };
2064    let probit_scale = probit_frailty_scale(initial_sigma);
2065    let (_raw_joint_designs, mut joint_specs) = build_term_collection_designs_and_freeze_joint(
2066        data_view,
2067        &[spec.marginalspec.clone(), spec.logslopespec.clone()],
2068    )
2069    .map_err(|e| e.to_string())?;
2070    let marginalspec_boot = joint_specs.remove(0);
2071    let logslopespec_boot = joint_specs.remove(0);
2072    // Rebuild the probe designs from the frozen `*_boot` specs so the probe's
2073    // penalty topology matches the topology produced by every other build path
2074    // in this optimization. The spatial optimizer's own bootstrap
2075    // (`build_term_collection_designs_and_freeze_joint(data, &[marginalspec_boot,
2076    // logslopespec_boot])` inside `optimize_spatial_length_scale_exact_joint`)
2077    // and every subsequent kappa-driven rebuild feed the basis builder the
2078    // captured `FrozenTransform` identifiability. Applying that captured
2079    // transform changes the exact coefficient chart of the penalty blocks.
2080    // Without this rebuild,
2081    // `marginal_penalty_count` / `logslope_design.penalties.len()` are taken
2082    // from the raw build but every subsequent evaluator measures the frozen
2083    // build, and `evaluate_custom_family_joint_hyper` refuses with a
2084    // "joint hyper rho dimension mismatch". Mirrors the CTN-side fix in
2085    // `fit_transformation_normal`.
2086    let (mut joint_designs, _) = build_term_collection_designs_and_freeze_joint(
2087        data_view,
2088        &[marginalspec_boot.clone(), logslopespec_boot.clone()],
2089    )
2090    .map_err(|e| format!("failed to rebuild frozen probe BMS joint designs: {e}"))?;
2091    let marginal_design = joint_designs.remove(0);
2092    let logslope_design = joint_designs.remove(0);
2093    spec.marginal_offset = marginal_design
2094        .compose_offset(spec.marginal_offset.view(), "BMS marginal block")
2095        .map_err(|error| error.to_string())?;
2096    spec.logslope_offset = logslope_design
2097        .compose_offset(spec.logslope_offset.view(), "BMS logslope block")
2098        .map_err(|error| error.to_string())?;
2099    // #905: the conditional `E[z|C]`/`Var(z|C)` Rao gate conditions on the
2100    // marginal-index span a(C) (= the marginal design columns), which is
2101    // exactly where the `b(C)·m(C)` leakage lives. It is engaged only on the
2102    // raw-z path (no CTN Stage-1 influence absorber); when an absorber is
2103    // active the conditional leakage is already absorbed (#461) and the
2104    // widened-marginal predict seam must not be perturbed by replacing z.
2105    let absorber_active = spec
2106        .score_influence_jacobian
2107        .as_ref()
2108        .is_some_and(|j| j.ncols() > 0);
2109    let conditioning_dense = if absorber_active {
2110        None
2111    } else {
2112        Some(
2113            marginal_design
2114                .design
2115                .try_to_dense_arc("bernoulli marginal-slope conditional latent-z gate")?,
2116        )
2117    };
2118    let (latent_measure, latent_z_calibration) = build_latent_measure_with_geometry(
2119        &spec.z,
2120        &spec.weights,
2121        &spec.latent_z_policy,
2122        conditioning_dense.as_ref().map(|d| d.view()),
2123    )?;
2124    if latent_measure.is_empirical() && sigma_learnable {
2125        return Err("empirical latent-measure marginal-slope calibration requires fixed GaussianShift sigma; learnable sigma derivatives must be fit under the standard-normal latent measure"
2126                    .to_string());
2127    }
2128
2129    let y = Arc::new(spec.y.clone());
2130    let weights = Arc::new(spec.weights.clone());
2131    // Apply rank-INT calibration to training z before any downstream
2132    // consumer (pooled probit baseline, term-collection designs, the
2133    // family's PIRLS loops) sees it. The calibration is persisted on the
2134    // fit result so prediction applies the identical monotone map.
2135    let z = match &latent_z_calibration {
2136        LatentMeasureCalibration::None => Arc::new(spec.z.clone()),
2137        LatentMeasureCalibration::RankInverseNormal(cal) => {
2138            Arc::new(cal.apply_to_training(&spec.z)?)
2139        }
2140        LatentMeasureCalibration::ConditionalLocationScale(cal) => {
2141            // ζ = (z − m(C))/√v(C) on the marginal-index span. The conditioning
2142            // block was built above (raw-z path only), so it is present here.
2143            let a_block = conditioning_dense.as_ref().ok_or_else(|| {
2144                "conditional latent calibration requires the marginal conditioning block"
2145                    .to_string()
2146            })?;
2147            Arc::new(cal.apply(spec.z.view(), a_block.view())?)
2148        }
2149    };
2150    let z_train = z.as_ref();
2151    let pilot_baseline = pooled_probit_baseline(&spec.y, z_train, &spec.weights)?;
2152    let baseline = (
2153        bernoulli_marginal_slope_eta_from_probability(
2154            &spec.base_link,
2155            normal_cdf(pilot_baseline.0),
2156            "bernoulli marginal-slope baseline link inversion",
2157        )?,
2158        pilot_baseline.1 / probit_scale,
2159    );
2160
2161    // Score-warp basis construction is β-independent (identifiability is
2162    // provided by the smoothness-null-space drop on the basis transform,
2163    // not by a data-distribution moment anchor at the rigid-pilot η₀), so
2164    // the standard-normal and empirical latent-measure branches build the
2165    // same block. There is no row-weight pilot to thread into the basis;
2166    // the latent-measure split is enforced upstream via the empirical
2167    // intercept solve in `build_row_exact_context_with_stats`, not in the
2168    // deviation basis.
2169    // Score-warp basis is built first, then immediately reparameterised
2170    // against the parametric span (marginal + logslope columns at the n
2171    // training rows) so its column span is orthogonal to span(X_marginal,
2172    // X_logslope) by construction. This is the first half of the joint-
2173    // design identifiability invariant; the second half (link-deviation
2174    // orthogonalised against parametric + the now-reparameterised score-
2175    // warp) runs inside the link-deviation closure below. Together they
2176    // ensure `[X_marginal | X_logslope | Φ_score_warp · T_sw |
2177    // Φ_link_dev · T_lw]` has full numerical column rank, structurally
2178    // bounding `σ_min(joint H + S) ≥ λ_min(S) > 0` regardless of how β
2179    // drifts the linear predictor distribution during PIRLS.
2180    // Cross-block W-metric pilot. The joint penalised Hessian during PIRLS
2181    // uses the probit-style data Hessian row metric
2182    //
2183    //   W_pirls[i] = spec.weights[i] · φ(η_i)² / (μ_i·(1−μ_i))
2184    //
2185    // which is the canonical IRLS row weight. The cross-block
2186    // orthogonalisation below must use this metric (not uniform
2187    // spec.weights) so that `Aᵀ W C̃ = 0` holds in the same inner product
2188    // the joint Hessian sees — otherwise A and C̃ are merely Euclidean-
2189    // orthogonal, `Aᵀ W_pirls C̃ ≠ 0`, the joint Hessian carries a near-
2190    // null direction along the W-metric alias, and REML can drive the
2191    // flex block's λ small enough that the alias direction's joint
2192    // Hessian eigenvalue collapses. β then runs away along the alias
2193    // (manifest as `rho≈2.0`, constant `step_inf`, growing `beta_inf`
2194    // during PIRLS, and the inner solve hitting `inner_max_cycles`
2195    // without satisfying the KKT residual).
2196    //
2197    // Use the rigid pooled-probit pilot η for score-warp (its basis is
2198    // β-independent in z, so the rigid pilot suffices) and the one-GN-
2199    // stepped pilot η for link-deviation (its basis is evaluated at the
2200    // same eta_pilot used here, so the orthogonalisation metric matches
2201    // the basis evaluation point exactly). Both are β-independent so the
2202    // orthogonalisation remains a one-shot construction-time step.
2203    let rigid_pilot_eta = rigid_pooled_probit_pilot_eta(
2204        &spec.base_link,
2205        z_train,
2206        &spec.marginal_offset,
2207        &spec.logslope_offset,
2208        baseline.0,
2209        baseline.1,
2210        probit_scale,
2211    )?;
2212    let cross_block_pilot_w_score_warp =
2213        pilot_irls_hessian_row_metric_at_eta(&rigid_pilot_eta, &spec.weights);
2214
2215    // Absorbed Stage-1 influence columns (#461, design §3). When the workflow
2216    // chained a CTN Stage-1 into this marginal-slope fit, `spec.score_influence_
2217    // jacobian` carries the out-of-fold `J = ∂z/∂θ₁`; the realized leakage
2218    // directions `Z_infl = diag(s_f·β̂₀)·J` are residualised against the fitted
2219    // marginal+logslope target span and appended to the additive marginal-index
2220    // block as a fixed-ridge absorber, so the joint penalised solve makes the
2221    // (α,β) score orthogonal to the remaining nuisance span without letting the
2222    // absorber compete for identifiable β(x) signal. `None` ⇒ raw z, and the
2223    // free score_warp spline below is the x-free-column fallback. β̂₀(x_i) is
2224    // the rigid-pilot logslope `baseline.1 + logslope_offset[i]`; s_f =
2225    // probit_scale.
2226    let influence_columns = if let Some(jac) = spec
2227        .score_influence_jacobian
2228        .as_ref()
2229        .filter(|j| j.ncols() > 0)
2230    {
2231        let protected_design = DesignMatrix::hstack(vec![
2232            marginal_design.design.clone(),
2233            logslope_design.design.clone(),
2234        ])
2235        .map_err(|e| {
2236            format!(
2237                "bernoulli marginal-slope influence-block protected projection stack failed to concatenate marginal + logslope design: {e}"
2238            )
2239        })?;
2240        let protected_dense_for_proj = protected_design
2241            .try_to_dense_arc("bernoulli marginal-slope influence-block protected projection")?;
2242        let protected_dense = protected_dense_for_proj.as_ref();
2243        if jac.nrows() != protected_dense.nrows() {
2244            return Err(format!(
2245                "influence block: Jacobian has {} rows, protected design has {}",
2246                jac.nrows(),
2247                protected_dense.nrows()
2248            ));
2249        }
2250        // Z̃ = residualize(diag(s_f·β̂₀)·J) against the fitted target span in
2251        // the rigid-pilot W-metric.  For BMS the absorbed columns are installed
2252        // in the same additive predictor as the marginal surface; if we protect
2253        // only M, any component of Z_infl aligned with the logslope design G can
2254        // be assigned to the fixed-ridge absorber by the joint solve, erasing
2255        // genuine β(x) heterogeneity.  Projecting out [M | G] keeps the nuisance
2256        // absorber orthogonal to both parametric target surfaces while still
2257        // absorbing Stage-1 leakage directions outside that identifiable target
2258        // span. β̂₀(x_i) = baseline.1 + logslope_offset[i]; s_f = probit_scale;
2259        // W = the rigid-pilot PIRLS row metric.
2260        let rigid_logslope_at_rows = &spec.logslope_offset + baseline.1;
2261        let residualized = crate::marginal_slope_orthogonal::residualized_influence_block(
2262            jac,
2263            z_train,
2264            &rigid_logslope_at_rows,
2265            probit_scale,
2266            protected_dense.view(),
2267            &cross_block_pilot_w_score_warp,
2268        )?;
2269        Some(residualized)
2270    } else {
2271        None
2272    };
2273    let mut cross_block_warnings: Vec<CrossBlockIdentifiabilityWarning> = Vec::new();
2274    let score_warp_prepared = if let Some(cfg) = spec.score_warp.as_ref() {
2275        use super::deviation_runtime::ParametricAnchorBlock;
2276        let mut prepared = build_score_warp_deviation_block_from_seed(z_train, cfg)?;
2277        // `install_compiled_flex_block_into_runtime` now delegates
2278        // its math body to `identifiability::families::compiler::compile` (commit
2279        // 4e20b8dc8); the prior Phase-4a shadow compile here was a
2280        // duplicate of that internal call and has been removed.
2281        let outcome = install_compiled_flex_block_into_runtime(
2282            &mut prepared,
2283            z_train,
2284            cfg,
2285            &[
2286                (&marginal_design.design, ParametricAnchorBlock::Marginal),
2287                (&logslope_design.design, ParametricAnchorBlock::Logslope),
2288            ],
2289            &[],
2290            &cross_block_pilot_w_score_warp,
2291        )?;
2292        match outcome {
2293            FlexCompileOutcome::Reparameterised => Some(prepared),
2294            FlexCompileOutcome::FullyAliased { reason } => {
2295                // Record via the structured channel. Keep the original
2296                // (non-compiled) design so the unified audit sees score_warp_dev
2297                // and attributes the drop via dropped_columns (gauge_priority=80
2298                // is below marginal=150 / logslope=120, so RRQR correctly
2299                // demotes score_warp_dev when it aliases those blocks).
2300                cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2301                    candidate_label: "score_warp",
2302                    anchor_summary: "marginal+logslope".to_string(),
2303                    reason,
2304                });
2305                Some(prepared)
2306            }
2307        }
2308    } else {
2309        None
2310    };
2311    // Build the link-deviation block. The basis lives in η-space, and at
2312    // PIRLS time `runtime.design(η_current)` is re-evaluated at the
2313    // current β-dependent η, so the basis is genuinely β-dependent during
2314    // optimisation. The construction-time seed is used only for (a) knot
2315    // placement in η-space and (b) the cross-block identifiability check
2316    // that computes the basis-space transform `T` orthogonalising the
2317    // candidate against the parametric and score-warp anchors at training
2318    // rows.
2319    //
2320    // Using the rigid pooled probit pilot directly (`q0 = a₀·√(…) + s_f·
2321    // b₀·z`) is structurally degenerate: with zero per-row offsets it is
2322    // affine in z, so a degree-3 I-spline of `q0` spans the same column
2323    // space at training rows as a degree-3 I-spline of z, and the cross-
2324    // block check finds the candidate fully aliased by the score-warp
2325    // anchor even though at any non-rigid β the link-deviation carries
2326    // PC/age structure the score-warp cannot represent.
2327    //
2328    // Instead, seed both knot placement and the orthogonalisation pivot at
2329    // a non-rigid pilot η computed via one probit Gauss-Newton step from
2330    // the rigid pilot onto the full marginal design (see
2331    // `pilot_eta_for_link_dev_orthogonalisation`). The pilot is row-varying
2332    // in PCs/age and the resulting `T` drops only directions aliased
2333    // across all β. The score-warp basis at training rows is also threaded
2334    // in as a flex anchor when active so the kept directions are jointly
2335    // orthogonal to parametric ⊕ score-warp.
2336    let link_dev_prepared = if let Some(cfg) = spec.link_dev.as_ref() {
2337        let eta_pilot = pilot_eta_for_link_dev_orthogonalisation(
2338            &spec.base_link,
2339            &spec.y,
2340            z_train,
2341            &spec.weights,
2342            &marginal_design.design,
2343            &spec.marginal_offset,
2344            &spec.logslope_offset,
2345            baseline.0,
2346            baseline.1,
2347            probit_scale,
2348        )?;
2349        let link_dev_seed = padded_deviation_seed(&eta_pilot, 1.0, 0.5);
2350        let mut prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
2351            &link_dev_seed,
2352            &eta_pilot,
2353            cfg,
2354        )?;
2355        // Cross-block identifiability for the link-deviation basis. The
2356        // anchor union covers BOTH possible aliasing channels:
2357        //
2358        //  - Parametric: location and logslope designs evaluated at the n
2359        //    training rows. Columns of `Φ_link_dev(q0)` that reproduce
2360        //    parametric features become null-direction targets in the
2361        //    joint penalised Hessian since `S_link_dev` has no mass on
2362        //    them.
2363        //
2364        //  - Score-warp (when active): the now-reparameterised score-warp
2365        //    basis, also evaluated at training rows. Both flex bases are
2366        //    cubic I-spline cubic combinations of an η-pilot scalar, and
2367        //    even with each block's own smoothness-null-space drop their
2368        //    column spans can still overlap inside the orthogonal
2369        //    complement of `{1, η_pilot}`.
2370        //
2371        // After the orthogonalisation, `[X_marginal | X_logslope |
2372        // Φ_score_warp · T_sw | Φ_link_dev · T_lw]` has full numerical
2373        // column rank at training rows, so `σ_min(joint H+S) ≥ λ_min(S)
2374        // > 0` for every β. This is the standard GAM `gam.side`
2375        // convention generalised to multi-anchor unions (mgcv applies it
2376        // sequentially across smooths sharing a covariate).
2377        // When `install_compiled_flex_block_into_runtime`
2378        // reparameterised the score-warp runtime against the parametric
2379        // anchor union (marginal + logslope), it installed an
2380        // `anchor_residual` and cached the training-row parametric
2381        // anchor matrix on the runtime. `runtime.design()` on a
2382        // residualised runtime returns the *raw* basis evaluation,
2383        // which `assert`s the caller hasn't conflated with the
2384        // reparameterised basis — we want the reparameterised one
2385        // here, so go through `design_at_training_with_residual` so
2386        // the cached anchor rows are folded in. For score-warp
2387        // configurations where reparameterisation was a no-op (no
2388        // residual installed) the same call falls back to the raw
2389        // `design()` path, so the residual-vs-no-residual branches
2390        // converge on the right matrix.
2391        let score_warp_anchor_design = score_warp_prepared
2392            .as_ref()
2393            .map(|sw| sw.runtime.design_at_training_with_residual(z_train))
2394            .transpose()?;
2395        use super::deviation_runtime::ParametricAnchorBlock;
2396        let parametric_anchors: [(&DesignMatrix, ParametricAnchorBlock); 2] = [
2397            (&marginal_design.design, ParametricAnchorBlock::Marginal),
2398            (&logslope_design.design, ParametricAnchorBlock::Logslope),
2399        ];
2400        let flex_anchor_slot: Option<&Array2<f64>> = score_warp_anchor_design.as_ref();
2401        let flex_anchors: Vec<&Array2<f64>> = flex_anchor_slot.into_iter().collect();
2402        // W-metric for link-deviation orthogonalisation: same IRLS-style
2403        // probit Hessian row weight as the score-warp path, but evaluated at
2404        // `eta_pilot` (the one-GN-stepped pilot at which the link-dev basis
2405        // itself is anchored).
2406        let cross_block_pilot_w_link_dev =
2407            pilot_irls_hessian_row_metric_at_eta(&eta_pilot, &spec.weights);
2408        let outcome = install_compiled_flex_block_into_runtime(
2409            &mut prepared,
2410            &eta_pilot,
2411            cfg,
2412            &parametric_anchors,
2413            &flex_anchors,
2414            &cross_block_pilot_w_link_dev,
2415        )?;
2416        match outcome {
2417            FlexCompileOutcome::Reparameterised => Some(prepared),
2418            FlexCompileOutcome::FullyAliased { reason } => {
2419                // Record via the structured channel. Keep the original
2420                // (non-compiled) design so the unified audit sees link_dev
2421                // and attributes the drop via dropped_columns (gauge_priority=60
2422                // is below all parametric blocks so RRQR correctly demotes
2423                // link_dev when it aliases marginal / logslope / score_warp).
2424                cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2425                    candidate_label: "link_deviation",
2426                    anchor_summary: "marginal+logslope+score_warp".to_string(),
2427                    reason,
2428                });
2429                Some(prepared)
2430            }
2431        }
2432    } else {
2433        None
2434    };
2435    let extra_rho0 = {
2436        let mut out = Vec::new();
2437        if let Some(ref prepared) = score_warp_prepared {
2438            out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2439        }
2440        if let Some(ref prepared) = link_dev_prepared {
2441            out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2442        }
2443        out
2444    };
2445    // Reduced-basis orthogonalisation of the logslope design through the BMS
2446    // family's OWN internal `logslope_design` geometry (robust cure for the
2447    // marginal↔logslope structural confound). Robustness is unconditional, so we
2448    // always reparameterize the logslope coordinate space to a full-rank reduced
2449    // basis `T` whose effective weighted columns are W-orthogonal to the marginal
2450    // span at the rigid pilot — removing the rank-soft confounded direction the
2451    // former pinned overlap ridge merely penalised. The transform is
2452    // β/ρ-independent (pilot geometry only), so it is a one-shot construction-
2453    // time map applied to every per-iteration logslope design inside
2454    // `build_blocks` / `make_family`, and inverted at fit-result assembly so the
2455    // reported logslope β is in the original basis. `None` ⇒ nothing to reduce
2456    // (no rank-soft confounded direction) ⇒ raw design used everywhere.
2457    let logslope_reduced_reparam: Option<ReducedLogslopeReparam> = build_reduced_logslope_reparam(
2458        &marginal_design,
2459        &logslope_design,
2460        z.as_ref(),
2461        &cross_block_pilot_w_score_warp,
2462        &spec.marginal_offset,
2463        &spec.logslope_offset,
2464        baseline.0,
2465        baseline.1,
2466        probit_scale,
2467    )?;
2468    // Apply the reduced reparam to a logslope `TermCollectionDesign`, or return
2469    // the raw design clone when the reparam is absent (flag off / nothing to
2470    // reduce). Used by both `build_blocks` and `make_family` so the family's
2471    // internal design, the block design, β width, jacobian, penalty, and the
2472    // `validate_exact_block_state_shapes` check all agree at the reduced width.
2473    let reduce_logslope_design =
2474        |logslope_design: &TermCollectionDesign| -> Result<TermCollectionDesign, String> {
2475            match logslope_reduced_reparam.as_ref() {
2476                Some(reparam) => reparameterize_logslope_design_reduced(logslope_design, reparam),
2477                None => Ok(logslope_design.clone()),
2478            }
2479        };
2480
2481    // With the #461 influence absorber active, the marginal block carries one
2482    // extra REML-learned ρ coordinate (the absorber ridge precision), seeded
2483    // at the ln(n) leakage scale and clamped into the outer ρ box.
2484    let absorber_slots = usize::from(influence_columns.is_some());
2485    let absorber_rho0 = influence_columns
2486        .as_ref()
2487        .map(|_| influence_absorber_log_lambda(spec.z.len()).clamp(-12.0, 12.0));
2488    let marginal_penalty_count = marginal_design.penalties.len() + absorber_slots;
2489    let setup = joint_setup(
2490        data_view,
2491        &marginalspec_boot,
2492        &logslopespec_boot,
2493        marginal_penalty_count,
2494        logslope_design.penalties.len(),
2495        absorber_rho0,
2496        &extra_rho0,
2497        &effective_kappa_options,
2498    )
2499    .map_err(|error| error.to_string())?;
2500    let setup = if sigma_learnable {
2501        setup.with_auxiliary(
2502            Array1::from_vec(vec![initial_sigma.expect("learnable sigma seed").ln()]),
2503            Array1::from_vec(vec![0.01_f64.ln()]),
2504            Array1::from_vec(vec![5.0_f64.ln()]),
2505        )
2506    } else {
2507        setup
2508    };
2509    let final_sigma_cell = std::cell::Cell::new(initial_sigma);
2510    let exact_mode_branch = RefCell::new(ExactCoefficientModeBranch::default());
2511    let runaway_error = RefCell::new(None::<String>);
2512    // Outer ρ-cache β-seed staging slot. On a cache hit the spatial-joint
2513    // optimizer invokes `seed_inner_beta_fn` before the first eval at the
2514    // restored ρ: per-block column widths aren't known until the first
2515    // `build_blocks(rho, …)` runs, so we stash the flat β here and the eval
2516    // closures promote it into the deterministic coefficient-mode branch on
2517    // their first invocation.
2518    let pending_beta_seed = RefCell::new(None::<Array1<f64>>);
2519    let hints = RefCell::new(ThetaHints::default());
2520    let score_warp_runtime = score_warp_prepared.as_ref().map(|p| p.runtime.clone());
2521    let link_dev_runtime = link_dev_prepared.as_ref().map(|p| p.runtime.clone());
2522
2523    let build_blocks = |rho: &Array1<f64>,
2524                        marginal_design: &TermCollectionDesign,
2525                        logslope_design: &TermCollectionDesign|
2526     -> Result<Vec<ParameterBlockSpec>, String> {
2527        let hints = hints.borrow();
2528        let mut cursor = 0usize;
2529        // Reduced-basis orthogonalisation: replace the per-iteration logslope
2530        // design with its full-rank reduced reparameterization `G·T` (flag ON);
2531        // a no-op clone when off. The reduced design carries the SAME number of
2532        // penalties (each S → Tᵀ S T), so the `rho_logslope` slice width below
2533        // is unchanged. Every consumer (marginal jacobian's c_i, logslope
2534        // blockspec design/β/penalty/jacobian) now agrees at the reduced width.
2535        let logslope_design_reduced = reduce_logslope_design(logslope_design)?;
2536        let logslope_design = &logslope_design_reduced;
2537        // The marginal slice carries the genuine smooth penalties plus, when
2538        // the #461 influence absorber is active, one TRAILING coordinate for
2539        // the absorber ridge — its precision is REML-learned like every other
2540        // penalty (seeded at the ln(n) leakage scale by `joint_setup`).
2541        let marginal_rho_len = marginal_design.penalties.len() + absorber_slots;
2542        let rho_marginal = rho.slice(s![cursor..cursor + marginal_rho_len]).to_owned();
2543        cursor += marginal_rho_len;
2544        let rho_logslope = rho
2545            .slice(s![cursor..cursor + logslope_design.penalties.len()])
2546            .to_owned();
2547        cursor += logslope_design.penalties.len();
2548        let p_m = marginal_design.design.ncols()
2549            + influence_columns.as_ref().map(|z| z.ncols()).unwrap_or(0);
2550        let mut blocks = vec![
2551            build_marginal_blockspec_bms(
2552                marginal_design,
2553                baseline.0,
2554                &spec.marginal_offset,
2555                rho_marginal,
2556                hints.marginal_beta.clone(),
2557                logslope_design,
2558                &spec.logslope_offset,
2559                baseline.1,
2560                p_m,
2561                influence_columns.as_ref(),
2562            )?,
2563            build_logslope_blockspec_bms(
2564                logslope_design,
2565                baseline.1,
2566                &spec.logslope_offset,
2567                rho_logslope,
2568                hints.logslope_beta.clone(),
2569                marginal_design,
2570                &spec.marginal_offset,
2571                baseline.0,
2572                Arc::clone(&z),
2573                p_m,
2574                influence_columns.as_ref(),
2575            )?,
2576        ];
2577        push_deviation_aux_blockspecs(
2578            &mut blocks,
2579            rho,
2580            &mut cursor,
2581            score_warp_prepared.as_ref(),
2582            link_dev_prepared.as_ref(),
2583            hints.score_warp_beta.clone(),
2584            hints.link_dev_beta.clone(),
2585        )?;
2586        Ok(blocks)
2587    };
2588
2589    let intercept_warm_starts = new_intercept_warm_start_cache(y.len());
2590    let cell_moment_lru = new_cell_moment_lru_cache(policy);
2591    let cell_moment_cache_stats = new_cell_moment_cache_stats();
2592    let make_family = |marginal_design: &TermCollectionDesign,
2593                       logslope_design: &TermCollectionDesign,
2594                       sigma: Option<f64>|
2595     -> BernoulliMarginalSlopeFamily {
2596        // The kernel reads the marginal index from a matched (self.marginal_
2597        // design, β_m) pair. When the Stage-1 influence absorber is active the
2598        // marginal β is widened to [β_m; γ], so the family's marginal design
2599        // MUST be the widened [M | Z̃] for every per-row projection to slice
2600        // correctly (#461). With no absorber it is the raw design unchanged.
2601        let kernel_marginal_design = match influence_columns.as_ref() {
2602            Some(z_infl) => {
2603                let raw = marginal_design
2604                    .design
2605                    .try_to_dense_arc("make_family::widened-marginal")
2606                    .expect("dense marginal design for influence widening");
2607                let widened = widen_marginal_dense_with_influence(&raw, Some(z_infl))
2608                    .expect("widen marginal design with influence columns");
2609                DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
2610                    (*widened).clone(),
2611                ))
2612            }
2613            None => marginal_design.design.clone(),
2614        };
2615        // The family's row kernel reconstructs η_logslope = G·β_s and the
2616        // logslope Jacobian factor_i·G_i from this matched (logslope_design,
2617        // β_s) pair, so it MUST be the SAME reduced design `G·T` the block specs
2618        // fit against — otherwise β_s (reduced width) and the family design
2619        // (full width) desync. A no-op clone when the reparam is absent.
2620        let kernel_logslope_design = reduce_logslope_design(logslope_design)
2621            .expect("reduce logslope design for family construction")
2622            .design;
2623        BernoulliMarginalSlopeFamily {
2624            y: Arc::clone(&y),
2625            weights: Arc::clone(&weights),
2626            z: Arc::clone(&z),
2627            latent_measure: latent_measure.clone(),
2628            gaussian_frailty_sd: sigma,
2629            base_link: spec.base_link.clone(),
2630            marginal_design: kernel_marginal_design,
2631            logslope_design: kernel_logslope_design,
2632            score_warp: score_warp_runtime.clone(),
2633            link_dev: link_dev_runtime.clone(),
2634            policy: policy.clone(),
2635            cell_moment_lru: Arc::clone(&cell_moment_lru),
2636            cell_moment_cache_stats: Arc::clone(&cell_moment_cache_stats),
2637            intercept_warm_starts: Some(Arc::clone(&intercept_warm_starts)),
2638            auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2639            auto_subsample_last_rho: Arc::new(Mutex::new(None)),
2640        }
2641    };
2642
2643    let marginal_terms = spatial_length_scale_term_indices(&marginalspec_boot);
2644    let logslope_terms = spatial_length_scale_term_indices(&logslopespec_boot);
2645    let marginal_has_spatial = !marginal_terms.is_empty();
2646    let logslope_has_spatial = !logslope_terms.is_empty();
2647    let analytic_joint_derivatives_available =
2648        marginal_has_spatial || logslope_has_spatial || setup.log_kappa_dim() == 0;
2649    if setup.log_kappa_dim() > 0 && !analytic_joint_derivatives_available {
2650        return Err("exact bernoulli marginal-slope spatial optimization requires analytic joint psi derivatives"
2651                    .to_string());
2652    }
2653    let initial_rho = setup.theta0().slice(s![..setup.rho_dim()]).to_owned();
2654    let initial_blocks = build_blocks(&initial_rho, &marginal_design, &logslope_design)?;
2655    let initial_family = make_family(&marginal_design, &logslope_design, initial_sigma);
2656    let (joint_gradient, joint_hessian) =
2657        custom_family_outer_derivatives(&initial_family, &initial_blocks, options);
2658    let analytic_joint_gradient_available = analytic_joint_derivatives_available
2659        && matches!(joint_gradient, gam_problem::Derivative::Analytic);
2660    // Keep the analytic outer Hessian advertised at large scale. The
2661    // row-tensor terms below are represented through block-local
2662    // `HyperOperator`s and cached exact-Hessian workspaces, so ARC/trust-region
2663    // can consume exact HVPs without falling back to BFGS merely because the
2664    // realized problem is large.
2665    let analytic_joint_hessian_available =
2666        analytic_joint_derivatives_available && joint_hessian.is_analytic();
2667    let kappa_options_ref: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
2668    let sigma_from_theta = |theta: &Array1<f64>| -> Option<f64> {
2669        if sigma_learnable {
2670            Some(theta[setup.rho_dim() + setup.log_kappa_dim()].exp())
2671        } else {
2672            initial_sigma
2673        }
2674    };
2675    let hyper_layout_cache = RefCell::new(
2676        None::<(
2677            Array1<f64>,
2678            crate::custom_family::SharedCustomFamilyHyperLayout,
2679        )>,
2680    );
2681    let theta_matches = |left: &Array1<f64>, right: &Array1<f64>| -> bool {
2682        left.len() == right.len()
2683            && left
2684                .iter()
2685                .zip(right.iter())
2686                .all(|(lhs, rhs)| lhs.to_bits() == rhs.to_bits())
2687    };
2688    let get_hyper_layout = |theta: &Array1<f64>,
2689                            specs: &[TermCollectionSpec],
2690                            designs: &[TermCollectionDesign]|
2691     -> Result<crate::custom_family::SharedCustomFamilyHyperLayout, String> {
2692        if let Some((cached_theta, cached_layout)) = hyper_layout_cache.borrow().as_ref()
2693            && theta_matches(cached_theta, theta)
2694        {
2695            return Ok(Arc::clone(cached_layout));
2696        }
2697
2698        let built = |specs: &[TermCollectionSpec],
2699                     designs: &[TermCollectionDesign]|
2700         -> Result<
2701            Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>>,
2702            String,
2703        > {
2704            let marginal_psi_derivs = if marginal_has_spatial {
2705                build_block_spatial_psi_derivatives(data_view, &specs[0], &designs[0])?.ok_or_else(
2706                    || {
2707                        "bernoulli marginal-slope: marginal block has spatial terms \
2708                         but spatial psi derivatives are unavailable"
2709                            .to_string()
2710                    },
2711                )?
2712            } else {
2713                Vec::new()
2714            };
2715            let logslope_psi_derivs = if logslope_has_spatial {
2716                let built = if let Some(reparam) = logslope_reduced_reparam.as_ref() {
2717                    let transform =
2718                        CoefficientSpatialPsiBlockTransform::new(&reparam.transform)?;
2719                    build_block_spatial_psi_derivatives_with_transform(
2720                        data_view,
2721                        &specs[1],
2722                        &designs[1],
2723                        &transform,
2724                    )?
2725                } else {
2726                    build_block_spatial_psi_derivatives(data_view, &specs[1], &designs[1])?
2727                };
2728                built.ok_or_else(|| {
2729                    "bernoulli marginal-slope: logslope block has spatial terms \
2730                     but spatial psi derivatives are unavailable"
2731                        .to_string()
2732                })?
2733            } else {
2734                Vec::new()
2735            };
2736            let mut derivative_blocks = vec![marginal_psi_derivs, logslope_psi_derivs];
2737            if score_warp_runtime.is_some() {
2738                derivative_blocks.push(Vec::new());
2739            }
2740            if link_dev_runtime.is_some() {
2741                derivative_blocks.push(Vec::new());
2742            }
2743            Ok(derivative_blocks)
2744        }(specs, designs)?;
2745        let family_axes = if sigma_learnable { vec![0] } else { Vec::new() };
2746        let hyper_values = theta.slice(s![setup.rho_dim()..]).to_owned();
2747        let layout = Arc::new(crate::custom_family::CustomFamilyHyperLayout::new(
2748            built,
2749            family_axes,
2750            hyper_values,
2751        )?);
2752        hyper_layout_cache.replace(Some((theta.clone(), Arc::clone(&layout))));
2753        Ok(layout)
2754    };
2755
2756    // Bernoulli marginal-slope is a multi-block family with β-dependent
2757    // joint Hessian: EFS/HybridEFS fixed-point structural invariant fails,
2758    // so we disable fixed-point at plan time rather than burning cycles on
2759    // a stalled first attempt that silently falls back.
2760    let outer_policy = {
2761        let psi_dim = setup.theta0().len() - setup.rho_dim();
2762        initial_family.outer_derivative_policy(&initial_blocks, psi_dim, options)
2763    };
2764    let exact_spatial_outer_tol = kappa_options_ref.rel_tol.max(EXACT_SPATIAL_OUTER_TOL_FLOOR);
2765    let solved = optimize_spatial_length_scale_exact_joint(
2766        data_view,
2767        &[marginalspec_boot.clone(), logslopespec_boot.clone()],
2768        &[marginal_terms.clone(), logslope_terms.clone()],
2769        kappa_options_ref,
2770        &setup,
2771        gam_solve::seeding::SeedRiskProfile::GeneralizedLinear,
2772        analytic_joint_gradient_available,
2773        analytic_joint_hessian_available,
2774        true,
2775        None,
2776        outer_policy,
2777        |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign], provenance| {
2778            if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2779                return Err(err);
2780            }
2781            assert_eq!(
2782                specs.len(),
2783                designs.len(),
2784                "spatial joint optimizer must supply one spec per design",
2785            );
2786            let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2787            let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2788            let sigma = sigma_from_theta(theta);
2789            final_sigma_cell.set(sigma);
2790            let family = make_family(&designs[0], &designs[1], sigma);
2791            let fit = match provenance {
2792                SpatialFitProvenance::NoOuterOptimization => {
2793                    inner_fit(&family, &blocks, options)?
2794                }
2795                SpatialFitProvenance::Certified { outer, mode } => {
2796                    inner_fit_from_certified_outer(
2797                        &family, &blocks, options, mode, theta, outer,
2798                    )?
2799                }
2800            };
2801            if let Some(block) = fit.block_states.first()
2802                && let Some(err) = bernoulli_marginal_slope_runaway_error_from_beta(
2803                    block.beta.view(),
2804                    &designs[0],
2805                    &specs[0],
2806                    true,
2807                    "final fit",
2808                )
2809            {
2810                runaway_error.replace(Some(err.clone()));
2811                return Err(err);
2812            }
2813            let mut hints_mut = hints.borrow_mut();
2814            let mut bidx = 0usize;
2815            if let Some(block) = fit.block_states.get(bidx) {
2816                hints_mut.marginal_beta = Some(block.beta.clone());
2817            }
2818            bidx += 1;
2819            if let Some(block) = fit.block_states.get(bidx) {
2820                hints_mut.logslope_beta = Some(block.beta.clone());
2821            }
2822            bidx += 1;
2823            if score_warp_prepared.is_some() {
2824                if let Some(block) = fit.block_states.get(bidx) {
2825                    hints_mut.score_warp_beta = Some(block.beta.clone());
2826                }
2827                bidx += 1;
2828            }
2829            if link_dev_prepared.is_some()
2830                && let Some(block) = fit.block_states.get(bidx)
2831            {
2832                hints_mut.link_dev_beta = Some(block.beta.clone());
2833            }
2834            Ok(fit)
2835        },
2836        |theta,
2837         specs: &[TermCollectionSpec],
2838         designs: &[TermCollectionDesign],
2839         eval_mode,
2840         row_set: &crate::row_kernel::RowSet| {
2841            if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2842                return Err(err);
2843            }
2844            use gam_problem::EvalMode;
2845            // One-shot row-measure waypoint. This closure runs on EVERY outer
2846            // objective evaluation (value/gradient/Hessian probes, line-search
2847            // cost-only probes, EFS evals), so an unconditional per-eval line
2848            // floods the biobank fit log with thousands of near-identical
2849            // entries. The bridge already emits a timed `[STAGE] outer eval`
2850            // marker per eval; this one records the row-measure exactly once.
2851            static BMS_OUTER_EVAL_ROWSET_LOGGED: std::sync::Once = std::sync::Once::new();
2852            BMS_OUTER_EVAL_ROWSET_LOGGED.call_once(|| {
2853                let row_set_rows = match row_set {
2854                    crate::row_kernel::RowSet::All => spec.y.len(),
2855                    crate::row_kernel::RowSet::Subsample { rows, .. } => rows.len(),
2856                };
2857                log::debug!(
2858                    "[BMS exact outer eval] mode={eval_mode:?} row_set_rows={row_set_rows}"
2859                );
2860            });
2861            let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2862            let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2863            // Promote a staged β seed (deposited by the outer ρ-cache hit
2864            // before any eval ran) into the deterministic mode branch now that
2865            // we know the per-block widths from the freshly built blocks.
2866            if let Some(beta_seed) = pending_beta_seed.borrow_mut().take() {
2867                let widths: Vec<usize> = blocks.iter().map(|b| b.design.ncols()).collect();
2868                match CustomFamilyWarmStart::from_cached_beta(&widths, &beta_seed) {
2869                    Ok(ws) => {
2870                        if !exact_mode_branch.borrow_mut().install_seed(ws) {
2871                            log::debug!(
2872                                "[BMS] ignored a late outer-cache coefficient seed after the exact mode branch froze"
2873                            );
2874                        }
2875                    }
2876                    Err(e) => {
2877                        log::warn!(
2878                            "[BMS] outer ρ-cache β-warm-start rejected: {e}; falling back to cold β"
2879                        );
2880                    }
2881                }
2882            }
2883            let sigma = sigma_from_theta(theta);
2884            final_sigma_cell.set(sigma);
2885            let family = make_family(&designs[0], &designs[1], sigma);
2886            let hyper_layout = get_hyper_layout(theta, specs, designs)?;
2887            // Downgrade to ValueAndGradient when the caller asks for a
2888            // Hessian we can't provide; preserve ValueOnly probes for
2889            // line-search cost-only evaluation.
2890            let effective_mode = match eval_mode {
2891                EvalMode::ValueGradientHessian if !analytic_joint_hessian_available => {
2892                    EvalMode::ValueAndGradient
2893                }
2894                other => other,
2895            };
2896            let tolerance_options =
2897                joint_hyper_options_for_outer_tolerance(options, exact_spatial_outer_tol);
2898            let eval_options = crate::outer_subsample::exact_outer_options_for_row_set(
2899                &tolerance_options,
2900                row_set,
2901            );
2902            let (froze, candidates) = exact_mode_branch
2903                .borrow_mut()
2904                .candidates(effective_mode, &rho);
2905            if froze {
2906                log::info!(
2907                    "[BMS] froze deterministic exact coefficient-mode branch at the first derivative-bearing outer evaluation"
2908                );
2909            }
2910            let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
2911                &family,
2912                &blocks,
2913                &eval_options,
2914                &rho,
2915                hyper_layout,
2916                &candidates,
2917                effective_mode,
2918            )?;
2919            if let Some(err) = bernoulli_marginal_slope_runaway_error(
2920                &selection.result.warm_start,
2921                &designs[0],
2922                &specs[0],
2923                selection.result.inner_converged,
2924                "exact outer evaluation",
2925            ) {
2926                runaway_error.replace(Some(err.clone()));
2927                return Err(err);
2928            }
2929            exact_mode_branch
2930                .borrow_mut()
2931                .record_value(eval_mode, selection.result.warm_start.clone());
2932            if !selection.result.inner_converged {
2933                return Err(
2934                    "exact bernoulli marginal-slope inner solve did not converge".to_string(),
2935                );
2936            }
2937            if matches!(eval_mode, EvalMode::ValueGradientHessian)
2938                && analytic_joint_hessian_available
2939                && !selection.result.outer_hessian.is_analytic()
2940            {
2941                return Err("exact bernoulli marginal-slope joint [rho, psi] objective did not return an outer Hessian"
2942                            .to_string());
2943            }
2944            Ok(ExactJointEvaluation {
2945                objective: selection.result.objective,
2946                gradient: selection.result.gradient.clone(),
2947                hessian: selection.result.outer_hessian.clone(),
2948                mode: selection,
2949            })
2950        },
2951        |_theta, _specs, _designs, _row_set| {
2952            Err::<ExactJointEfsEvaluation<CustomFamilyJointHyperModeSelection>, String>(
2953                "bernoulli marginal-slope EFS callback invoked even though fixed-point optimization is disabled for beta-dependent exact curvature".to_string(),
2954            )
2955        },
2956        crate::marginal_slope_shared::make_beta_seed_validator(&pending_beta_seed),
2957    )?;
2958
2959    let mut resolved_specs = solved.resolved_specs;
2960    let mut designs = solved.designs;
2961    let mut solved_fit = solved.fit;
2962    // #905 GENERATED-REGRESSOR (Murphy–Topel) SEAM. When the conditional
2963    // location-scale gate fired, the slope fit above treated the calibrated
2964    // score `ζ = (z − m̂(C))/√v̂(C)` as KNOWN, so `solved_fit.beta_covariance()`
2965    // is the naive second-stage covariance `V_β^naive = H_β⁻¹` that ignores the
2966    // first-stage estimation error in `θ₁ = (mean_coeffs, var_coeffs)`. The
2967    // honest two-stage covariance is
2968    //   `V_β = V_β^naive + (H_β⁻¹ G) V₁ (H_β⁻¹ G)ᵀ`,  `G = ∂(score_β)/∂θ₁`.
2969    // The closed-form first-stage covariance `V₁` and the per-row chain-rule
2970    // sensitivity `∂ζ_i/∂θ₁` are computed and stored on the calibration at fit
2971    // time (see `LatentZConditionalCalibration::{theta1_covariance,
2972    // zeta_theta1_jacobian_row, generated_regressor_term}`), so the correction
2973    // is consumable wherever the slope information `G` (the per-row
2974    // `∂score_β/∂ζ_i` of the marginal/logslope blocks) is available.
2975    //
2976    // The full correction is assembled by
2977    // `LatentZConditionalCalibration::generated_regressor_correction` (mod.rs):
2978    // given the per-row reduced-frame slope-score sensitivity to the
2979    // calibrated score `s_i = ∂score_β,i/∂ζ_i` (an `n × p_β` matrix), it
2980    //   1. builds `J_zeta` row-by-row via `zeta_theta1_jacobian_row` (exact-zero
2981    //      on floored rows, so `G`'s support is the gate-fired rows),
2982    //   2. accumulates `G = Σ_i s_i ⊗ (∂ζ_i/∂θ₁)` (`p_β × dim θ₁`),
2983    //   3. forms `Vb·G = solved_fit.beta_covariance()·G` (the naive reduced-frame
2984    //      covariance IS `H_β⁻¹`, so `H_β⁻¹ G = Vb·G`), and
2985    //   4. returns `(Vb·G)·V₁·(Vb·G)ᵀ` (PSD ⇒ corrected slope SE strictly ≥
2986    //      naive whenever the gate fires).
2987    // `V₁`, `∂ζ/∂θ₁`, the `Vb` frame, and the whole congruence therefore meet
2988    // here with the exact row-kernel `s_i` channel.
2989    //
2990    // `s_i = ∂²ℓ_i/∂β∂ζ_i = J_iᵀ·(∂²ℓ_i/∂η_i∂ζ_i)` is the mixed `(β, ζ)` second
2991    // derivative of the row kernel contracted through the slope Jacobian `J_i`.
2992    // For the rigid StandardNormal kernel the mixed 2-vector is read from the
2993    // same row tower by seeding ζ as a third axis. For score_warp/link_dev flex
2994    // kernels, the span-local cubic row calculus differentiates its observed-z
2995    // coefficient jets and scatters every active primary into the same reduced
2996    // full-beta frame as `covariance_conditional` (#2303).
2997    let (latent_z_rank_int_calibration, latent_z_conditional_calibration) =
2998        match latent_z_calibration {
2999            LatentMeasureCalibration::None => (None, None),
3000            LatentMeasureCalibration::RankInverseNormal(cal) => (Some(cal), None),
3001            LatentMeasureCalibration::ConditionalLocationScale(cal) => (None, Some(cal)),
3002        };
3003    // #905/#1028/#2303: apply the Murphy–Topel correction while every block is
3004    // still in the exact reduced covariance frame. The rigid kernel uses its
3005    // three-axis tower; flex fits use the observed-z derivative channel from
3006    // the same span-local cubic row calculus that produced their score and
3007    // Hessian, including every score_warp/link_dev coefficient. No active block
3008    // may be padded with zero sensitivity or skipped.
3009    if let Some(cal) = latent_z_conditional_calibration.as_ref()
3010        && let Some(vb) = solved_fit.covariance_conditional.clone()
3011    {
3012        if !matches!(latent_measure, LatentMeasureKind::StandardNormal) {
3013            return Err(
3014                "BMS Murphy-Topel generated-regressor covariance is unavailable for a conditional latent-z calibration whose second-stage latent measure is not StandardNormal"
3015                    .to_string(),
3016            );
3017        }
3018        let p_beta = vb.nrows();
3019        let calibration_marginal_dense = marginal_design
3020            .design
3021            .try_to_dense_arc("bms generated-regressor marginal design")?;
3022        if p_beta != vb.ncols() {
3023            return Err(format!(
3024                "bms generated-regressor: covariance_conditional must be square, got {}×{}",
3025                vb.nrows(),
3026                vb.ncols()
3027            ));
3028        }
3029        let correction_family =
3030            make_family(&marginal_design, &logslope_design, final_sigma_cell.get());
3031        let flex_active = score_warp_runtime.is_some() || link_dev_runtime.is_some();
3032        let s = if flex_active {
3033            correction_family.flex_score_zeta_sensitivity(
3034                &solved_fit.block_states,
3035                options,
3036                p_beta,
3037            )?
3038        } else {
3039            // Use the FAMILY designs here, not the raw reporting designs: they
3040            // are exactly the widened-marginal/reduced-logslope coefficient
3041            // frame carried by Vb (including an active influence absorber).
3042            let score_marginal_dense = correction_family
3043                .marginal_design
3044                .try_to_dense_arc("bms generated-regressor fitted marginal design")?;
3045            let score_logslope_dense = correction_family
3046                .logslope_design
3047                .try_to_dense_arc("bms generated-regressor fitted logslope design")?;
3048            let p_score = score_marginal_dense.ncols() + score_logslope_dense.ncols();
3049            if p_beta != p_score {
3050                return Err(format!(
3051                    "bms generated-regressor rigid covariance/frame mismatch: covariance width {p_beta} != marginal({}) + logslope({})",
3052                    score_marginal_dense.ncols(),
3053                    score_logslope_dense.ncols()
3054                ));
3055            }
3056            let marginal_eta = &solved_fit.block_states[0].eta;
3057            let slope_eta = &solved_fit.block_states[1].eta;
3058            let probit_scale = probit_frailty_scale(final_sigma_cell.get());
3059            rigid_standard_normal_score_zeta_sensitivity(
3060                &spec.base_link,
3061                marginal_eta,
3062                slope_eta,
3063                z.as_ref(),
3064                y.as_ref(),
3065                weights.as_ref(),
3066                probit_scale,
3067                score_marginal_dense.view(),
3068                score_logslope_dense.view(),
3069                p_beta,
3070            )?
3071        };
3072        // The first-stage Jacobian expects the RAW normalized score and the
3073        // raw calibration design, whereas `s` above is evaluated at the
3074        // calibrated score consumed by the second-stage kernel.
3075        let correction = cal.generated_regressor_correction(
3076            s.view(),
3077            spec.z.view(),
3078            calibration_marginal_dense.view(),
3079            vb.view(),
3080        )?;
3081        if let Some(cov) = solved_fit.covariance_conditional.as_mut() {
3082            *cov = &*cov + &correction;
3083        }
3084        if let Some(cov) = solved_fit.covariance_corrected.as_mut() {
3085            *cov = &*cov + &correction;
3086        }
3087        log::info!(
3088            "[BMS latent-z] Murphy–Topel generated-regressor SE correction applied: \
3089             p_beta={p_beta} flex_active={flex_active} theta1_dim={} max_diag_inflation={:.3e}",
3090            cal.theta1_dim(),
3091            (0..p_beta)
3092                .map(|i| correction[[i, i]])
3093                .fold(0.0_f64, f64::max),
3094        );
3095    }
3096    // Only after covariance correction is complete may the reported logslope
3097    // coefficients be lifted from fitted G*T coordinates into the original
3098    // full-width G frame used by prediction and reporting.
3099    if let Some(reparam) = logslope_reduced_reparam.as_ref() {
3100        let r = reparam.reduced_cols();
3101        if let Some(block) = solved_fit.blocks.get_mut(1)
3102            && block.beta.len() == r
3103        {
3104            block.beta = reparam.recover_original_logslope_beta(&block.beta)?;
3105        }
3106        if let Some(state) = solved_fit.block_states.get_mut(1)
3107            && state.beta.len() == r
3108        {
3109            state.beta = reparam.recover_original_logslope_beta(&state.beta)?;
3110        }
3111    }
3112    // #461: PREDICT SEAM — when the Stage-1 influence absorber is active
3113    // (spec.score_influence_jacobian.is_some()), `fit.block_states[0].beta` is
3114    // the WIDENED marginal coefficient `[β_m; γ]` (length p_m + p₁), but
3115    // `marginal_design` below is the RAW term-collection design (p_m columns):
3116    // the absorbed influence columns Z̃_infl are a TRAINING-only leakage
3117    // absorber and do NOT exist at predict rows (no Stage-1 fold there). The
3118    // orthogonalized β̂_m is a property of the training fit, so prediction must
3119    // use ONLY the first p_m entries of block_states[0].beta against this raw
3120    // marginal_design and DROP the trailing γ. The model-payload / predict
3121    // builder (src/main.rs run_fit_bernoulli_marginal_slope → inference) owns
3122    // that truncation; it must record p_m (= marginal_design.design.ncols())
3123    // and slice the persisted marginal β to it. Survival mirrors this seam.
3124    Ok(BernoulliMarginalSlopeFitResult {
3125        fit: solved_fit,
3126        marginalspec_resolved: resolved_specs.remove(0),
3127        logslopespec_resolved: resolved_specs.remove(0),
3128        marginal_design: designs.remove(0),
3129        logslope_design: designs.remove(0),
3130        baseline_marginal: baseline.0,
3131        baseline_logslope: baseline.1,
3132        z_normalization,
3133        latent_measure,
3134        score_warp_runtime,
3135        link_dev_runtime,
3136        gaussian_frailty_sd: final_sigma_cell.get(),
3137        cross_block_warnings,
3138        latent_z_rank_int_calibration,
3139        latent_z_conditional_calibration,
3140    })
3141}