Skip to main content

gam_models/bms/
mod.rs

1use crate::cubic_cell_kernel as exact_kernel;
2use crate::custom_family::{
3    BatchedOuterGradientTerms, BlockEffectiveJacobian, BlockWorkingSet, BlockwiseFitOptions,
4    CustomFamily, CustomFamilyJointHyperModeSelection, CustomFamilyWarmStart, EvalMode,
5    ExactNewtonJointGradientEvaluation, ExactNewtonJointHessianWorkspace, FamilyEvaluation,
6    FamilyLinearizationState, ParameterBlockSpec, ParameterBlockState, PenaltyMatrix,
7    custom_family_outer_derivatives, evaluate_custom_family_joint_hyper_best_mode_shared,
8    fit_custom_family, fit_custom_family_fixed_log_lambdas_from_mode_selection,
9    joint_hyper_options_for_outer_tolerance,
10};
11use crate::exact_mode_branch::ExactCoefficientModeBranch;
12use crate::fit_orchestration::drivers::{
13    ExactJointEfsEvaluation, ExactJointEvaluation, ExactJointHyperSetup, SpatialFitProvenance,
14    apply_spatial_anisotropy_pilot_initializer, build_term_collection_designs_and_freeze_joint,
15    optimize_spatial_length_scale_exact_joint, spatial_length_scale_term_indices,
16};
17use crate::marginal_slope_shared::{
18    CoeffSupport, ObservedDenestedCellPartials, SparsePrimaryCoeffJetView, add_optional_matrix,
19    add_optional_vector, add_two_surface_psi_outer,
20    build_denested_partition_cells as shared_denested_partition_cells, chunked_row_reduction,
21    eval_coeff4_at, first_parameter_directional_order2_terms, first_parameter_order2_terms,
22    observed_denested_cell_partials as shared_observed_denested_cell_partials, outer_row_indices,
23    outer_weighted_rows, parameter_block_specs_match_rows, probit_frailty_scale,
24    psi_derivative_location, scale_coeff4, second_parameter_order2_terms,
25};
26use crate::model_types::UnifiedFitResult;
27use crate::outer_subsample::WeightedOuterRow;
28use crate::parameter_block::ParameterBlockInput;
29use crate::probability::{
30    chi_square_sf, normal_cdf, normal_logcdf, normal_pdf, normal_two_sided_probability,
31    signed_probit_logcdf_and_mills_ratio, standard_normal_quantile,
32};
33use crate::row_kernel::{
34    RowKernel, RowKernelHessianWorkspace, build_row_kernel_cache, row_kernel_gradient,
35    row_kernel_hessian_dense, row_kernel_log_likelihood,
36};
37use crate::spatial_psi_bridge::{
38    CoefficientSpatialPsiBlockTransform, build_block_spatial_psi_derivatives,
39    build_block_spatial_psi_derivatives_with_transform,
40};
41use crate::survival::lognormal_kernel::{FrailtyScale, FrailtySpec};
42use gam_linalg::matrix::{DesignMatrix, SymmetricMatrix};
43use gam_problem::{
44    ExactNewtonJointPsiSecondOrderTerms, ExactNewtonJointPsiTerms, ExactNewtonJointPsiWorkspace,
45    HyperOperator, InverseLink, StandardLink, WigglePenaltyConfig,
46};
47use gam_solve::estimate::reml::reml_outer_engine::{DenseSpectralOperator, HessianFactorization};
48use gam_solve::pirls::LinearInequalityConstraints;
49use gam_terms::smooth::{
50    SpatialLengthScaleOptimizationOptions, SpatialLogKappaCoords, TermCollectionDesign,
51    TermCollectionSpec,
52};
53use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayViewMut1, s};
54use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
55use serde::{Deserialize, Serialize};
56use std::cell::RefCell;
57use std::collections::HashMap;
58use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
59use std::sync::{Arc, Mutex, OnceLock};
60
61mod alo_replay;
62pub mod deviation_runtime;
63pub mod gpu;
64pub(crate) use alo_replay::exact_runtime_from_saved;
65pub use alo_replay::{
66    BernoulliMarginalSlopeAloRowGeometry, BernoulliMarginalSlopeAloRowInput,
67    BernoulliMarginalSlopeSavedAloReplay, BernoulliMarginalSlopeSavedAloRowGeometry,
68    bernoulli_marginal_slope_alo_row_geometry,
69};
70pub(crate) use alo_replay::{
71    BernoulliMarginalSlopeSavedAloReplayInput, replay_saved_bernoulli_marginal_slope_alo,
72};
73pub use deviation_runtime::DeviationRuntime;
74pub use deviation_runtime::ParametricAnchorBlock;
75
76/// Above this size, FLEX spatial length-scale optimization uses the pilot
77/// geometry initializer and skips the iterative joint κ/ψ outer loop. This is
78/// a spatial-optimizer policy only; it must not gate exact outer Hessian
79/// capability or row-cell moment materialization.
80pub(crate) const BMS_FLEX_SPATIAL_OUTER_PILOT_ROW_THRESHOLD: usize = 50_000;
81
82#[derive(Clone, Debug)]
83pub struct DeviationBlockConfig {
84    pub degree: usize,
85    pub num_internal_knots: usize,
86    pub penalty_order: usize,
87    pub penalty_orders: Vec<usize>,
88    pub double_penalty: bool,
89    pub monotonicity_eps: f64,
90}
91
92impl Default for DeviationBlockConfig {
93    fn default() -> Self {
94        WigglePenaltyConfig::cubic_triple_operator_default().into()
95    }
96}
97
98impl DeviationBlockConfig {
99    pub fn triple_penalty_default() -> Self {
100        Self::default()
101    }
102}
103
104impl From<WigglePenaltyConfig> for DeviationBlockConfig {
105    fn from(cfg: WigglePenaltyConfig) -> Self {
106        let penalty_order = *cfg.penalty_orders.iter().max().unwrap_or(&2);
107        Self {
108            degree: cfg.degree,
109            num_internal_knots: cfg.num_internal_knots,
110            penalty_order,
111            penalty_orders: cfg.penalty_orders,
112            double_penalty: cfg.double_penalty,
113            monotonicity_eps: cfg.monotonicity_eps,
114        }
115    }
116}
117
118#[derive(Clone)]
119pub(crate) struct DeviationPrepared {
120    pub(crate) block: ParameterBlockInput,
121    pub(crate) runtime: DeviationRuntime,
122}
123
124impl std::fmt::Debug for DeviationPrepared {
125    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
126        f.debug_struct("DeviationPrepared").finish_non_exhaustive()
127    }
128}
129
130#[derive(Clone)]
131pub struct BernoulliMarginalSlopeTermSpec {
132    pub y: Array1<f64>,
133    pub weights: Array1<f64>,
134    pub z: Array1<f64>,
135    pub base_link: InverseLink,
136    pub marginalspec: TermCollectionSpec,
137    pub logslopespec: TermCollectionSpec,
138    pub marginal_offset: Array1<f64>,
139    pub logslope_offset: Array1<f64>,
140    /// GaussianShift frailty on the final probit index: U ~ N(0, σ²) added
141    /// to the scalar argument of Φ.  This is exact because the sextic
142    /// microcell kernel is preserved — the Gaussian-decoupling identity
143    /// E[Φ(η + U)] = Φ(η / √(1+σ²)) rescales the index by 1/τ where
144    /// τ = √(1+σ²), and every derivative chain rule factor is polynomial
145    /// in τ, so all six kernel derivatives remain closed-form.
146    ///
147    /// **HazardMultiplier frailty is NOT supported in this family.**
148    /// HazardMultiplier frailty + score_warp/linkwiggle cubic marginal-slope
149    /// is not finite-state exact.  For hazard-multiplier frailty, use the
150    /// standalone LatentCloglogBinomial / LatentSurvival families instead.
151    pub frailty: FrailtySpec,
152    pub score_warp: Option<DeviationBlockConfig>,
153    pub link_dev: Option<DeviationBlockConfig>,
154    pub latent_z_policy: LatentZPolicy,
155    /// Out-of-fold Stage-1 score-influence Jacobian `J = ∂z/∂θ₁` (n × p₁)
156    /// from cross-fitting a CTN transformation-normal Stage-1 model (#461).
157    /// When `Some`, the realized leakage directions `Z_infl = diag(s_f·β̂₀)·J`
158    /// are absorbed as a null-penalized block so the joint solve makes the
159    /// β estimating equation orthogonal to `span(Z_infl)` — the x-dependent
160    /// realization of `ψ − Π_η[ψ]`. `None` ⇒ raw `--z-column` with no CTN
161    /// Stage-1, in which case the free 1-D `score_warp` spline is the
162    /// fallback basis (it spans only the x-free leakage column).
163    pub score_influence_jacobian: Option<Array2<f64>>,
164}
165
166pub struct BernoulliMarginalSlopeFitResult {
167    pub fit: UnifiedFitResult,
168    pub marginalspec_resolved: TermCollectionSpec,
169    pub logslopespec_resolved: TermCollectionSpec,
170    pub marginal_design: TermCollectionDesign,
171    pub logslope_design: TermCollectionDesign,
172    pub baseline_marginal: f64,
173    pub baseline_logslope: f64,
174    pub z_normalization: LatentZNormalization,
175    pub latent_measure: LatentMeasureKind,
176    pub score_warp_runtime: Option<DeviationRuntime>,
177    pub link_dev_runtime: Option<DeviationRuntime>,
178    /// Learned or fixed Gaussian-shift frailty SD.  `None` = no frailty.
179    pub gaussian_frailty_sd: Option<f64>,
180    /// Structured warnings emitted during fit-time setup when a flex
181    /// block was fully aliased by its anchor union and got dropped. The
182    /// fit proceeds without the dropped block (its contribution to the
183    /// joint design was numerically reproducible by the anchor span, so
184    /// keeping it would leave the joint Hessian rank-deficient). Empty
185    /// for fits where every flex block carried independent directions.
186    pub cross_block_warnings: Vec<CrossBlockIdentifiabilityWarning>,
187    /// Optional weighted rank inverse-normal (Blom rankit) calibration
188    /// installed at fit time when the auto latent-z normality check
189    /// failed. `Some(_)` ⇒ the training z was transformed in place via
190    /// [`LatentZRankIntCalibration::apply_to_training`] before any
191    /// downstream consumer (pooled probit baseline, term-collection
192    /// designs, family PIRLS loops) saw it, and the rigid kernel
193    /// routes through the standard-normal closed-form path on the
194    /// calibrated scale. `None` ⇒ no calibration was applied (training
195    /// z already passed the standard-normal diagnostics, or the caller
196    /// explicitly selected a non-Auto `LatentMeasureSpec`).
197    ///
198    /// Persisted to disk so prediction applies the same monotone map
199    /// via [`LatentZRankIntCalibration::apply_at_predict`] to incoming
200    /// z before the standard-normal kernel runs. The public field name
201    /// is `latent_z_rank_int_calibration` — Agent D's persistence
202    /// pipeline reads it under that exact identifier.
203    pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
204    /// Optional conditional location-scale calibration of the latent score
205    /// (#905). `Some(_)` ⇒ the Auto path's conditional `E[z|C]`/`Var(z|C)` Rao
206    /// gate detected PC/grouping-dependence that the pooled-marginal gate
207    /// cannot see, so the training z was replaced in place by
208    /// `ζ = (z − m(C))/√v(C)` (via [`LatentZConditionalCalibration::apply`])
209    /// before any downstream consumer saw it. Mutually exclusive with
210    /// `latent_z_rank_int_calibration`: rank-INT fixes a pooled-marginal
211    /// defect, the conditional correction fixes a conditional-shift defect that
212    /// rank-INT provably cannot. Persisted so prediction rebuilds `a(C)` from
213    /// the (reproducible) marginal design and applies the identical map.
214    pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
215}
216
217#[derive(Clone, Debug)]
218pub enum LatentZCheckMode {
219    Strict,
220    WarnOnly,
221    Off,
222}
223
224#[derive(Clone, Debug)]
225pub enum LatentZNormalizationMode {
226    None,
227    FitWeighted,
228    Frozen { mean: f64, sd: f64 },
229}
230
231pub const DEFAULT_EMPIRICAL_LATENT_GRID_SIZE: usize = 65;
232pub(crate) const AUTO_Z_NORMAL_SKEW_TOL: f64 = 0.10;
233pub(crate) const AUTO_Z_NORMAL_KURT_TOL: f64 = 0.25;
234pub(crate) const AUTO_Z_NORMAL_KS_TOL: f64 = 0.025;
235pub(crate) const AUTO_Z_NORMAL_MAX_ABS: f64 = 8.0;
236/// Inner σ level at which the empirical tail mass of latent z is compared
237/// against the standard normal's theoretical two-sided tail in the auto
238/// normality gate. Chosen well inside `AUTO_Z_NORMAL_MAX_ABS` so a fat inner
239/// tail is caught before any single observation trips the hard `max |z|` bound.
240pub(crate) const AUTO_Z_NORMAL_TAIL_SIGMA_INNER: f64 = 4.0;
241/// Outer σ level for the same tail-mass comparison; catches heavier far-tail
242/// excess that the inner level can miss.
243pub(crate) const AUTO_Z_NORMAL_TAIL_SIGMA_OUTER: f64 = 6.0;
244/// Multiplier applied to the normal's theoretical tail mass before comparison:
245/// the empirical tail may be up to this many times the Gaussian tail at the
246/// same σ before the gate fails, allowing for finite-sample sampling noise.
247pub(crate) const AUTO_Z_NORMAL_TAIL_MASS_SLACK: f64 = 2.0;
248/// Absolute additive floor on the inner-σ tail comparison, so the gate does
249/// not fail on round-off when the Gaussian tail itself is already tiny.
250pub(crate) const AUTO_Z_NORMAL_TAIL_FLOOR_INNER: f64 = 1e-5;
251/// Absolute additive floor on the outer-σ tail comparison; smaller than the
252/// inner floor because the 6σ Gaussian tail is many orders smaller than 4σ.
253pub(crate) const AUTO_Z_NORMAL_TAIL_FLOOR_OUTER: f64 = 1e-8;
254/// Significance level for the conditional `E[z|C]` / `Var(z|C)` Rao gate in the
255/// core Auto path (#905). When the latent score's conditional mean or variance
256/// on the marginal-index span `a(C)` is significant at this level, the Auto
257/// path escalates from the pooled-marginal rank-INT to a conditional
258/// location-scale correction. Chosen small (0.1%) so the escalation fires only
259/// on clear conditional structure, not finite-sample noise — the gate runs once
260/// over the whole training sample, so a per-test α this tight still has ample
261/// power against the grouping mean-shift the issue names.
262pub(crate) const AUTO_Z_CONDITIONAL_RAO_ALPHA: f64 = 1.0e-3;
263/// Relative ridge added to the weighted normal equations when regressing the
264/// latent score on the marginal-index span for the conditional correction.
265/// Stabilizes the solve when `a(C)` is rank-deficient or collinear (penalized
266/// spline marginal indices routinely are) without materially biasing the
267/// conditional mean/variance fit.
268pub(crate) const AUTO_Z_CONDITIONAL_RIDGE_REL: f64 = 1.0e-8;
269/// Floor on the fitted conditional variance `v(C)`, as a fraction of the global
270/// weighted variance of the latent score. Keeps `ζ = (z−m)/√v` finite and
271/// well-scaled where the linear variance model would otherwise fit a
272/// non-positive or vanishing conditional variance.
273pub(crate) const AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC: f64 = 1.0e-3;
274
275#[derive(Clone, Copy, Debug, PartialEq, Eq)]
276pub enum LatentMeasureSpec {
277    Auto { grid_size: usize },
278    StandardNormal,
279    GlobalEmpirical { grid_size: usize },
280}
281
282impl LatentMeasureSpec {
283    pub fn auto_default() -> Self {
284        Self::Auto {
285            grid_size: DEFAULT_EMPIRICAL_LATENT_GRID_SIZE,
286        }
287    }
288}
289
290impl Default for LatentMeasureSpec {
291    fn default() -> Self {
292        Self::auto_default()
293    }
294}
295
296#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
297pub struct EmpiricalZGrid {
298    pub nodes: Vec<f64>,
299    pub weights: Vec<f64>,
300}
301
302impl EmpiricalZGrid {
303    /// Construct a grid whose node/weight invariants (equal length ≥ 2, finite
304    /// ascending nodes, finite positive weights, weights summing to 1 within
305    /// 1e-8) are enforced up-front. Sorted order is part of the input contract
306    /// so hot denested-cell kernels can consume contiguous buckets without a
307    /// constructor-side reorder or allocation. Prefer this over building the
308    /// struct literally; every code path that goes through `new` satisfies the
309    /// same contract that `validate_empirical_z_grid` checks on read.
310    pub fn new(nodes: Vec<f64>, weights: Vec<f64>, context: &str) -> Result<Self, String> {
311        validate_empirical_z_grid(&nodes, &weights, context)?;
312        Ok(Self { nodes, weights })
313    }
314
315    /// Iterate over co-indexed `(node, weight)` pairs. Use this instead of
316    /// reading `.nodes`/`.weights` separately whenever a loop wants both
317    /// arrays in lockstep — eliminates the chance of mismatched indexing.
318    #[inline]
319    pub fn pairs(&self) -> impl Iterator<Item = (f64, f64)> + '_ {
320        self.nodes.iter().copied().zip(self.weights.iter().copied())
321    }
322}
323
324#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
325#[serde(tag = "kind", rename_all = "kebab-case")]
326#[derive(Default)]
327pub enum LatentMeasureKind {
328    #[default]
329    StandardNormal,
330    GlobalEmpirical {
331        grid: EmpiricalZGrid,
332    },
333    LocalEmpirical {
334        feature_cols: Vec<usize>,
335        #[serde(default)]
336        input_scales: Option<Vec<f64>>,
337        centers: Vec<Vec<f64>>,
338        grids: Vec<EmpiricalZGrid>,
339        top_k: usize,
340        bandwidth: f64,
341        #[serde(skip)]
342        train_row_mixtures: Arc<Vec<Vec<(usize, f64)>>>,
343    },
344}
345
346impl LatentMeasureKind {
347    pub fn validate(&self, context: &str) -> Result<(), String> {
348        match self {
349            Self::StandardNormal => Ok(()),
350            Self::GlobalEmpirical { grid } => {
351                validate_empirical_z_grid(&grid.nodes, &grid.weights, context)
352            }
353            Self::LocalEmpirical {
354                feature_cols,
355                input_scales,
356                centers,
357                grids,
358                top_k,
359                bandwidth,
360                ..
361            } => {
362                if feature_cols.is_empty() {
363                    return Err(format!(
364                        "{context} local empirical latent measure needs feature columns"
365                    ));
366                }
367                if centers.is_empty() {
368                    return Err(format!(
369                        "{context} local empirical latent measure needs centers"
370                    ));
371                }
372                if centers.len() != grids.len() {
373                    return Err(format!(
374                        "{context} local empirical latent measure center/grid length mismatch: centers={}, grids={}",
375                        centers.len(),
376                        grids.len()
377                    ));
378                }
379                if *top_k == 0 || *top_k > centers.len() {
380                    return Err(format!(
381                        "{context} local empirical latent measure top_k must be in 1..={}, got {top_k}",
382                        centers.len()
383                    ));
384                }
385                if !(*bandwidth).is_finite() || *bandwidth <= 0.0 {
386                    return Err(format!(
387                        "{context} local empirical latent measure bandwidth must be finite and positive, got {bandwidth}"
388                    ));
389                }
390                if let Some(scales) = input_scales.as_ref() {
391                    if scales.len() != feature_cols.len() {
392                        return Err(format!(
393                            "{context} local empirical latent measure input scale dimension mismatch: scales={}, features={}",
394                            scales.len(),
395                            feature_cols.len()
396                        ));
397                    }
398                    for (scale_idx, scale) in scales.iter().enumerate() {
399                        if !(scale.is_finite() && *scale > 0.0) {
400                            return Err(format!(
401                                "{context} local empirical latent measure input scale {scale_idx} must be finite and positive, got {scale}"
402                            ));
403                        }
404                    }
405                }
406                for (center_idx, center) in centers.iter().enumerate() {
407                    if center.len() != feature_cols.len() {
408                        return Err(format!(
409                            "{context} local empirical latent center {center_idx} dimension mismatch: got {}, expected {}",
410                            center.len(),
411                            feature_cols.len()
412                        ));
413                    }
414                    if center.iter().any(|value| !value.is_finite()) {
415                        return Err(format!(
416                            "{context} local empirical latent center {center_idx} has non-finite coordinates"
417                        ));
418                    }
419                }
420                for (grid_idx, grid) in grids.iter().enumerate() {
421                    validate_empirical_z_grid(
422                        &grid.nodes,
423                        &grid.weights,
424                        &format!("{context} local empirical grid {grid_idx}"),
425                    )?;
426                }
427                Ok(())
428            }
429        }
430    }
431
432    pub(crate) fn is_empirical(&self) -> bool {
433        matches!(
434            self,
435            Self::GlobalEmpirical { .. } | Self::LocalEmpirical { .. }
436        )
437    }
438
439    /// Per-row empirical latent grid, borrowed where possible. This sits in
440    /// the innermost per-row loops of every criterion/gradient/Hessian
441    /// evaluation, so the global grid MUST come back as a borrow — the old
442    /// `grid.clone()` here allocated two `grid_size`-length vectors per row
443    /// per evaluation across the whole fit. Only the local-mixture path,
444    /// which genuinely synthesizes a new grid per row, returns an owned
445    /// value.
446    pub(crate) fn empirical_grid_for_training_row(
447        &self,
448        row: usize,
449    ) -> Result<Option<std::borrow::Cow<'_, EmpiricalZGrid>>, String> {
450        match self {
451            Self::StandardNormal => Ok(None),
452            Self::GlobalEmpirical { grid } => Ok(Some(std::borrow::Cow::Borrowed(grid))),
453            Self::LocalEmpirical {
454                grids,
455                train_row_mixtures,
456                ..
457            } => {
458                let mixture = train_row_mixtures.get(row).ok_or_else(|| {
459                    format!(
460                        "local empirical latent measure is missing training mixture for row {row}"
461                    )
462                })?;
463                Ok(Some(std::borrow::Cow::Owned(combine_empirical_grids(
464                    grids, mixture,
465                )?)))
466            }
467        }
468    }
469}
470
471/// Allocation-free heapsort of parallel empirical node/weight storage.
472/// Used by the local-mixture constructor, whose concatenated sorted component
473/// grids are not globally ordered. Moving pairs in place avoids the third
474/// temporary allocation that a `Vec<(node, weight)>` canonicalization would
475/// add to that per-row path.
476fn sort_empirical_node_weight_pairs(nodes: &mut [f64], weights: &mut [f64]) {
477    assert_eq!(
478        nodes.len(),
479        weights.len(),
480        "empirical grid nodes and weights must remain parallel"
481    );
482    fn sift_down(nodes: &mut [f64], weights: &mut [f64], mut root: usize, end: usize) {
483        loop {
484            let mut child = 2 * root + 1;
485            if child >= end {
486                return;
487            }
488            if child + 1 < end && nodes[child].total_cmp(&nodes[child + 1]).is_lt() {
489                child += 1;
490            }
491            if !nodes[root].total_cmp(&nodes[child]).is_lt() {
492                return;
493            }
494            nodes.swap(root, child);
495            weights.swap(root, child);
496            root = child;
497        }
498    }
499
500    let len = nodes.len();
501    for root in (0..len / 2).rev() {
502        sift_down(nodes, weights, root, len);
503    }
504    for end in (1..len).rev() {
505        nodes.swap(0, end);
506        weights.swap(0, end);
507        sift_down(nodes, weights, 0, end);
508    }
509}
510
511pub(crate) fn validate_empirical_z_grid(
512    nodes: &[f64],
513    weights: &[f64],
514    context: &str,
515) -> Result<(), String> {
516    if nodes.len() != weights.len() {
517        return Err(format!(
518            "{context} empirical latent measure node/weight length mismatch: nodes={}, weights={}",
519            nodes.len(),
520            weights.len()
521        ));
522    }
523    if nodes.len() < 2 {
524        return Err(format!(
525            "{context} empirical latent measure requires at least two nodes"
526        ));
527    }
528    let mut total = 0.0;
529    let mut previous_node = f64::NEG_INFINITY;
530    for (idx, (&node, &weight)) in nodes.iter().zip(weights.iter()).enumerate() {
531        if !node.is_finite() {
532            return Err(format!(
533                "{context} empirical latent measure node {idx} is non-finite ({node})"
534            ));
535        }
536        if !(weight.is_finite() && weight > 0.0) {
537            return Err(format!(
538                "{context} empirical latent measure weight {idx} must be finite and positive, got {weight}"
539            ));
540        }
541        if node < previous_node {
542            return Err(format!(
543                "{context} empirical latent measure nodes must be sorted ascending, but node {idx} ({node}) is below node {} ({previous_node})",
544                idx - 1
545            ));
546        }
547        previous_node = node;
548        total += weight;
549    }
550    if !(total.is_finite() && (total - 1.0).abs() <= 1e-8) {
551        return Err(format!(
552            "{context} empirical latent measure weights must sum to 1, got {total}"
553        ));
554    }
555    Ok(())
556}
557
558pub(crate) fn combine_empirical_grids(
559    grids: &[EmpiricalZGrid],
560    mixture: &[(usize, f64)],
561) -> Result<EmpiricalZGrid, String> {
562    if mixture.is_empty() {
563        return Err("local empirical latent measure row mixture is empty".to_string());
564    }
565    let mut nodes = Vec::new();
566    let mut weights = Vec::new();
567    for &(grid_idx, grid_weight) in mixture {
568        if !grid_weight.is_finite() || grid_weight <= 0.0 {
569            return Err(format!(
570                "local empirical latent mixture weight must be finite and positive, got {grid_weight}"
571            ));
572        }
573        let grid = grids.get(grid_idx).ok_or_else(|| {
574            format!("local empirical latent mixture references missing grid {grid_idx}")
575        })?;
576        for (node, weight) in grid.pairs() {
577            nodes.push(node);
578            weights.push(grid_weight * weight);
579        }
580    }
581    let total = weights.iter().copied().sum::<f64>();
582    if !(total.is_finite() && total > 0.0) {
583        return Err(
584            "local empirical latent combined grid has non-positive total weight".to_string(),
585        );
586    }
587    for weight in &mut weights {
588        *weight /= total;
589    }
590    sort_empirical_node_weight_pairs(&mut nodes, &mut weights);
591    validate_empirical_z_grid(&nodes, &weights, "local empirical latent combined grid")?;
592    Ok(EmpiricalZGrid { nodes, weights })
593}
594
595#[derive(Clone, Debug)]
596pub struct LatentZPolicy {
597    pub check_mode: LatentZCheckMode,
598    pub normalization: LatentZNormalizationMode,
599    pub latent_measure: LatentMeasureSpec,
600    pub mean_tol_multiplier: f64,
601    pub sd_tol_multiplier: f64,
602    pub max_abs_skew: f64,
603    pub max_abs_excess_kurtosis: f64,
604}
605
606impl LatentZPolicy {
607    pub fn frozen_transformation_normal() -> Self {
608        // Defaults relaxed to `WarnOnly` with the same thresholds the
609        // exploratory-weighted preset uses (skew ≤ 4.0, |excess kurt| ≤ 20.0).
610        // Rationale: the upstream conditional transformation-normal
611        // preprocessor may be fit isotropically (no per-axis κ). At large-scale
612        // dimensionality (16 PCs, 15 ancestries) an isotropic fit can leave
613        // the global latent-z distribution mildly heavy-tailed (skew ≈ 4,
614        // excess kurt ≈ 30–40 in synthetic studies) without violating per-
615        // grouping mean/variance calibration. The downstream marginal-slope
616        // model still uses the latent-Gaussian probit/score-warp link; the
617        // emitted warning makes the deviation visible without aborting the
618        // fit. Callers that need strict enforcement can construct a custom
619        // `LatentZPolicy` with `check_mode: LatentZCheckMode::Strict`.
620        Self {
621            check_mode: LatentZCheckMode::WarnOnly,
622            normalization: LatentZNormalizationMode::Frozen { mean: 0.0, sd: 1.0 },
623            latent_measure: LatentMeasureSpec::auto_default(),
624            mean_tol_multiplier: 4.0,
625            sd_tol_multiplier: 4.0,
626            max_abs_skew: 4.0,
627            max_abs_excess_kurtosis: 20.0,
628        }
629    }
630
631    pub fn exploratory_fit_weighted() -> Self {
632        Self {
633            check_mode: LatentZCheckMode::WarnOnly,
634            normalization: LatentZNormalizationMode::FitWeighted,
635            latent_measure: LatentMeasureSpec::auto_default(),
636            mean_tol_multiplier: 8.0,
637            sd_tol_multiplier: 8.0,
638            max_abs_skew: 4.0,
639            max_abs_excess_kurtosis: 20.0,
640        }
641    }
642}
643
644impl Default for LatentZPolicy {
645    fn default() -> Self {
646        Self::frozen_transformation_normal()
647    }
648}
649
650#[derive(Clone, Copy, Debug, PartialEq)]
651pub struct LatentZNormalization {
652    pub mean: f64,
653    pub sd: f64,
654}
655
656impl LatentZNormalization {
657    pub fn apply(&self, z: &Array1<f64>, context: &str) -> Result<Array1<f64>, String> {
658        if !(self.mean.is_finite() && self.sd.is_finite() && self.sd > BMS_VARIANCE_FLOOR) {
659            return Err(format!(
660                "{context} requires finite latent z normalization with sd > {BMS_VARIANCE_FLOOR:e}; got mean={} sd={}",
661                self.mean, self.sd
662            ));
663        }
664        if z.iter().any(|value| !value.is_finite()) {
665            return Err(format!("{context} requires finite z values"));
666        }
667        Ok(z.mapv(|zi| (zi - self.mean) / self.sd))
668    }
669}
670
671/// Weighted mid-distribution rank inverse-normal transform for the
672/// latent score.
673///
674/// When the latent z fails the standard-normal auto-detection
675/// (`latent_z_normal_adequacy`), the BMS family applied to
676/// pretend the score is N(0,1) anyway would distort the closed-form
677/// probit log-CDF kernel. The historical fallback (local- or
678/// global-empirical latent measure) is *mathematically correct* but
679/// triggers the per-row intercept Newton solve in the empirical-grid
680/// closed-form kernels (`empirical_rigid_primary_grad_hess_closed_form`
681/// and its higher-order siblings); at large scale that is the dominant
682/// cost.
683///
684/// **Rank-INT is a modeling choice, not a reparameterisation.** The
685/// rigid BMS predictor is *affine* in the latent score
686/// (`η = q·√(1+b²) + b·z`), so a nonlinear monotone map of `z` changes
687/// the model class and its likelihood — it does not leave them
688/// invariant. Applying the calibration redefines the latent axis: the
689/// affine model is *specified on the calibrated score* `T(z)`. Nor is
690/// the calibrated training sample exactly N(0,1): a finite set of
691/// normal scores is discrete, and with heavy ties it can stay far from
692/// Gaussian. The closed-form standard-normal kernel is therefore
693/// adequate only when the calibrated sample itself passes the same
694/// standard-normal adequacy gate applied to raw z
695/// (`latent_z_normal_adequacy`);
696/// `build_latent_measure_with_geometry` re-checks the calibrated
697/// sample and falls back to the mathematically exact global-empirical
698/// latent measure when that re-check fails. On the passing path the
699/// kept work is the same closed-form
700/// `signed_probit_logcdf_and_mills_ratio` evaluation as the
701/// no-calibration path; the dropped work is the empirical-grid jet
702/// machinery. Persisted to disk so prediction applies the same
703/// monotone map to incoming z and re-routes through the closed-form
704/// kernel.
705#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
706pub struct LatentZRankIntCalibration {
707    /// Sorted unique positive-mass z values seen during training, ascending.
708    /// Knot table for `apply_to_training` / `apply_at_predict`. Zero-weight
709    /// knots carry no probability mass and are not stored.
710    pub sorted_z: Vec<f64>,
711    /// Weighted mid-distribution rank `(W_before + w_knot/2) / W_total` at
712    /// each `sorted_z` knot. Strictly increasing, strictly inside `(0, 1)`
713    /// (each knot is bounded away from the endpoints by half its own mass),
714    /// and invariant to a common rescaling of the weights.
715    pub weighted_cdf: Vec<f64>,
716    /// Weighted mean of the calibrated training sample. Used as a
717    /// sanity-check value on `fit`; should be very close to zero.
718    pub post_mean: f64,
719    /// Weighted SD of the calibrated training sample. Used as a
720    /// sanity-check value on `fit`; should be very close to one.
721    pub post_sd: f64,
722}
723
724impl LatentZRankIntCalibration {
725    /// Fit the weighted rank-INT calibration from training z and weights.
726    ///
727    /// Algorithm:
728    /// 1. Sort rows by ascending z and merge ties into one knot per unique
729    ///    z with the tie group's total weight `w_g`; discard zero-mass
730    ///    knots.
731    /// 2. Weighted mid-distribution rank at each knot:
732    ///    `p_g = (W_before + w_g/2) / W_total`,
733    ///    with `W_before` the cumulative weight strictly below the knot.
734    /// 3. Store `(sorted_z, weighted_cdf = p_g)`.
735    ///
736    /// The mid-rank depends only on *relative* weights (rescaling every
737    /// weight by a common factor leaves every `p_g` unchanged), is strictly
738    /// increasing across knots, and lies strictly inside `(0, 1)` — each
739    /// knot is separated from the endpoints by half its own mass, so no
740    /// ad-hoc clamp is needed and `Φ⁻¹(p_g)` is always finite.
741    ///
742    /// Returns the calibration plus the post-transform sample's weighted
743    /// mean / SD for sanity-check logging.
744    pub fn fit(z: &Array1<f64>, weights: &Array1<f64>) -> Result<Self, String> {
745        if z.len() != weights.len() {
746            return Err(format!(
747                "rank-INT calibration: z length {} != weights length {}",
748                z.len(),
749                weights.len()
750            ));
751        }
752        if z.is_empty() {
753            return Err("rank-INT calibration requires at least one observation".to_string());
754        }
755        let w_total = weights.iter().copied().sum::<f64>();
756        if !(w_total.is_finite() && w_total > 0.0) {
757            return Err(format!(
758                "rank-INT calibration requires positive finite total weight, got {w_total}"
759            ));
760        }
761        for (idx, value) in z.iter().enumerate() {
762            if !value.is_finite() {
763                return Err(format!(
764                    "rank-INT calibration: z[{idx}] = {value} not finite"
765                ));
766            }
767        }
768        for (idx, weight) in weights.iter().enumerate() {
769            if !(weight.is_finite() && *weight >= 0.0) {
770                return Err(format!(
771                    "rank-INT calibration: weight[{idx}] = {weight} not finite/non-negative"
772                ));
773            }
774        }
775        let mut order: Vec<usize> = (0..z.len()).collect();
776        order.sort_by(|&a, &b| z[a].partial_cmp(&z[b]).unwrap_or(std::cmp::Ordering::Equal));
777
778        let mut sorted_z: Vec<f64> = Vec::with_capacity(z.len());
779        let mut weighted_cdf: Vec<f64> = Vec::with_capacity(z.len());
780        // Merge ties into one knot per unique z, then assign the weighted
781        // mid-distribution rank p_g = (W_before + w_g/2) / W_total. This is
782        // the mid-point of the tie group's probability mass, so it depends
783        // only on relative weights, is strictly increasing, and sits
784        // strictly inside (0, 1) without any clamp. Zero-mass tie groups
785        // are not knots of the weighted empirical distribution and are
786        // dropped.
787        let mut cum_before = 0.0_f64;
788        let mut pos = 0usize;
789        while pos < order.len() {
790            let zi = z[order[pos]];
791            let mut w_group = 0.0_f64;
792            let mut end = pos;
793            while end < order.len() && z[order[end]] == zi {
794                w_group += weights[order[end]];
795                end += 1;
796            }
797            if w_group > 0.0 {
798                sorted_z.push(zi);
799                weighted_cdf.push((cum_before + 0.5 * w_group) / w_total);
800                cum_before += w_group;
801            }
802            pos = end;
803        }
804        if sorted_z.is_empty() {
805            return Err(
806                "rank-INT calibration requires at least one positive-weight observation"
807                    .to_string(),
808            );
809        }
810
811        // Compute sanity-check post-mean and post-sd on the transformed
812        // sample, weighted by the original weights.
813        let mut sum_wz = 0.0_f64;
814        let mut sum_w = 0.0_f64;
815        for &idx in &order {
816            let zi = z[idx];
817            let calibrated = Self::apply_with_knots(zi, &sorted_z, &weighted_cdf);
818            sum_wz += weights[idx] * calibrated;
819            sum_w += weights[idx];
820        }
821        let post_mean = if sum_w > 0.0 { sum_wz / sum_w } else { 0.0 };
822        let mut sum_w_dev = 0.0_f64;
823        for &idx in &order {
824            let zi = z[idx];
825            let calibrated = Self::apply_with_knots(zi, &sorted_z, &weighted_cdf);
826            let d = calibrated - post_mean;
827            sum_w_dev += weights[idx] * d * d;
828        }
829        let post_sd = if sum_w > 0.0 {
830            (sum_w_dev / sum_w).sqrt()
831        } else {
832            1.0
833        };
834
835        Ok(Self {
836            sorted_z,
837            weighted_cdf,
838            post_mean,
839            post_sd,
840        })
841    }
842
843    /// Apply the calibration to the full training z vector, returning the
844    /// calibrated sample. Equivalent to mapping each row's z through
845    /// [`Self::apply_at_predict`], but vectorised.
846    pub fn apply_to_training(&self, z: &Array1<f64>) -> Result<Array1<f64>, String> {
847        if self.sorted_z.is_empty() {
848            return Err("rank-INT calibration has no knots".to_string());
849        }
850        let mut out = Array1::<f64>::zeros(z.len());
851        for (idx, &zi) in z.iter().enumerate() {
852            if !zi.is_finite() {
853                return Err(format!(
854                    "rank-INT calibration apply: z[{idx}] = {zi} not finite"
855                ));
856            }
857            out[idx] = self.apply_at_predict(zi);
858        }
859        Ok(out)
860    }
861
862    /// Apply the calibration to a single z at predict time.
863    ///
864    /// Linear interpolation on `(sorted_z, weighted_cdf)` to obtain
865    /// `p ∈ [eps, 1 − eps]`, then `Φ⁻¹(p)` via
866    /// [`standard_normal_quantile`]. Out-of-range z's clip to the
867    /// boundary CDF before the quantile, so the calibration extrapolates
868    /// monotonically beyond the training support.
869    pub fn apply_at_predict(&self, z: f64) -> f64 {
870        Self::apply_with_knots(z, &self.sorted_z, &self.weighted_cdf)
871    }
872
873    pub(crate) fn apply_with_knots(z: f64, sorted_z: &[f64], weighted_cdf: &[f64]) -> f64 {
874        assert_eq!(sorted_z.len(), weighted_cdf.len());
875        assert!(!sorted_z.is_empty());
876        let n = sorted_z.len();
877        let p = if z <= sorted_z[0] {
878            weighted_cdf[0]
879        } else if z >= sorted_z[n - 1] {
880            weighted_cdf[n - 1]
881        } else {
882            // Binary search for the right knot.
883            let mut lo = 0usize;
884            let mut hi = n - 1;
885            while hi - lo > 1 {
886                let mid = (lo + hi) / 2;
887                if sorted_z[mid] <= z {
888                    lo = mid;
889                } else {
890                    hi = mid;
891                }
892            }
893            let z_lo = sorted_z[lo];
894            let z_hi = sorted_z[hi];
895            let p_lo = weighted_cdf[lo];
896            let p_hi = weighted_cdf[hi];
897            if z_hi == z_lo {
898                p_hi
899            } else {
900                let t = (z - z_lo) / (z_hi - z_lo);
901                p_lo + t * (p_hi - p_lo)
902            }
903        };
904        // Φ⁻¹(p); clip away from {0, 1} to keep the quantile finite.
905        standard_normal_quantile(p).unwrap_or_else(|err| {
906            let clipped = if p < 0.5 { -8.0 } else { 8.0 };
907            log::debug!(
908                "standard_normal_quantile({p}) failed ({err}); clipping the latent score to {clipped}"
909            );
910            clipped
911        })
912    }
913}
914
915/// Optional calibration applied to the latent score before the BMS
916/// kernel runs. When `RankInverseNormal`, both the training and predict
917/// paths route the input z through `LatentZRankIntCalibration::apply_*`
918/// before the standard-normal closed-form kernel is invoked.
919#[derive(Clone, Debug)]
920pub enum LatentMeasureCalibration {
921    None,
922    RankInverseNormal(LatentZRankIntCalibration),
923    ConditionalLocationScale(LatentZConditionalCalibration),
924}
925
926/// Conditional location-scale calibration of the latent score (#905).
927///
928/// The marginal-slope Auto trigger's pooled-z gate (KS / skewness / kurtosis +
929/// the rank inverse-normal transform) only inspects the **marginal** law of
930/// `z`. A conditional shift `E[z | C] = m(C) ≠ 0` — the allele-frequency-driven
931/// grouping mean shift — passes the marginal gate while leaving `z | C`
932/// off-center, so the slope contribution `b(C)·m(C)` leaks into the influence
933/// channel `q`. Rank-INT provably cannot fix this: no transform `T` depending
934/// only on the marginal `F_Z` can enforce `E[T(Z) | C] ≡ const` for all joint
935/// laws.
936///
937/// The unique Fisher-orthogonal location-scale correction (for the Gaussian
938/// working metric the closed-form probit kernel assumes) is
939/// `ζ = (z − m(C)) / √v(C)`, where `m(C) = E[z|C]` and `v(C) = Var(z|C)` are
940/// estimated by weighted ridge regression of `z` (and its squared residual) on
941/// the marginal-index span `a(C) = [1 | X_marginal]`. The corrected `ζ` is
942/// conditionally centered (and homoskedastic when the variance block is
943/// active) by construction, so the `b(C)·m(C)` leakage vanishes. Matching the
944/// first two conditional moments does **not** by itself make `ζ` standard
945/// normal (a two-point residual law survives location-scale correction
946/// unchanged in shape), so `build_latent_measure_with_geometry` re-checks
947/// the calibrated sample against the standard-normal adequacy gate and
948/// retains an empirical latent measure for the residual distribution when
949/// that re-check fails; only a passing `ζ` uses the closed-form
950/// standard-normal kernel. Persisted so prediction rebuilds `a(C)` from the
951/// (reproducible) marginal design and applies the identical map to incoming
952/// z.
953#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
954pub struct LatentZConditionalCalibration {
955    /// Coefficients for the conditional mean `m(C) = β_m·[1 | a(C)]` over the
956    /// basis `[1 | marginal-design row]`. Length `1 + basis_ncols` (leading
957    /// entry is the intercept).
958    pub mean_coeffs: Vec<f64>,
959    /// Coefficients for the conditional variance
960    /// `v(C) = max(β_v·[1 | a(C)], var_floor)`. Length `1 + basis_ncols`, or
961    /// empty when the conditional-variance block of the Rao gate was not
962    /// significant (mean-only correction); then `v(C) ≡ homoskedastic_var`.
963    pub var_coeffs: Vec<f64>,
964    /// Number of marginal-design columns in the basis (excludes the leading
965    /// intercept). The predict-time marginal design must present exactly this
966    /// many columns.
967    pub basis_ncols: usize,
968    /// Floor on the fitted conditional variance, in the (normalized)
969    /// latent-score scale (= `AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC ·` the global
970    /// weighted variance of the training score).
971    pub var_floor: f64,
972    /// The homoskedastic conditional variance `v(C) ≡ Var(z | C)`, used when the
973    /// Breusch-Pagan stage did not fire and `v` is therefore constant in `C`.
974    ///
975    /// This is the *residual* variance of the conditional-mean regression,
976    /// `Σ w (z − m̂(C))² / Σ w`, and NOT the global (marginal) variance of z —
977    /// which is what it used to hold, and what its old on-disk name still says
978    /// (gam#2768).
979    ///
980    /// The distinction is the whole correction. `ζ = (z − m(C))/√v` is supposed
981    /// to be conditionally standard normal, because the marginal-slope identity
982    /// that makes `q` the MARGINAL index,
983    /// `E_ζ[Φ(q√(1+b²) + bζ)] = Φ(q)`, holds only at `Var(ζ|C) = 1`; at
984    /// `Var(ζ|C) = v` it becomes `Φ(q√(1+b²)/√(1+b²v))`, a multiplicative
985    /// distortion of every marginal coefficient. Dividing by the marginal
986    /// variance guaranteed `v ≠ 1`: with z standardised,
987    /// `1 = Var(m(C)) + E[Var(z|C)]`, so the residual variance is `1 − R²` and
988    /// is strictly below the marginal variance *whenever the gate fires at all*.
989    /// The bug was therefore not a corner case — it was on every fired gate, and
990    /// it grew with exactly the conditional structure the correction exists to
991    /// remove. At `R² = 0.25` it left `sd(ζ) = 0.87` against the `post_sd ≈ 1`
992    /// this struct documents, distorted the marginal coefficients by ~4%, and
993    /// (worse) made the calibrated residual fail the standard-normal adequacy
994    /// re-check on the SD clause alone at any appreciable n — sending BMS to the
995    /// empirical measure and, per gam#2718, withholding the covariance.
996    ///
997    /// Kept under the on-disk name `global_var` so a model saved before the fix
998    /// still deserializes AND still applies the map it was fitted with: the
999    /// stored number is whatever that fit divided by, and predict must reproduce
1000    /// the fit, not the current formula.
1001    #[serde(rename = "global_var")]
1002    pub homoskedastic_var: f64,
1003    /// Weighted mean of the calibrated training sample (sanity-check, ≈ 0).
1004    pub post_mean: f64,
1005    /// Weighted SD of the calibrated training sample (sanity-check, ≈ 1).
1006    pub post_sd: f64,
1007    /// Joint first-stage (generated-regressor) sandwich covariance of
1008    /// `θ₁ = (mean_coeffs, var_coeffs)`, shape `dim θ₁ × dim θ₁`.
1009    ///
1010    /// Replaced a stored PAIR of per-stage sandwiches that
1011    /// [`Self::theta1_covariance`] assembled block-diagonally -- i.e. that
1012    /// asserted `Cov(mean_coeffs, var_coeffs) = 0`. The stages are not
1013    /// independent: stage B regresses the squared MEAN residual on the same
1014    /// basis, so the stacked bread is block lower-triangular and the meat has a
1015    /// cross-block proportional to the residual's third moment. Both vanish
1016    /// under a Gaussian residual, and neither vanishes on the branch this
1017    /// covariance serves (gam#2484). Built by
1018    /// `stacked_first_stage_sandwich_cov`; the two retired fields were its
1019    /// diagonal blocks.
1020    ///
1021    /// Fit-time only: predict applies the map from `mean_coeffs`/`var_coeffs`
1022    /// and never reads their uncertainty. Verified rather than assumed -- the
1023    /// only production consumer of this matrix is the Murphy-Topel assembly in
1024    /// `block_specs.rs`, which runs at fit.
1025    ///
1026    /// `#[serde(default)]` so a model saved before gam#2484 -- carrying the two
1027    /// retired per-stage blocks and no joint -- still deserializes on the
1028    /// predict path. An empty default is NOT a silent zero: it is refused at the
1029    /// single point of consumption (see
1030    /// [`Self::generated_regressor_correction`]), because a covariance that
1031    /// quietly vanishes is exactly the failure this issue exists to prevent.
1032    #[serde(default)]
1033    pub theta1_cov: Array2<f64>,
1034}
1035
1036impl LatentZConditionalCalibration {
1037    #[inline]
1038    pub(crate) fn affine(coeffs: &[f64], a_row: ArrayView1<'_, f64>) -> f64 {
1039        let mut acc = coeffs[0];
1040        for (c, &x) in coeffs[1..].iter().zip(a_row.iter()) {
1041            acc += c * x;
1042        }
1043        acc
1044    }
1045
1046    pub(crate) fn conditional_mean(&self, a_row: ArrayView1<'_, f64>) -> f64 {
1047        Self::affine(&self.mean_coeffs, a_row)
1048    }
1049
1050    pub(crate) fn conditional_var(&self, a_row: ArrayView1<'_, f64>) -> f64 {
1051        if self.var_coeffs.is_empty() {
1052            self.homoskedastic_var.max(self.var_floor)
1053        } else {
1054            Self::affine(&self.var_coeffs, a_row).max(self.var_floor)
1055        }
1056    }
1057
1058    /// Apply `ζ = (z − m(C))/√v(C)` to a batch. `a_block` is the marginal
1059    /// design (`n × basis_ncols`); `z` is the (normalized) latent score. Used
1060    /// at both training and predict time, so the map is identical.
1061    pub fn apply(
1062        &self,
1063        z: ArrayView1<'_, f64>,
1064        a_block: ArrayView2<'_, f64>,
1065    ) -> Result<Array1<f64>, String> {
1066        if a_block.ncols() != self.basis_ncols {
1067            return Err(format!(
1068                "conditional latent calibration expects {} basis columns, got {}",
1069                self.basis_ncols,
1070                a_block.ncols()
1071            ));
1072        }
1073        if a_block.nrows() != z.len() {
1074            return Err(format!(
1075                "conditional latent calibration row mismatch: z={}, basis rows={}",
1076                z.len(),
1077                a_block.nrows()
1078            ));
1079        }
1080        if self.mean_coeffs.len() != self.basis_ncols + 1 {
1081            return Err(format!(
1082                "conditional latent calibration mean coefficient length {} != basis_ncols+1 ({})",
1083                self.mean_coeffs.len(),
1084                self.basis_ncols + 1
1085            ));
1086        }
1087        let mut out = Array1::<f64>::zeros(z.len());
1088        for i in 0..z.len() {
1089            let a_row = a_block.row(i);
1090            if !z[i].is_finite() {
1091                return Err(format!(
1092                    "conditional latent calibration: z[{i}] = {} not finite",
1093                    z[i]
1094                ));
1095            }
1096            let m = self.conditional_mean(a_row);
1097            let v = self.conditional_var(a_row);
1098            if !(v.is_finite() && v > 0.0) {
1099                return Err(format!(
1100                    "conditional latent calibration produced non-positive variance {v} at row {i}"
1101                ));
1102            }
1103            let zeta = (z[i] - m) / v.sqrt();
1104            if !zeta.is_finite() {
1105                return Err(format!(
1106                    "conditional latent calibration produced non-finite zeta at row {i}"
1107                ));
1108            }
1109            out[i] = zeta;
1110        }
1111        Ok(out)
1112    }
1113
1114    /// Dimension of the first-stage parameter vector `θ₁ = (mean_coeffs,
1115    /// var_coeffs)` whose estimation uncertainty the generated-regressor
1116    /// correction propagates. Equals `len(mean_coeffs)` when the variance block
1117    /// is inactive, otherwise `len(mean_coeffs) + len(var_coeffs)`.
1118    pub fn theta1_dim(&self) -> usize {
1119        self.mean_coeffs.len() + self.var_coeffs.len()
1120    }
1121
1122    /// Per-row sensitivity `∂ζ_i/∂θ₁` of the calibrated score to the first-stage
1123    /// calibration coefficients, stacked as `[∂ζ/∂mean_coeffs | ∂ζ/∂var_coeffs]`
1124    /// (length [`Self::theta1_dim`]). With `ζ = (z − m(C))/√v(C)`,
1125    /// `A_i = [1 | a(C_i)]`, `m = A_iᵀ·mean_coeffs`, `v = A_iᵀ·var_coeffs`:
1126    ///
1127    ///   `∂ζ/∂m = −1/√v`,  `∂ζ/∂v = −(z − m)/(2 v^{3/2}) = −ζ/(2v)`,
1128    ///
1129    /// and by the chain rule through the affine basis
1130    /// `∂ζ/∂mean_coeffs = (∂ζ/∂m)·A_i`, `∂ζ/∂var_coeffs = (∂ζ/∂v)·A_i`. The
1131    /// variance block contributes only when `var_coeffs` is active AND the
1132    /// fitted `v(C_i)` is above the floor (a floored row has `∂v/∂var_coeffs = 0`
1133    /// in the applied map). `z` is the (normalized) raw latent score at this row.
1134    pub fn zeta_theta1_jacobian_row(&self, z: f64, a_row: ArrayView1<'_, f64>) -> Vec<f64> {
1135        let m = self.conditional_mean(a_row);
1136        let v = self.conditional_var(a_row);
1137        let inv_sqrt_v = 1.0 / v.sqrt();
1138        // Intercept-augmented basis row A_i = [1 | a(C_i)].
1139        let mut out = Vec::with_capacity(self.theta1_dim());
1140        let dzeta_dm = -inv_sqrt_v;
1141        out.push(dzeta_dm); // intercept column of A
1142        for &x in a_row.iter() {
1143            out.push(dzeta_dm * x);
1144        }
1145        if !self.var_coeffs.is_empty() {
1146            // ∂ζ/∂v active only off the floor; on the floor the applied v(C) is
1147            // constant in var_coeffs, so the variance sensitivity is exactly 0.
1148            let raw_v = Self::affine(&self.var_coeffs, a_row);
1149            let dzeta_dv = if raw_v > self.var_floor {
1150                let zeta = (z - m) * inv_sqrt_v;
1151                -zeta / (2.0 * v)
1152            } else {
1153                0.0
1154            };
1155            out.push(dzeta_dv);
1156            for &x in a_row.iter() {
1157                out.push(dzeta_dv * x);
1158            }
1159        }
1160        out
1161    }
1162
1163    /// Joint first-stage covariance `V₁` of `θ₁ = (mean_coeffs, var_coeffs)`
1164    /// of `θ₁`, ordered to match [`Self::zeta_theta1_jacobian_row`]. The two
1165    /// stages are fit on (asymptotically) uncorrelated estimating equations
1166    /// (the mean score `Σ w û A` and the Breusch–Pagan variance score
1167    /// `Σ w (û² − v) A` are orthogonal under the Gaussian working model), so the
1168    /// joint first-stage covariance is block-diagonal to first order — the same
1169    /// approximation the Rao gate above uses.
1170    pub fn theta1_covariance(&self) -> Array2<f64> {
1171        self.theta1_cov.clone()
1172    }
1173
1174    /// Murphy–Topel generated-regressor correction term for the second-stage
1175    /// slope covariance. Given the second-stage information `H_β` (the penalized
1176    /// joint Hessian of the slope fit, whose inverse is the naive `V_β`) and the
1177    /// cross-derivative `G = ∂(score_β)/∂θ₁` (`p_β × dim θ₁`), the corrected
1178    /// covariance is
1179    ///
1180    ///   `V_β = V_β^naive + (H_β⁻¹ G) V₁ (H_β⁻¹ G)ᵀ`.
1181    ///
1182    /// This returns the additive rank-`dim θ₁` term `(H_β⁻¹ G) V₁ (H_β⁻¹ G)ᵀ`
1183    /// given the already-formed `hbeta_inv_g = H_β⁻¹ G` (`p_β × dim θ₁`). The
1184    /// caller forms `G` by accumulating the per-row slope-score sensitivity to
1185    /// `ζ_i` times [`Self::zeta_theta1_jacobian_row`] (chain rule
1186    /// `∂score_β/∂θ₁ = Σ_i (∂score_β/∂ζ_i) (∂ζ_i/∂θ₁)`).
1187    pub fn generated_regressor_term(&self, hbeta_inv_g: ArrayView2<'_, f64>) -> Array2<f64> {
1188        let v1 = self.theta1_covariance();
1189        hbeta_inv_g.dot(&v1).dot(&hbeta_inv_g.t())
1190    }
1191
1192    /// Assemble the full Murphy–Topel generated-regressor correction
1193    /// `(Vb·G)·V₁·(Vb·G)ᵀ` for the second-stage slope covariance, given the ONE
1194    /// engine-side quantity it cannot reconstruct post-fit: the per-row
1195    /// reduced-frame slope-score sensitivity to the calibrated score,
1196    /// `s_i = ∂score_β,i/∂ζ_i` (a `p_β`-vector in the joint flat-β reduced frame
1197    /// `solved_fit.beta_covariance()` lives in). With `score_β,i = ∂ℓ_i/∂β`,
1198    /// `s_i = ∂²ℓ_i/∂β∂ζ_i = J_iᵀ·(∂²ℓ_i/∂η_i∂ζ_i)` is the mixed `(β, ζ)`
1199    /// second derivative of the warped row kernel contracted through the slope
1200    /// design Jacobian `J_i` — exactly the #932 `RowProgram` z-jet
1201    /// channel (`z` is already a row-program input; one extra mixed `(β, z)` jet
1202    /// channel reads off `∂²ℓ/∂β∂z`). It must be evaluated at the converged `β̂`
1203    /// in the SAME reduced frame as `vb`.
1204    ///
1205    /// Everything else is built here from the stored first-stage quantities and
1206    /// the second-stage fit, dissolving the post-fit-reconstruction blocker:
1207    ///   - `G = Σ_i s_i · (∂ζ_i/∂θ₁)ᵀ` (`p_β × dim θ₁`), the chain-rule outer
1208    ///     product accumulated row-by-row with `∂ζ_i/∂θ₁ =
1209    ///     `[`Self::zeta_theta1_jacobian_row`]`(z_i, a_row_i)` (exact-zero on
1210    ///     floored rows, so floored rows contribute nothing — `G`'s support is
1211    ///     the gate-fired rows);
1212    ///   - `Vb·G = vb·G` since the naive second-stage covariance `vb` IS
1213    ///     `H_β⁻¹` (the coordinator's `H_β⁻¹ G = Vb.dot(G)`);
1214    ///   - the term `(Vb·G)·V₁·(Vb·G)ᵀ` via [`Self::generated_regressor_term`].
1215    ///
1216    /// `score_zeta_sensitivity` is `n × p_β` (row `i` = `s_i`); `z` is the
1217    /// per-row normalized latent score (`n`); `a_block` is the marginal design
1218    /// `n × basis_ncols` whose rows feed `zeta_theta1_jacobian_row`; `vb` is the
1219    /// naive reduced-frame slope covariance `n_β × n_β`. The returned term is
1220    /// PSD (a congruence of the PSD `V₁`), so adding it to `vb` makes the
1221    /// corrected slope SE strictly ≥ the naive SE whenever the gate fires
1222    /// (`G ≠ 0`) and exactly equal when every row is floored (`G = 0`).
1223    pub fn generated_regressor_correction(
1224        &self,
1225        score_zeta_sensitivity: ArrayView2<'_, f64>,
1226        z: ArrayView1<'_, f64>,
1227        a_block: ArrayView2<'_, f64>,
1228        vb: ArrayView2<'_, f64>,
1229    ) -> Result<Array2<f64>, String> {
1230        let n = score_zeta_sensitivity.nrows();
1231        let p_beta = score_zeta_sensitivity.ncols();
1232        if z.len() != n || a_block.nrows() != n {
1233            return Err(format!(
1234                "generated_regressor_correction row mismatch: score_zeta_sensitivity rows={n}, \
1235                 z={}, a_block rows={}",
1236                z.len(),
1237                a_block.nrows()
1238            ));
1239        }
1240        if a_block.ncols() != self.basis_ncols {
1241            return Err(format!(
1242                "generated_regressor_correction expects {} basis columns, got {}",
1243                self.basis_ncols,
1244                a_block.ncols()
1245            ));
1246        }
1247        if vb.nrows() != p_beta || vb.ncols() != p_beta {
1248            return Err(format!(
1249                "generated_regressor_correction: vb must be {p_beta}×{p_beta}, got {}×{}",
1250                vb.nrows(),
1251                vb.ncols()
1252            ));
1253        }
1254        // G = Σ_i s_i ⊗ (∂ζ_i/∂θ₁)  (p_β × dim θ₁). Each row contributes the
1255        // rank-1 outer product `s_i ⊗ J_zeta_i`, so summed over the n rows this
1256        // is exactly the cross product `G = Sᵀ·J` of the score-sensitivity
1257        // matrix `S` (`n × p_β`, supplied) and the per-row ζ-Jacobian matrix
1258        // `J` (`n × dim θ₁`). Forming `J` row-by-row is O(n·dim θ₁); the cross
1259        // product is then a single BLAS-3 GEMM rather than the O(n·p_β·dim θ₁)
1260        // scalar triple loop (≈1.5e9 FMA at biobank scale, n≈194k, the dominant
1261        // ~13s/disease cost of the SE correction). Floored rows yield an exact
1262        // all-zero `J` row, so they contribute zero to the GEMM — bit-identical
1263        // to skipping them, no approximation.
1264        // gam#2484: refuse an absent or ill-shaped first-stage covariance rather
1265        // than multiplying by it. This fires for a payload written before the
1266        // joint covariance existed (`theta1_cov` defaults to `0×0` there); such
1267        // a model cannot supply the term, and silently contributing nothing
1268        // would understate the interval by exactly the amount the correction
1269        // exists to add.
1270        let dim_theta1 = self.mean_coeffs.len() + self.var_coeffs.len();
1271        if self.theta1_cov.nrows() != dim_theta1 || self.theta1_cov.ncols() != dim_theta1 {
1272            return Err(format!(
1273                "generated_regressor_correction: the first-stage covariance is {}×{} but \
1274                 theta1 has dimension {dim_theta1}. A calibration deserialized from a payload \
1275                 written before gam#2484 carries no joint first-stage covariance; refit rather \
1276                 than publishing an uncorrected interval.",
1277                self.theta1_cov.nrows(),
1278                self.theta1_cov.ncols()
1279            ));
1280        }
1281        let j_mat = self.build_zeta_theta1_jacobian(z, a_block);
1282        let vb_g = self.beta_theta1_sensitivity(score_zeta_sensitivity, j_mat.view(), vb)?;
1283        Ok(self.generated_regressor_term(vb_g.view()))
1284    }
1285
1286    /// Per-row ζ-Jacobian matrix `J` (`n × dim θ₁`, row `i` = `∂ζ_i/∂θ₁`) built
1287    /// row-by-row from [`Self::zeta_theta1_jacobian_row`]. Floored rows yield an
1288    /// exact all-zero row, so they contribute nothing to the `G = Sᵀ·J` cross
1289    /// product (bit-identical to skipping them).
1290    fn build_zeta_theta1_jacobian(
1291        &self,
1292        z: ArrayView1<'_, f64>,
1293        a_block: ArrayView2<'_, f64>,
1294    ) -> Array2<f64> {
1295        let n = a_block.nrows();
1296        let dim_theta1 = self.theta1_dim();
1297        let mut j_mat = Array2::<f64>::zeros((n, dim_theta1));
1298        for i in 0..n {
1299            let j_zeta_row = self.zeta_theta1_jacobian_row(z[i], a_block.row(i));
1300            assert_eq!(
1301                j_zeta_row.len(),
1302                dim_theta1,
1303                "J_zeta row width must match the first-stage hyperparameter dimension"
1304            );
1305            let mut dst = j_mat.row_mut(i);
1306            for (slot, jz) in dst.iter_mut().zip(j_zeta_row.into_iter()) {
1307                *slot = jz;
1308            }
1309        }
1310        j_mat
1311    }
1312
1313    /// Signed first-order sensitivity `∂β̂/∂θ₁ = Vb·G` (`p_β × dim θ₁`) of the
1314    /// converged second-stage slope to the first-stage calibration parameters,
1315    /// the SIGNED quantity the Murphy–Topel correction is built from.
1316    ///
1317    /// `G = Sᵀ·J = Σ_i s_i ⊗ (∂ζ_i/∂θ₁)` with `s_i = ∂score_β,i/∂ζ_i` the
1318    /// LOG-LIKELIHOOD-score sensitivity (the sign convention #1131 fixes at the
1319    /// source in [`gradient_paths::rigid_standard_normal_mixed_z_sensitivity`]),
1320    /// and `Vb = H_β⁻¹` the NLL-Hessian inverse. Under this convention the
1321    /// implicit-function theorem on `∂(log L)/∂β = 0` gives
1322    /// `∂β̂/∂θ₁ = +H_β⁻¹·G = +Vb·G`, so the returned matrix matches the finite
1323    /// difference of the refit slope in θ₁ in BOTH sign and magnitude — unlike
1324    /// the PSD correction term [`Self::generated_regressor_correction`], which is
1325    /// invariant to this sign. `j_zeta` is the per-row ζ-Jacobian matrix
1326    /// (`n × dim θ₁`, row `i` = `∂ζ_i/∂θ₁`).
1327    fn beta_theta1_sensitivity(
1328        &self,
1329        score_zeta_sensitivity: ArrayView2<'_, f64>,
1330        j_zeta: ArrayView2<'_, f64>,
1331        vb: ArrayView2<'_, f64>,
1332    ) -> Result<Array2<f64>, String> {
1333        // G = Sᵀ·J (p_β × dim θ₁) via the SIMD/GPU-routed cross product.
1334        let g = gam_linalg::faer_ndarray::fast_atb(&score_zeta_sensitivity, &j_zeta);
1335        // Vb·G = H_β⁻¹·G (vb is the naive reduced-frame covariance the fit
1336        // already produced — reused, never recomputed).
1337        Ok(vb.dot(&g))
1338    }
1339}
1340
1341/// First-stage robust (HC0) sandwich covariance of a weighted-ridge coefficient
1342/// vector: `V₁ = M⁺ (Σ_i w_i² û_i² A_i A_iᵀ) M⁺` with `M = AᵀWA + λR` the
1343/// ridge normal matrix that produced the coefficients, `W = diag(weights)`,
1344/// `û_i` the per-row residual, and `A` the regression basis (here `[1 | a(C)]`).
1345/// `M⁺` is the Moore–Penrose pseudo-inverse via eigendecomposition with a
1346/// relative tolerance: identifiable directions get the usual `(λ_eff)⁻¹` weight,
1347/// and rank-deficient directions (where some θ₁ components are not identified
1348/// by `A`) are zeroed — they carry no asymptotic distribution, so V₁ in those
1349/// directions is zero, and the Murphy–Topel propagation through identifiable
1350/// functionals of β remains finite and consistent. Using the ordinary inverse
1351/// here let the unregularized direction's `1/ε` blow `M⁻¹·meat·M⁻¹` through
1352/// the f64 range whenever the wide marginal-index span had a near-null
1353/// direction (the bug behind "conditional latent calibration sandwich
1354/// covariance is non-finite" on wide rank-deficient duchon/spline conditioning).
1355/// The meat is formed as `BᵀB` with `B_i = w_i û_i A_iᵀ` (signed) so the
1356/// fused-multiply GEMM is the same SIMD path used everywhere else in the
1357/// codebase, instead of a hand-rolled triple loop whose partial sums could
1358/// overflow on a single pathological row of the basis.
1359/// `M⁺` for a weighted-ridge normal matrix, in RAW coordinates, computed on the
1360/// Jacobi-preconditioned matrix for the conditioning reasons spelled out in
1361/// [`weighted_ridge_sandwich_cov`].
1362///
1363/// `M̃ = S·M·S` with `S = diag(1/√M_jj)`, so `M⁻¹ = S·M̃⁻¹·S`. Returned rather
1364/// than folded into a sandwich because the stacked system needs the SAME `M⁺`
1365/// in three different products.
1366///
1367/// `weighted_ridge_sandwich_cov` deliberately keeps its own inlined copy of this
1368/// arithmetic: extracting it there would regroup its floating-point operations
1369/// and move numbers on a currently-passing path for no benefit.
1370pub(crate) fn preconditioned_normal_pseudoinverse(
1371    normal_matrix: &Array2<f64>,
1372) -> Result<Array2<f64>, String> {
1373    let p = normal_matrix.nrows();
1374    if normal_matrix.ncols() != p {
1375        return Err(format!(
1376            "stacked first-stage sandwich needs a square normal matrix, got {}x{}",
1377            normal_matrix.nrows(),
1378            normal_matrix.ncols()
1379        ));
1380    }
1381    let mut m_sym = normal_matrix.clone();
1382    gam_linalg::matrix::symmetrize_in_place(&mut m_sym);
1383    let scale: Vec<f64> = (0..p)
1384        .map(|j| 1.0 / m_sym[[j, j]].max(f64::MIN_POSITIVE).sqrt())
1385        .collect();
1386    let mut m_scaled = m_sym;
1387    for i in 0..p {
1388        for j in 0..p {
1389            m_scaled[[i, j]] *= scale[i] * scale[j];
1390        }
1391    }
1392    let mut pinv = gam_linalg::utils::rank_certified_psd_pseudoinverse(&m_scaled, 1.0e-10)
1393        .map_err(|e| format!("stacked first-stage sandwich pseudo-inverse failed: {e}"))?
1394        .into_pseudoinverse();
1395    for i in 0..p {
1396        for j in 0..p {
1397            pinv[[i, j]] *= scale[i] * scale[j];
1398        }
1399    }
1400    Ok(pinv)
1401}
1402
1403/// Weighted Gram `Σ_i s_i · A_i A_iᵀ` for SIGNED per-row scalars `s`.
1404///
1405/// Not expressible as `BᵀB` (that forces a non-negative weight), so this is the
1406/// explicit `Aᵀ diag(s) A`.
1407fn signed_weighted_gram(basis: ArrayView2<'_, f64>, s: &[f64]) -> Array2<f64> {
1408    let mut scaled = basis.to_owned();
1409    for (mut row, &value) in scaled.rows_mut().into_iter().zip(s.iter()) {
1410        row.iter_mut().for_each(|entry| *entry *= value);
1411    }
1412    basis.t().dot(&scaled)
1413}
1414
1415/// The joint first-stage covariance `V₁` of `θ₁ = (mean_coeffs, var_coeffs)`,
1416/// as the sandwich of the STACKED estimating system (gam#2484).
1417///
1418/// With `A = [1 | a(C)]`, `W = diag(w)`, `M = AᵀWA + λR`,
1419/// `û_i = z_i − A_iᵀβ_m` and `r_i = û_i² − A_iᵀβ_v`:
1420///
1421/// ```text
1422/// ψ^m_i = w_i A_i û_i          ψ^v_i = w_i A_i r_i
1423///
1424/// B = [ M      0 ]   with  M_vm = ∂ψ^v/∂β_m = −2·Σ_i w_i û_i A_i A_iᵀ
1425///     [ M_vm   M ]
1426///
1427/// Ω = Σ_i w_i² [A_i û_i ; A_i r_i][·]ᵀ      V₁ = B⁻¹ Ω B⁻ᵀ
1428/// ```
1429///
1430/// The previous form kept only `blkdiag(M⁻¹Ω_mm M⁻¹, M⁻¹Ω_vv M⁻¹)`, which is
1431/// this expression with `M_vm` and `Ω_mv` set to zero. Both are zero in
1432/// expectation for a Gaussian residual (`E[û|A] = 0` and `E[û³] = 0`) and
1433/// neither is zero on the branch this covariance serves.
1434///
1435/// `B⁻¹ = [[M⁻¹, 0], [K, M⁻¹]]` with `K = −M⁻¹·M_vm·M⁻¹`, so no `2p × 2p`
1436/// factorization is formed: one `M⁺` is reused. `V₁` stays PSD -- it is a
1437/// congruence of the PSD Gram `Ω` -- so the caller's PSD guarantee on the
1438/// generated-regressor term is untouched.
1439pub(crate) fn stacked_first_stage_sandwich_cov(
1440    basis: ArrayView2<'_, f64>,
1441    weights: ArrayView1<'_, f64>,
1442    mean_residuals: &[f64],
1443    var_residuals: &[f64],
1444    normal_matrix: &Array2<f64>,
1445) -> Result<Array2<f64>, String> {
1446    let n = basis.nrows();
1447    let p = basis.ncols();
1448    if mean_residuals.len() != n || var_residuals.len() != n || weights.len() != n {
1449        return Err(format!(
1450            "stacked first-stage sandwich length mismatch: rows={n}, mean_residuals={}, \
1451             var_residuals={}, weights={}",
1452            mean_residuals.len(),
1453            var_residuals.len(),
1454            weights.len()
1455        ));
1456    }
1457    let m_pinv = preconditioned_normal_pseudoinverse(normal_matrix)?;
1458    if m_pinv.nrows() != p {
1459        return Err(format!(
1460            "stacked first-stage sandwich normal-matrix shape mismatch: basis cols={p}, normal {}",
1461            m_pinv.nrows()
1462        ));
1463    }
1464
1465    // Meat blocks as Grams of the score-scaled basis.
1466    let scaled = |values: &[f64]| -> Array2<f64> {
1467        let mut b = basis.to_owned();
1468        for (i, mut row) in b.rows_mut().into_iter().enumerate() {
1469            let scale = weights[i] * values[i];
1470            if scale == 0.0 {
1471                row.fill(0.0);
1472            } else {
1473                row.iter_mut().for_each(|entry| *entry *= scale);
1474            }
1475        }
1476        b
1477    };
1478    let bm = scaled(mean_residuals);
1479    let bv = scaled(var_residuals);
1480    let omega_mm = gam_linalg::faer_ndarray::fast_ata(&bm);
1481    let omega_vv = gam_linalg::faer_ndarray::fast_ata(&bv);
1482    let omega_mv = bm.t().dot(&bv);
1483
1484    // Bread cross-block and the resulting inverse-bread row.
1485    let wu: Vec<f64> = (0..n).map(|i| weights[i] * mean_residuals[i]).collect();
1486    let m_vm = signed_weighted_gram(basis, &wu).mapv(|value| -2.0 * value);
1487    let k = m_pinv.dot(&m_vm).dot(&m_pinv).mapv(|value| -value);
1488
1489    let v11 = m_pinv.dot(&omega_mm).dot(&m_pinv);
1490    let v12 = m_pinv.dot(&omega_mm).dot(&k.t()) + m_pinv.dot(&omega_mv).dot(&m_pinv);
1491    let v22 = k.dot(&omega_mm).dot(&k.t())
1492        + k.dot(&omega_mv).dot(&m_pinv)
1493        + m_pinv.dot(&omega_mv.t()).dot(&k.t())
1494        + m_pinv.dot(&omega_vv).dot(&m_pinv);
1495
1496    let mut v1 = Array2::<f64>::zeros((2 * p, 2 * p));
1497    v1.slice_mut(s![..p, ..p]).assign(&v11);
1498    v1.slice_mut(s![..p, p..]).assign(&v12);
1499    v1.slice_mut(s![p.., ..p]).assign(&v12.t());
1500    v1.slice_mut(s![p.., p..]).assign(&v22);
1501    gam_linalg::matrix::symmetrize_in_place(&mut v1);
1502    if v1.iter().any(|value| !value.is_finite()) {
1503        return Err("stacked first-stage sandwich covariance is non-finite".to_string());
1504    }
1505    Ok(v1)
1506}
1507
1508pub(crate) fn weighted_ridge_sandwich_cov(
1509    basis: ArrayView2<'_, f64>,
1510    residuals: &[f64],
1511    weights: ArrayView1<'_, f64>,
1512    normal_matrix: &Array2<f64>,
1513) -> Result<Array2<f64>, String> {
1514    let n = basis.nrows();
1515    let p = basis.ncols();
1516    if residuals.len() != n || weights.len() != n {
1517        return Err(format!(
1518            "weighted ridge sandwich length mismatch: rows={n}, residuals={}, weights={}",
1519            residuals.len(),
1520            weights.len()
1521        ));
1522    }
1523    if normal_matrix.nrows() != p || normal_matrix.ncols() != p {
1524        return Err(format!(
1525            "weighted ridge sandwich normal-matrix shape mismatch: basis cols={p}, normal {}x{}",
1526            normal_matrix.nrows(),
1527            normal_matrix.ncols()
1528        ));
1529    }
1530    // Robust HC0 meat as a Gram: build `B` with `B_i = (w_i û_i) A_iᵀ` (rows of
1531    // basis scaled by `w_i û_i`, sign carried), so `meat = BᵀB = Σ_i w_i² û_i²
1532    // A_i A_iᵀ` from one BLAS Gramian. Identical math to the per-row outer-
1533    // product accumulation, but the GEMM path keeps partial sums vectorized
1534    // and is less sensitive to a single pathological row producing an
1535    // intermediate that overflows f64 before the column-wise reduction cancels.
1536    let mut b = basis.to_owned();
1537    for i in 0..n {
1538        let wi = weights[i];
1539        let ri = residuals[i];
1540        let scale = wi * ri;
1541        if scale == 0.0 {
1542            b.row_mut(i).fill(0.0);
1543            continue;
1544        }
1545        b.row_mut(i).iter_mut().for_each(|value| *value *= scale);
1546    }
1547    let meat = gam_linalg::faer_ndarray::fast_ata(&b);
1548    // SPD pseudo-inverse of `M = AᵀWA + λR` via eigendecomposition with a
1549    // relative tolerance; symmetrize first to absorb floating-point asymmetry
1550    // accumulated in the AᵀWA assembly.
1551    let mut m_sym = normal_matrix.clone();
1552    gam_linalg::matrix::symmetrize_in_place(&mut m_sym);
1553    // Jacobi (symmetric diagonal) preconditioning. When the conditioning basis
1554    // spans many orders of magnitude — a power-9 Duchon RBF over 16 standardized
1555    // PCs produces columns differing by ~30 decades — `M` and `meat` live on
1556    // wildly different per-column scales, and the eigendecomposition behind
1557    // `M⁺ meat M⁺` loses all accuracy: the relative truncation tolerance is set
1558    // by `λ_max(M)` (dominated by the largest-scale column), so a genuinely
1559    // identified small-scale direction can be dropped while a near-null one is
1560    // kept, and the surviving `1/λ` then multiplies the huge `meat` straight
1561    // through the f64 range. Precondition by `D = diag(√M_jj)`. Because the ridge
1562    // penalty diagonal is built as the weighted Gram diagonal itself
1563    // (`penalty_jj = Σ_i w_i a_ij² = (AᵀWA)_jj`), `M_jj = (1+ρ)(AᵀWA)_jj`, so
1564    // `M̃ = D⁻¹ M D⁻¹` has EXACT unit diagonal and `M̃ = C + (ρ/(1+ρ))·I` with
1565    // `C` the basis correlation matrix (PSD). Hence `λ_min(M̃) ≥ ρ/(1+ρ) ≈ 1e-8`
1566    // even for a fully collinear basis, which clears the pseudo-inverse's
1567    // relative tolerance `≈ 1e-10·λ_max(M̃)` for the conditioning widths that
1568    // occur here: no direction is spuriously dropped, so `M̃⁺ = M̃⁻¹ = D M⁻¹ D`
1569    // and `cov = D⁻¹ (M̃⁻¹ meat̃ M̃⁻¹) D⁻¹ = M⁻¹ meat M⁻¹` EXACTLY — the scaling
1570    // cancels, this is the same sandwich, only computed on a well-conditioned
1571    // matrix. (Should a pure-ridge direction ever fall under tolerance at very
1572    // large width, dropping it is the correct scale-invariant identifiability
1573    // call.) `meat̃ = D⁻¹ meat D⁻¹`; `M_jj > 0` (Gram diagonal floored positive)
1574    // so `D` is always finite and invertible.
1575    let scale: Vec<f64> = (0..p)
1576        .map(|j| 1.0 / m_sym[[j, j]].max(f64::MIN_POSITIVE).sqrt())
1577        .collect();
1578    let mut m_scaled = m_sym;
1579    let mut meat_scaled = meat;
1580    for i in 0..p {
1581        for j in 0..p {
1582            let s = scale[i] * scale[j];
1583            m_scaled[[i, j]] *= s;
1584            meat_scaled[[i, j]] *= s;
1585        }
1586    }
1587    let m_pinv = gam_linalg::utils::rank_certified_psd_pseudoinverse(&m_scaled, 1.0e-10)
1588        .map_err(|e| format!("conditional latent calibration sandwich pseudo-inverse failed: {e}"))?
1589        .into_pseudoinverse();
1590    let mut cov = m_pinv.dot(&meat_scaled).dot(&m_pinv);
1591    // Undo the symmetric scaling: cov_raw = D⁻¹ cov_scaled D⁻¹.
1592    for i in 0..p {
1593        for j in 0..p {
1594            cov[[i, j]] *= scale[i] * scale[j];
1595        }
1596    }
1597    if cov.iter().any(|v| !v.is_finite()) {
1598        return Err("conditional latent calibration sandwich covariance is non-finite".to_string());
1599    }
1600    Ok(cov)
1601}
1602
1603/// Weighted mean of a slice of values.
1604pub(crate) fn weighted_mean(
1605    values: &[f64],
1606    weights: ArrayView1<'_, f64>,
1607    total_weight: f64,
1608) -> f64 {
1609    values
1610        .iter()
1611        .zip(weights.iter())
1612        .map(|(&v, &w)| w * v)
1613        .sum::<f64>()
1614        / total_weight
1615}
1616
1617/// Robust (heteroskedasticity-consistent) Rao/LM score-test p-value for the
1618/// null that the centered basis columns `ã(C)` carry no information about the
1619/// centered response `u`. This is the LAN locally-optimal statistic the issue
1620/// names: `s = Σ_i w_i u_i ã(C_i)`, `Ω̂ = Σ_i w_i² u_i² ã(C_i)ã(C_i)ᵀ`,
1621/// `D = sᵀ Ω̂⁺ s ⟶ χ²_{rank Ω̂}`. Both the conditional-mean test
1622/// (`u_i = z_i − z̄`) and the conditional-variance / Breusch-Pagan test
1623/// (`u_i = (z_i − z̄)² − σ̂²`) are this statistic with the same centered basis.
1624///
1625/// Returns `None` when the test is degenerate (no usable basis directions),
1626/// otherwise the asymptotic p-value.
1627pub(crate) fn robust_conditional_score_pvalue(
1628    a_centered: ArrayView2<'_, f64>,
1629    u: &[f64],
1630    weights: ArrayView1<'_, f64>,
1631) -> Result<Option<f64>, String> {
1632    let n = a_centered.nrows();
1633    let r = a_centered.ncols();
1634    if r == 0 || n == 0 {
1635        return Ok(None);
1636    }
1637    if u.len() != n || weights.len() != n {
1638        return Err(format!(
1639            "conditional score test length mismatch: rows={n}, u={}, weights={}",
1640            u.len(),
1641            weights.len()
1642        ));
1643    }
1644    // Build the per-row scaled basis `B` with `B_i = (w_i u_i) ã_i` once, then
1645    // recover both the score and the HC0 robust meat from it with two BLAS-3
1646    // GEMMs over chunked row-blocks instead of an `O(n · r²)` per-row scatter:
1647    //   • score  `s   = ãᵀ (w ∘ u) = Bᵀ 1`     (column sums of `B`),
1648    //   • meat   `Ω̂  = Σ_i w_i² u_i² ã_i ã_iᵀ = BᵀB` since `(w_i u_i)² = w_i² u_i²`.
1649    // A non-positive weight zeroes that row of `B` (its score and meat
1650    // contributions both vanish), reproducing the `wi <= 0.0` skip EXACTLY.
1651    // `fast_ata` is the same parallel Gramian the second-stage sandwich uses, so
1652    // the statistic is numerically identical to the row-accumulated form up to
1653    // the deterministic GEMM reduction order.
1654    let mut b = a_centered.to_owned();
1655    for i in 0..n {
1656        let wi = weights[i];
1657        let scale = if wi > 0.0 { wi * u[i] } else { 0.0 };
1658        if scale == 0.0 {
1659            b.row_mut(i).fill(0.0);
1660            continue;
1661        }
1662        b.row_mut(i).iter_mut().for_each(|value| *value *= scale);
1663    }
1664    let s = b.sum_axis(ndarray::Axis(0));
1665    let omega = gam_linalg::faer_ndarray::fast_ata(&b);
1666    if !s.iter().all(|v| v.is_finite()) || !omega.iter().all(|v| v.is_finite()) {
1667        return Ok(None);
1668    }
1669    let omega_geometry = gam_linalg::utils::rank_certified_psd_pseudoinverse(&omega, 1.0e-10)
1670        .map_err(|e| format!("conditional score test pseudo-inverse failed: {e}"))?;
1671    let rank = omega_geometry.rank();
1672    let omega_pinv = omega_geometry.into_pseudoinverse();
1673    if rank == 0 {
1674        return Ok(None);
1675    }
1676    let d_stat = s.dot(&omega_pinv.dot(&s));
1677    if !(d_stat.is_finite() && d_stat >= 0.0) {
1678        return Ok(None);
1679    }
1680    // The shared survival primitive owns both the direct upper-gamma identity
1681    // and its exact `Q(a, 0) = 1` boundary.
1682    let p_value = chi_square_sf(d_stat, rank as f64);
1683    Ok(p_value.is_finite().then_some(p_value))
1684}
1685
1686/// Fit the conditional location-scale calibration (#905) if the conditional
1687/// `E[z|C]`/`Var(z|C)` Rao gate fires on the marginal-index basis `a_block`.
1688///
1689/// Returns `None` when there is no conditional structure to correct (the gate
1690/// does not fire, or the basis is degenerate) — in that case the caller falls
1691/// back to the existing pooled-marginal gate (rank-INT or no calibration).
1692pub(crate) fn fit_conditional_latent_calibration_if_needed(
1693    z: &Array1<f64>,
1694    weights: &Array1<f64>,
1695    a_block: ArrayView2<'_, f64>,
1696) -> Result<Option<LatentZConditionalCalibration>, String> {
1697    let n = z.len();
1698    let p = a_block.ncols();
1699    if n != weights.len() {
1700        return Err(format!(
1701            "conditional latent gate length mismatch: z={n}, weights={}",
1702            weights.len()
1703        ));
1704    }
1705    if a_block.nrows() != n {
1706        return Err(format!(
1707            "conditional latent gate row mismatch: z={n}, basis rows={}",
1708            a_block.nrows()
1709        ));
1710    }
1711    if p == 0 {
1712        return Ok(None);
1713    }
1714    let total_weight = weights.iter().copied().sum::<f64>();
1715    if !(total_weight.is_finite() && total_weight > 0.0) {
1716        return Ok(None);
1717    }
1718    if z.iter().any(|v| !v.is_finite()) || a_block.iter().any(|v| !v.is_finite()) {
1719        return Ok(None);
1720    }
1721
1722    let z_mean = z
1723        .iter()
1724        .zip(weights.iter())
1725        .map(|(&zi, &wi)| wi * zi)
1726        .sum::<f64>()
1727        / total_weight;
1728    let global_var = z
1729        .iter()
1730        .zip(weights.iter())
1731        .map(|(&zi, &wi)| wi * (zi - z_mean) * (zi - z_mean))
1732        .sum::<f64>()
1733        / total_weight;
1734    if !(global_var.is_finite() && global_var > 0.0) {
1735        return Ok(None);
1736    }
1737
1738    // Center each basis column by its weighted mean so the score test is about
1739    // conditional structure *beyond* the global level (the intercept nuisance).
1740    // A constant marginal-design column collapses to ~0 and is dropped by the
1741    // pseudo-inverse rank, so an intercept already present in a(C) is harmless.
1742    let mut a_centered = a_block.to_owned();
1743    for j in 0..p {
1744        let col = a_block.column(j);
1745        let col_mean = col
1746            .iter()
1747            .zip(weights.iter())
1748            .map(|(&v, &w)| w * v)
1749            .sum::<f64>()
1750            / total_weight;
1751        a_centered.column_mut(j).mapv_inplace(|v| v - col_mean);
1752    }
1753
1754    // Conditional-mean Rao test: u = z − z̄.
1755    let u_mean: Vec<f64> = z.iter().map(|&zi| zi - z_mean).collect();
1756    let p_mean = robust_conditional_score_pvalue(a_centered.view(), &u_mean, weights.view())?;
1757    // Conditional-variance (Breusch-Pagan) Rao test: u = (z − z̄)² − σ̂².
1758    let u_var: Vec<f64> = u_mean.iter().map(|&e| e * e - global_var).collect();
1759    let p_var = robust_conditional_score_pvalue(a_centered.view(), &u_var, weights.view())?;
1760
1761    let mean_fires = p_mean.is_some_and(|p| p < AUTO_Z_CONDITIONAL_RAO_ALPHA);
1762    let var_fires = p_var.is_some_and(|p| p < AUTO_Z_CONDITIONAL_RAO_ALPHA);
1763    if !mean_fires && !var_fires {
1764        return Ok(None);
1765    }
1766
1767    // Escalation fires. Fit the conditional mean over the full basis
1768    // [1 | a(C)] via a weighted ridge (the ridge stabilizes a rank-deficient
1769    // marginal-index span; it does not meaningfully shrink the few directions
1770    // that triggered the gate). The conditional-mean correction is applied
1771    // whenever the gate fires (a pure-variance trigger leaves the C-slopes of
1772    // m(C) ≈ 0, so it reduces to harmless global centering).
1773    let basis = build_intercept_basis(a_block);
1774    // Per-column Tikhonov penalty scaled by the weighted Gram diagonal, so the
1775    // ridge is *relative* to each column's scale (a 1e-8 absolute ridge would
1776    // be negligible against an O(n) Gram and would not stabilize a
1777    // rank-deficient penalized-spline marginal index). `diag_jj = Σ_i w_i a_ij²`;
1778    // floored positive so the all-zero (already-dropped) directions still
1779    // receive a finite ridge and the factorization cannot fail.
1780    let mut penalty = Array2::<f64>::zeros((basis.ncols(), basis.ncols()));
1781    for j in 0..basis.ncols() {
1782        let diag_jj = basis
1783            .column(j)
1784            .iter()
1785            .zip(weights.iter())
1786            .map(|(&x, &w)| w * x * x)
1787            .sum::<f64>()
1788            .max(f64::MIN_POSITIVE);
1789        penalty[[j, j]] = diag_jj;
1790    }
1791    let z_col = z.view().insert_axis(ndarray::Axis(1));
1792    let (mean_coeffs_mat, mean_fitted) = gam_linalg::utils::gaussian_weighted_ridge(
1793        basis.view(),
1794        z_col,
1795        penalty.view(),
1796        weights.view(),
1797        AUTO_Z_CONDITIONAL_RIDGE_REL,
1798    )?;
1799    let mean_coeffs: Vec<f64> = mean_coeffs_mat.column(0).to_vec();
1800
1801    // First-stage (generated-regressor) normal matrix `M = AᵀWA + λR`, the same
1802    // weighted-ridge system `gaussian_weighted_ridge` factorizes internally;
1803    // rebuilt here so its inverse can form the closed-form coefficient sandwich
1804    // `V₁` that the second-stage Murphy–Topel correction consumes. `p` is the
1805    // marginal-index width (small), so this is a cheap dense `(p+1)²` form.
1806    let normal_matrix = {
1807        let mut wa = basis.to_owned();
1808        for i in 0..wa.nrows() {
1809            let wi = weights[i];
1810            wa.row_mut(i).iter_mut().for_each(|value| *value *= wi);
1811        }
1812        let mut m = basis.t().dot(&wa);
1813        m += &(penalty.to_owned() * AUTO_Z_CONDITIONAL_RIDGE_REL);
1814        m
1815    };
1816    let mean_residuals: Vec<f64> = z
1817        .iter()
1818        .zip(mean_fitted.column(0).iter())
1819        .map(|(&zi, &mi)| zi - mi)
1820        .collect();
1821    let mean_cov = weighted_ridge_sandwich_cov(
1822        basis.view(),
1823        &mean_residuals,
1824        weights.view(),
1825        &normal_matrix,
1826    )?;
1827
1828    let var_floor = (AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC * global_var).max(f64::MIN_POSITIVE);
1829    // gam#2484: the per-stage variance sandwich is gone -- it was one diagonal
1830    // block of a matrix that is now built jointly, and computing it here only to
1831    // discard it would be dead work. Its Breusch-Pagan residual is kept, because
1832    // the stacked meat's cross-block needs it.
1833    let mut var_residuals: Option<Vec<f64>> = None;
1834    let var_coeffs: Vec<f64> = if var_fires {
1835        // Conditional-variance correction: regress the squared mean-residual on
1836        // the same basis. Fitted values are floored at `var_floor` when applied.
1837        let resid_sq: Array1<f64> = mean_residuals.iter().map(|&e| e * e).collect();
1838        let resid_col = resid_sq.view().insert_axis(ndarray::Axis(1));
1839        let (var_coeffs_mat, var_fitted) = gam_linalg::utils::gaussian_weighted_ridge(
1840            basis.view(),
1841            resid_col,
1842            penalty.view(),
1843            weights.view(),
1844            AUTO_Z_CONDITIONAL_RIDGE_REL,
1845        )?;
1846        // `r_i = (z−m̂)²_i − v̂_i`, the Breusch-Pagan residual of stage B.
1847        var_residuals = Some(
1848            resid_sq
1849                .iter()
1850                .zip(var_fitted.column(0).iter())
1851                .map(|(&si, &vi)| si - vi)
1852                .collect(),
1853        );
1854        var_coeffs_mat.column(0).to_vec()
1855    } else {
1856        Vec::new()
1857    };
1858
1859    // gam#2484: the joint sandwich when the variance stage fired, otherwise the
1860    // mean-only sandwich -- with no variance coefficients there is no second
1861    // block and no cross-term to get wrong.
1862    let theta1_cov = match var_residuals.as_ref() {
1863        Some(residuals) => stacked_first_stage_sandwich_cov(
1864            basis.view(),
1865            weights.view(),
1866            &mean_residuals,
1867            residuals,
1868            &normal_matrix,
1869        )?,
1870        None => mean_cov,
1871    };
1872    // gam#2768: the homoskedastic branch's `v(C)` is the RESIDUAL variance of
1873    // the conditional-mean regression, not the marginal variance of z. See the
1874    // field doc for why the difference is the correction rather than a detail:
1875    // with z standardised, `1 = Var(m(C)) + E[Var(z|C)]`, so the marginal
1876    // variance overstates `Var(z|C)` by exactly the structure the gate just
1877    // detected, and dividing by it leaves `ζ` at `sd = √(1−R²)`.
1878    //
1879    // Its own estimation uncertainty is not propagated into `theta1_cov`, for
1880    // the same reason the marginal variance's never was: the second-stage
1881    // Murphy-Topel correction treats a CONSTANT scale as known, and a plug-in
1882    // variance's contribution is one order down in `n` from the mean
1883    // coefficients' (`θ₁` carries the mean block, whose sensitivity
1884    // `∂ζ/∂m = −1/√v` is O(1)). When the Breusch-Pagan stage fires, `v(C)` is a
1885    // fitted function of the basis and IS carried in `θ₁`.
1886    let homoskedastic_var = {
1887        let residual_sum = mean_residuals
1888            .iter()
1889            .zip(weights.iter())
1890            .map(|(&e, &w)| w * e * e)
1891            .sum::<f64>();
1892        (residual_sum / total_weight).max(var_floor)
1893    };
1894    let mut calibration = LatentZConditionalCalibration {
1895        mean_coeffs,
1896        var_coeffs,
1897        basis_ncols: p,
1898        var_floor,
1899        homoskedastic_var,
1900        post_mean: 0.0,
1901        post_sd: 1.0,
1902        theta1_cov,
1903    };
1904
1905    // Sanity-check post-correction moments on the training sample.
1906    let calibrated = calibration.apply(z.view(), a_block)?;
1907    let post_mean = weighted_mean(
1908        calibrated
1909            .as_slice()
1910            .expect("calibration.apply returns an owned standard-layout 1-D array"),
1911        weights.view(),
1912        total_weight,
1913    );
1914    let post_var = calibrated
1915        .iter()
1916        .zip(weights.iter())
1917        .map(|(&zi, &wi)| wi * (zi - post_mean) * (zi - post_mean))
1918        .sum::<f64>()
1919        / total_weight;
1920    calibration.post_mean = post_mean;
1921    calibration.post_sd = post_var.max(0.0).sqrt();
1922
1923    Ok(Some(calibration))
1924}
1925
1926/// Prepend a column of ones to `a_block`, producing the `[1 | a(C)]` regression
1927/// basis used by the conditional location-scale fit.
1928pub(crate) fn build_intercept_basis(a_block: ArrayView2<'_, f64>) -> Array2<f64> {
1929    let n = a_block.nrows();
1930    let p = a_block.ncols();
1931    let mut basis = Array2::<f64>::ones((n, p + 1));
1932    basis.slice_mut(s![.., 1..]).assign(&a_block);
1933    basis
1934}
1935
1936/// Which latent measures the *calling family's row kernel* can actually
1937/// evaluate.
1938///
1939/// This is the only thing that differs between the two marginal-slope families'
1940/// latent-measure decisions, so it is the only argument
1941/// [`build_latent_measure_decision`] takes to serve both. The Bernoulli kernel
1942/// owns an empirical-grid branch (`empirical_rigid_primary_grad_hess_closed_form`
1943/// and its higher-order siblings, driven by a per-row intercept Newton solve);
1944/// the survival marginal-slope kernel does not — its row program is the
1945/// closed-form standard-normal probit lowering and nothing else. A decision that
1946/// handed the survival family a `GlobalEmpirical` measure would be a measure it
1947/// cannot evaluate, so the two families must reach *different* terminal states
1948/// from the *same* gate. Making the capability an argument is what keeps the
1949/// gate itself a single object (gam#2768).
1950#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1951pub(crate) enum EmpiricalLatentMeasureSupport {
1952    /// The caller can evaluate `LatentMeasureKind::GlobalEmpirical`.
1953    Available,
1954    /// The caller can only evaluate `LatentMeasureKind::StandardNormal`.
1955    StandardNormalOnly,
1956}
1957
1958/// The latent-measure gate's verdict: the measure the kernel will integrate
1959/// against, the pre-transform applied to z before it reaches that kernel, and —
1960/// for a [`EmpiricalLatentMeasureSupport::StandardNormalOnly`] caller — the
1961/// adequacy ledger of the sample the standard-normal kernel is actually being
1962/// handed when that sample failed the gate.
1963pub(crate) struct LatentMeasureDecision {
1964    pub(crate) kind: LatentMeasureKind,
1965    pub(crate) calibration: LatentMeasureCalibration,
1966    /// `Some` only for a `StandardNormalOnly` caller, and only when the sample
1967    /// the kernel sees failed the standard-normal adequacy gate with no
1968    /// empirical measure available to carry the residual law. Never a silent
1969    /// state: the caller must route it through its own [`LatentZCheckMode`].
1970    pub(crate) unmodelled_residual: Option<LatentNormalAdequacy>,
1971    /// `Some` exactly when `kind` is a freshly built `GlobalEmpirical`: the
1972    /// record of the equal-mass compression that produced it, which is what
1973    /// makes the measure differentiable in the sample it was built from.
1974    ///
1975    /// Fit-time only and deliberately not part of `kind`: the measure is on the
1976    /// persistence wire and its identity is its nodes and weights, while this is
1977    /// provenance about the rows behind them. The gam#2484 Murphy–Topel
1978    /// correction is its only consumer.
1979    pub(crate) empirical_build: Option<empirical_measure_sensitivity::EmpiricalZGridBuild>,
1980}
1981
1982impl LatentMeasureDecision {
1983    fn standard_normal(calibration: LatentMeasureCalibration) -> Self {
1984        Self {
1985            kind: LatentMeasureKind::StandardNormal,
1986            calibration,
1987            unmodelled_residual: None,
1988            empirical_build: None,
1989        }
1990    }
1991}
1992
1993pub(crate) fn build_latent_measure_with_geometry(
1994    z: &Array1<f64>,
1995    weights: &Array1<f64>,
1996    policy: &LatentZPolicy,
1997    conditioning: Option<ArrayView2<'_, f64>>,
1998) -> Result<
1999    (
2000        LatentMeasureKind,
2001        LatentMeasureCalibration,
2002        Option<empirical_measure_sensitivity::EmpiricalZGridBuild>,
2003    ),
2004    String,
2005> {
2006    let decision = build_latent_measure_decision(
2007        z,
2008        weights,
2009        policy,
2010        conditioning,
2011        EmpiricalLatentMeasureSupport::Available,
2012        "BMS",
2013    )?;
2014    if decision.unmodelled_residual.is_some() {
2015        // Structurally unreachable: an `Available` caller always has an
2016        // empirical measure to carry a failing residual law, so the gate never
2017        // hands one back unmodelled. Checked rather than assumed, because the
2018        // silent alternative is publishing a standard-normal kernel over a
2019        // sample the gate rejected — the exact failure this decision exists to
2020        // make impossible.
2021        return Err(
2022            "BMS latent-measure gate returned an unmodelled residual law even though the \
2023             empirical latent measure is available to this kernel"
2024                .to_string(),
2025        );
2026    }
2027    Ok((
2028        decision.kind,
2029        decision.calibration,
2030        decision.empirical_build,
2031    ))
2032}
2033
2034/// The latent-measure gate, shared by both marginal-slope families.
2035///
2036/// The gate order is fixed and family-independent:
2037///
2038/// 1. the conditional `E[z|C]` / `Var(z|C)` Rao gate on the marginal-index span
2039///    `a(C)` (#905) — the `b(C)·m(C)` leakage the pooled gate cannot see and
2040///    that rank-INT provably cannot fix, so it takes precedence;
2041/// 2. the pooled standard-normal adequacy gate on raw z;
2042/// 3. the weighted mid-rank inverse-normal transform, re-gated on its own
2043///    output.
2044///
2045/// What differs between families is only the *terminal* state when the sample a
2046/// kernel would see fails the adequacy gate, and that is exactly what `support`
2047/// selects. With [`EmpiricalLatentMeasureSupport::Available`] the decision falls
2048/// back to the mathematically exact empirical latent measure. With
2049/// [`EmpiricalLatentMeasureSupport::StandardNormalOnly`] there is no such
2050/// fallback, so the decision instead keeps the *best available pre-transform* —
2051/// the one whose sample is closest to the kernel's own assumption — and returns
2052/// the failing ledger rather than dropping it.
2053///
2054/// Keeping the pre-transform in that case is not a smaller version of the
2055/// empirical branch, it is the correct choice among the two axes a
2056/// standard-normal-only kernel can be given: the kernel assumes N(0,1) by
2057/// construction, the mid-rank transform matches every quantile of that law up to
2058/// the sample's own discreteness, and raw z can be arbitrarily far from it. The
2059/// residual inadequacy is reported, never absorbed.
2060pub(crate) fn build_latent_measure_decision(
2061    z: &Array1<f64>,
2062    weights: &Array1<f64>,
2063    policy: &LatentZPolicy,
2064    conditioning: Option<ArrayView2<'_, f64>>,
2065    support: EmpiricalLatentMeasureSupport,
2066    context: &str,
2067) -> Result<LatentMeasureDecision, String> {
2068    match policy.latent_measure {
2069        LatentMeasureSpec::Auto { grid_size } => {
2070            // #905: conditional `E[z|C]`/`Var(z|C)` Rao gate. Inspect the latent
2071            // score's conditional moments on the marginal-index span a(C)
2072            // BEFORE the pooled-marginal gate. A significant conditional shift
2073            // is the `b(C)·m(C)` leakage the pooled gate cannot see and that
2074            // rank-INT provably cannot fix, so it takes precedence: route to the
2075            // conditional location-scale correction `ζ = (z−m(C))/√v(C)`.
2076            if let Some(a_block) = conditioning
2077                && let Some(cal) =
2078                    fit_conditional_latent_calibration_if_needed(z, weights, a_block)?
2079            {
2080                // Matching the first two conditional moments does not
2081                // establish Gaussianity of the residual ζ (a two-point
2082                // residual law survives location-scale correction unchanged
2083                // in shape). The closed-form standard-normal kernel is only
2084                // admissible when the calibrated sample passes the same
2085                // pooled adequacy gate raw z faces; otherwise retain an
2086                // empirical latent measure built from ζ, so the residual
2087                // distribution stays the one the data show.
2088                let zeta = cal.apply(z.view(), a_block)?;
2089                let residual_adequacy = latent_z_normal_adequacy(&zeta, weights, policy)?;
2090                let residual_is_standard_normal = residual_adequacy.passes();
2091                let (kind, empirical_build) = match (residual_is_standard_normal, support) {
2092                    (true, _) | (false, EmpiricalLatentMeasureSupport::StandardNormalOnly) => {
2093                        (LatentMeasureKind::StandardNormal, None)
2094                    }
2095                    (false, EmpiricalLatentMeasureSupport::Available) => {
2096                        let (kind, build) =
2097                            build_global_empirical_latent_measure(&zeta, weights, grid_size)?;
2098                        (kind, Some(build))
2099                    }
2100                };
2101                log::info!(
2102                    "[{context} latent-z] conditional location-scale calibrated: basis_ncols={} var_active={} post_mean={:.3e} post_sd={:.3e} residual_measure={} (E[z|C]/Var(z|C) Rao gate fired)",
2103                    cal.basis_ncols,
2104                    !cal.var_coeffs.is_empty(),
2105                    cal.post_mean,
2106                    cal.post_sd,
2107                    if matches!(kind, LatentMeasureKind::StandardNormal) {
2108                        "standard-normal"
2109                    } else {
2110                        "global-empirical"
2111                    },
2112                );
2113                if !residual_is_standard_normal
2114                    && support == EmpiricalLatentMeasureSupport::StandardNormalOnly
2115                {
2116                    // No empirical measure exists for this kernel, so the
2117                    // conditional correction is kept (it removes the first-order
2118                    // `b(C)·m(C)` leakage regardless of the residual's shape)
2119                    // and the residual's distance from the kernel's assumption
2120                    // is handed back to the caller's `LatentZCheckMode`.
2121                    return Ok(LatentMeasureDecision {
2122                        kind,
2123                        calibration: LatentMeasureCalibration::ConditionalLocationScale(cal),
2124                        unmodelled_residual: Some(residual_adequacy),
2125                        empirical_build,
2126                    });
2127                }
2128                if !residual_is_standard_normal {
2129                    // gam#2484: this pair is a legitimate POINT-ESTIMATION
2130                    // state, so it is minted rather than refused here -- but a
2131                    // later Murphy-Topel generated-regressor covariance request
2132                    // is already determined to fail, and it used to fail three
2133                    // stages away with no reference back to the decision that
2134                    // caused it. Say so at the decision, with the evidence.
2135                    log::warn!(
2136                        "[{context} latent-z] the calibrated residual FAILED the standard-normal \
2137                         adequacy gate, so the second-stage latent measure is global-empirical. \
2138                         Point estimation is unaffected; a Murphy-Topel generated-regressor \
2139                         covariance will be REFUSED for this fit, because that correction needs \
2140                         a per-row mixed derivative and an empirical measure built from the \
2141                         whole calibrated-residual vector does not have one (gam#2484). \
2142                         Adequacy ledger (x = statistic / bound, x<=1 passed): {}",
2143                        residual_adequacy.ledger(),
2144                    );
2145                }
2146                return Ok(LatentMeasureDecision {
2147                    kind,
2148                    calibration: LatentMeasureCalibration::ConditionalLocationScale(cal),
2149                    unmodelled_residual: None,
2150                    empirical_build,
2151                });
2152            }
2153            let pooled_adequacy = latent_z_normal_adequacy(z, weights, policy)?;
2154            if pooled_adequacy.passes() {
2155                Ok(LatentMeasureDecision::standard_normal(
2156                    LatentMeasureCalibration::None,
2157                ))
2158            } else {
2159                // P4: route bad-normal latent z through a weighted
2160                // mid-distribution-rank inverse-normal transform. Rank-INT
2161                // redefines the latent axis (the affine rigid model is
2162                // specified on the calibrated score); it makes the calibrated
2163                // sample approximately — not exactly — N(0,1), so the
2164                // closed-form standard-normal kernel is admitted only when
2165                // the calibrated sample itself passes the adequacy gate.
2166                // When it cannot (heavy ties leave the calibrated law
2167                // discrete), fall back to the mathematically exact
2168                // global-empirical latent measure on the raw score.
2169                let calibration = LatentZRankIntCalibration::fit(z, weights)?;
2170                let calibrated = calibration.apply_to_training(z)?;
2171                let calibrated_adequacy = latent_z_normal_adequacy(&calibrated, weights, policy)?;
2172                if calibrated_adequacy.passes() {
2173                    log::info!(
2174                        "[{context} latent-z] rank-INT calibrated: post_mean={:.3e} post_sd={:.3e} knots={}",
2175                        calibration.post_mean,
2176                        calibration.post_sd,
2177                        calibration.sorted_z.len(),
2178                    );
2179                    Ok(LatentMeasureDecision::standard_normal(
2180                        LatentMeasureCalibration::RankInverseNormal(calibration),
2181                    ))
2182                } else {
2183                    match support {
2184                        EmpiricalLatentMeasureSupport::Available => {
2185                            log::info!(
2186                                "[{context} latent-z] rank-INT output failed the standard-normal adequacy gate (post_mean={:.3e} post_sd={:.3e} knots={}); using the global-empirical latent measure",
2187                                calibration.post_mean,
2188                                calibration.post_sd,
2189                                calibration.sorted_z.len(),
2190                            );
2191                            let (kind, build) =
2192                                build_global_empirical_latent_measure(z, weights, grid_size)?;
2193                            Ok(LatentMeasureDecision {
2194                                kind,
2195                                calibration: LatentMeasureCalibration::None,
2196                                unmodelled_residual: None,
2197                                empirical_build: Some(build),
2198                            })
2199                        }
2200                        EmpiricalLatentMeasureSupport::StandardNormalOnly => {
2201                            // Both candidate axes are inadequate and there is no
2202                            // empirical measure to carry either law. Take the
2203                            // one the gate itself measures as closer to the
2204                            // kernel's assumption -- by construction the mid-rank
2205                            // transform matches every quantile of N(0,1) up to
2206                            // the sample's discreteness, which raw z need not do
2207                            // at all -- and report the residual gap.
2208                            Ok(LatentMeasureDecision {
2209                                kind: LatentMeasureKind::StandardNormal,
2210                                calibration: LatentMeasureCalibration::RankInverseNormal(
2211                                    calibration,
2212                                ),
2213                                unmodelled_residual: Some(calibrated_adequacy),
2214                                empirical_build: None,
2215                            })
2216                        }
2217                    }
2218                }
2219            }
2220        }
2221        LatentMeasureSpec::StandardNormal => Ok(LatentMeasureDecision::standard_normal(
2222            LatentMeasureCalibration::None,
2223        )),
2224        LatentMeasureSpec::GlobalEmpirical { grid_size } => match support {
2225            EmpiricalLatentMeasureSupport::Available => {
2226                let (kind, build) = build_global_empirical_latent_measure(z, weights, grid_size)?;
2227                Ok(LatentMeasureDecision {
2228                    kind,
2229                    calibration: LatentMeasureCalibration::None,
2230                    unmodelled_residual: None,
2231                    empirical_build: Some(build),
2232                })
2233            }
2234            EmpiricalLatentMeasureSupport::StandardNormalOnly => Err(format!(
2235                "{context} was asked for a global-empirical latent measure, but its row kernel \
2236                 exists only in the closed-form standard-normal branch: there is no empirical-grid \
2237                 lowering to integrate against. Use the Auto latent measure (which will apply the \
2238                 conditional location-scale or rank inverse-normal pre-transform when the data \
2239                 need one) or fit this data with the Bernoulli marginal-slope family, whose kernel \
2240                 owns the empirical branch"
2241            )),
2242        },
2243    }
2244}
2245
2246/// The standard-normal adequacy verdict on a latent-z sample: every statistic
2247/// the gate forms, beside the bound it was judged against.
2248///
2249/// The gate used to return a bare `bool`, and all three of its call sites threw
2250/// the evidence away. That is why "how far is the calibrated residual from
2251/// standard normal when the gate trips?" could not be answered from a fit --
2252/// the failing clause and its margin existed only inside the conjunction and
2253/// were discarded at the `&&`. The conjunction now lives in [`Self::passes`]
2254/// and the evidence survives it, which is what lets the conditional
2255/// location-scale branch say what the data did (gam#2484).
2256///
2257/// Every field is in the units of its own clause; no field is a ratio, so a
2258/// consumer can report either the raw statistic or its margin.
2259#[derive(Clone, Debug)]
2260pub(crate) struct LatentNormalAdequacy {
2261    /// Kish effective sample size `(Σw)² / Σw²`, which sets the moment bounds.
2262    pub(crate) effective_n: f64,
2263    pub(crate) mean: f64,
2264    pub(crate) mean_tol: f64,
2265    pub(crate) sd: f64,
2266    pub(crate) sd_tol: f64,
2267    pub(crate) skew: f64,
2268    pub(crate) skew_tol: f64,
2269    pub(crate) excess_kurtosis: f64,
2270    pub(crate) excess_kurtosis_tol: f64,
2271    pub(crate) ks: f64,
2272    pub(crate) ks_tol: f64,
2273    pub(crate) tail_mass_inner: f64,
2274    pub(crate) tail_bound_inner: f64,
2275    pub(crate) tail_mass_outer: f64,
2276    pub(crate) tail_bound_outer: f64,
2277    pub(crate) max_abs: f64,
2278    pub(crate) max_abs_tol: f64,
2279}
2280
2281impl LatentNormalAdequacy {
2282    /// The nine-clause conjunction that admits the closed-form standard-normal
2283    /// kernel. A `NaN` statistic never passes: the `is_finite` clauses and the
2284    /// comparisons both reject it, which is how a degenerate sample (constant
2285    /// or non-finite) is refused without being reported as a large deviation it
2286    /// was never measured to have.
2287    pub(crate) fn passes(&self) -> bool {
2288        self.mean.abs() <= self.mean_tol
2289            && (self.sd - 1.0).abs() <= self.sd_tol
2290            && self.skew.is_finite()
2291            && self.skew.abs() <= self.skew_tol
2292            && self.excess_kurtosis.is_finite()
2293            && self.excess_kurtosis.abs() <= self.excess_kurtosis_tol
2294            && self.ks.is_finite()
2295            && self.ks <= self.ks_tol
2296            && self.tail_mass_inner <= self.tail_bound_inner
2297            && self.tail_mass_outer <= self.tail_bound_outer
2298            && self.max_abs < self.max_abs_tol
2299    }
2300
2301    /// One line naming every clause, its statistic, its bound, and the factor by
2302    /// which it missed -- so a reader can tell a marginal failure from a
2303    /// structural one without re-running the fit. `x` is the ratio of the
2304    /// statistic to its bound; a clause with `x <= 1` passed.
2305    pub(crate) fn ledger(&self) -> String {
2306        fn clause(name: &str, value: f64, bound: f64) -> String {
2307            format!(
2308                "{name}={value:.4e}/{bound:.4e}(x{:.2})",
2309                (value / bound).abs()
2310            )
2311        }
2312        format!(
2313            "n_eff={:.1} {} {} {} {} {} {} {} {}",
2314            self.effective_n,
2315            clause("|mean|", self.mean, self.mean_tol),
2316            clause("|sd-1|", self.sd - 1.0, self.sd_tol),
2317            clause("|skew|", self.skew, self.skew_tol),
2318            clause(
2319                "|excess_kurtosis|",
2320                self.excess_kurtosis,
2321                self.excess_kurtosis_tol
2322            ),
2323            clause("ks", self.ks, self.ks_tol),
2324            clause(
2325                "tail_mass_inner",
2326                self.tail_mass_inner,
2327                self.tail_bound_inner
2328            ),
2329            clause(
2330                "tail_mass_outer",
2331                self.tail_mass_outer,
2332                self.tail_bound_outer
2333            ),
2334            clause("max_abs", self.max_abs, self.max_abs_tol),
2335        )
2336    }
2337}
2338
2339/// Measure a latent-z sample against the standard-normal adequacy gate,
2340/// returning every statistic and bound rather than only the verdict.
2341pub(crate) fn latent_z_normal_adequacy(
2342    z: &Array1<f64>,
2343    weights: &Array1<f64>,
2344    policy: &LatentZPolicy,
2345) -> Result<LatentNormalAdequacy, String> {
2346    if z.len() != weights.len() {
2347        return Err(format!(
2348            "latent-measure auto-detection length mismatch: z={}, weights={}",
2349            z.len(),
2350            weights.len()
2351        ));
2352    }
2353    let weight_sum = weights.iter().copied().sum::<f64>();
2354    let weight_sq_sum = weights.iter().map(|&w| w * w).sum::<f64>();
2355    if !(weight_sum.is_finite()
2356        && weight_sum > 0.0
2357        && weight_sq_sum.is_finite()
2358        && weight_sq_sum > 0.0)
2359    {
2360        return Err("latent-measure auto-detection requires positive finite weights".to_string());
2361    }
2362    let effective_n = weight_sum * weight_sum / weight_sq_sum;
2363    if !(effective_n.is_finite() && effective_n > 1.0) {
2364        return Err(
2365            "latent-measure auto-detection requires at least two effective observations"
2366                .to_string(),
2367        );
2368    }
2369    let mean = z
2370        .iter()
2371        .zip(weights.iter())
2372        .map(|(&zi, &wi)| wi * zi)
2373        .sum::<f64>()
2374        / weight_sum;
2375    let var = z
2376        .iter()
2377        .zip(weights.iter())
2378        .map(|(&zi, &wi)| wi * (zi - mean) * (zi - mean))
2379        .sum::<f64>()
2380        / weight_sum;
2381    let sd = var.sqrt();
2382    // The two moment bounds depend only on the effective sample size, so they
2383    // are formed before the degeneracy check and stated once for both exits.
2384    let mean_tol = policy.mean_tol_multiplier / effective_n.sqrt();
2385    let sd_tol = policy.sd_tol_multiplier / (2.0 * (effective_n - 1.0).max(1.0)).sqrt();
2386    if !(mean.is_finite() && sd.is_finite() && sd > 0.0) {
2387        // A constant or non-finite sample has no shape: every standardized
2388        // statistic divides by `sd`, so skewness, kurtosis and the tail masses
2389        // are UNDEFINED here rather than merely large. `NaN` is the honest
2390        // entry -- `passes` rejects it through the same `is_finite` clauses the
2391        // conjunction always had, so this exit refuses the standard-normal
2392        // kernel exactly as the previous `Ok(false)` did, without reporting a
2393        // deviation that was never measured.
2394        return Ok(LatentNormalAdequacy {
2395            effective_n,
2396            mean,
2397            mean_tol,
2398            sd,
2399            sd_tol,
2400            skew: f64::NAN,
2401            skew_tol: policy.max_abs_skew.min(AUTO_Z_NORMAL_SKEW_TOL),
2402            excess_kurtosis: f64::NAN,
2403            excess_kurtosis_tol: policy.max_abs_excess_kurtosis.min(AUTO_Z_NORMAL_KURT_TOL),
2404            ks: f64::NAN,
2405            ks_tol: AUTO_Z_NORMAL_KS_TOL,
2406            tail_mass_inner: f64::NAN,
2407            tail_bound_inner: AUTO_Z_NORMAL_TAIL_MASS_SLACK
2408                * normal_two_sided_probability(AUTO_Z_NORMAL_TAIL_SIGMA_INNER)
2409                + AUTO_Z_NORMAL_TAIL_FLOOR_INNER,
2410            tail_mass_outer: f64::NAN,
2411            tail_bound_outer: AUTO_Z_NORMAL_TAIL_MASS_SLACK
2412                * normal_two_sided_probability(AUTO_Z_NORMAL_TAIL_SIGMA_OUTER)
2413                + AUTO_Z_NORMAL_TAIL_FLOOR_OUTER,
2414            max_abs: f64::NAN,
2415            max_abs_tol: AUTO_Z_NORMAL_MAX_ABS,
2416        });
2417    }
2418    let skew = z
2419        .iter()
2420        .zip(weights.iter())
2421        .map(|(&zi, &wi)| {
2422            let centered = (zi - mean) / sd;
2423            wi * centered.powi(3)
2424        })
2425        .sum::<f64>()
2426        / weight_sum;
2427    let excess_kurtosis = z
2428        .iter()
2429        .zip(weights.iter())
2430        .map(|(&zi, &wi)| {
2431            let centered = (zi - mean) / sd;
2432            wi * centered.powi(4)
2433        })
2434        .sum::<f64>()
2435        / weight_sum
2436        - 3.0;
2437    let ks_to_normal = weighted_ks_to_standard_normal(z, weights, weight_sum)?;
2438    let tail_mass_4 = weighted_tail_mass(z, weights, weight_sum, AUTO_Z_NORMAL_TAIL_SIGMA_INNER);
2439    let tail_mass_6 = weighted_tail_mass(z, weights, weight_sum, AUTO_Z_NORMAL_TAIL_SIGMA_OUTER);
2440    let max_abs_z = z.iter().fold(0.0_f64, |acc, &zi| acc.max(zi.abs()));
2441    let normal_tail_4 = normal_two_sided_probability(AUTO_Z_NORMAL_TAIL_SIGMA_INNER);
2442    let normal_tail_6 = normal_two_sided_probability(AUTO_Z_NORMAL_TAIL_SIGMA_OUTER);
2443    Ok(LatentNormalAdequacy {
2444        effective_n,
2445        mean,
2446        mean_tol,
2447        sd,
2448        sd_tol,
2449        skew,
2450        skew_tol: policy.max_abs_skew.min(AUTO_Z_NORMAL_SKEW_TOL),
2451        excess_kurtosis,
2452        excess_kurtosis_tol: policy.max_abs_excess_kurtosis.min(AUTO_Z_NORMAL_KURT_TOL),
2453        ks: ks_to_normal,
2454        ks_tol: AUTO_Z_NORMAL_KS_TOL,
2455        tail_mass_inner: tail_mass_4,
2456        tail_bound_inner: AUTO_Z_NORMAL_TAIL_MASS_SLACK * normal_tail_4
2457            + AUTO_Z_NORMAL_TAIL_FLOOR_INNER,
2458        tail_mass_outer: tail_mass_6,
2459        tail_bound_outer: AUTO_Z_NORMAL_TAIL_MASS_SLACK * normal_tail_6
2460            + AUTO_Z_NORMAL_TAIL_FLOOR_OUTER,
2461        max_abs: max_abs_z,
2462        max_abs_tol: AUTO_Z_NORMAL_MAX_ABS,
2463    })
2464}
2465
2466/// The global-empirical latent measure AND the fit-time record of how it was
2467/// built.
2468///
2469/// The two are returned together, never separately: the record is what makes
2470/// the measure differentiable in the sample it was compressed from (gam#2484),
2471/// and a record obtained from a second call to the builder would be a record of
2472/// a second compression.
2473pub(crate) fn build_global_empirical_latent_measure(
2474    z: &Array1<f64>,
2475    weights: &Array1<f64>,
2476    grid_size: usize,
2477) -> Result<
2478    (
2479        LatentMeasureKind,
2480        empirical_measure_sensitivity::EmpiricalZGridBuild,
2481    ),
2482    String,
2483> {
2484    let build = empirical_measure_sensitivity::build_empirical_z_grid_with_alpha(
2485        z.view(),
2486        weights.view(),
2487        grid_size,
2488        "empirical latent measure",
2489    )?;
2490    let measure = LatentMeasureKind::GlobalEmpirical {
2491        grid: build.grid.clone(),
2492    };
2493    measure.validate("empirical latent measure")?;
2494    Ok((measure, build))
2495}
2496
2497pub(crate) fn weighted_ks_to_standard_normal(
2498    z: &Array1<f64>,
2499    weights: &Array1<f64>,
2500    total_weight: f64,
2501) -> Result<f64, String> {
2502    let mut pairs = Vec::<(f64, f64)>::with_capacity(z.len());
2503    for (&zi, &wi) in z.iter().zip(weights.iter()) {
2504        if !zi.is_finite() || !wi.is_finite() || wi < 0.0 {
2505            return Err(
2506                "latent-measure KS diagnostic requires finite z and non-negative finite weights"
2507                    .to_string(),
2508            );
2509        }
2510        if wi > 0.0 {
2511            pairs.push((zi, wi));
2512        }
2513    }
2514    pairs.sort_by(|left, right| {
2515        left.0
2516            .partial_cmp(&right.0)
2517            .expect("validated latent z values are finite")
2518    });
2519    let mut prev = 0.0;
2520    let mut ks = 0.0_f64;
2521    for (zi, wi) in pairs {
2522        let cdf = normal_cdf(zi);
2523        let next = prev + wi / total_weight;
2524        ks = ks.max((cdf - prev).abs()).max((cdf - next).abs());
2525        prev = next;
2526    }
2527    Ok(ks)
2528}
2529
2530pub(crate) fn weighted_tail_mass(
2531    z: &Array1<f64>,
2532    weights: &Array1<f64>,
2533    total_weight: f64,
2534    cutoff: f64,
2535) -> f64 {
2536    z.iter()
2537        .zip(weights.iter())
2538        .filter(|&(&zi, _)| zi.abs() > cutoff)
2539        .map(|(_, &wi)| wi)
2540        .sum::<f64>()
2541        / total_weight
2542}
2543
2544// ---------------------------------------------------------------------------
2545// Cross-module constants — declared here so all submodules can reach them
2546// via `use super::*` without promoting implementation details to pub(crate).
2547// ---------------------------------------------------------------------------
2548pub(super) const BMS_AUTO_SUBSAMPLE_PHASE1_BUDGET: usize = 12;
2549pub(super) const BERNOULLI_LINK_PROBABILITY_EPS: f64 = 1e-12;
2550pub(super) const BMS_VARIANCE_FLOOR: f64 = 1e-12;
2551pub(super) const BMS_DERIV_TOL: f64 = 1e-8;
2552/// Relative tolerance below which a residual weight is treated as exhausted in
2553/// the equal-mass empirical-grid compression loop. Used both for the per-bin
2554/// "need" remaining (relative to the target bin weight) and for the per-pair
2555/// remainder (relative to that pair's weight), so a pair/bin that is filled to
2556/// within a few ulps advances the cursor instead of spinning on round-off.
2557pub(super) const EMPIRICAL_GRID_WEIGHT_EXHAUSTED_REL_TOL: f64 = 1e-14;
2558/// Upper bound (and large-`n` default) for rows-per-chunk in the parallel
2559/// row-accumulation phases.
2560///
2561/// This is also a hard *ceiling* the [`bms_row_chunk_size`] chunk sizing must
2562/// respect: several per-chunk fast paths (block-Hessian / block-gradient
2563/// assembly) allocate fixed `[0.0f64; ROW_CHUNK_SIZE]` stack buffers and index
2564/// them by the chunk's local row position, so a chunk may never carry more than
2565/// `ROW_CHUNK_SIZE` rows.
2566pub(super) const ROW_CHUNK_SIZE: usize = 1024;
2567/// Floor for rows-per-chunk: below it the per-chunk scratch allocation +
2568/// scheduler hand-off cost dominates the row arithmetic. Small enough that a
2569/// moderate `n` on a many-core box still carves several chunks per worker.
2570pub(super) const ROW_CHUNK_MIN: usize = 64;
2571/// Target number of row-chunks per rayon worker for the BMS exact-Newton
2572/// row-fan-out phases (gradient / HVP / diagonal directional-derivative sweeps).
2573///
2574/// Several chunks per worker keeps the pool load-balanced across the uneven
2575/// per-row cost tail (work-stealing moves whole chunks, never partial sums) so
2576/// the heavy coord-corrections / row-stream phases saturate the cores instead
2577/// of stranding the tail on one worker.
2578pub(super) const ROW_CHUNKS_PER_WORKER: usize = 4;
2579
2580/// Pool-aware rows-per-chunk for the BMS exact-Newton row fan-outs.
2581///
2582/// A *fixed* `ROW_CHUNK_SIZE` divisor makes the chunk **count** scale with `n`,
2583/// so at moderate `n` (e.g. `n = 10·ROW_CHUNK_SIZE` on a 64-core box) the
2584/// `into_par_iter` over `⌈n/ROW_CHUNK_SIZE⌉` chunks has far fewer tasks than
2585/// workers and most cores idle — the measured ~30-90% core utilization on the
2586/// biobank coord-corrections / row-stream phases. This sizes the chunk so the
2587/// chunk count targets `ROW_CHUNKS_PER_WORKER × worker_count` (the same policy
2588/// `chunked_row_reduction` uses), clamped to `[ROW_CHUNK_MIN, ROW_CHUNK_SIZE]`:
2589///
2590/// * the `ROW_CHUNK_SIZE` ceiling is mandatory — the block-assembly fast paths
2591///   index fixed `[…; ROW_CHUNK_SIZE]` stack buffers by local row, so a chunk
2592///   can never exceed it. At large `n` the per-1024-row count already exceeds
2593///   the worker count, so the clamp costs nothing there;
2594/// * the `ROW_CHUNK_MIN` floor stops sub-floor fan-out at tiny `n`.
2595///
2596/// Reproducibility contract (#1045): the worker count used here is the
2597/// process-stable machine parallelism (`reproducible_chunk_parallelism`), NOT
2598/// the live `rayon::current_num_threads()` of the executing (possibly scoped,
2599/// possibly shrunk) pool. Keying the chunk *count* — and hence the chunk
2600/// boundaries `chunk_idx·chunk → (chunk_idx+1)·chunk` — to the transient pool
2601/// size made the per-chunk row sums regroup when the pool was narrowed, so a
2602/// fit reduced over these chunks and fed into the iterative REML optimizer moved
2603/// its `(ρ, λ)` selection with the pool size. Anchoring to a process constant
2604/// makes the boundaries — and therefore the `try_fold`/`try_reduce` reduction
2605/// tree that round-trips through them — invariant to how many workers run the
2606/// fit, while rayon still fans the chunks across whatever workers exist. For a
2607/// given `n` the returned chunk size is stable across calls and pool sizes.
2608#[inline]
2609pub(super) fn bms_row_chunk_size(n: usize) -> usize {
2610    if n == 0 {
2611        return ROW_CHUNK_SIZE;
2612    }
2613    let workers = crate::marginal_slope_shared::reproducible_chunk_parallelism();
2614    let target_chunks = workers.saturating_mul(ROW_CHUNKS_PER_WORKER).max(1);
2615    // Rows per chunk that yields ≈ `target_chunks` chunks, clamped into
2616    // `[ROW_CHUNK_MIN, ROW_CHUNK_SIZE]`.
2617    n.div_ceil(target_chunks)
2618        .clamp(ROW_CHUNK_MIN, ROW_CHUNK_SIZE)
2619}
2620pub(super) const EXACT_WORK_LOG_MIN_ROWS: usize = 50_000;
2621pub(super) const BMS_ROW_PRIMARY_HESSIAN_EXPECTED_REUSE_PASSES: usize = 3;
2622pub(super) const BMS_ROW_PRIMARY_HESSIAN_MIN_REUSE_PASSES: usize = 2;
2623pub(super) const BMS_ROW_PRIMARY_HESSIAN_TILE_ROWS: usize = 8192;
2624pub(super) const BMS_ROW_PRIMARY_HESSIAN_SINGLE_FRACTION_NUM: u64 = 1;
2625pub(super) const BMS_ROW_PRIMARY_HESSIAN_SINGLE_FRACTION_DEN: u64 = 4;
2626pub(super) const BMS_ROW_PRIMARY_HESSIAN_GLOBAL_FRACTION_NUM: u64 = 1;
2627pub(super) const BMS_ROW_PRIMARY_HESSIAN_GLOBAL_FRACTION_DEN: u64 = 2;
2628pub(super) const BERNOULLI_MARGSLOPE_LINE_SEARCH_EARLY_EXIT_CHUNK_ROWS: usize = 10_000;
2629
2630// ---------------------------------------------------------------------------
2631// Submodule declarations
2632// ---------------------------------------------------------------------------
2633pub(crate) mod block_specs;
2634pub mod conditional_score_covariance;
2635pub(crate) mod exact_eval_cache;
2636pub(crate) mod family;
2637pub(crate) mod flex_row_program;
2638pub(crate) mod gradient_paths;
2639pub(crate) mod hessian_paths;
2640pub(crate) mod install_flex;
2641pub(crate) mod row_kernel;
2642#[cfg(test)]
2643mod tests {
2644    include!("../../../../tests/src_modules/misc/families_bms_identifiability_rigid_tests.rs");
2645    include!(
2646        "../../../../tests/src_modules/optimization/families_bms_joint_hessian_hvp_correction_tests.rs"
2647    );
2648
2649    #[test]
2650    fn empirical_grid_constructor_preserves_canonical_node_order() {
2651        let grid = EmpiricalZGrid::new(
2652            vec![-2.0, 0.5, 1.0],
2653            vec![0.3, 0.5, 0.2],
2654            "sorted-grid invariant",
2655        )
2656        .expect("canonical sorted grid");
2657        assert_eq!(grid.nodes, vec![-2.0, 0.5, 1.0]);
2658        assert_eq!(grid.weights, vec![0.3, 0.5, 0.2]);
2659    }
2660
2661    #[test]
2662    fn empirical_grid_constructor_rejects_noncanonical_node_order() {
2663        let err = EmpiricalZGrid::new(vec![0.0, -1.0], vec![0.5, 0.5], "sorted-grid invariant")
2664            .expect_err("constructed grids must already be canonical");
2665        assert!(err.contains("nodes must be sorted ascending"), "{err}");
2666    }
2667}
2668
2669#[cfg(test)]
2670mod stacked_first_stage_sandwich_2484_tests {
2671    use super::{stacked_first_stage_sandwich_cov, weighted_ridge_sandwich_cov};
2672    use ndarray::{Array1, Array2, array};
2673
2674    /// `A`, weights, and a normal matrix `M = AᵀWA + λR` built exactly the way
2675    /// the calibration builds it, so the sandwich is exercised on a realistic
2676    /// bread rather than on an identity.
2677    fn system(basis: &Array2<f64>, weights: &Array1<f64>) -> Array2<f64> {
2678        let mut wa = basis.clone();
2679        for (mut row, &w) in wa.rows_mut().into_iter().zip(weights.iter()) {
2680            row.iter_mut().for_each(|value| *value *= w);
2681        }
2682        let mut m = basis.t().dot(&wa);
2683        let diag: Vec<f64> = (0..m.nrows()).map(|j| m[[j, j]]).collect();
2684        for (j, value) in diag.iter().enumerate() {
2685            m[[j, j]] += value * 1.0e-8;
2686        }
2687        m
2688    }
2689
2690    /// PAIRED fixture: every conditioning row appears twice with equal weight
2691    /// and mean residuals `+c` / `−c`.
2692    ///
2693    /// That makes both cross-terms cancel EXACTLY rather than approximately —
2694    /// the bread's `Σ w·û·A Aᵀ` cancels because `û` flips sign while `A Aᵀ`
2695    /// does not, and the meat's `Σ w²·û·r·A Aᵀ` cancels because `r` depends on
2696    /// `û²` and is therefore equal across the pair. So this arm asserts an
2697    /// identity, not a tolerance: with no third moment, the joint sandwich must
2698    /// reproduce the block-diagonal form the code used to assume.
2699    #[test]
2700    fn symmetric_residuals_reproduce_the_block_diagonal_form_2484() {
2701        let basis = array![
2702            [1.0, 0.4],
2703            [1.0, 0.4],
2704            [1.0, -0.7],
2705            [1.0, -0.7],
2706            [1.0, 1.3],
2707            [1.0, 1.3],
2708        ];
2709        let weights = Array1::from(vec![1.0, 1.0, 0.5, 0.5, 2.0, 2.0]);
2710        let mean_residuals = vec![0.6, -0.6, 0.9, -0.9, 0.3, -0.3];
2711        // r depends on û only through û², so it is equal within each pair.
2712        let var_residuals: Vec<f64> = mean_residuals.iter().map(|&u| u * u - 0.5).collect();
2713        let m = system(&basis, &weights);
2714
2715        let joint = stacked_first_stage_sandwich_cov(
2716            basis.view(),
2717            weights.view(),
2718            &mean_residuals,
2719            &var_residuals,
2720            &m,
2721        )
2722        .expect("joint sandwich");
2723        let mean_block =
2724            weighted_ridge_sandwich_cov(basis.view(), &mean_residuals, weights.view(), &m)
2725                .expect("mean sandwich");
2726        let var_block =
2727            weighted_ridge_sandwich_cov(basis.view(), &var_residuals, weights.view(), &m)
2728                .expect("var sandwich");
2729
2730        let p = basis.ncols();
2731        for i in 0..p {
2732            for j in 0..p {
2733                let tol = 1.0e-9 * (1.0 + mean_block[[i, j]].abs());
2734                assert!(
2735                    (joint[[i, j]] - mean_block[[i, j]]).abs() <= tol,
2736                    "gam#2484: with no third moment the (mean,mean) block must reproduce the \
2737                     standalone sandwich; got {} vs {}",
2738                    joint[[i, j]],
2739                    mean_block[[i, j]]
2740                );
2741                let tol_v = 1.0e-9 * (1.0 + var_block[[i, j]].abs());
2742                assert!(
2743                    (joint[[p + i, p + j]] - var_block[[i, j]]).abs() <= tol_v,
2744                    "gam#2484: with no third moment the (var,var) block must reproduce the \
2745                     standalone sandwich; got {} vs {}",
2746                    joint[[p + i, p + j]],
2747                    var_block[[i, j]]
2748                );
2749                assert!(
2750                    joint[[i, p + j]].abs() <= 1.0e-9,
2751                    "gam#2484: the cross-block must vanish when the residual third moment does; \
2752                     got {}",
2753                    joint[[i, p + j]]
2754                );
2755            }
2756        }
2757    }
2758
2759    /// The arm that asserts the change DOES something. Break the pairing so the
2760    /// residual carries a third moment, and the cross-block must be measurably
2761    /// non-zero — otherwise an implementation that quietly returns the old
2762    /// block-diagonal form passes the test above and ships.
2763    #[test]
2764    fn skewed_residuals_make_the_cross_block_nonzero_2484() {
2765        let basis = array![
2766            [1.0, 0.4],
2767            [1.0, 0.9],
2768            [1.0, -0.7],
2769            [1.0, 0.2],
2770            [1.0, 1.3],
2771            [1.0, -1.1],
2772        ];
2773        let weights = Array1::from(vec![1.0, 1.0, 0.5, 0.5, 2.0, 2.0]);
2774        // Strongly right-skewed mean residuals: one large positive, the rest
2775        // small negative. This is the shape the adequacy gate rejects.
2776        let mean_residuals = vec![-0.2, -0.3, -0.25, -0.15, 2.4, -0.35];
2777        let var_residuals: Vec<f64> = mean_residuals.iter().map(|&u| u * u - 0.5).collect();
2778        let m = system(&basis, &weights);
2779
2780        let joint = stacked_first_stage_sandwich_cov(
2781            basis.view(),
2782            weights.view(),
2783            &mean_residuals,
2784            &var_residuals,
2785            &m,
2786        )
2787        .expect("joint sandwich");
2788        let var_block =
2789            weighted_ridge_sandwich_cov(basis.view(), &var_residuals, weights.view(), &m)
2790                .expect("var sandwich");
2791
2792        let p = basis.ncols();
2793        let cross = (0..p)
2794            .flat_map(|i| (0..p).map(move |j| (i, j)))
2795            .map(|(i, j)| joint[[i, p + j]].abs())
2796            .fold(0.0_f64, f64::max);
2797        let scale = (0..p)
2798            .map(|i| joint[[i, i]].abs().max(joint[[p + i, p + i]].abs()))
2799            .fold(0.0_f64, f64::max);
2800        assert!(
2801            cross > 1.0e-6 * scale.max(1.0e-12),
2802            "gam#2484: a skewed residual must produce a NON-ZERO first-stage cross-block. \
2803             max|cross|={cross:e} against block scale {scale:e}. If this fails, the \
2804             implementation is still returning a block-diagonal V1 and the whole change is \
2805             inert."
2806        );
2807
2808        // And the (var,var) block must differ from the standalone sandwich: the
2809        // triangular bread feeds the mean-stage uncertainty into it.
2810        let var_shift = (0..p)
2811            .flat_map(|i| (0..p).map(move |j| (i, j)))
2812            .map(|(i, j)| (joint[[p + i, p + j]] - var_block[[i, j]]).abs())
2813            .fold(0.0_f64, f64::max);
2814        assert!(
2815            var_shift > 1.0e-6 * scale.max(1.0e-12),
2816            "gam#2484: the (var,var) block must absorb the mean-stage uncertainty through the \
2817             lower-triangular bread; max shift {var_shift:e} is indistinguishable from the old \
2818             standalone block"
2819        );
2820    }
2821}
2822
2823pub(crate) mod axis_direction_search;
2824pub(crate) mod cell_moment_assembly;
2825#[cfg(test)]
2826mod empirical_measure_2484_tests;
2827pub(crate) mod empirical_measure_sensitivity;
2828// #932 BMS flex single-source jet substrate (runtime-dimension `Jet2` + IFT
2829// lift + cell base-moment jets). A bare `#[cfg(test)] mod` with an allowed name
2830// so the build.rs ban-scanner exempts it; shared by its own FD gates and the
2831// `cell_moment_assembly` flex-fixture oracle gate as a private child of `bms`.
2832#[cfg(test)]
2833mod test_support;
2834// #932 INDEPENDENT adversarial verifier (bms-flex-verify): a high-order
2835// finite-difference oracle on the production compiled lowering
2836// `lower_bms_flex_row_order2_from_parts` of the canonical BMS FLEX program,
2837// plus a moving-edge Leibniz
2838// cross-check + a planted-corruption tripwire. Bare `#[cfg(test)] mod` with the
2839// allowed `*_tests` name so the build.rs ban-scanner exempts it; owned solely by
2840// the verifier (never edits the implementer's row_primary_hessian /
2841// gradient_paths / cell_moment_assembly).
2842pub(crate) mod custom_family_impl;
2843#[cfg(test)]
2844mod flex_verify_932_tests;
2845// #932 direct production-path measurement: forced 65-node empirical grid,
2846// warmed/cold row-op allocation counting + ns/row diagnostics for the MSI
2847// A/B ledger. The asserted gate is per-row allocation calls (deterministic);
2848// timing is eprintln-only per the SPEC ban on wall-clock correctness budgets.
2849#[cfg(test)]
2850mod flex_measure_932_tests;
2851// gam#2768 unit gates on the shared latent-measure decision and the conditional
2852// location-scale calibration it escalates to. Bare `#[cfg(test)] mod` with the
2853// allowed `*_tests` name so the build.rs ban-scanner exempts it.
2854#[cfg(test)]
2855mod latent_measure_2768_tests;
2856pub(crate) mod row_primary_hessian;
2857
2858pub use block_specs::fit_bernoulli_marginal_slope_terms;
2859pub use conditional_score_covariance::{
2860    ConditionalScoreCoordinate, ConditionalScoreCovariance, ScoreCovarianceField,
2861};
2862pub use gradient_paths::{
2863    MarginalSlopeCovariance, MarginalSlopeCovarianceShape, marginal_slope_covariance_from_scores,
2864    marginal_slope_preserving_scale, marginal_slope_probit_eta, padded_deviation_seed,
2865};
2866pub use install_flex::CrossBlockIdentifiabilityWarning;
2867pub(crate) use install_flex::FlexCompileOutcome;
2868
2869// pub(crate) re-exports for internal callers:
2870pub(crate) use block_specs::push_deviation_aux_blockspecs;
2871pub use block_specs::{BmsLogslopeJacobian, BmsMarginalJacobian};
2872pub(crate) use family::{
2873    BernoulliMarginalLinkMap, bernoulli_marginal_link_map,
2874    build_link_deviation_block_from_knots_design_seed_and_weights,
2875    build_score_warp_deviation_block_from_seed,
2876};
2877pub(crate) use gradient_paths::MarginalSlopeCovarianceRef;
2878pub(crate) use gradient_paths::signed_probit_neglog_unary_stack;
2879pub(crate) use gradient_paths::standardize_latent_z_with_policy;
2880pub(crate) use gradient_paths::{
2881    empirical_intercept_from_marginal, signed_probit_neglog_derivatives_up_to_fourth,
2882    unary_derivatives_inverse_sqrt, unary_derivatives_log, unary_derivatives_log_normal_pdf,
2883    unary_derivatives_neglog_phi, unary_derivatives_sqrt,
2884};
2885pub(crate) use install_flex::{
2886    install_compiled_flex_block_into_runtime, project_monotone_feasible_beta,
2887};