gam_models/bms/mod.rs
1use crate::cubic_cell_kernel as exact_kernel;
2use crate::custom_family::{
3 BatchedOuterGradientTerms, BlockEffectiveJacobian, BlockWorkingSet, BlockwiseFitOptions,
4 CustomFamily, CustomFamilyWarmStart, EvalMode, ExactNewtonJointGradientEvaluation,
5 ExactNewtonJointHessianWorkspace, FamilyEvaluation, FamilyLinearizationState,
6 ParameterBlockSpec, ParameterBlockState, PenaltyMatrix, custom_family_outer_derivatives,
7 evaluate_custom_family_joint_hyper_efs_shared, evaluate_custom_family_joint_hyper_shared,
8 fit_custom_family, joint_hyper_options_for_outer_tolerance,
9};
10use crate::fit_orchestration::drivers::{
11 ExactJointHyperSetup, apply_spatial_anisotropy_pilot_initializer,
12 build_term_collection_designs_and_freeze_joint, optimize_spatial_length_scale_exact_joint,
13 spatial_length_scale_term_indices,
14};
15use crate::marginal_slope_shared::{
16 CoeffSupport, DirectionalScaleJets, ObservedDenestedCellPartials, SparsePrimaryCoeffJetView,
17 add_optional_matrix, add_optional_vector, add_two_surface_psi_outer,
18 build_denested_partition_cells as shared_denested_partition_cells, chunked_row_reduction,
19 directional_obj_grad_hess, eval_coeff4_at, is_sigma_aux_index as shared_is_sigma_aux_index,
20 observed_denested_cell_partials as shared_observed_denested_cell_partials, outer_row_indices,
21 outer_weighted_rows, parameter_block_specs_match_rows, probit_frailty_scale,
22 probit_frailty_scale_multi_dir_jet, psi_derivative_location, scale_coeff4,
23};
24use crate::model_types::UnifiedFitResult;
25use crate::outer_subsample::WeightedOuterRow;
26use crate::parameter_block::ParameterBlockInput;
27use crate::probability::{
28 normal_cdf, normal_logcdf, normal_pdf, signed_probit_logcdf_and_mills_ratio,
29 standard_normal_quantile,
30};
31use crate::row_kernel::{
32 RowKernel, RowKernelHessianWorkspace, build_row_kernel_cache, row_kernel_gradient,
33 row_kernel_hessian_dense, row_kernel_log_likelihood,
34};
35use crate::spatial_psi_bridge::build_block_spatial_psi_derivatives;
36use crate::survival::lognormal_kernel::FrailtySpec;
37use crate::wiggle::initializewiggle_knots_from_seed;
38use gam_linalg::matrix::{DesignMatrix, SymmetricMatrix};
39use gam_math::jet_partitions::MultiDirJet;
40use gam_problem::{
41 ExactNewtonJointPsiSecondOrderTerms, ExactNewtonJointPsiTerms, ExactNewtonJointPsiWorkspace,
42 HyperOperator, InverseLink, StandardLink, WigglePenaltyConfig,
43};
44use gam_solve::estimate::reml::reml_outer_engine::{DenseSpectralOperator, HessianFactorization};
45use gam_solve::pirls::LinearInequalityConstraints;
46use gam_terms::smooth::{
47 SpatialLengthScaleOptimizationOptions, SpatialLogKappaCoords, TermCollectionDesign,
48 TermCollectionSpec,
49};
50use ndarray::{Array1, Array2, ArrayView1, ArrayView2, ArrayViewMut1, s};
51use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
52use serde::{Deserialize, Serialize};
53use std::cell::RefCell;
54use std::collections::HashMap;
55use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
56use std::sync::{Arc, Mutex, OnceLock};
57
58pub mod deviation_runtime;
59pub mod gpu;
60pub use deviation_runtime::DeviationRuntime;
61pub use deviation_runtime::ParametricAnchorBlock;
62
63/// Above this size, FLEX spatial length-scale optimization uses the pilot
64/// geometry initializer and skips the iterative joint κ/ψ outer loop. This is
65/// a spatial-optimizer policy only; it must not gate exact outer Hessian
66/// capability or row-cell moment materialization.
67pub(crate) const BMS_FLEX_SPATIAL_OUTER_PILOT_ROW_THRESHOLD: usize = 50_000;
68
69#[derive(Clone, Debug)]
70pub struct DeviationBlockConfig {
71 pub degree: usize,
72 pub num_internal_knots: usize,
73 pub penalty_order: usize,
74 pub penalty_orders: Vec<usize>,
75 pub double_penalty: bool,
76 pub monotonicity_eps: f64,
77}
78
79impl Default for DeviationBlockConfig {
80 fn default() -> Self {
81 WigglePenaltyConfig::cubic_triple_operator_default().into()
82 }
83}
84
85impl DeviationBlockConfig {
86 pub fn triple_penalty_default() -> Self {
87 Self::default()
88 }
89}
90
91impl From<WigglePenaltyConfig> for DeviationBlockConfig {
92 fn from(cfg: WigglePenaltyConfig) -> Self {
93 let penalty_order = *cfg.penalty_orders.iter().max().unwrap_or(&2);
94 Self {
95 degree: cfg.degree,
96 num_internal_knots: cfg.num_internal_knots,
97 penalty_order,
98 penalty_orders: cfg.penalty_orders,
99 double_penalty: cfg.double_penalty,
100 monotonicity_eps: cfg.monotonicity_eps,
101 }
102 }
103}
104
105#[derive(Clone)]
106pub(crate) struct DeviationPrepared {
107 pub(crate) block: ParameterBlockInput,
108 pub(crate) runtime: DeviationRuntime,
109}
110
111impl std::fmt::Debug for DeviationPrepared {
112 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113 f.debug_struct("DeviationPrepared").finish_non_exhaustive()
114 }
115}
116
117#[derive(Clone)]
118pub struct BernoulliMarginalSlopeTermSpec {
119 pub y: Array1<f64>,
120 pub weights: Array1<f64>,
121 pub z: Array1<f64>,
122 pub base_link: InverseLink,
123 pub marginalspec: TermCollectionSpec,
124 pub logslopespec: TermCollectionSpec,
125 pub marginal_offset: Array1<f64>,
126 pub logslope_offset: Array1<f64>,
127 /// GaussianShift frailty on the final probit index: U ~ N(0, σ²) added
128 /// to the scalar argument of Φ. This is exact because the sextic
129 /// microcell kernel is preserved — the Gaussian-decoupling identity
130 /// E[Φ(η + U)] = Φ(η / √(1+σ²)) rescales the index by 1/τ where
131 /// τ = √(1+σ²), and every derivative chain rule factor is polynomial
132 /// in τ, so all six kernel derivatives remain closed-form.
133 ///
134 /// **HazardMultiplier frailty is NOT supported in this family.**
135 /// HazardMultiplier frailty + score_warp/linkwiggle cubic marginal-slope
136 /// is not finite-state exact. For hazard-multiplier frailty, use the
137 /// standalone LatentCloglogBinomial / LatentSurvival families instead.
138 pub frailty: FrailtySpec,
139 pub score_warp: Option<DeviationBlockConfig>,
140 pub link_dev: Option<DeviationBlockConfig>,
141 pub latent_z_policy: LatentZPolicy,
142 /// Out-of-fold Stage-1 score-influence Jacobian `J = ∂z/∂θ₁` (n × p₁)
143 /// from cross-fitting a CTN transformation-normal Stage-1 model (#461).
144 /// When `Some`, the realized leakage directions `Z_infl = diag(s_f·β̂₀)·J`
145 /// are absorbed as a null-penalized block so the joint solve makes the
146 /// β estimating equation orthogonal to `span(Z_infl)` — the x-dependent
147 /// realization of `ψ − Π_η[ψ]`. `None` ⇒ raw `--z-column` with no CTN
148 /// Stage-1, in which case the free 1-D `score_warp` spline is the
149 /// fallback basis (it spans only the x-free leakage column).
150 pub score_influence_jacobian: Option<Array2<f64>>,
151}
152
153pub struct BernoulliMarginalSlopeFitResult {
154 pub fit: UnifiedFitResult,
155 pub marginalspec_resolved: TermCollectionSpec,
156 pub logslopespec_resolved: TermCollectionSpec,
157 pub marginal_design: TermCollectionDesign,
158 pub logslope_design: TermCollectionDesign,
159 pub baseline_marginal: f64,
160 pub baseline_logslope: f64,
161 pub z_normalization: LatentZNormalization,
162 pub latent_measure: LatentMeasureKind,
163 pub score_warp_runtime: Option<DeviationRuntime>,
164 pub link_dev_runtime: Option<DeviationRuntime>,
165 /// Learned or fixed Gaussian-shift frailty SD. `None` = no frailty.
166 pub gaussian_frailty_sd: Option<f64>,
167 /// Structured warnings emitted during fit-time setup when a flex
168 /// block was fully aliased by its anchor union and got dropped. The
169 /// fit proceeds without the dropped block (its contribution to the
170 /// joint design was numerically reproducible by the anchor span, so
171 /// keeping it would leave the joint Hessian rank-deficient). Empty
172 /// for fits where every flex block carried independent directions.
173 pub cross_block_warnings: Vec<CrossBlockIdentifiabilityWarning>,
174 /// Optional weighted rank inverse-normal (Blom rankit) calibration
175 /// installed at fit time when the auto latent-z normality check
176 /// failed. `Some(_)` ⇒ the training z was transformed in place via
177 /// [`LatentZRankIntCalibration::apply_to_training`] before any
178 /// downstream consumer (pooled probit baseline, term-collection
179 /// designs, family PIRLS loops) saw it, and the rigid kernel
180 /// routes through the standard-normal closed-form path on the
181 /// calibrated scale. `None` ⇒ no calibration was applied (training
182 /// z already passed the standard-normal diagnostics, or the caller
183 /// explicitly selected a non-Auto `LatentMeasureSpec`).
184 ///
185 /// Persisted to disk so prediction applies the same monotone map
186 /// via [`LatentZRankIntCalibration::apply_at_predict`] to incoming
187 /// z before the standard-normal kernel runs. The public field name
188 /// is `latent_z_rank_int_calibration` — Agent D's persistence
189 /// pipeline reads it under that exact identifier.
190 pub latent_z_rank_int_calibration: Option<LatentZRankIntCalibration>,
191 /// Optional conditional location-scale calibration of the latent score
192 /// (#905). `Some(_)` ⇒ the Auto path's conditional `E[z|C]`/`Var(z|C)` Rao
193 /// gate detected PC/grouping-dependence that the pooled-marginal gate
194 /// cannot see, so the training z was replaced in place by
195 /// `ζ = (z − m(C))/√v(C)` (via [`LatentZConditionalCalibration::apply`])
196 /// before any downstream consumer saw it. Mutually exclusive with
197 /// `latent_z_rank_int_calibration`: rank-INT fixes a pooled-marginal
198 /// defect, the conditional correction fixes a conditional-shift defect that
199 /// rank-INT provably cannot. Persisted so prediction rebuilds `a(C)` from
200 /// the (reproducible) marginal design and applies the identical map.
201 pub latent_z_conditional_calibration: Option<LatentZConditionalCalibration>,
202}
203
204#[derive(Clone, Debug)]
205pub enum LatentZCheckMode {
206 Strict,
207 WarnOnly,
208 Off,
209}
210
211#[derive(Clone, Debug)]
212pub enum LatentZNormalizationMode {
213 None,
214 FitWeighted,
215 Frozen { mean: f64, sd: f64 },
216}
217
218pub const DEFAULT_EMPIRICAL_LATENT_GRID_SIZE: usize = 65;
219pub(crate) const AUTO_Z_NORMAL_SKEW_TOL: f64 = 0.10;
220pub(crate) const AUTO_Z_NORMAL_KURT_TOL: f64 = 0.25;
221pub(crate) const AUTO_Z_NORMAL_KS_TOL: f64 = 0.025;
222pub(crate) const AUTO_Z_NORMAL_MAX_ABS: f64 = 8.0;
223/// Inner σ level at which the empirical tail mass of latent z is compared
224/// against the standard normal's theoretical two-sided tail in the auto
225/// normality gate. Chosen well inside `AUTO_Z_NORMAL_MAX_ABS` so a fat inner
226/// tail is caught before any single observation trips the hard `max |z|` bound.
227pub(crate) const AUTO_Z_NORMAL_TAIL_SIGMA_INNER: f64 = 4.0;
228/// Outer σ level for the same tail-mass comparison; catches heavier far-tail
229/// excess that the inner level can miss.
230pub(crate) const AUTO_Z_NORMAL_TAIL_SIGMA_OUTER: f64 = 6.0;
231/// Multiplier applied to the normal's theoretical tail mass before comparison:
232/// the empirical tail may be up to this many times the Gaussian tail at the
233/// same σ before the gate fails, allowing for finite-sample sampling noise.
234pub(crate) const AUTO_Z_NORMAL_TAIL_MASS_SLACK: f64 = 2.0;
235/// Absolute additive floor on the inner-σ tail comparison, so the gate does
236/// not fail on round-off when the Gaussian tail itself is already tiny.
237pub(crate) const AUTO_Z_NORMAL_TAIL_FLOOR_INNER: f64 = 1e-5;
238/// Absolute additive floor on the outer-σ tail comparison; smaller than the
239/// inner floor because the 6σ Gaussian tail is many orders smaller than 4σ.
240pub(crate) const AUTO_Z_NORMAL_TAIL_FLOOR_OUTER: f64 = 1e-8;
241/// Significance level for the conditional `E[z|C]` / `Var(z|C)` Rao gate in the
242/// core Auto path (#905). When the latent score's conditional mean or variance
243/// on the marginal-index span `a(C)` is significant at this level, the Auto
244/// path escalates from the pooled-marginal rank-INT to a conditional
245/// location-scale correction. Chosen small (0.1%) so the escalation fires only
246/// on clear conditional structure, not finite-sample noise — the gate runs once
247/// over the whole training sample, so a per-test α this tight still has ample
248/// power against the grouping mean-shift the issue names.
249pub(crate) const AUTO_Z_CONDITIONAL_RAO_ALPHA: f64 = 1.0e-3;
250/// Relative ridge added to the weighted normal equations when regressing the
251/// latent score on the marginal-index span for the conditional correction.
252/// Stabilizes the solve when `a(C)` is rank-deficient or collinear (penalized
253/// spline marginal indices routinely are) without materially biasing the
254/// conditional mean/variance fit.
255pub(crate) const AUTO_Z_CONDITIONAL_RIDGE_REL: f64 = 1.0e-8;
256/// Floor on the fitted conditional variance `v(C)`, as a fraction of the global
257/// weighted variance of the latent score. Keeps `ζ = (z−m)/√v` finite and
258/// well-scaled where the linear variance model would otherwise fit a
259/// non-positive or vanishing conditional variance.
260pub(crate) const AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC: f64 = 1.0e-3;
261
262#[derive(Clone, Copy, Debug, PartialEq, Eq)]
263pub enum LatentMeasureSpec {
264 Auto { grid_size: usize },
265 StandardNormal,
266 GlobalEmpirical { grid_size: usize },
267}
268
269impl LatentMeasureSpec {
270 pub fn auto_default() -> Self {
271 Self::Auto {
272 grid_size: DEFAULT_EMPIRICAL_LATENT_GRID_SIZE,
273 }
274 }
275}
276
277impl Default for LatentMeasureSpec {
278 fn default() -> Self {
279 Self::auto_default()
280 }
281}
282
283#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
284pub struct EmpiricalZGrid {
285 pub nodes: Vec<f64>,
286 pub weights: Vec<f64>,
287}
288
289impl EmpiricalZGrid {
290 /// Construct a grid whose node/weight invariants (equal length ≥ 2, finite
291 /// nodes, finite positive weights, weights summing to 1 within 1e-8) are
292 /// enforced up-front. Prefer this over building the struct literally;
293 /// every code path that goes through `new` is guaranteed to satisfy the
294 /// same contract that `validate_empirical_z_grid` checks on read.
295 pub fn new(nodes: Vec<f64>, weights: Vec<f64>, context: &str) -> Result<Self, String> {
296 validate_empirical_z_grid(&nodes, &weights, context)?;
297 Ok(Self { nodes, weights })
298 }
299
300 /// Iterate over co-indexed `(node, weight)` pairs. Use this instead of
301 /// reading `.nodes`/`.weights` separately whenever a loop wants both
302 /// arrays in lockstep — eliminates the chance of mismatched indexing.
303 #[inline]
304 pub fn pairs(&self) -> impl Iterator<Item = (f64, f64)> + '_ {
305 self.nodes.iter().copied().zip(self.weights.iter().copied())
306 }
307}
308
309#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
310#[serde(tag = "kind", rename_all = "kebab-case")]
311#[derive(Default)]
312pub enum LatentMeasureKind {
313 #[default]
314 StandardNormal,
315 GlobalEmpirical {
316 grid: EmpiricalZGrid,
317 },
318 LocalEmpirical {
319 feature_cols: Vec<usize>,
320 #[serde(default)]
321 input_scales: Option<Vec<f64>>,
322 centers: Vec<Vec<f64>>,
323 grids: Vec<EmpiricalZGrid>,
324 top_k: usize,
325 bandwidth: f64,
326 #[serde(skip)]
327 train_row_mixtures: Arc<Vec<Vec<(usize, f64)>>>,
328 },
329}
330
331impl LatentMeasureKind {
332 pub fn validate(&self, context: &str) -> Result<(), String> {
333 match self {
334 Self::StandardNormal => Ok(()),
335 Self::GlobalEmpirical { grid } => {
336 validate_empirical_z_grid(&grid.nodes, &grid.weights, context)
337 }
338 Self::LocalEmpirical {
339 feature_cols,
340 input_scales,
341 centers,
342 grids,
343 top_k,
344 bandwidth,
345 ..
346 } => {
347 if feature_cols.is_empty() {
348 return Err(format!(
349 "{context} local empirical latent measure needs feature columns"
350 ));
351 }
352 if centers.is_empty() {
353 return Err(format!(
354 "{context} local empirical latent measure needs centers"
355 ));
356 }
357 if centers.len() != grids.len() {
358 return Err(format!(
359 "{context} local empirical latent measure center/grid length mismatch: centers={}, grids={}",
360 centers.len(),
361 grids.len()
362 ));
363 }
364 if *top_k == 0 || *top_k > centers.len() {
365 return Err(format!(
366 "{context} local empirical latent measure top_k must be in 1..={}, got {top_k}",
367 centers.len()
368 ));
369 }
370 if !(*bandwidth).is_finite() || *bandwidth <= 0.0 {
371 return Err(format!(
372 "{context} local empirical latent measure bandwidth must be finite and positive, got {bandwidth}"
373 ));
374 }
375 if let Some(scales) = input_scales.as_ref() {
376 if scales.len() != feature_cols.len() {
377 return Err(format!(
378 "{context} local empirical latent measure input scale dimension mismatch: scales={}, features={}",
379 scales.len(),
380 feature_cols.len()
381 ));
382 }
383 for (scale_idx, scale) in scales.iter().enumerate() {
384 if !(scale.is_finite() && *scale > 0.0) {
385 return Err(format!(
386 "{context} local empirical latent measure input scale {scale_idx} must be finite and positive, got {scale}"
387 ));
388 }
389 }
390 }
391 for (center_idx, center) in centers.iter().enumerate() {
392 if center.len() != feature_cols.len() {
393 return Err(format!(
394 "{context} local empirical latent center {center_idx} dimension mismatch: got {}, expected {}",
395 center.len(),
396 feature_cols.len()
397 ));
398 }
399 if center.iter().any(|value| !value.is_finite()) {
400 return Err(format!(
401 "{context} local empirical latent center {center_idx} has non-finite coordinates"
402 ));
403 }
404 }
405 for (grid_idx, grid) in grids.iter().enumerate() {
406 validate_empirical_z_grid(
407 &grid.nodes,
408 &grid.weights,
409 &format!("{context} local empirical grid {grid_idx}"),
410 )?;
411 }
412 Ok(())
413 }
414 }
415 }
416
417 pub(crate) fn is_empirical(&self) -> bool {
418 matches!(
419 self,
420 Self::GlobalEmpirical { .. } | Self::LocalEmpirical { .. }
421 )
422 }
423
424 /// Per-row empirical latent grid, borrowed where possible. This sits in
425 /// the innermost per-row loops of every criterion/gradient/Hessian
426 /// evaluation, so the global grid MUST come back as a borrow — the old
427 /// `grid.clone()` here allocated two `grid_size`-length vectors per row
428 /// per evaluation across the whole fit. Only the local-mixture path,
429 /// which genuinely synthesizes a new grid per row, returns an owned
430 /// value.
431 pub(crate) fn empirical_grid_for_training_row(
432 &self,
433 row: usize,
434 ) -> Result<Option<std::borrow::Cow<'_, EmpiricalZGrid>>, String> {
435 match self {
436 Self::StandardNormal => Ok(None),
437 Self::GlobalEmpirical { grid } => Ok(Some(std::borrow::Cow::Borrowed(grid))),
438 Self::LocalEmpirical {
439 grids,
440 train_row_mixtures,
441 ..
442 } => {
443 let mixture = train_row_mixtures.get(row).ok_or_else(|| {
444 format!(
445 "local empirical latent measure is missing training mixture for row {row}"
446 )
447 })?;
448 Ok(Some(std::borrow::Cow::Owned(combine_empirical_grids(
449 grids, mixture,
450 )?)))
451 }
452 }
453 }
454}
455
456pub(crate) fn validate_empirical_z_grid(
457 nodes: &[f64],
458 weights: &[f64],
459 context: &str,
460) -> Result<(), String> {
461 if nodes.len() != weights.len() {
462 return Err(format!(
463 "{context} empirical latent measure node/weight length mismatch: nodes={}, weights={}",
464 nodes.len(),
465 weights.len()
466 ));
467 }
468 if nodes.len() < 2 {
469 return Err(format!(
470 "{context} empirical latent measure requires at least two nodes"
471 ));
472 }
473 let mut total = 0.0;
474 for (idx, (&node, &weight)) in nodes.iter().zip(weights.iter()).enumerate() {
475 if !node.is_finite() {
476 return Err(format!(
477 "{context} empirical latent measure node {idx} is non-finite ({node})"
478 ));
479 }
480 if !(weight.is_finite() && weight > 0.0) {
481 return Err(format!(
482 "{context} empirical latent measure weight {idx} must be finite and positive, got {weight}"
483 ));
484 }
485 total += weight;
486 }
487 if !(total.is_finite() && (total - 1.0).abs() <= 1e-8) {
488 return Err(format!(
489 "{context} empirical latent measure weights must sum to 1, got {total}"
490 ));
491 }
492 Ok(())
493}
494
495pub(crate) fn combine_empirical_grids(
496 grids: &[EmpiricalZGrid],
497 mixture: &[(usize, f64)],
498) -> Result<EmpiricalZGrid, String> {
499 if mixture.is_empty() {
500 return Err("local empirical latent measure row mixture is empty".to_string());
501 }
502 let mut nodes = Vec::new();
503 let mut weights = Vec::new();
504 for &(grid_idx, grid_weight) in mixture {
505 if !grid_weight.is_finite() || grid_weight <= 0.0 {
506 return Err(format!(
507 "local empirical latent mixture weight must be finite and positive, got {grid_weight}"
508 ));
509 }
510 let grid = grids.get(grid_idx).ok_or_else(|| {
511 format!("local empirical latent mixture references missing grid {grid_idx}")
512 })?;
513 for (node, weight) in grid.pairs() {
514 nodes.push(node);
515 weights.push(grid_weight * weight);
516 }
517 }
518 let total = weights.iter().copied().sum::<f64>();
519 if !(total.is_finite() && total > 0.0) {
520 return Err(
521 "local empirical latent combined grid has non-positive total weight".to_string(),
522 );
523 }
524 for weight in &mut weights {
525 *weight /= total;
526 }
527 EmpiricalZGrid::new(nodes, weights, "local empirical latent combined grid")
528}
529
530#[derive(Clone, Debug)]
531pub struct LatentZPolicy {
532 pub check_mode: LatentZCheckMode,
533 pub normalization: LatentZNormalizationMode,
534 pub latent_measure: LatentMeasureSpec,
535 pub mean_tol_multiplier: f64,
536 pub sd_tol_multiplier: f64,
537 pub max_abs_skew: f64,
538 pub max_abs_excess_kurtosis: f64,
539}
540
541impl LatentZPolicy {
542 pub fn frozen_transformation_normal() -> Self {
543 // Defaults relaxed to `WarnOnly` with the same thresholds the
544 // exploratory-weighted preset uses (skew ≤ 4.0, |excess kurt| ≤ 20.0).
545 // Rationale: the upstream conditional transformation-normal
546 // preprocessor may be fit isotropically (no per-axis κ). At large-scale
547 // dimensionality (16 PCs, 15 ancestries) an isotropic fit can leave
548 // the global latent-z distribution mildly heavy-tailed (skew ≈ 4,
549 // excess kurt ≈ 30–40 in synthetic studies) without violating per-
550 // grouping mean/variance calibration. The downstream marginal-slope
551 // model still uses the latent-Gaussian probit/score-warp link; the
552 // emitted warning makes the deviation visible without aborting the
553 // fit. Callers that need strict enforcement can construct a custom
554 // `LatentZPolicy` with `check_mode: LatentZCheckMode::Strict`.
555 Self {
556 check_mode: LatentZCheckMode::WarnOnly,
557 normalization: LatentZNormalizationMode::Frozen { mean: 0.0, sd: 1.0 },
558 latent_measure: LatentMeasureSpec::auto_default(),
559 mean_tol_multiplier: 4.0,
560 sd_tol_multiplier: 4.0,
561 max_abs_skew: 4.0,
562 max_abs_excess_kurtosis: 20.0,
563 }
564 }
565
566 pub fn exploratory_fit_weighted() -> Self {
567 Self {
568 check_mode: LatentZCheckMode::WarnOnly,
569 normalization: LatentZNormalizationMode::FitWeighted,
570 latent_measure: LatentMeasureSpec::auto_default(),
571 mean_tol_multiplier: 8.0,
572 sd_tol_multiplier: 8.0,
573 max_abs_skew: 4.0,
574 max_abs_excess_kurtosis: 20.0,
575 }
576 }
577}
578
579impl Default for LatentZPolicy {
580 fn default() -> Self {
581 Self::frozen_transformation_normal()
582 }
583}
584
585#[derive(Clone, Copy, Debug, PartialEq)]
586pub struct LatentZNormalization {
587 pub mean: f64,
588 pub sd: f64,
589}
590
591impl LatentZNormalization {
592 pub fn apply(&self, z: &Array1<f64>, context: &str) -> Result<Array1<f64>, String> {
593 if !(self.mean.is_finite() && self.sd.is_finite() && self.sd > BMS_VARIANCE_FLOOR) {
594 return Err(format!(
595 "{context} requires finite latent z normalization with sd > {BMS_VARIANCE_FLOOR:e}; got mean={} sd={}",
596 self.mean, self.sd
597 ));
598 }
599 if z.iter().any(|value| !value.is_finite()) {
600 return Err(format!("{context} requires finite z values"));
601 }
602 Ok(z.mapv(|zi| (zi - self.mean) / self.sd))
603 }
604}
605
606/// Weighted mid-distribution rank inverse-normal transform for the
607/// latent score.
608///
609/// When the latent z fails the standard-normal auto-detection
610/// ([`latent_z_is_standard_normal_enough`]), the BMS family applied to
611/// pretend the score is N(0,1) anyway would distort the closed-form
612/// probit log-CDF kernel. The historical fallback (local- or
613/// global-empirical latent measure) is *mathematically correct* but
614/// triggers the per-row intercept Newton solve in the empirical-grid
615/// closed-form kernels (`empirical_rigid_primary_grad_hess_closed_form`
616/// and its higher-order siblings); at large scale that is the dominant
617/// cost.
618///
619/// **Rank-INT is a modeling choice, not a reparameterisation.** The
620/// rigid BMS predictor is *affine* in the latent score
621/// (`η = q·√(1+b²) + b·z`), so a nonlinear monotone map of `z` changes
622/// the model class and its likelihood — it does not leave them
623/// invariant. Applying the calibration redefines the latent axis: the
624/// affine model is *specified on the calibrated score* `T(z)`. Nor is
625/// the calibrated training sample exactly N(0,1): a finite set of
626/// normal scores is discrete, and with heavy ties it can stay far from
627/// Gaussian. The closed-form standard-normal kernel is therefore
628/// adequate only when the calibrated sample itself passes the same
629/// standard-normal adequacy gate applied to raw z
630/// ([`latent_z_is_standard_normal_enough`]);
631/// [`build_latent_measure_with_geometry`] re-checks the calibrated
632/// sample and falls back to the mathematically exact global-empirical
633/// latent measure when that re-check fails. On the passing path the
634/// kept work is the same closed-form
635/// `signed_probit_logcdf_and_mills_ratio` evaluation as the
636/// no-calibration path; the dropped work is the empirical-grid jet
637/// machinery. Persisted to disk so prediction applies the same
638/// monotone map to incoming z and re-routes through the closed-form
639/// kernel.
640#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
641pub struct LatentZRankIntCalibration {
642 /// Sorted unique positive-mass z values seen during training, ascending.
643 /// Knot table for `apply_to_training` / `apply_at_predict`. Zero-weight
644 /// knots carry no probability mass and are not stored.
645 pub sorted_z: Vec<f64>,
646 /// Weighted mid-distribution rank `(W_before + w_knot/2) / W_total` at
647 /// each `sorted_z` knot. Strictly increasing, strictly inside `(0, 1)`
648 /// (each knot is bounded away from the endpoints by half its own mass),
649 /// and invariant to a common rescaling of the weights.
650 pub weighted_cdf: Vec<f64>,
651 /// Weighted mean of the calibrated training sample. Used as a
652 /// sanity-check value on `fit`; should be very close to zero.
653 pub post_mean: f64,
654 /// Weighted SD of the calibrated training sample. Used as a
655 /// sanity-check value on `fit`; should be very close to one.
656 pub post_sd: f64,
657}
658
659impl LatentZRankIntCalibration {
660 /// Fit the weighted rank-INT calibration from training z and weights.
661 ///
662 /// Algorithm:
663 /// 1. Sort rows by ascending z and merge ties into one knot per unique
664 /// z with the tie group's total weight `w_g`; discard zero-mass
665 /// knots.
666 /// 2. Weighted mid-distribution rank at each knot:
667 /// `p_g = (W_before + w_g/2) / W_total`,
668 /// with `W_before` the cumulative weight strictly below the knot.
669 /// 3. Store `(sorted_z, weighted_cdf = p_g)`.
670 ///
671 /// The mid-rank depends only on *relative* weights (rescaling every
672 /// weight by a common factor leaves every `p_g` unchanged), is strictly
673 /// increasing across knots, and lies strictly inside `(0, 1)` — each
674 /// knot is separated from the endpoints by half its own mass, so no
675 /// ad-hoc clamp is needed and `Φ⁻¹(p_g)` is always finite.
676 ///
677 /// Returns the calibration plus the post-transform sample's weighted
678 /// mean / SD for sanity-check logging.
679 pub fn fit(z: &Array1<f64>, weights: &Array1<f64>) -> Result<Self, String> {
680 if z.len() != weights.len() {
681 return Err(format!(
682 "rank-INT calibration: z length {} != weights length {}",
683 z.len(),
684 weights.len()
685 ));
686 }
687 if z.is_empty() {
688 return Err("rank-INT calibration requires at least one observation".to_string());
689 }
690 let w_total = weights.iter().copied().sum::<f64>();
691 if !(w_total.is_finite() && w_total > 0.0) {
692 return Err(format!(
693 "rank-INT calibration requires positive finite total weight, got {w_total}"
694 ));
695 }
696 for (idx, value) in z.iter().enumerate() {
697 if !value.is_finite() {
698 return Err(format!(
699 "rank-INT calibration: z[{idx}] = {value} not finite"
700 ));
701 }
702 }
703 for (idx, weight) in weights.iter().enumerate() {
704 if !(weight.is_finite() && *weight >= 0.0) {
705 return Err(format!(
706 "rank-INT calibration: weight[{idx}] = {weight} not finite/non-negative"
707 ));
708 }
709 }
710 let mut order: Vec<usize> = (0..z.len()).collect();
711 order.sort_by(|&a, &b| z[a].partial_cmp(&z[b]).unwrap_or(std::cmp::Ordering::Equal));
712
713 let mut sorted_z: Vec<f64> = Vec::with_capacity(z.len());
714 let mut weighted_cdf: Vec<f64> = Vec::with_capacity(z.len());
715 // Merge ties into one knot per unique z, then assign the weighted
716 // mid-distribution rank p_g = (W_before + w_g/2) / W_total. This is
717 // the mid-point of the tie group's probability mass, so it depends
718 // only on relative weights, is strictly increasing, and sits
719 // strictly inside (0, 1) without any clamp. Zero-mass tie groups
720 // are not knots of the weighted empirical distribution and are
721 // dropped.
722 let mut cum_before = 0.0_f64;
723 let mut pos = 0usize;
724 while pos < order.len() {
725 let zi = z[order[pos]];
726 let mut w_group = 0.0_f64;
727 let mut end = pos;
728 while end < order.len() && z[order[end]] == zi {
729 w_group += weights[order[end]];
730 end += 1;
731 }
732 if w_group > 0.0 {
733 sorted_z.push(zi);
734 weighted_cdf.push((cum_before + 0.5 * w_group) / w_total);
735 cum_before += w_group;
736 }
737 pos = end;
738 }
739 if sorted_z.is_empty() {
740 return Err(
741 "rank-INT calibration requires at least one positive-weight observation"
742 .to_string(),
743 );
744 }
745
746 // Compute sanity-check post-mean and post-sd on the transformed
747 // sample, weighted by the original weights.
748 let mut sum_wz = 0.0_f64;
749 let mut sum_w = 0.0_f64;
750 for &idx in &order {
751 let zi = z[idx];
752 let calibrated = Self::apply_with_knots(zi, &sorted_z, &weighted_cdf);
753 sum_wz += weights[idx] * calibrated;
754 sum_w += weights[idx];
755 }
756 let post_mean = if sum_w > 0.0 { sum_wz / sum_w } else { 0.0 };
757 let mut sum_w_dev = 0.0_f64;
758 for &idx in &order {
759 let zi = z[idx];
760 let calibrated = Self::apply_with_knots(zi, &sorted_z, &weighted_cdf);
761 let d = calibrated - post_mean;
762 sum_w_dev += weights[idx] * d * d;
763 }
764 let post_sd = if sum_w > 0.0 {
765 (sum_w_dev / sum_w).sqrt()
766 } else {
767 1.0
768 };
769
770 Ok(Self {
771 sorted_z,
772 weighted_cdf,
773 post_mean,
774 post_sd,
775 })
776 }
777
778 /// Apply the calibration to the full training z vector, returning the
779 /// calibrated sample. Equivalent to mapping each row's z through
780 /// [`Self::apply_at_predict`], but vectorised.
781 pub fn apply_to_training(&self, z: &Array1<f64>) -> Result<Array1<f64>, String> {
782 if self.sorted_z.is_empty() {
783 return Err("rank-INT calibration has no knots".to_string());
784 }
785 let mut out = Array1::<f64>::zeros(z.len());
786 for (idx, &zi) in z.iter().enumerate() {
787 if !zi.is_finite() {
788 return Err(format!(
789 "rank-INT calibration apply: z[{idx}] = {zi} not finite"
790 ));
791 }
792 out[idx] = self.apply_at_predict(zi);
793 }
794 Ok(out)
795 }
796
797 /// Apply the calibration to a single z at predict time.
798 ///
799 /// Linear interpolation on `(sorted_z, weighted_cdf)` to obtain
800 /// `p ∈ [eps, 1 − eps]`, then `Φ⁻¹(p)` via
801 /// [`standard_normal_quantile`]. Out-of-range z's clip to the
802 /// boundary CDF before the quantile, so the calibration extrapolates
803 /// monotonically beyond the training support.
804 pub fn apply_at_predict(&self, z: f64) -> f64 {
805 Self::apply_with_knots(z, &self.sorted_z, &self.weighted_cdf)
806 }
807
808 pub(crate) fn apply_with_knots(z: f64, sorted_z: &[f64], weighted_cdf: &[f64]) -> f64 {
809 assert_eq!(sorted_z.len(), weighted_cdf.len());
810 assert!(!sorted_z.is_empty());
811 let n = sorted_z.len();
812 let p = if z <= sorted_z[0] {
813 weighted_cdf[0]
814 } else if z >= sorted_z[n - 1] {
815 weighted_cdf[n - 1]
816 } else {
817 // Binary search for the right knot.
818 let mut lo = 0usize;
819 let mut hi = n - 1;
820 while hi - lo > 1 {
821 let mid = (lo + hi) / 2;
822 if sorted_z[mid] <= z {
823 lo = mid;
824 } else {
825 hi = mid;
826 }
827 }
828 let z_lo = sorted_z[lo];
829 let z_hi = sorted_z[hi];
830 let p_lo = weighted_cdf[lo];
831 let p_hi = weighted_cdf[hi];
832 if z_hi == z_lo {
833 p_hi
834 } else {
835 let t = (z - z_lo) / (z_hi - z_lo);
836 p_lo + t * (p_hi - p_lo)
837 }
838 };
839 // Φ⁻¹(p); clip away from {0, 1} to keep the quantile finite.
840 standard_normal_quantile(p).unwrap_or_else(|_| if p < 0.5 { -8.0 } else { 8.0 })
841 }
842}
843
844/// Optional calibration applied to the latent score before the BMS
845/// kernel runs. When `RankInverseNormal`, both the training and predict
846/// paths route the input z through [`LatentZRankIntCalibration::apply_*`]
847/// before the standard-normal closed-form kernel is invoked.
848#[derive(Clone, Debug)]
849pub enum LatentMeasureCalibration {
850 None,
851 RankInverseNormal(LatentZRankIntCalibration),
852 ConditionalLocationScale(LatentZConditionalCalibration),
853}
854
855/// Conditional location-scale calibration of the latent score (#905).
856///
857/// The marginal-slope Auto trigger's pooled-z gate (KS / skewness / kurtosis +
858/// the rank inverse-normal transform) only inspects the **marginal** law of
859/// `z`. A conditional shift `E[z | C] = m(C) ≠ 0` — the allele-frequency-driven
860/// grouping mean shift — passes the marginal gate while leaving `z | C`
861/// off-center, so the slope contribution `b(C)·m(C)` leaks into the influence
862/// channel `q`. Rank-INT provably cannot fix this: no transform `T` depending
863/// only on the marginal `F_Z` can enforce `E[T(Z) | C] ≡ const` for all joint
864/// laws.
865///
866/// The unique Fisher-orthogonal location-scale correction (for the Gaussian
867/// working metric the closed-form probit kernel assumes) is
868/// `ζ = (z − m(C)) / √v(C)`, where `m(C) = E[z|C]` and `v(C) = Var(z|C)` are
869/// estimated by weighted ridge regression of `z` (and its squared residual) on
870/// the marginal-index span `a(C) = [1 | X_marginal]`. The corrected `ζ` is
871/// conditionally centered (and homoskedastic when the variance block is
872/// active) by construction, so the `b(C)·m(C)` leakage vanishes. Matching the
873/// first two conditional moments does **not** by itself make `ζ` standard
874/// normal (a two-point residual law survives location-scale correction
875/// unchanged in shape), so [`build_latent_measure_with_geometry`] re-checks
876/// the calibrated sample against the standard-normal adequacy gate and
877/// retains an empirical latent measure for the residual distribution when
878/// that re-check fails; only a passing `ζ` uses the closed-form
879/// standard-normal kernel. Persisted so prediction rebuilds `a(C)` from the
880/// (reproducible) marginal design and applies the identical map to incoming
881/// z.
882#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
883pub struct LatentZConditionalCalibration {
884 /// Coefficients for the conditional mean `m(C) = β_m·[1 | a(C)]` over the
885 /// basis `[1 | marginal-design row]`. Length `1 + basis_ncols` (leading
886 /// entry is the intercept).
887 pub mean_coeffs: Vec<f64>,
888 /// Coefficients for the conditional variance
889 /// `v(C) = max(β_v·[1 | a(C)], var_floor)`. Length `1 + basis_ncols`, or
890 /// empty when the conditional-variance block of the Rao gate was not
891 /// significant (mean-only correction); then `v(C) ≡ global_var`.
892 pub var_coeffs: Vec<f64>,
893 /// Number of marginal-design columns in the basis (excludes the leading
894 /// intercept). The predict-time marginal design must present exactly this
895 /// many columns.
896 pub basis_ncols: usize,
897 /// Floor on the fitted conditional variance, in the (normalized)
898 /// latent-score scale (= `AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC · global_var`).
899 pub var_floor: f64,
900 /// Global weighted variance of the (normalized) training latent score. Used
901 /// as `v(C)` when `var_coeffs` is empty.
902 pub global_var: f64,
903 /// Weighted mean of the calibrated training sample (sanity-check, ≈ 0).
904 pub post_mean: f64,
905 /// Weighted SD of the calibrated training sample (sanity-check, ≈ 1).
906 pub post_sd: f64,
907 /// First-stage (generated-regressor) sandwich covariance of `mean_coeffs`,
908 /// `V₁ᵐ = M⁻¹ (Σ_i w_i² û_i² A_i A_iᵀ) M⁻¹` with `A = [1 | a(C)]`,
909 /// `M = AᵀWA + λR` (the same weighted-ridge normal matrix that produced
910 /// `mean_coeffs`), `û_i = z_i − m̂(C_i)` the HC0 mean residual, and
911 /// `W = diag(w_i)`. Shape `(1+basis_ncols) × (1+basis_ncols)`. This is the
912 /// closed-form estimation uncertainty of `m(C)` that the second stage
913 /// (Murphy–Topel) needs; see [`Self::generated_regressor_term`].
914 pub mean_cov: Array2<f64>,
915 /// First-stage sandwich covariance of `var_coeffs`, computed identically on
916 /// the squared-mean-residual response. Empty (`0 × 0`) exactly when
917 /// `var_coeffs` is empty (mean-only correction; `v(C) ≡ global_var` is a
918 /// constant carrying no first-stage slope uncertainty).
919 pub var_cov: Array2<f64>,
920}
921
922impl LatentZConditionalCalibration {
923 #[inline]
924 pub(crate) fn affine(coeffs: &[f64], a_row: ArrayView1<'_, f64>) -> f64 {
925 let mut acc = coeffs[0];
926 for (c, &x) in coeffs[1..].iter().zip(a_row.iter()) {
927 acc += c * x;
928 }
929 acc
930 }
931
932 pub(crate) fn conditional_mean(&self, a_row: ArrayView1<'_, f64>) -> f64 {
933 Self::affine(&self.mean_coeffs, a_row)
934 }
935
936 pub(crate) fn conditional_var(&self, a_row: ArrayView1<'_, f64>) -> f64 {
937 if self.var_coeffs.is_empty() {
938 self.global_var.max(self.var_floor)
939 } else {
940 Self::affine(&self.var_coeffs, a_row).max(self.var_floor)
941 }
942 }
943
944 /// Apply `ζ = (z − m(C))/√v(C)` to a batch. `a_block` is the marginal
945 /// design (`n × basis_ncols`); `z` is the (normalized) latent score. Used
946 /// at both training and predict time, so the map is identical.
947 pub fn apply(
948 &self,
949 z: ArrayView1<'_, f64>,
950 a_block: ArrayView2<'_, f64>,
951 ) -> Result<Array1<f64>, String> {
952 if a_block.ncols() != self.basis_ncols {
953 return Err(format!(
954 "conditional latent calibration expects {} basis columns, got {}",
955 self.basis_ncols,
956 a_block.ncols()
957 ));
958 }
959 if a_block.nrows() != z.len() {
960 return Err(format!(
961 "conditional latent calibration row mismatch: z={}, basis rows={}",
962 z.len(),
963 a_block.nrows()
964 ));
965 }
966 if self.mean_coeffs.len() != self.basis_ncols + 1 {
967 return Err(format!(
968 "conditional latent calibration mean coefficient length {} != basis_ncols+1 ({})",
969 self.mean_coeffs.len(),
970 self.basis_ncols + 1
971 ));
972 }
973 let mut out = Array1::<f64>::zeros(z.len());
974 for i in 0..z.len() {
975 let a_row = a_block.row(i);
976 if !z[i].is_finite() {
977 return Err(format!(
978 "conditional latent calibration: z[{i}] = {} not finite",
979 z[i]
980 ));
981 }
982 let m = self.conditional_mean(a_row);
983 let v = self.conditional_var(a_row);
984 if !(v.is_finite() && v > 0.0) {
985 return Err(format!(
986 "conditional latent calibration produced non-positive variance {v} at row {i}"
987 ));
988 }
989 let zeta = (z[i] - m) / v.sqrt();
990 if !zeta.is_finite() {
991 return Err(format!(
992 "conditional latent calibration produced non-finite zeta at row {i}"
993 ));
994 }
995 out[i] = zeta;
996 }
997 Ok(out)
998 }
999
1000 /// Dimension of the first-stage parameter vector `θ₁ = (mean_coeffs,
1001 /// var_coeffs)` whose estimation uncertainty the generated-regressor
1002 /// correction propagates. Equals `len(mean_coeffs)` when the variance block
1003 /// is inactive, otherwise `len(mean_coeffs) + len(var_coeffs)`.
1004 pub fn theta1_dim(&self) -> usize {
1005 self.mean_coeffs.len() + self.var_coeffs.len()
1006 }
1007
1008 /// Per-row sensitivity `∂ζ_i/∂θ₁` of the calibrated score to the first-stage
1009 /// calibration coefficients, stacked as `[∂ζ/∂mean_coeffs | ∂ζ/∂var_coeffs]`
1010 /// (length [`Self::theta1_dim`]). With `ζ = (z − m(C))/√v(C)`,
1011 /// `A_i = [1 | a(C_i)]`, `m = A_iᵀ·mean_coeffs`, `v = A_iᵀ·var_coeffs`:
1012 ///
1013 /// `∂ζ/∂m = −1/√v`, `∂ζ/∂v = −(z − m)/(2 v^{3/2}) = −ζ/(2v)`,
1014 ///
1015 /// and by the chain rule through the affine basis
1016 /// `∂ζ/∂mean_coeffs = (∂ζ/∂m)·A_i`, `∂ζ/∂var_coeffs = (∂ζ/∂v)·A_i`. The
1017 /// variance block contributes only when `var_coeffs` is active AND the
1018 /// fitted `v(C_i)` is above the floor (a floored row has `∂v/∂var_coeffs = 0`
1019 /// in the applied map). `z` is the (normalized) raw latent score at this row.
1020 pub fn zeta_theta1_jacobian_row(&self, z: f64, a_row: ArrayView1<'_, f64>) -> Vec<f64> {
1021 let m = self.conditional_mean(a_row);
1022 let v = self.conditional_var(a_row);
1023 let inv_sqrt_v = 1.0 / v.sqrt();
1024 // Intercept-augmented basis row A_i = [1 | a(C_i)].
1025 let mut out = Vec::with_capacity(self.theta1_dim());
1026 let dzeta_dm = -inv_sqrt_v;
1027 out.push(dzeta_dm); // intercept column of A
1028 for &x in a_row.iter() {
1029 out.push(dzeta_dm * x);
1030 }
1031 if !self.var_coeffs.is_empty() {
1032 // ∂ζ/∂v active only off the floor; on the floor the applied v(C) is
1033 // constant in var_coeffs, so the variance sensitivity is exactly 0.
1034 let raw_v = Self::affine(&self.var_coeffs, a_row);
1035 let dzeta_dv = if raw_v > self.var_floor {
1036 let zeta = (z - m) * inv_sqrt_v;
1037 -zeta / (2.0 * v)
1038 } else {
1039 0.0
1040 };
1041 out.push(dzeta_dv);
1042 for &x in a_row.iter() {
1043 out.push(dzeta_dv * x);
1044 }
1045 }
1046 out
1047 }
1048
1049 /// Block-diagonal first-stage covariance `V₁ = blkdiag(mean_cov, var_cov)`
1050 /// of `θ₁`, ordered to match [`Self::zeta_theta1_jacobian_row`]. The two
1051 /// stages are fit on (asymptotically) uncorrelated estimating equations
1052 /// (the mean score `Σ w û A` and the Breusch–Pagan variance score
1053 /// `Σ w (û² − v) A` are orthogonal under the Gaussian working model), so the
1054 /// joint first-stage covariance is block-diagonal to first order — the same
1055 /// approximation the Rao gate above uses.
1056 pub fn theta1_covariance(&self) -> Array2<f64> {
1057 let dm = self.mean_coeffs.len();
1058 let dv = self.var_coeffs.len();
1059 let mut v1 = Array2::<f64>::zeros((dm + dv, dm + dv));
1060 v1.slice_mut(s![..dm, ..dm]).assign(&self.mean_cov);
1061 if dv > 0 {
1062 v1.slice_mut(s![dm.., dm..]).assign(&self.var_cov);
1063 }
1064 v1
1065 }
1066
1067 /// Murphy–Topel generated-regressor correction term for the second-stage
1068 /// slope covariance. Given the second-stage information `H_β` (the penalized
1069 /// joint Hessian of the slope fit, whose inverse is the naive `V_β`) and the
1070 /// cross-derivative `G = ∂(score_β)/∂θ₁` (`p_β × dim θ₁`), the corrected
1071 /// covariance is
1072 ///
1073 /// `V_β = V_β^naive + (H_β⁻¹ G) V₁ (H_β⁻¹ G)ᵀ`.
1074 ///
1075 /// This returns the additive rank-`dim θ₁` term `(H_β⁻¹ G) V₁ (H_β⁻¹ G)ᵀ`
1076 /// given the already-formed `hbeta_inv_g = H_β⁻¹ G` (`p_β × dim θ₁`). The
1077 /// caller forms `G` by accumulating the per-row slope-score sensitivity to
1078 /// `ζ_i` times [`Self::zeta_theta1_jacobian_row`] (chain rule
1079 /// `∂score_β/∂θ₁ = Σ_i (∂score_β/∂ζ_i) (∂ζ_i/∂θ₁)`).
1080 pub fn generated_regressor_term(&self, hbeta_inv_g: ArrayView2<'_, f64>) -> Array2<f64> {
1081 let v1 = self.theta1_covariance();
1082 hbeta_inv_g.dot(&v1).dot(&hbeta_inv_g.t())
1083 }
1084
1085 /// Assemble the full Murphy–Topel generated-regressor correction
1086 /// `(Vb·G)·V₁·(Vb·G)ᵀ` for the second-stage slope covariance, given the ONE
1087 /// engine-side quantity it cannot reconstruct post-fit: the per-row
1088 /// reduced-frame slope-score sensitivity to the calibrated score,
1089 /// `s_i = ∂score_β,i/∂ζ_i` (a `p_β`-vector in the joint flat-β reduced frame
1090 /// `solved_fit.beta_covariance()` lives in). With `score_β,i = ∂ℓ_i/∂β`,
1091 /// `s_i = ∂²ℓ_i/∂β∂ζ_i = J_iᵀ·(∂²ℓ_i/∂η_i∂ζ_i)` is the mixed `(β, ζ)`
1092 /// second derivative of the warped row kernel contracted through the slope
1093 /// design Jacobian `J_i` — exactly the #932 RowNllProgram/Tower4 z-jet
1094 /// channel (`z` is already a row-program input; one extra mixed `(β, z)` jet
1095 /// channel reads off `∂²ℓ/∂β∂z`). It must be evaluated at the converged `β̂`
1096 /// in the SAME reduced frame as `vb`.
1097 ///
1098 /// Everything else is built here from the stored first-stage quantities and
1099 /// the second-stage fit, dissolving the post-fit-reconstruction blocker:
1100 /// - `G = Σ_i s_i · (∂ζ_i/∂θ₁)ᵀ` (`p_β × dim θ₁`), the chain-rule outer
1101 /// product accumulated row-by-row with `∂ζ_i/∂θ₁ =
1102 /// `[`Self::zeta_theta1_jacobian_row`]`(z_i, a_row_i)` (exact-zero on
1103 /// floored rows, so floored rows contribute nothing — `G`'s support is
1104 /// the gate-fired rows);
1105 /// - `Vb·G = vb·G` since the naive second-stage covariance `vb` IS
1106 /// `H_β⁻¹` (the coordinator's `H_β⁻¹ G = Vb.dot(G)`);
1107 /// - the term `(Vb·G)·V₁·(Vb·G)ᵀ` via [`Self::generated_regressor_term`].
1108 ///
1109 /// `score_zeta_sensitivity` is `n × p_β` (row `i` = `s_i`); `z` is the
1110 /// per-row normalized latent score (`n`); `a_block` is the marginal design
1111 /// `n × basis_ncols` whose rows feed `zeta_theta1_jacobian_row`; `vb` is the
1112 /// naive reduced-frame slope covariance `n_β × n_β`. The returned term is
1113 /// PSD (a congruence of the PSD `V₁`), so adding it to `vb` makes the
1114 /// corrected slope SE strictly ≥ the naive SE whenever the gate fires
1115 /// (`G ≠ 0`) and exactly equal when every row is floored (`G = 0`).
1116 pub fn generated_regressor_correction(
1117 &self,
1118 score_zeta_sensitivity: ArrayView2<'_, f64>,
1119 z: ArrayView1<'_, f64>,
1120 a_block: ArrayView2<'_, f64>,
1121 vb: ArrayView2<'_, f64>,
1122 ) -> Result<Array2<f64>, String> {
1123 let n = score_zeta_sensitivity.nrows();
1124 let p_beta = score_zeta_sensitivity.ncols();
1125 if z.len() != n || a_block.nrows() != n {
1126 return Err(format!(
1127 "generated_regressor_correction row mismatch: score_zeta_sensitivity rows={n}, \
1128 z={}, a_block rows={}",
1129 z.len(),
1130 a_block.nrows()
1131 ));
1132 }
1133 if a_block.ncols() != self.basis_ncols {
1134 return Err(format!(
1135 "generated_regressor_correction expects {} basis columns, got {}",
1136 self.basis_ncols,
1137 a_block.ncols()
1138 ));
1139 }
1140 if vb.nrows() != p_beta || vb.ncols() != p_beta {
1141 return Err(format!(
1142 "generated_regressor_correction: vb must be {p_beta}×{p_beta}, got {}×{}",
1143 vb.nrows(),
1144 vb.ncols()
1145 ));
1146 }
1147 // G = Σ_i s_i ⊗ (∂ζ_i/∂θ₁) (p_β × dim θ₁). Each row contributes the
1148 // rank-1 outer product `s_i ⊗ J_zeta_i`, so summed over the n rows this
1149 // is exactly the cross product `G = Sᵀ·J` of the score-sensitivity
1150 // matrix `S` (`n × p_β`, supplied) and the per-row ζ-Jacobian matrix
1151 // `J` (`n × dim θ₁`). Forming `J` row-by-row is O(n·dim θ₁); the cross
1152 // product is then a single BLAS-3 GEMM rather than the O(n·p_β·dim θ₁)
1153 // scalar triple loop (≈1.5e9 FMA at biobank scale, n≈194k, the dominant
1154 // ~13s/disease cost of the SE correction). Floored rows yield an exact
1155 // all-zero `J` row, so they contribute zero to the GEMM — bit-identical
1156 // to skipping them, no approximation.
1157 let j_mat = self.build_zeta_theta1_jacobian(z, a_block);
1158 let vb_g = self.beta_theta1_sensitivity(score_zeta_sensitivity, j_mat.view(), vb)?;
1159 Ok(self.generated_regressor_term(vb_g.view()))
1160 }
1161
1162 /// Per-row ζ-Jacobian matrix `J` (`n × dim θ₁`, row `i` = `∂ζ_i/∂θ₁`) built
1163 /// row-by-row from [`Self::zeta_theta1_jacobian_row`]. Floored rows yield an
1164 /// exact all-zero row, so they contribute nothing to the `G = Sᵀ·J` cross
1165 /// product (bit-identical to skipping them).
1166 fn build_zeta_theta1_jacobian(
1167 &self,
1168 z: ArrayView1<'_, f64>,
1169 a_block: ArrayView2<'_, f64>,
1170 ) -> Array2<f64> {
1171 let n = a_block.nrows();
1172 let dim_theta1 = self.theta1_dim();
1173 let mut j_mat = Array2::<f64>::zeros((n, dim_theta1));
1174 for i in 0..n {
1175 let j_zeta_row = self.zeta_theta1_jacobian_row(z[i], a_block.row(i));
1176 assert_eq!(
1177 j_zeta_row.len(),
1178 dim_theta1,
1179 "J_zeta row width must match the first-stage hyperparameter dimension"
1180 );
1181 let mut dst = j_mat.row_mut(i);
1182 for (slot, jz) in dst.iter_mut().zip(j_zeta_row.into_iter()) {
1183 *slot = jz;
1184 }
1185 }
1186 j_mat
1187 }
1188
1189 /// Signed first-order sensitivity `∂β̂/∂θ₁ = Vb·G` (`p_β × dim θ₁`) of the
1190 /// converged second-stage slope to the first-stage calibration parameters,
1191 /// the SIGNED quantity the Murphy–Topel correction is built from.
1192 ///
1193 /// `G = Sᵀ·J = Σ_i s_i ⊗ (∂ζ_i/∂θ₁)` with `s_i = ∂score_β,i/∂ζ_i` the
1194 /// LOG-LIKELIHOOD-score sensitivity (the sign convention #1131 fixes at the
1195 /// source in [`gradient_paths::rigid_standard_normal_mixed_z_sensitivity`]),
1196 /// and `Vb = H_β⁻¹` the NLL-Hessian inverse. Under this convention the
1197 /// implicit-function theorem on `∂(log L)/∂β = 0` gives
1198 /// `∂β̂/∂θ₁ = +H_β⁻¹·G = +Vb·G`, so the returned matrix matches the finite
1199 /// difference of the refit slope in θ₁ in BOTH sign and magnitude — unlike
1200 /// the PSD correction term [`Self::generated_regressor_correction`], which is
1201 /// invariant to this sign. `j_zeta` is the per-row ζ-Jacobian matrix
1202 /// (`n × dim θ₁`, row `i` = `∂ζ_i/∂θ₁`).
1203 fn beta_theta1_sensitivity(
1204 &self,
1205 score_zeta_sensitivity: ArrayView2<'_, f64>,
1206 j_zeta: ArrayView2<'_, f64>,
1207 vb: ArrayView2<'_, f64>,
1208 ) -> Result<Array2<f64>, String> {
1209 // G = Sᵀ·J (p_β × dim θ₁) via the SIMD/GPU-routed cross product.
1210 let g = gam_linalg::faer_ndarray::fast_atb(&score_zeta_sensitivity, &j_zeta);
1211 // Vb·G = H_β⁻¹·G (vb is the naive reduced-frame covariance the fit
1212 // already produced — reused, never recomputed).
1213 Ok(vb.dot(&g))
1214 }
1215}
1216
1217/// First-stage robust (HC0) sandwich covariance of a weighted-ridge coefficient
1218/// vector: `V₁ = M⁺ (Σ_i w_i² û_i² A_i A_iᵀ) M⁺` with `M = AᵀWA + λR` the
1219/// ridge normal matrix that produced the coefficients, `W = diag(weights)`,
1220/// `û_i` the per-row residual, and `A` the regression basis (here `[1 | a(C)]`).
1221/// `M⁺` is the Moore–Penrose pseudo-inverse via eigendecomposition with a
1222/// relative tolerance: identifiable directions get the usual `(λ_eff)⁻¹` weight,
1223/// and rank-deficient directions (where some θ₁ components are not identified
1224/// by `A`) are zeroed — they carry no asymptotic distribution, so V₁ in those
1225/// directions is zero, and the Murphy–Topel propagation through identifiable
1226/// functionals of β remains finite and consistent. Using the ordinary inverse
1227/// here let the unregularized direction's `1/ε` blow `M⁻¹·meat·M⁻¹` through
1228/// the f64 range whenever the wide marginal-index span had a near-null
1229/// direction (the bug behind "conditional latent calibration sandwich
1230/// covariance is non-finite" on wide rank-deficient duchon/spline conditioning).
1231/// The meat is formed as `BᵀB` with `B_i = w_i û_i A_iᵀ` (signed) so the
1232/// fused-multiply GEMM is the same SIMD path used everywhere else in the
1233/// codebase, instead of a hand-rolled triple loop whose partial sums could
1234/// overflow on a single pathological row of the basis.
1235pub(crate) fn weighted_ridge_sandwich_cov(
1236 basis: ArrayView2<'_, f64>,
1237 residuals: &[f64],
1238 weights: ArrayView1<'_, f64>,
1239 normal_matrix: &Array2<f64>,
1240) -> Result<Array2<f64>, String> {
1241 let n = basis.nrows();
1242 let p = basis.ncols();
1243 if residuals.len() != n || weights.len() != n {
1244 return Err(format!(
1245 "weighted ridge sandwich length mismatch: rows={n}, residuals={}, weights={}",
1246 residuals.len(),
1247 weights.len()
1248 ));
1249 }
1250 if normal_matrix.nrows() != p || normal_matrix.ncols() != p {
1251 return Err(format!(
1252 "weighted ridge sandwich normal-matrix shape mismatch: basis cols={p}, normal {}x{}",
1253 normal_matrix.nrows(),
1254 normal_matrix.ncols()
1255 ));
1256 }
1257 // Robust HC0 meat as a Gram: build `B` with `B_i = (w_i û_i) A_iᵀ` (rows of
1258 // basis scaled by `w_i û_i`, sign carried), so `meat = BᵀB = Σ_i w_i² û_i²
1259 // A_i A_iᵀ` from one BLAS Gramian. Identical math to the per-row outer-
1260 // product accumulation, but the GEMM path keeps partial sums vectorized
1261 // and is less sensitive to a single pathological row producing an
1262 // intermediate that overflows f64 before the column-wise reduction cancels.
1263 let mut b = basis.to_owned();
1264 for i in 0..n {
1265 let wi = weights[i];
1266 let ri = residuals[i];
1267 let scale = wi * ri;
1268 if scale == 0.0 {
1269 b.row_mut(i).fill(0.0);
1270 continue;
1271 }
1272 b.row_mut(i).iter_mut().for_each(|value| *value *= scale);
1273 }
1274 let meat = gam_linalg::faer_ndarray::fast_ata(&b);
1275 // SPD pseudo-inverse of `M = AᵀWA + λR` via eigendecomposition with a
1276 // relative tolerance; symmetrize first to absorb floating-point asymmetry
1277 // accumulated in the AᵀWA assembly.
1278 let mut m_sym = normal_matrix.clone();
1279 gam_linalg::matrix::symmetrize_in_place(&mut m_sym);
1280 // Jacobi (symmetric diagonal) preconditioning. When the conditioning basis
1281 // spans many orders of magnitude — a power-9 Duchon RBF over 16 standardized
1282 // PCs produces columns differing by ~30 decades — `M` and `meat` live on
1283 // wildly different per-column scales, and the eigendecomposition behind
1284 // `M⁺ meat M⁺` loses all accuracy: the relative truncation tolerance is set
1285 // by `λ_max(M)` (dominated by the largest-scale column), so a genuinely
1286 // identified small-scale direction can be dropped while a near-null one is
1287 // kept, and the surviving `1/λ` then multiplies the huge `meat` straight
1288 // through the f64 range. Precondition by `D = diag(√M_jj)`. Because the ridge
1289 // penalty diagonal is built as the weighted Gram diagonal itself
1290 // (`penalty_jj = Σ_i w_i a_ij² = (AᵀWA)_jj`), `M_jj = (1+ρ)(AᵀWA)_jj`, so
1291 // `M̃ = D⁻¹ M D⁻¹` has EXACT unit diagonal and `M̃ = C + (ρ/(1+ρ))·I` with
1292 // `C` the basis correlation matrix (PSD). Hence `λ_min(M̃) ≥ ρ/(1+ρ) ≈ 1e-8`
1293 // even for a fully collinear basis, which clears the pseudo-inverse's
1294 // relative tolerance `≈ 1e-10·λ_max(M̃)` for the conditioning widths that
1295 // occur here: no direction is spuriously dropped, so `M̃⁺ = M̃⁻¹ = D M⁻¹ D`
1296 // and `cov = D⁻¹ (M̃⁻¹ meat̃ M̃⁻¹) D⁻¹ = M⁻¹ meat M⁻¹` EXACTLY — the scaling
1297 // cancels, this is the same sandwich, only computed on a well-conditioned
1298 // matrix. (Should a pure-ridge direction ever fall under tolerance at very
1299 // large width, dropping it is the correct scale-invariant identifiability
1300 // call.) `meat̃ = D⁻¹ meat D⁻¹`; `M_jj > 0` (Gram diagonal floored positive)
1301 // so `D` is always finite and invertible.
1302 let scale: Vec<f64> = (0..p)
1303 .map(|j| 1.0 / m_sym[[j, j]].max(f64::MIN_POSITIVE).sqrt())
1304 .collect();
1305 let mut m_scaled = m_sym;
1306 let mut meat_scaled = meat;
1307 for i in 0..p {
1308 for j in 0..p {
1309 let s = scale[i] * scale[j];
1310 m_scaled[[i, j]] *= s;
1311 meat_scaled[[i, j]] *= s;
1312 }
1313 }
1314 let (_rank, m_pinv) =
1315 gam_linalg::utils::block_penalty_rank_and_pinv(&m_scaled).map_err(|e| {
1316 format!("conditional latent calibration sandwich pseudo-inverse failed: {e}")
1317 })?;
1318 let mut cov = m_pinv.dot(&meat_scaled).dot(&m_pinv);
1319 // Undo the symmetric scaling: cov_raw = D⁻¹ cov_scaled D⁻¹.
1320 for i in 0..p {
1321 for j in 0..p {
1322 cov[[i, j]] *= scale[i] * scale[j];
1323 }
1324 }
1325 if cov.iter().any(|v| !v.is_finite()) {
1326 return Err("conditional latent calibration sandwich covariance is non-finite".to_string());
1327 }
1328 Ok(cov)
1329}
1330
1331/// Weighted mean of a slice of values.
1332pub(crate) fn weighted_mean(
1333 values: &[f64],
1334 weights: ArrayView1<'_, f64>,
1335 total_weight: f64,
1336) -> f64 {
1337 values
1338 .iter()
1339 .zip(weights.iter())
1340 .map(|(&v, &w)| w * v)
1341 .sum::<f64>()
1342 / total_weight
1343}
1344
1345/// Robust (heteroskedasticity-consistent) Rao/LM score-test p-value for the
1346/// null that the centered basis columns `ã(C)` carry no information about the
1347/// centered response `u`. This is the LAN locally-optimal statistic the issue
1348/// names: `s = Σ_i w_i u_i ã(C_i)`, `Ω̂ = Σ_i w_i² u_i² ã(C_i)ã(C_i)ᵀ`,
1349/// `D = sᵀ Ω̂⁺ s ⟶ χ²_{rank Ω̂}`. Both the conditional-mean test
1350/// (`u_i = z_i − z̄`) and the conditional-variance / Breusch-Pagan test
1351/// (`u_i = (z_i − z̄)² − σ̂²`) are this statistic with the same centered basis.
1352///
1353/// Returns `None` when the test is degenerate (no usable basis directions),
1354/// otherwise the asymptotic p-value.
1355pub(crate) fn robust_conditional_score_pvalue(
1356 a_centered: ArrayView2<'_, f64>,
1357 u: &[f64],
1358 weights: ArrayView1<'_, f64>,
1359) -> Result<Option<f64>, String> {
1360 let n = a_centered.nrows();
1361 let r = a_centered.ncols();
1362 if r == 0 || n == 0 {
1363 return Ok(None);
1364 }
1365 if u.len() != n || weights.len() != n {
1366 return Err(format!(
1367 "conditional score test length mismatch: rows={n}, u={}, weights={}",
1368 u.len(),
1369 weights.len()
1370 ));
1371 }
1372 // Build the per-row scaled basis `B` with `B_i = (w_i u_i) ã_i` once, then
1373 // recover both the score and the HC0 robust meat from it with two BLAS-3
1374 // GEMMs over chunked row-blocks instead of an `O(n · r²)` per-row scatter:
1375 // • score `s = ãᵀ (w ∘ u) = Bᵀ 1` (column sums of `B`),
1376 // • meat `Ω̂ = Σ_i w_i² u_i² ã_i ã_iᵀ = BᵀB` since `(w_i u_i)² = w_i² u_i²`.
1377 // A non-positive weight zeroes that row of `B` (its score and meat
1378 // contributions both vanish), reproducing the `wi <= 0.0` skip EXACTLY.
1379 // `fast_ata` is the same parallel Gramian the second-stage sandwich uses, so
1380 // the statistic is numerically identical to the row-accumulated form up to
1381 // the deterministic GEMM reduction order.
1382 let mut b = a_centered.to_owned();
1383 for i in 0..n {
1384 let wi = weights[i];
1385 let scale = if wi > 0.0 { wi * u[i] } else { 0.0 };
1386 if scale == 0.0 {
1387 b.row_mut(i).fill(0.0);
1388 continue;
1389 }
1390 b.row_mut(i).iter_mut().for_each(|value| *value *= scale);
1391 }
1392 let s = b.sum_axis(ndarray::Axis(0));
1393 let omega = gam_linalg::faer_ndarray::fast_ata(&b);
1394 if !s.iter().all(|v| v.is_finite()) || !omega.iter().all(|v| v.is_finite()) {
1395 return Ok(None);
1396 }
1397 let (rank, omega_pinv) = gam_linalg::utils::block_penalty_rank_and_pinv(&omega)
1398 .map_err(|e| format!("conditional score test pseudo-inverse failed: {e}"))?;
1399 if rank == 0 {
1400 return Ok(None);
1401 }
1402 let d_stat = s.dot(&omega_pinv.dot(&s));
1403 if !(d_stat.is_finite() && d_stat >= 0.0) {
1404 return Ok(None);
1405 }
1406 // p = 1 − CDF_{χ²_rank}(D) = 1 − P(rank/2, D/2) (regularized lower gamma).
1407 let p_lower = statrs::function::gamma::gamma_lr(rank as f64 / 2.0, d_stat / 2.0);
1408 let p_value = (1.0 - p_lower).clamp(0.0, 1.0);
1409 Ok(Some(p_value))
1410}
1411
1412/// Fit the conditional location-scale calibration (#905) if the conditional
1413/// `E[z|C]`/`Var(z|C)` Rao gate fires on the marginal-index basis `a_block`.
1414///
1415/// Returns `None` when there is no conditional structure to correct (the gate
1416/// does not fire, or the basis is degenerate) — in that case the caller falls
1417/// back to the existing pooled-marginal gate (rank-INT or no calibration).
1418pub(crate) fn fit_conditional_latent_calibration_if_needed(
1419 z: &Array1<f64>,
1420 weights: &Array1<f64>,
1421 a_block: ArrayView2<'_, f64>,
1422) -> Result<Option<LatentZConditionalCalibration>, String> {
1423 let n = z.len();
1424 let p = a_block.ncols();
1425 if n != weights.len() {
1426 return Err(format!(
1427 "conditional latent gate length mismatch: z={n}, weights={}",
1428 weights.len()
1429 ));
1430 }
1431 if a_block.nrows() != n {
1432 return Err(format!(
1433 "conditional latent gate row mismatch: z={n}, basis rows={}",
1434 a_block.nrows()
1435 ));
1436 }
1437 if p == 0 {
1438 return Ok(None);
1439 }
1440 let total_weight = weights.iter().copied().sum::<f64>();
1441 if !(total_weight.is_finite() && total_weight > 0.0) {
1442 return Ok(None);
1443 }
1444 if z.iter().any(|v| !v.is_finite()) || a_block.iter().any(|v| !v.is_finite()) {
1445 return Ok(None);
1446 }
1447
1448 let z_mean = z
1449 .iter()
1450 .zip(weights.iter())
1451 .map(|(&zi, &wi)| wi * zi)
1452 .sum::<f64>()
1453 / total_weight;
1454 let global_var = z
1455 .iter()
1456 .zip(weights.iter())
1457 .map(|(&zi, &wi)| wi * (zi - z_mean) * (zi - z_mean))
1458 .sum::<f64>()
1459 / total_weight;
1460 if !(global_var.is_finite() && global_var > 0.0) {
1461 return Ok(None);
1462 }
1463
1464 // Center each basis column by its weighted mean so the score test is about
1465 // conditional structure *beyond* the global level (the intercept nuisance).
1466 // A constant marginal-design column collapses to ~0 and is dropped by the
1467 // pseudo-inverse rank, so an intercept already present in a(C) is harmless.
1468 let mut a_centered = a_block.to_owned();
1469 for j in 0..p {
1470 let col = a_block.column(j);
1471 let col_mean = col
1472 .iter()
1473 .zip(weights.iter())
1474 .map(|(&v, &w)| w * v)
1475 .sum::<f64>()
1476 / total_weight;
1477 a_centered.column_mut(j).mapv_inplace(|v| v - col_mean);
1478 }
1479
1480 // Conditional-mean Rao test: u = z − z̄.
1481 let u_mean: Vec<f64> = z.iter().map(|&zi| zi - z_mean).collect();
1482 let p_mean = robust_conditional_score_pvalue(a_centered.view(), &u_mean, weights.view())?;
1483 // Conditional-variance (Breusch-Pagan) Rao test: u = (z − z̄)² − σ̂².
1484 let u_var: Vec<f64> = u_mean.iter().map(|&e| e * e - global_var).collect();
1485 let p_var = robust_conditional_score_pvalue(a_centered.view(), &u_var, weights.view())?;
1486
1487 let mean_fires = p_mean.is_some_and(|p| p < AUTO_Z_CONDITIONAL_RAO_ALPHA);
1488 let var_fires = p_var.is_some_and(|p| p < AUTO_Z_CONDITIONAL_RAO_ALPHA);
1489 if !mean_fires && !var_fires {
1490 return Ok(None);
1491 }
1492
1493 // Escalation fires. Fit the conditional mean over the full basis
1494 // [1 | a(C)] via a weighted ridge (the ridge stabilizes a rank-deficient
1495 // marginal-index span; it does not meaningfully shrink the few directions
1496 // that triggered the gate). The conditional-mean correction is applied
1497 // whenever the gate fires (a pure-variance trigger leaves the C-slopes of
1498 // m(C) ≈ 0, so it reduces to harmless global centering).
1499 let basis = build_intercept_basis(a_block);
1500 // Per-column Tikhonov penalty scaled by the weighted Gram diagonal, so the
1501 // ridge is *relative* to each column's scale (a 1e-8 absolute ridge would
1502 // be negligible against an O(n) Gram and would not stabilize a
1503 // rank-deficient penalized-spline marginal index). `diag_jj = Σ_i w_i a_ij²`;
1504 // floored positive so the all-zero (already-dropped) directions still
1505 // receive a finite ridge and the factorization cannot fail.
1506 let mut penalty = Array2::<f64>::zeros((basis.ncols(), basis.ncols()));
1507 for j in 0..basis.ncols() {
1508 let diag_jj = basis
1509 .column(j)
1510 .iter()
1511 .zip(weights.iter())
1512 .map(|(&x, &w)| w * x * x)
1513 .sum::<f64>()
1514 .max(f64::MIN_POSITIVE);
1515 penalty[[j, j]] = diag_jj;
1516 }
1517 let z_col = z.view().insert_axis(ndarray::Axis(1));
1518 let (mean_coeffs_mat, mean_fitted) = gam_linalg::utils::gaussian_weighted_ridge(
1519 basis.view(),
1520 z_col,
1521 penalty.view(),
1522 weights.view(),
1523 AUTO_Z_CONDITIONAL_RIDGE_REL,
1524 )?;
1525 let mean_coeffs: Vec<f64> = mean_coeffs_mat.column(0).to_vec();
1526
1527 // First-stage (generated-regressor) normal matrix `M = AᵀWA + λR`, the same
1528 // weighted-ridge system `gaussian_weighted_ridge` factorizes internally;
1529 // rebuilt here so its inverse can form the closed-form coefficient sandwich
1530 // `V₁` that the second-stage Murphy–Topel correction consumes. `p` is the
1531 // marginal-index width (small), so this is a cheap dense `(p+1)²` form.
1532 let normal_matrix = {
1533 let mut wa = basis.to_owned();
1534 for i in 0..wa.nrows() {
1535 let wi = weights[i];
1536 wa.row_mut(i).iter_mut().for_each(|value| *value *= wi);
1537 }
1538 let mut m = basis.t().dot(&wa);
1539 m += &(penalty.to_owned() * AUTO_Z_CONDITIONAL_RIDGE_REL);
1540 m
1541 };
1542 let mean_residuals: Vec<f64> = z
1543 .iter()
1544 .zip(mean_fitted.column(0).iter())
1545 .map(|(&zi, &mi)| zi - mi)
1546 .collect();
1547 let mean_cov = weighted_ridge_sandwich_cov(
1548 basis.view(),
1549 &mean_residuals,
1550 weights.view(),
1551 &normal_matrix,
1552 )?;
1553
1554 let var_floor = (AUTO_Z_CONDITIONAL_VAR_FLOOR_FRAC * global_var).max(f64::MIN_POSITIVE);
1555 let (var_coeffs, var_cov): (Vec<f64>, Array2<f64>) = if var_fires {
1556 // Conditional-variance correction: regress the squared mean-residual on
1557 // the same basis. Fitted values are floored at `var_floor` when applied.
1558 let resid_sq: Array1<f64> = mean_residuals.iter().map(|&e| e * e).collect();
1559 let resid_col = resid_sq.view().insert_axis(ndarray::Axis(1));
1560 let (var_coeffs_mat, var_fitted) = gam_linalg::utils::gaussian_weighted_ridge(
1561 basis.view(),
1562 resid_col,
1563 penalty.view(),
1564 weights.view(),
1565 AUTO_Z_CONDITIONAL_RIDGE_REL,
1566 )?;
1567 // First-stage sandwich for the variance coefficients on the same ridge
1568 // normal matrix `M` (the basis and weights are identical; only the
1569 // response — and hence the residual — differs). `û_i = (z−m̂)²_i − v̂_i`
1570 // is the Breusch–Pagan residual.
1571 let var_residuals: Vec<f64> = resid_sq
1572 .iter()
1573 .zip(var_fitted.column(0).iter())
1574 .map(|(&si, &vi)| si - vi)
1575 .collect();
1576 let cov = weighted_ridge_sandwich_cov(
1577 basis.view(),
1578 &var_residuals,
1579 weights.view(),
1580 &normal_matrix,
1581 )?;
1582 (var_coeffs_mat.column(0).to_vec(), cov)
1583 } else {
1584 (Vec::new(), Array2::<f64>::zeros((0, 0)))
1585 };
1586
1587 let mut calibration = LatentZConditionalCalibration {
1588 mean_coeffs,
1589 var_coeffs,
1590 basis_ncols: p,
1591 var_floor,
1592 global_var,
1593 post_mean: 0.0,
1594 post_sd: 1.0,
1595 mean_cov,
1596 var_cov,
1597 };
1598
1599 // Sanity-check post-correction moments on the training sample.
1600 let calibrated = calibration.apply(z.view(), a_block)?;
1601 let post_mean = weighted_mean(calibrated.as_slice().unwrap(), weights.view(), total_weight);
1602 let post_var = calibrated
1603 .iter()
1604 .zip(weights.iter())
1605 .map(|(&zi, &wi)| wi * (zi - post_mean) * (zi - post_mean))
1606 .sum::<f64>()
1607 / total_weight;
1608 calibration.post_mean = post_mean;
1609 calibration.post_sd = post_var.max(0.0).sqrt();
1610
1611 Ok(Some(calibration))
1612}
1613
1614/// Prepend a column of ones to `a_block`, producing the `[1 | a(C)]` regression
1615/// basis used by the conditional location-scale fit.
1616pub(crate) fn build_intercept_basis(a_block: ArrayView2<'_, f64>) -> Array2<f64> {
1617 let n = a_block.nrows();
1618 let p = a_block.ncols();
1619 let mut basis = Array2::<f64>::ones((n, p + 1));
1620 basis.slice_mut(s![.., 1..]).assign(&a_block);
1621 basis
1622}
1623
1624pub(crate) fn build_latent_measure_with_geometry(
1625 z: &Array1<f64>,
1626 weights: &Array1<f64>,
1627 policy: &LatentZPolicy,
1628 conditioning: Option<ArrayView2<'_, f64>>,
1629) -> Result<(LatentMeasureKind, LatentMeasureCalibration), String> {
1630 match policy.latent_measure {
1631 LatentMeasureSpec::Auto { grid_size } => {
1632 // #905: conditional `E[z|C]`/`Var(z|C)` Rao gate. Inspect the latent
1633 // score's conditional moments on the marginal-index span a(C)
1634 // BEFORE the pooled-marginal gate. A significant conditional shift
1635 // is the `b(C)·m(C)` leakage the pooled gate cannot see and that
1636 // rank-INT provably cannot fix, so it takes precedence: route to the
1637 // conditional location-scale correction `ζ = (z−m(C))/√v(C)`.
1638 if let Some(a_block) = conditioning
1639 && let Some(cal) =
1640 fit_conditional_latent_calibration_if_needed(z, weights, a_block)?
1641 {
1642 // Matching the first two conditional moments does not
1643 // establish Gaussianity of the residual ζ (a two-point
1644 // residual law survives location-scale correction unchanged
1645 // in shape). The closed-form standard-normal kernel is only
1646 // admissible when the calibrated sample passes the same
1647 // pooled adequacy gate raw z faces; otherwise retain an
1648 // empirical latent measure built from ζ, so the residual
1649 // distribution stays the one the data show.
1650 let zeta = cal.apply(z.view(), a_block)?;
1651 let residual_is_standard_normal =
1652 latent_z_is_standard_normal_enough(&zeta, weights, policy)?;
1653 let kind = if residual_is_standard_normal {
1654 LatentMeasureKind::StandardNormal
1655 } else {
1656 build_global_empirical_latent_measure(&zeta, weights, grid_size)?
1657 };
1658 log::info!(
1659 "[BMS 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)",
1660 cal.basis_ncols,
1661 !cal.var_coeffs.is_empty(),
1662 cal.post_mean,
1663 cal.post_sd,
1664 if residual_is_standard_normal {
1665 "standard-normal"
1666 } else {
1667 "global-empirical"
1668 },
1669 );
1670 return Ok((
1671 kind,
1672 LatentMeasureCalibration::ConditionalLocationScale(cal),
1673 ));
1674 }
1675 if latent_z_is_standard_normal_enough(z, weights, policy)? {
1676 Ok((
1677 LatentMeasureKind::StandardNormal,
1678 LatentMeasureCalibration::None,
1679 ))
1680 } else {
1681 // P4: route bad-normal latent z through a weighted
1682 // mid-distribution-rank inverse-normal transform. Rank-INT
1683 // redefines the latent axis (the affine rigid model is
1684 // specified on the calibrated score); it makes the calibrated
1685 // sample approximately — not exactly — N(0,1), so the
1686 // closed-form standard-normal kernel is admitted only when
1687 // the calibrated sample itself passes the adequacy gate.
1688 // When it cannot (heavy ties leave the calibrated law
1689 // discrete), fall back to the mathematically exact
1690 // global-empirical latent measure on the raw score.
1691 let calibration = LatentZRankIntCalibration::fit(z, weights)?;
1692 let calibrated = calibration.apply_to_training(z)?;
1693 if latent_z_is_standard_normal_enough(&calibrated, weights, policy)? {
1694 log::info!(
1695 "[BMS latent-z] rank-INT calibrated: post_mean={:.3e} post_sd={:.3e} knots={}",
1696 calibration.post_mean,
1697 calibration.post_sd,
1698 calibration.sorted_z.len(),
1699 );
1700 Ok((
1701 LatentMeasureKind::StandardNormal,
1702 LatentMeasureCalibration::RankInverseNormal(calibration),
1703 ))
1704 } else {
1705 log::info!(
1706 "[BMS latent-z] rank-INT output failed the standard-normal adequacy gate (post_mean={:.3e} post_sd={:.3e} knots={}); using the global-empirical latent measure",
1707 calibration.post_mean,
1708 calibration.post_sd,
1709 calibration.sorted_z.len(),
1710 );
1711 Ok((
1712 build_global_empirical_latent_measure(z, weights, grid_size)?,
1713 LatentMeasureCalibration::None,
1714 ))
1715 }
1716 }
1717 }
1718 LatentMeasureSpec::StandardNormal => Ok((
1719 LatentMeasureKind::StandardNormal,
1720 LatentMeasureCalibration::None,
1721 )),
1722 LatentMeasureSpec::GlobalEmpirical { grid_size } => {
1723 let kind = build_global_empirical_latent_measure(z, weights, grid_size)?;
1724 Ok((kind, LatentMeasureCalibration::None))
1725 }
1726 }
1727}
1728
1729pub(crate) fn latent_z_is_standard_normal_enough(
1730 z: &Array1<f64>,
1731 weights: &Array1<f64>,
1732 policy: &LatentZPolicy,
1733) -> Result<bool, String> {
1734 if z.len() != weights.len() {
1735 return Err(format!(
1736 "latent-measure auto-detection length mismatch: z={}, weights={}",
1737 z.len(),
1738 weights.len()
1739 ));
1740 }
1741 let weight_sum = weights.iter().copied().sum::<f64>();
1742 let weight_sq_sum = weights.iter().map(|&w| w * w).sum::<f64>();
1743 if !(weight_sum.is_finite()
1744 && weight_sum > 0.0
1745 && weight_sq_sum.is_finite()
1746 && weight_sq_sum > 0.0)
1747 {
1748 return Err("latent-measure auto-detection requires positive finite weights".to_string());
1749 }
1750 let effective_n = weight_sum * weight_sum / weight_sq_sum;
1751 if !(effective_n.is_finite() && effective_n > 1.0) {
1752 return Err(
1753 "latent-measure auto-detection requires at least two effective observations"
1754 .to_string(),
1755 );
1756 }
1757 let mean = z
1758 .iter()
1759 .zip(weights.iter())
1760 .map(|(&zi, &wi)| wi * zi)
1761 .sum::<f64>()
1762 / weight_sum;
1763 let var = z
1764 .iter()
1765 .zip(weights.iter())
1766 .map(|(&zi, &wi)| wi * (zi - mean) * (zi - mean))
1767 .sum::<f64>()
1768 / weight_sum;
1769 let sd = var.sqrt();
1770 if !(mean.is_finite() && sd.is_finite() && sd > 0.0) {
1771 return Ok(false);
1772 }
1773 let skew = z
1774 .iter()
1775 .zip(weights.iter())
1776 .map(|(&zi, &wi)| {
1777 let centered = (zi - mean) / sd;
1778 wi * centered.powi(3)
1779 })
1780 .sum::<f64>()
1781 / weight_sum;
1782 let excess_kurtosis = z
1783 .iter()
1784 .zip(weights.iter())
1785 .map(|(&zi, &wi)| {
1786 let centered = (zi - mean) / sd;
1787 wi * centered.powi(4)
1788 })
1789 .sum::<f64>()
1790 / weight_sum
1791 - 3.0;
1792 let mean_tol = policy.mean_tol_multiplier / effective_n.sqrt();
1793 let sd_tol = policy.sd_tol_multiplier / (2.0 * (effective_n - 1.0).max(1.0)).sqrt();
1794 let ks_to_normal = weighted_ks_to_standard_normal(z, weights, weight_sum)?;
1795 let tail_mass_4 = weighted_tail_mass(z, weights, weight_sum, AUTO_Z_NORMAL_TAIL_SIGMA_INNER);
1796 let tail_mass_6 = weighted_tail_mass(z, weights, weight_sum, AUTO_Z_NORMAL_TAIL_SIGMA_OUTER);
1797 let max_abs_z = z.iter().fold(0.0_f64, |acc, &zi| acc.max(zi.abs()));
1798 let normal_tail_4 = 2.0 * (1.0 - normal_cdf(AUTO_Z_NORMAL_TAIL_SIGMA_INNER));
1799 let normal_tail_6 = 2.0 * (1.0 - normal_cdf(AUTO_Z_NORMAL_TAIL_SIGMA_OUTER));
1800 Ok(mean.abs() <= mean_tol
1801 && (sd - 1.0).abs() <= sd_tol
1802 && skew.is_finite()
1803 && skew.abs() <= policy.max_abs_skew.min(AUTO_Z_NORMAL_SKEW_TOL)
1804 && excess_kurtosis.is_finite()
1805 && excess_kurtosis.abs() <= policy.max_abs_excess_kurtosis.min(AUTO_Z_NORMAL_KURT_TOL)
1806 && ks_to_normal.is_finite()
1807 && ks_to_normal <= AUTO_Z_NORMAL_KS_TOL
1808 && tail_mass_4
1809 <= AUTO_Z_NORMAL_TAIL_MASS_SLACK * normal_tail_4 + AUTO_Z_NORMAL_TAIL_FLOOR_INNER
1810 && tail_mass_6
1811 <= AUTO_Z_NORMAL_TAIL_MASS_SLACK * normal_tail_6 + AUTO_Z_NORMAL_TAIL_FLOOR_OUTER
1812 && max_abs_z < AUTO_Z_NORMAL_MAX_ABS)
1813}
1814
1815pub(crate) fn build_global_empirical_latent_measure(
1816 z: &Array1<f64>,
1817 weights: &Array1<f64>,
1818 grid_size: usize,
1819) -> Result<LatentMeasureKind, String> {
1820 let grid = build_empirical_z_grid(z, weights, grid_size, "empirical latent measure")?;
1821 let measure = LatentMeasureKind::GlobalEmpirical { grid };
1822 measure.validate("empirical latent measure")?;
1823 Ok(measure)
1824}
1825
1826pub(crate) fn weighted_ks_to_standard_normal(
1827 z: &Array1<f64>,
1828 weights: &Array1<f64>,
1829 total_weight: f64,
1830) -> Result<f64, String> {
1831 let mut pairs = Vec::<(f64, f64)>::with_capacity(z.len());
1832 for (&zi, &wi) in z.iter().zip(weights.iter()) {
1833 if !zi.is_finite() || !wi.is_finite() || wi < 0.0 {
1834 return Err(
1835 "latent-measure KS diagnostic requires finite z and non-negative finite weights"
1836 .to_string(),
1837 );
1838 }
1839 if wi > 0.0 {
1840 pairs.push((zi, wi));
1841 }
1842 }
1843 pairs.sort_by(|left, right| {
1844 left.0
1845 .partial_cmp(&right.0)
1846 .expect("validated latent z values are finite")
1847 });
1848 let mut prev = 0.0;
1849 let mut ks = 0.0_f64;
1850 for (zi, wi) in pairs {
1851 let cdf = normal_cdf(zi);
1852 let next = prev + wi / total_weight;
1853 ks = ks.max((cdf - prev).abs()).max((cdf - next).abs());
1854 prev = next;
1855 }
1856 Ok(ks)
1857}
1858
1859pub(crate) fn weighted_tail_mass(
1860 z: &Array1<f64>,
1861 weights: &Array1<f64>,
1862 total_weight: f64,
1863 cutoff: f64,
1864) -> f64 {
1865 z.iter()
1866 .zip(weights.iter())
1867 .filter(|&(&zi, _)| zi.abs() > cutoff)
1868 .map(|(_, &wi)| wi)
1869 .sum::<f64>()
1870 / total_weight
1871}
1872
1873pub(crate) fn build_empirical_z_grid(
1874 z: &Array1<f64>,
1875 weights: &Array1<f64>,
1876 grid_size: usize,
1877 context: &str,
1878) -> Result<EmpiricalZGrid, String> {
1879 if grid_size < 3 {
1880 return Err(format!(
1881 "empirical latent measure grid_size must be at least 3, got {grid_size}"
1882 ));
1883 }
1884 if z.len() != weights.len() {
1885 return Err(format!(
1886 "{context} length mismatch: z={}, weights={}",
1887 z.len(),
1888 weights.len()
1889 ));
1890 }
1891 let mut pairs = Vec::<(f64, f64)>::with_capacity(z.len());
1892 for (idx, (&zi, &wi)) in z.iter().zip(weights.iter()).enumerate() {
1893 if !zi.is_finite() {
1894 return Err(format!(
1895 "{context} z value at row {idx} is non-finite ({zi})"
1896 ));
1897 }
1898 if !wi.is_finite() || wi < 0.0 {
1899 return Err(format!(
1900 "{context} weight at row {idx} must be finite and non-negative, got {wi}"
1901 ));
1902 }
1903 if wi > 0.0 {
1904 pairs.push((zi, wi));
1905 }
1906 }
1907 if pairs.len() < 2 {
1908 return Err(format!(
1909 "{context} requires at least two positive-weight rows"
1910 ));
1911 }
1912 pairs.sort_by(|left, right| {
1913 left.0
1914 .partial_cmp(&right.0)
1915 .expect("validated empirical latent z values are finite")
1916 });
1917 let total_weight = pairs.iter().map(|(_, weight)| *weight).sum::<f64>();
1918 if !(total_weight.is_finite() && total_weight > 0.0) {
1919 return Err(format!("{context} requires positive finite total weight"));
1920 }
1921
1922 let m = grid_size.min(pairs.len());
1923 let mut nodes = Vec::with_capacity(m);
1924 let mut out_weights = Vec::with_capacity(m);
1925 let bin_weight_target = total_weight / (m as f64);
1926 let mut cursor = 0usize;
1927 let mut remaining = pairs[0].1;
1928 for _ in 0..m {
1929 let mut need = bin_weight_target;
1930 let mut bin_weight = 0.0;
1931 let mut bin_sum = 0.0;
1932 while need > EMPIRICAL_GRID_WEIGHT_EXHAUSTED_REL_TOL * bin_weight_target
1933 && cursor < pairs.len()
1934 {
1935 let take = remaining.min(need);
1936 bin_sum += take * pairs[cursor].0;
1937 bin_weight += take;
1938 need -= take;
1939 remaining -= take;
1940 if remaining <= EMPIRICAL_GRID_WEIGHT_EXHAUSTED_REL_TOL * pairs[cursor].1 {
1941 cursor += 1;
1942 if cursor < pairs.len() {
1943 remaining = pairs[cursor].1;
1944 }
1945 }
1946 }
1947 if bin_weight > 0.0 {
1948 nodes.push(bin_sum / bin_weight);
1949 out_weights.push(bin_weight / total_weight);
1950 }
1951 }
1952 if nodes.len() < 2 {
1953 return Err(format!(
1954 "{context} compression produced fewer than two nodes"
1955 ));
1956 }
1957 recenter_rescale_empirical_grid(&mut nodes, &out_weights);
1958 let total = out_weights.iter().sum::<f64>();
1959 if total.is_finite() && total > 0.0 {
1960 for weight in &mut out_weights {
1961 *weight /= total;
1962 }
1963 }
1964 validate_empirical_z_grid(&nodes, &out_weights, context)?;
1965 Ok(EmpiricalZGrid {
1966 nodes,
1967 weights: out_weights,
1968 })
1969}
1970
1971pub(crate) fn recenter_rescale_empirical_grid(nodes: &mut [f64], weights: &[f64]) {
1972 let total = weights.iter().sum::<f64>();
1973 if !(total.is_finite() && total > 0.0) {
1974 return;
1975 }
1976 let mean = nodes
1977 .iter()
1978 .zip(weights.iter())
1979 .map(|(&node, &weight)| weight * node)
1980 .sum::<f64>()
1981 / total;
1982 let var = nodes
1983 .iter()
1984 .zip(weights.iter())
1985 .map(|(&node, &weight)| weight * (node - mean).powi(2))
1986 .sum::<f64>()
1987 / total;
1988 let sd = var.sqrt();
1989 if sd.is_finite() && sd > BMS_VARIANCE_FLOOR {
1990 for node in nodes {
1991 *node = (*node - mean) / sd;
1992 }
1993 }
1994}
1995
1996// ---------------------------------------------------------------------------
1997// Cross-module constants — declared here so all submodules can reach them
1998// via `use super::*` without promoting implementation details to pub(crate).
1999// ---------------------------------------------------------------------------
2000pub(super) const BMS_AUTO_SUBSAMPLE_PHASE1_BUDGET: usize = 12;
2001pub(super) const BERNOULLI_LINK_PROBABILITY_EPS: f64 = 1e-12;
2002pub(super) const BMS_VARIANCE_FLOOR: f64 = 1e-12;
2003pub(super) const BMS_DERIV_TOL: f64 = 1e-8;
2004/// Relative tolerance below which a residual weight is treated as exhausted in
2005/// the equal-mass empirical-grid compression loop. Used both for the per-bin
2006/// "need" remaining (relative to the target bin weight) and for the per-pair
2007/// remainder (relative to that pair's weight), so a pair/bin that is filled to
2008/// within a few ulps advances the cursor instead of spinning on round-off.
2009pub(super) const EMPIRICAL_GRID_WEIGHT_EXHAUSTED_REL_TOL: f64 = 1e-14;
2010/// Upper bound (and large-`n` default) for rows-per-chunk in the parallel
2011/// row-accumulation phases.
2012///
2013/// This is also a hard *ceiling* the [`bms_row_chunk_size`] chunk sizing must
2014/// respect: several per-chunk fast paths (block-Hessian / block-gradient
2015/// assembly) allocate fixed `[0.0f64; ROW_CHUNK_SIZE]` stack buffers and index
2016/// them by the chunk's local row position, so a chunk may never carry more than
2017/// `ROW_CHUNK_SIZE` rows.
2018pub(super) const ROW_CHUNK_SIZE: usize = 1024;
2019/// Floor for rows-per-chunk: below it the per-chunk scratch allocation +
2020/// scheduler hand-off cost dominates the row arithmetic. Small enough that a
2021/// moderate `n` on a many-core box still carves several chunks per worker.
2022pub(super) const ROW_CHUNK_MIN: usize = 64;
2023/// Target number of row-chunks per rayon worker for the BMS exact-Newton
2024/// row-fan-out phases (gradient / HVP / diagonal directional-derivative sweeps).
2025///
2026/// Several chunks per worker keeps the pool load-balanced across the uneven
2027/// per-row cost tail (work-stealing moves whole chunks, never partial sums) so
2028/// the heavy coord-corrections / row-stream phases saturate the cores instead
2029/// of stranding the tail on one worker.
2030pub(super) const ROW_CHUNKS_PER_WORKER: usize = 4;
2031
2032/// Pool-aware rows-per-chunk for the BMS exact-Newton row fan-outs.
2033///
2034/// A *fixed* `ROW_CHUNK_SIZE` divisor makes the chunk **count** scale with `n`,
2035/// so at moderate `n` (e.g. `n = 10·ROW_CHUNK_SIZE` on a 64-core box) the
2036/// `into_par_iter` over `⌈n/ROW_CHUNK_SIZE⌉` chunks has far fewer tasks than
2037/// workers and most cores idle — the measured ~30-90% core utilization on the
2038/// biobank coord-corrections / row-stream phases. This sizes the chunk so the
2039/// chunk count targets `ROW_CHUNKS_PER_WORKER × worker_count` (the same policy
2040/// `chunked_row_reduction` uses), clamped to `[ROW_CHUNK_MIN, ROW_CHUNK_SIZE]`:
2041///
2042/// * the `ROW_CHUNK_SIZE` ceiling is mandatory — the block-assembly fast paths
2043/// index fixed `[…; ROW_CHUNK_SIZE]` stack buffers by local row, so a chunk
2044/// can never exceed it. At large `n` the per-1024-row count already exceeds
2045/// the worker count, so the clamp costs nothing there;
2046/// * the `ROW_CHUNK_MIN` floor stops sub-floor fan-out at tiny `n`.
2047///
2048/// Reproducibility contract (#1045): the worker count used here is the
2049/// process-stable machine parallelism (`reproducible_chunk_parallelism`), NOT
2050/// the live `rayon::current_num_threads()` of the executing (possibly scoped,
2051/// possibly shrunk) pool. Keying the chunk *count* — and hence the chunk
2052/// boundaries `chunk_idx·chunk → (chunk_idx+1)·chunk` — to the transient pool
2053/// size made the per-chunk row sums regroup when the pool was narrowed, so a
2054/// fit reduced over these chunks and fed into the iterative REML optimizer moved
2055/// its `(ρ, λ)` selection with the pool size. Anchoring to a process constant
2056/// makes the boundaries — and therefore the `try_fold`/`try_reduce` reduction
2057/// tree that round-trips through them — invariant to how many workers run the
2058/// fit, while rayon still fans the chunks across whatever workers exist. For a
2059/// given `n` the returned chunk size is stable across calls and pool sizes.
2060#[inline]
2061pub(super) fn bms_row_chunk_size(n: usize) -> usize {
2062 if n == 0 {
2063 return ROW_CHUNK_SIZE;
2064 }
2065 let workers = crate::marginal_slope_shared::reproducible_chunk_parallelism();
2066 let target_chunks = workers.saturating_mul(ROW_CHUNKS_PER_WORKER).max(1);
2067 // Rows per chunk that yields ≈ `target_chunks` chunks, clamped into
2068 // `[ROW_CHUNK_MIN, ROW_CHUNK_SIZE]`.
2069 n.div_ceil(target_chunks)
2070 .clamp(ROW_CHUNK_MIN, ROW_CHUNK_SIZE)
2071}
2072pub(super) const EXACT_WORK_LOG_MIN_ROWS: usize = 50_000;
2073pub(super) const BMS_ROW_PRIMARY_HESSIAN_EXPECTED_REUSE_PASSES: usize = 3;
2074pub(super) const BMS_ROW_PRIMARY_HESSIAN_MIN_REUSE_PASSES: usize = 2;
2075pub(super) const BMS_ROW_PRIMARY_HESSIAN_TILE_ROWS: usize = 8192;
2076pub(super) const BMS_ROW_PRIMARY_HESSIAN_SINGLE_FRACTION_NUM: u64 = 1;
2077pub(super) const BMS_ROW_PRIMARY_HESSIAN_SINGLE_FRACTION_DEN: u64 = 4;
2078pub(super) const BMS_ROW_PRIMARY_HESSIAN_GLOBAL_FRACTION_NUM: u64 = 1;
2079pub(super) const BMS_ROW_PRIMARY_HESSIAN_GLOBAL_FRACTION_DEN: u64 = 2;
2080pub(super) const BERNOULLI_MARGSLOPE_LINE_SEARCH_EARLY_EXIT_CHUNK_ROWS: usize = 10_000;
2081
2082// ---------------------------------------------------------------------------
2083// Submodule declarations
2084// ---------------------------------------------------------------------------
2085pub(crate) mod block_specs;
2086pub(crate) mod exact_eval_cache;
2087pub(crate) mod family;
2088pub(crate) mod gradient_paths;
2089pub(crate) mod hessian_paths;
2090pub(crate) mod install_flex;
2091pub(crate) mod row_kernel;
2092#[cfg(test)]
2093mod tests {
2094 include!("../../../../tests/src_modules/misc/families_bms_identifiability_rigid_tests.rs");
2095 include!(
2096 "../../../../tests/src_modules/optimization/families_bms_joint_hessian_hvp_correction_tests.rs"
2097 );
2098}
2099pub(crate) mod axis_direction_search;
2100pub(crate) mod cell_moment_assembly;
2101// #932 BMS flex single-source jet substrate (runtime-dimension `Jet2` + IFT
2102// lift + cell base-moment jets). A bare `#[cfg(test)] mod` with an allowed name
2103// so the build.rs ban-scanner exempts it; shared by its own FD gates and the
2104// `cell_moment_assembly` flex-fixture oracle gate as a private child of `bms`.
2105#[cfg(test)]
2106mod test_support;
2107// #932 INDEPENDENT adversarial verifier (bms-flex-verify): a high-order
2108// finite-difference oracle on the production hand path
2109// `compute_row_analytic_flex_from_parts_into` + a moving-edge Leibniz
2110// cross-check + a planted-corruption tripwire. Bare `#[cfg(test)] mod` with the
2111// allowed `*_tests` name so the build.rs ban-scanner exempts it; owned solely by
2112// the verifier (never edits the implementer's row_primary_hessian /
2113// gradient_paths / cell_moment_assembly).
2114pub(crate) mod custom_family_impl;
2115#[cfg(test)]
2116mod flex_verify_932_tests;
2117// #932 direct production-path measurement: forced 65-node empirical grid,
2118// warmed/cold row-op allocation counting + ns/row diagnostics for the MSI
2119// A/B ledger. The asserted gate is per-row allocation calls (deterministic);
2120// timing is eprintln-only per the SPEC ban on wall-clock correctness budgets.
2121#[cfg(test)]
2122mod flex_measure_932_tests;
2123pub(crate) mod row_primary_hessian;
2124
2125pub use block_specs::fit_bernoulli_marginal_slope_terms;
2126pub use gradient_paths::{
2127 MarginalSlopeCovariance, MarginalSlopeCovarianceShape, marginal_slope_covariance_from_scores,
2128 marginal_slope_preserving_scale, marginal_slope_probit_eta, padded_deviation_seed,
2129};
2130pub use install_flex::CrossBlockIdentifiabilityWarning;
2131pub(crate) use install_flex::FlexCompileOutcome;
2132
2133// pub(crate) re-exports for internal callers:
2134pub(crate) use block_specs::push_deviation_aux_blockspecs;
2135pub use block_specs::{BmsLogslopeJacobian, BmsMarginalJacobian};
2136pub(crate) use family::{
2137 BernoulliMarginalLinkMap, bernoulli_marginal_link_map,
2138 build_link_deviation_block_from_knots_design_seed_and_weights,
2139 build_score_warp_deviation_block_from_seed,
2140};
2141pub(crate) use gradient_paths::standardize_latent_z_with_policy;
2142pub(crate) use gradient_paths::{
2143 empirical_intercept_from_marginal, signed_probit_neglog_derivatives_up_to_fourth,
2144 unary_derivatives_log, unary_derivatives_log_normal_pdf, unary_derivatives_neglog_phi,
2145 unary_derivatives_sqrt,
2146};
2147pub(crate) use install_flex::{
2148 install_compiled_flex_block_into_runtime, project_monotone_feasible_beta,
2149};