1use super::family::*;
2use super::gradient_paths::*;
3use super::hessian_paths::{new_cell_moment_cache_stats, new_cell_moment_lru_cache};
4use super::install_flex::validate_spec;
5use super::*;
6use crate::marginal_slope_orthogonal::influence_absorber_log_lambda;
7use faer::Side;
8use gam_linalg::faer_ndarray::{FaerEigh, fast_ab, fast_atb, fast_xt_diag_x};
9
10pub(crate) const BMS_PROBIT_SEPARATION_ETA_INF: f64 = 35.0;
25
26pub(super) const GAUGE_PRIORITY_ANCHOR: u8 = 200;
41pub(super) const GAUGE_PRIORITY_MARGINAL: u8 = 150;
44pub(super) const GAUGE_PRIORITY_LOGSLOPE: u8 = 120;
46pub(super) const GAUGE_PRIORITY_CANDIDATE_FLEX: u8 = 100;
49pub(super) const GAUGE_PRIORITY_SCORE_WARP_DEV: u8 = 80;
52pub(super) const GAUGE_PRIORITY_DEVIATION_DEFAULT: u8 = 70;
56pub(super) const GAUGE_PRIORITY_LINK_DEV: u8 = 60;
58
59pub(crate) const EXACT_SPATIAL_OUTER_TOL_FLOOR: f64 = 1e-6;
65
66pub struct BmsMarginalJacobian {
105 pub marginal_dense: Arc<Array2<f64>>,
107 pub logslope_dense: Arc<Array2<f64>>,
109 pub offset_m: Array1<f64>,
110 pub offset_s: Array1<f64>,
111 pub p_marginal: usize,
113}
114
115impl BmsMarginalJacobian {
116 pub fn new(
117 marginal_dense: Arc<Array2<f64>>,
118 logslope_dense: Arc<Array2<f64>>,
119 offset_m: Array1<f64>,
120 offset_s: Array1<f64>,
121 p_marginal: usize,
122 ) -> Self {
123 Self {
124 marginal_dense,
125 logslope_dense,
126 offset_m,
127 offset_s,
128 p_marginal,
129 }
130 }
131}
132
133impl BlockEffectiveJacobian for BmsMarginalJacobian {
134 fn effective_jacobian_rows(
135 &self,
136 state: &FamilyLinearizationState<'_>,
137 rows: std::ops::Range<usize>,
138 ) -> Result<Array2<f64>, String> {
139 let beta = state.beta;
140 let s = state.probit_frailty_scale;
141 let p_m = self.p_marginal;
142 let p_s_block = self.logslope_dense.ncols();
143 let beta_s_raw = if beta.len() > p_m {
144 &beta[p_m..]
145 } else {
146 &[][..]
147 };
148 let p_s_use = p_s_block.min(beta_s_raw.len());
149 let beta_s = &beta_s_raw[..p_s_use];
150 let n = self.marginal_dense.nrows();
151 let rows = rows.start.min(n)..rows.end.min(n);
152 let p_block = self.marginal_dense.ncols();
153
154 let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_block));
164 for i in rows.clone() {
165 let g_i = self.offset_s[i]
166 + self
167 .logslope_dense
168 .row(i)
169 .slice(ndarray::s![..p_s_use])
170 .dot(&ArrayView1::from(beta_s));
171 let sg = s * g_i;
172 let c_i = (1.0 + sg * sg).sqrt();
173 let m_row = self.marginal_dense.row(i);
175 out.row_mut(i - rows.start).assign(&m_row.mapv(|x| c_i * x));
176 }
177 Ok(out)
178 }
179
180 fn n_outputs(&self) -> usize {
181 1
182 }
183
184 fn locks_raw_width_reduction(&self) -> bool {
185 true
195 }
196}
197
198pub struct BmsLogslopeJacobian {
210 pub marginal_dense: Arc<Array2<f64>>,
212 pub logslope_dense: Arc<Array2<f64>>,
214 pub offset_m: Array1<f64>,
215 pub offset_s: Array1<f64>,
216 pub z: Arc<Array1<f64>>,
217 pub p_marginal: usize,
219}
220
221impl BmsLogslopeJacobian {
222 pub fn new(
223 marginal_dense: Arc<Array2<f64>>,
224 logslope_dense: Arc<Array2<f64>>,
225 offset_m: Array1<f64>,
226 offset_s: Array1<f64>,
227 z: Arc<Array1<f64>>,
228 p_marginal: usize,
229 ) -> Self {
230 Self {
231 marginal_dense,
232 logslope_dense,
233 offset_m,
234 offset_s,
235 z,
236 p_marginal,
237 }
238 }
239}
240
241impl BlockEffectiveJacobian for BmsLogslopeJacobian {
242 fn effective_jacobian_rows(
243 &self,
244 state: &FamilyLinearizationState<'_>,
245 rows: std::ops::Range<usize>,
246 ) -> Result<Array2<f64>, String> {
247 let beta = state.beta;
248 let s = state.probit_frailty_scale;
249 let p_m = self.p_marginal;
250 let p_m_use = p_m.min(beta.len());
251 let beta_m = &beta[..p_m_use];
252 let beta_s_raw = if beta.len() > p_m {
253 &beta[p_m..]
254 } else {
255 &[][..]
256 };
257 let p_s_block = self.logslope_dense.ncols();
258 let p_s_use = p_s_block.min(beta_s_raw.len());
259 let beta_s = &beta_s_raw[..p_s_use];
260 let n = self.logslope_dense.nrows();
261 let rows = rows.start.min(n)..rows.end.min(n);
262
263 let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_s_block));
276 for i in rows.clone() {
277 let q_i = self.offset_m[i]
278 + self
279 .marginal_dense
280 .row(i)
281 .slice(ndarray::s![..p_m_use])
282 .dot(&ArrayView1::from(beta_m));
283 let g_i = self.offset_s[i]
284 + self
285 .logslope_dense
286 .row(i)
287 .slice(ndarray::s![..p_s_use])
288 .dot(&ArrayView1::from(beta_s));
289 let sg = s * g_i;
290 let c_i = (1.0 + sg * sg).sqrt();
291 let z_i = self.z[i];
292 let factor = q_i * s * s * g_i / c_i + s * z_i;
294 let g_row = self.logslope_dense.row(i);
296 out.row_mut(i - rows.start)
297 .assign(&g_row.mapv(|x| factor * x));
298 }
299 Ok(out)
300 }
301
302 fn n_outputs(&self) -> usize {
303 1
304 }
305
306 fn locks_raw_width_reduction(&self) -> bool {
307 true
318 }
319}
320
321pub(crate) fn widen_marginal_dense_with_influence(
331 marginal_dense: &Arc<Array2<f64>>,
332 influence_columns: Option<&Array2<f64>>,
333) -> Result<Arc<Array2<f64>>, String> {
334 let Some(z_infl) = influence_columns else {
335 return Ok(Arc::clone(marginal_dense));
336 };
337 let n = marginal_dense.nrows();
338 if z_infl.nrows() != n {
339 return Err(format!(
340 "influence block: residualised columns have {} rows, marginal design has {n}",
341 z_infl.nrows()
342 ));
343 }
344 let p_m = marginal_dense.ncols();
345 let p1 = z_infl.ncols();
346 let mut widened = Array2::<f64>::zeros((n, p_m + p1));
347 widened
348 .slice_mut(s![.., ..p_m])
349 .assign(marginal_dense.as_ref());
350 widened.slice_mut(s![.., p_m..]).assign(z_infl);
351 Ok(Arc::new(widened))
352}
353
354pub(crate) const LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL: f64 = 1.0e-6;
362
363#[derive(Debug, Clone)]
415pub(super) struct ReducedLogslopeReparam {
416 transform: Array2<f64>,
419}
420
421impl ReducedLogslopeReparam {
422 #[inline]
424 pub(super) fn original_cols(&self) -> usize {
425 self.transform.nrows()
426 }
427
428 #[inline]
430 pub(super) fn reduced_cols(&self) -> usize {
431 self.transform.ncols()
432 }
433
434 pub(super) fn recover_original_logslope_beta(
438 &self,
439 beta_reduced: &Array1<f64>,
440 ) -> Result<Array1<f64>, String> {
441 if beta_reduced.len() != self.reduced_cols() {
442 return Err(format!(
443 "reduced logslope reparam: β' length ({}) != reduced width ({})",
444 beta_reduced.len(),
445 self.reduced_cols()
446 ));
447 }
448 Ok(self.transform.dot(beta_reduced))
449 }
450}
451
452fn build_reduced_logslope_reparam(
463 marginal_design: &TermCollectionDesign,
464 logslope_design: &TermCollectionDesign,
465 z: &Array1<f64>,
466 row_metric: &Array1<f64>,
467 marginal_offset: &Array1<f64>,
468 logslope_offset: &Array1<f64>,
469 marginal_baseline: f64,
470 logslope_baseline: f64,
471 probit_scale: f64,
472) -> Result<Option<ReducedLogslopeReparam>, String> {
473 let marginal = marginal_design
474 .design
475 .try_to_dense_arc("build_reduced_logslope_reparam::marginal")?;
476 let logslope = logslope_design
477 .design
478 .try_to_dense_arc("build_reduced_logslope_reparam::logslope")?;
479 let n = marginal.nrows();
480 if logslope.nrows() != n
481 || z.len() != n
482 || row_metric.len() != n
483 || marginal_offset.len() != n
484 || logslope_offset.len() != n
485 {
486 return Err(format!(
487 "reduced logslope reparam row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
488 marginal.nrows(),
489 logslope.nrows(),
490 z.len(),
491 row_metric.len(),
492 marginal_offset.len(),
493 logslope_offset.len(),
494 ));
495 }
496 let p_m = marginal.ncols();
497 let p_g = logslope.ncols();
498 if p_m == 0 || p_g == 0 {
499 return Ok(None);
500 }
501 if !marginal_baseline.is_finite()
502 || !logslope_baseline.is_finite()
503 || !probit_scale.is_finite()
504 || probit_scale <= 0.0
505 || z.iter().any(|v| !v.is_finite())
506 || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
507 || marginal_offset.iter().any(|v| !v.is_finite())
508 || logslope_offset.iter().any(|v| !v.is_finite())
509 {
510 return Err(
511 "reduced logslope reparam requires finite pilot geometry and finite non-negative row metric"
512 .to_string(),
513 );
514 }
515
516 match reduced_logslope_transform_effective(
527 marginal.view(),
528 logslope.view(),
529 z,
530 row_metric,
531 marginal_offset,
532 logslope_offset,
533 marginal_baseline,
534 logslope_baseline,
535 probit_scale,
536 )? {
537 ReducedLogslopeOutcome::Reduced(transform) => {
538 Ok(Some(ReducedLogslopeReparam { transform }))
539 }
540 ReducedLogslopeOutcome::FullRank => Ok(None),
541 ReducedLogslopeOutcome::FullyConfounded => Err(
542 "BMS score-slope block is fully confounded with the marginal index: every \
543 effective logslope direction diag(f)·G·v is W-explained by the effective \
544 marginal span at the rigid pilot, so the data identify only the sum of the \
545 marginal and score-slope surfaces and the smoothing penalty would select an \
546 arbitrary decomposition between them. Refusing to fit; remove the score-slope \
547 terms or supply covariates that separate them from the marginal index."
548 .to_string(),
549 ),
550 }
551}
552
553#[derive(Debug)]
560pub(crate) enum ReducedLogslopeOutcome {
561 FullRank,
564 Reduced(Array2<f64>),
567 FullyConfounded,
570}
571
572pub(crate) fn reduced_logslope_transform_effective(
592 marginal: ArrayView2<'_, f64>,
593 logslope: ArrayView2<'_, f64>,
594 z: &Array1<f64>,
595 row_metric: &Array1<f64>,
596 marginal_offset: &Array1<f64>,
597 logslope_offset: &Array1<f64>,
598 marginal_baseline: f64,
599 logslope_baseline: f64,
600 probit_scale: f64,
601) -> Result<ReducedLogslopeOutcome, String> {
602 let n = marginal.nrows();
603 let p_m = marginal.ncols();
604 let p_g = logslope.ncols();
605 if p_m == 0 || p_g == 0 {
606 return Ok(ReducedLogslopeOutcome::FullRank);
607 }
608
609 let mut m_eff = Array2::<f64>::zeros((n, p_m));
611 let mut g_eff = Array2::<f64>::zeros((n, p_g));
612 for i in 0..n {
613 let q_i = marginal_offset[i] + marginal_baseline;
614 let g_i = logslope_offset[i] + logslope_baseline;
615 let sg = probit_scale * g_i;
616 let c_i = (1.0 + sg * sg).sqrt();
617 let f_i = q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
618 for j in 0..p_m {
619 m_eff[[i, j]] = c_i * marginal[[i, j]];
620 }
621 for j in 0..p_g {
622 g_eff[[i, j]] = f_i * logslope[[i, j]];
623 }
624 }
625
626 let c_gram = fast_xt_diag_x(&g_eff, row_metric);
629 let energy_scale = (0..p_g).map(|i| c_gram[[i, i]]).fold(0.0_f64, f64::max);
630 if !energy_scale.is_finite() {
631 return Err(
632 "reduced logslope reparam: effective logslope Gram produced non-finite energy"
633 .to_string(),
634 );
635 }
636 if energy_scale <= 0.0 {
637 return Ok(ReducedLogslopeOutcome::FullyConfounded);
640 }
641
642 let mut a_gram = fast_xt_diag_x(&m_eff, row_metric);
646 let a_scale = (0..p_m).map(|i| a_gram[[i, i]]).fold(0.0_f64, f64::max);
647 let a_ridge = (a_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL).max(f64::EPSILON);
648 for i in 0..p_m {
649 a_gram[[i, i]] += a_ridge;
650 }
651
652 let b_cross = gam_linalg::faer_ndarray::fast_xt_diag_y(&m_eff, row_metric, &g_eff);
654 let a_view = gam_linalg::faer_ndarray::FaerArrayView::new(&a_gram);
655 let a_factor =
656 gam_linalg::faer_ndarray::factorize_symmetricwith_fallback(a_view.as_ref(), Side::Lower)
657 .map_err(|e| {
658 format!(
659 "reduced logslope reparam: effective marginal Gram factorization failed: {e}"
660 )
661 })?;
662 let b_view = gam_linalg::faer_ndarray::FaerArrayView::new(&b_cross);
663 let solved = a_factor.solve(b_view.as_ref()); let a_inv_b = Array2::from_shape_fn((p_m, p_g), |(i, j)| solved[(i, j)]);
665 let schur = fast_atb(&b_cross, &a_inv_b); let mut stt = &c_gram - &schur;
667 stt = (&stt + &stt.t()) * 0.5;
668 if stt.iter().any(|v| !v.is_finite()) {
669 return Err(
670 "reduced logslope reparam: effective Schur Gram produced non-finite entries"
671 .to_string(),
672 );
673 }
674
675 let (evals, evecs) = stt
676 .eigh(Side::Lower)
677 .map_err(|e| format!("reduced logslope reparam: eigendecomposition failed: {e:?}"))?;
678 let tol = energy_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL;
682 let mut kept: Vec<usize> = (0..evals.len()).filter(|&i| evals[i] > tol).collect();
683 kept.sort_by(|&a, &b| {
684 evals[b]
685 .partial_cmp(&evals[a])
686 .unwrap_or(std::cmp::Ordering::Equal)
687 });
688 let r = kept.len();
689 if r == p_g {
695 return Ok(ReducedLogslopeOutcome::FullRank);
696 }
697 if r == 0 {
698 return Ok(ReducedLogslopeOutcome::FullyConfounded);
699 }
700 let mut transform = Array2::<f64>::zeros((p_g, r));
701 for (out_col, &src) in kept.iter().enumerate() {
702 transform.column_mut(out_col).assign(&evecs.column(src));
703 }
704 if transform.iter().any(|v| !v.is_finite()) {
705 return Err(
706 "reduced logslope reparam: reduced transform produced non-finite entries".to_string(),
707 );
708 }
709 Ok(ReducedLogslopeOutcome::Reduced(transform))
710}
711
712fn reparameterize_logslope_design_reduced(
719 logslope_design: &TermCollectionDesign,
720 reparam: &ReducedLogslopeReparam,
721) -> Result<TermCollectionDesign, String> {
722 let g = logslope_design
723 .design
724 .try_to_dense_arc("reparameterize_logslope_design_reduced::logslope")?;
725 let p_g = g.ncols();
726 if p_g != reparam.original_cols() {
727 return Err(format!(
728 "reduced logslope reparam width mismatch: design has {p_g} cols, transform expects {}",
729 reparam.original_cols()
730 ));
731 }
732 let t = &reparam.transform;
733 let r = reparam.reduced_cols();
734 let g_reduced = fast_ab(&g, t);
736
737 let mut new_penalties: Vec<gam_terms::smooth::BlockwisePenalty> =
740 Vec::with_capacity(logslope_design.penalties.len());
741 let mut new_nullspace_dims: Vec<usize> = Vec::with_capacity(logslope_design.penalties.len());
742 for bp in &logslope_design.penalties {
743 let mut full = Array2::<f64>::zeros((p_g, p_g));
744 full.slice_mut(s![bp.col_range.clone(), bp.col_range.clone()])
745 .assign(&bp.local);
746 let st = fast_ab(&full, t); let mut s_reduced = fast_atb(t, &st); s_reduced = (&s_reduced + &s_reduced.t()) * 0.5;
750 let (evals, _) = s_reduced
752 .eigh(Side::Lower)
753 .map_err(|e| format!("reduced logslope penalty eigendecomposition failed: {e:?}"))?;
754 let max_eval = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
755 let pen_tol = (max_eval * 1.0e-12).max(f64::EPSILON);
756 let rank = evals.iter().filter(|&&v| v.abs() > pen_tol).count();
757 let nullspace_dim = r.saturating_sub(rank);
758 new_penalties.push(gam_terms::smooth::BlockwisePenalty::new(0..r, s_reduced));
759 new_nullspace_dims.push(nullspace_dim);
760 }
761
762 let new_design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(g_reduced));
763 Ok(TermCollectionDesign {
769 design: new_design,
770 penalties: new_penalties,
771 nullspace_dims: new_nullspace_dims,
772 penaltyinfo: Vec::new(),
773 dropped_penaltyinfo: Vec::new(),
774 coefficient_lower_bounds: None,
775 linear_constraints: None,
776 intercept_range: 0..0,
777 linear_ranges: Vec::new(),
778 random_effect_ranges: Vec::new(),
779 random_effect_levels: Vec::new(),
780 smooth: gam_terms::smooth::SmoothDesign {
781 term_designs: Vec::new(),
782 penalties: Vec::new(),
783 nullspace_dims: Vec::new(),
784 penaltyinfo: Vec::new(),
785 dropped_penaltyinfo: Vec::new(),
786 terms: Vec::new(),
787 coefficient_lower_bounds: None,
788 linear_constraints: None,
789 },
790 })
791}
792
793pub(crate) fn marginal_penalties_with_influence_ridge(
816 design: &TermCollectionDesign,
817 rho_marginal: &Array1<f64>,
818 influence_columns: Option<&Array2<f64>>,
819) -> Result<(Vec<PenaltyMatrix>, Vec<usize>, Array1<f64>), String> {
820 let p_m = design.design.ncols();
821 let p1 = influence_columns.map(|z| z.ncols()).unwrap_or(0);
822 let total_dim = p_m + p1;
823 let expected_rho = design.penalties.len() + usize::from(p1 > 0);
824 if rho_marginal.len() != expected_rho {
825 return Err(format!(
826 "marginal rho width {} != smooth penalties {} + absorber slot {}",
827 rho_marginal.len(),
828 design.penalties.len(),
829 usize::from(p1 > 0),
830 ));
831 }
832 let mut penalties: Vec<PenaltyMatrix> = design
835 .penalties
836 .iter()
837 .map(|bp| bp.to_penalty_matrix(total_dim))
838 .collect();
839 let mut nullspace_dims = design.nullspace_dims.clone();
840 let log_lambdas = rho_marginal.to_vec();
841
842 if p1 > 0 {
848 penalties.push(PenaltyMatrix::Blockwise {
849 local: Array2::<f64>::eye(p1),
850 col_range: p_m..total_dim,
851 total_dim,
852 });
853 nullspace_dims.push(0);
854 }
855
856 Ok((penalties, nullspace_dims, Array1::from_vec(log_lambdas)))
857}
858
859pub(crate) fn widen_marginal_beta_hint(
862 beta_hint: Option<Array1<f64>>,
863 p_marginal_widened: usize,
864) -> Option<Array1<f64>> {
865 beta_hint.map(|hint| {
866 if hint.len() == p_marginal_widened {
867 hint
868 } else {
869 let mut widened = Array1::<f64>::zeros(p_marginal_widened);
870 let copy = hint.len().min(p_marginal_widened);
871 widened
872 .slice_mut(s![..copy])
873 .assign(&hint.slice(s![..copy]));
874 widened
875 }
876 })
877}
878
879fn marginal_fitted_eta_sup_norm(design: &TermCollectionDesign, masked_beta: &Array1<f64>) -> f64 {
889 let x = &design.design;
890 let n = x.nrows();
891 if n == 0 || x.ncols() == 0 {
892 return 0.0;
893 }
894 let mut sup = 0.0_f64;
895 for row in 0..n {
896 let eta = x.dot_row_view(row, masked_beta.view());
897 if eta.is_finite() {
898 sup = sup.max(eta.abs());
899 }
900 }
901 sup
902}
903
904fn marginal_design_beta(
907 design: &TermCollectionDesign,
908 block_beta: ArrayView1<'_, f64>,
909) -> Array1<f64> {
910 let ncols = design.design.ncols();
911 let mut masked = Array1::<f64>::zeros(ncols);
912 let copy = ncols.min(block_beta.len());
913 masked
914 .slice_mut(s![..copy])
915 .assign(&block_beta.slice(s![..copy]));
916 masked
917}
918
919fn mask_parametric_columns(
925 design: &TermCollectionDesign,
926 spec: &TermCollectionSpec,
927 full: &Array1<f64>,
928) -> Array1<f64> {
929 let ncols = design.design.ncols();
930 let mut masked = Array1::<f64>::zeros(ncols);
931 if design.intercept_range.len() == 1 {
932 let idx = design.intercept_range.start;
933 if idx < ncols {
934 masked[idx] = full[idx];
935 }
936 }
937 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
938 if linear.double_penalty {
939 continue;
940 }
941 for col in range.clone() {
942 if col < ncols {
943 masked[col] = full[col];
944 }
945 }
946 }
947 masked
948}
949
950pub(crate) fn bernoulli_marginal_slope_runaway_error_from_beta(
961 block_beta: ArrayView1<'_, f64>,
962 design: &TermCollectionDesign,
963 spec: &TermCollectionSpec,
964 inner_converged: bool,
965 eval_label: &str,
966) -> Option<String> {
967 let full_beta = marginal_design_beta(design, block_beta);
968 let parametric_beta = mask_parametric_columns(design, spec, &full_beta);
969
970 let eta_parametric = marginal_fitted_eta_sup_norm(design, ¶metric_beta);
971 let eta_full = marginal_fitted_eta_sup_norm(design, &full_beta);
972
973 let (eta_inf, explanation) = if eta_parametric >= BMS_PROBIT_SEPARATION_ETA_INF {
974 (
975 eta_parametric,
976 "an unpenalized parametric marginal direction has no stable finite probit optimum and its fitted predictor has run to the probit underflow scale",
977 )
978 } else if eta_full >= BMS_PROBIT_SEPARATION_ETA_INF {
979 (
980 eta_full,
981 "a marginal direction is trading off against the logslope surface; this is the under-constrained marginal/logslope coupling that appears when the score is correlated with the shared surface covariates",
982 )
983 } else {
984 return None;
988 };
989
990 let inner_status = if inner_converged {
991 "the inner solve reached a KKT certificate at this separation-scale predictor"
992 } else {
993 "the inner solve failed while already carrying a separation-scale predictor"
994 };
995 let beta_abs = full_beta
997 .iter()
998 .copied()
999 .filter(|v| v.is_finite())
1000 .fold(0.0_f64, |acc, v| acc.max(v.abs()));
1001
1002 Some(format!(
1003 "bernoulli marginal-slope probit marginal/logslope runaway detected in block \
1004 'marginal_surface' during {eval_label}: the fitted marginal predictor has \
1005 |η|∞={eta_inf:.3e} (numerical-degeneracy threshold \
1006 {BMS_PROBIT_SEPARATION_ETA_INF:.1}; raw |β|∞={beta_abs:.3e} is reported for \
1007 context only and does not gate this diagnostic). The joint design is \
1008 identifiable; {explanation}. {inner_status}. The robust Jeffreys curvature \
1009 path is already installed for this fit, so this diagnostic means the current \
1010 coupled surface still drives the linear predictor to the probit underflow \
1011 scale rather than a request for an external bias-reduction prior. Reduce or \
1012 reparameterize the coupled marginal/logslope surface, or use a \
1013 lower-dimensional logslope interaction. This is not a \
1014 Matérn/Duchon polynomial-nullspace or cross-block gauge-priority \
1015 failure."
1016 ))
1017}
1018
1019pub(crate) fn bernoulli_marginal_slope_runaway_error(
1020 warm_start: &CustomFamilyWarmStart,
1021 design: &TermCollectionDesign,
1022 spec: &TermCollectionSpec,
1023 inner_converged: bool,
1024 eval_label: &str,
1025) -> Option<String> {
1026 let block_beta = warm_start.block_beta_view(0)?;
1027 bernoulli_marginal_slope_runaway_error_from_beta(
1028 block_beta,
1029 design,
1030 spec,
1031 inner_converged,
1032 eval_label,
1033 )
1034}
1035
1036#[cfg(test)]
1037mod runaway_tests {
1038 use super::*;
1039 use gam_linalg::faer_ndarray::{
1040 FaerArrayView, factorize_symmetricwith_fallback, fast_xt_diag_y,
1041 };
1042 use gam_terms::smooth::{LinearCoefficientGeometry, LinearTermSpec};
1043
1044 pub(crate) fn marginal_logslope_overlap_penalty(
1050 marginal_design: &DesignMatrix,
1051 logslope_design: &DesignMatrix,
1052 z: &Array1<f64>,
1053 row_metric: &Array1<f64>,
1054 marginal_offset: &Array1<f64>,
1055 logslope_offset: &Array1<f64>,
1056 marginal_baseline: f64,
1057 logslope_baseline: f64,
1058 probit_scale: f64,
1059 ) -> Result<Option<Array2<f64>>, String> {
1060 let marginal =
1061 marginal_design.try_to_dense_arc("marginal_logslope_overlap_penalty::marginal")?;
1062 let logslope =
1063 logslope_design.try_to_dense_arc("marginal_logslope_overlap_penalty::logslope")?;
1064 let n = marginal.nrows();
1065 if logslope.nrows() != n
1066 || z.len() != n
1067 || row_metric.len() != n
1068 || marginal_offset.len() != n
1069 || logslope_offset.len() != n
1070 {
1071 return Err(format!(
1072 "marginal/logslope overlap penalty row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
1073 marginal.nrows(),
1074 logslope.nrows(),
1075 z.len(),
1076 row_metric.len(),
1077 marginal_offset.len(),
1078 logslope_offset.len(),
1079 ));
1080 }
1081 let p_m = marginal.ncols();
1082 let p_g = logslope.ncols();
1083 if p_m == 0 || p_g == 0 {
1084 return Ok(None);
1085 }
1086 if !marginal_baseline.is_finite()
1087 || !logslope_baseline.is_finite()
1088 || !probit_scale.is_finite()
1089 || probit_scale <= 0.0
1090 || z.iter().any(|v| !v.is_finite())
1091 || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
1092 || marginal_offset.iter().any(|v| !v.is_finite())
1093 || logslope_offset.iter().any(|v| !v.is_finite())
1094 {
1095 return Err(
1096 "marginal/logslope overlap penalty requires finite pilot geometry and finite non-negative row metric"
1097 .to_string(),
1098 );
1099 }
1100
1101 let mut marginal_effective = Array2::<f64>::zeros((n, p_m));
1102 let mut effective_logslope = Array2::<f64>::zeros((n, p_g));
1103 for i in 0..n {
1104 let q_i = marginal_offset[i] + marginal_baseline;
1105 let g_i = logslope_offset[i] + logslope_baseline;
1106 let sg = probit_scale * g_i;
1107 let c_i = (1.0 + sg * sg).sqrt();
1108 let logslope_factor =
1109 q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
1110 for j in 0..p_m {
1111 marginal_effective[[i, j]] = c_i * marginal[[i, j]];
1112 }
1113 for j in 0..p_g {
1114 effective_logslope[[i, j]] = logslope_factor * logslope[[i, j]];
1115 }
1116 }
1117 if effective_logslope.iter().all(|v| v.abs() <= f64::EPSILON) {
1118 return Ok(None);
1119 }
1120
1121 let mut gram = fast_xt_diag_x(&effective_logslope, row_metric);
1122 let gram_scale = gram.diag().iter().copied().fold(0.0_f64, f64::max);
1123 if !gram_scale.is_finite() || gram_scale <= 0.0 {
1124 return Ok(None);
1125 }
1126 let projection_ridge = (gram_scale * 1.0e-10).max(f64::EPSILON);
1127 for i in 0..p_g {
1128 gram[[i, i]] += projection_ridge;
1129 }
1130 let cross = fast_xt_diag_y(&effective_logslope, row_metric, &marginal_effective);
1131 let gram_view = FaerArrayView::new(&gram);
1132 let factor = factorize_symmetricwith_fallback(gram_view.as_ref(), Side::Lower)
1133 .map_err(|e| format!("marginal/logslope overlap Gram factorization failed: {e}"))?;
1134 let rhsview = FaerArrayView::new(&cross);
1135 let coeffs_mat = factor.solve(rhsview.as_ref());
1136 let coeffs = Array2::from_shape_fn((p_g, p_m), |(i, j)| coeffs_mat[(i, j)]);
1137 let projected_marginal = fast_ab(&effective_logslope, &coeffs);
1138 let mut penalty = fast_xt_diag_y(&marginal_effective, row_metric, &projected_marginal);
1139 penalty = (&penalty + &penalty.t()) * 0.5;
1140 let max_abs = penalty.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1141 if !max_abs.is_finite() || max_abs <= 1.0e-12 {
1142 return Ok(None);
1143 }
1144 Ok(Some(penalty))
1145 }
1146
1147 #[test]
1155 pub(crate) fn effective_reduction_drops_score_weighted_confound_raw_audit_misses() {
1156 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1158 let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 1.0, 2.0, 2.0, 3.0, 9.0]).unwrap();
1159 let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1160 let w = Array1::<f64>::ones(3);
1161 let zero = Array1::<f64>::zeros(3);
1162
1163 let reparam = match reduced_logslope_transform_effective(
1166 m.view(),
1167 g.view(),
1168 &z,
1169 &w,
1170 &zero,
1171 &zero,
1172 0.0,
1173 0.0,
1174 1.0,
1175 )
1176 .expect("effective reduction must succeed")
1177 {
1178 ReducedLogslopeOutcome::Reduced(t) => t,
1179 other => panic!(
1180 "effective audit must reduce the score-weighted confound (raw audit would not), got {}",
1181 match other {
1182 ReducedLogslopeOutcome::FullRank => "FullRank",
1183 ReducedLogslopeOutcome::FullyConfounded => "FullyConfounded",
1184 ReducedLogslopeOutcome::Reduced(_) => unreachable!(),
1185 }
1186 ),
1187 };
1188 assert_eq!(
1189 reparam.ncols(),
1190 1,
1191 "exactly one effective-identifiable logslope direction should survive"
1192 );
1193
1194 let g_eff = {
1198 let mut e = Array2::<f64>::zeros((3, 2));
1199 for i in 0..3 {
1200 for j in 0..2 {
1201 e[[i, j]] = z[i] * g[[i, j]];
1202 }
1203 }
1204 e
1205 };
1206 let img = g_eff.dot(&reparam.column(0));
1207 let mean = img.iter().sum::<f64>() / 3.0;
1208 let var = img.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / 3.0;
1209 assert!(
1210 var > 1.0e-6,
1211 "kept direction must be the identifiable (non-constant) effective column, var={var}"
1212 );
1213 }
1214
1215 #[test]
1222 pub(crate) fn effective_reduction_fully_confounded_single_column_is_distinct_outcome() {
1223 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1224 let g = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
1225 let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1226 let w = Array1::<f64>::ones(3);
1227 let zero = Array1::<f64>::zeros(3);
1228 let outcome = reduced_logslope_transform_effective(
1229 m.view(),
1230 g.view(),
1231 &z,
1232 &w,
1233 &zero,
1234 &zero,
1235 0.0,
1236 0.0,
1237 1.0,
1238 )
1239 .expect("effective reduction must succeed");
1240 assert!(
1241 matches!(outcome, ReducedLogslopeOutcome::FullyConfounded),
1242 "fully effective-confounded logslope must surface the distinct FullyConfounded outcome"
1243 );
1244 }
1245
1246 #[test]
1249 pub(crate) fn effective_reduction_no_confound_returns_none() {
1250 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1251 let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]).unwrap();
1253 let z = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1254 let w = Array1::<f64>::ones(3);
1255 let zero = Array1::<f64>::zeros(3);
1256 let outcome = reduced_logslope_transform_effective(
1257 m.view(),
1258 g.view(),
1259 &z,
1260 &w,
1261 &zero,
1262 &zero,
1263 0.0,
1264 0.0,
1265 1.0,
1266 )
1267 .expect("effective reduction must succeed");
1268 assert!(
1269 matches!(outcome, ReducedLogslopeOutcome::FullRank),
1270 "no effective confound ⇒ FullRank (raw design kept unchanged)"
1271 );
1272 }
1273
1274 #[test]
1275 pub(crate) fn spatial_joint_setup_counts_only_learned_penalties_in_rho() {
1276 let data = Array2::<f64>::zeros((3, 1));
1277 let empty_terms = TermCollectionSpec {
1278 linear_terms: Vec::new(),
1279 random_effect_terms: Vec::new(),
1280 smooth_terms: Vec::new(),
1281 };
1282 let setup = joint_setup(
1283 data.view(),
1284 &empty_terms,
1285 &empty_terms,
1286 2,
1287 3,
1288 Some(2.5),
1289 &[0.4],
1290 &SpatialLengthScaleOptimizationOptions::default(),
1291 );
1292
1293 assert_eq!(
1294 setup.rho_dim(),
1295 6,
1296 "BMS spatial setup rho holds every learned marginal/logslope/auxiliary penalty; the #461 absorber ridge occupies the trailing marginal slot"
1297 );
1298 assert_eq!(
1299 setup.theta0()[1], 2.5,
1300 "absorber ridge seeds the trailing marginal rho coordinate at the ln(n) leakage scale"
1301 );
1302 }
1303
1304 #[test]
1305 pub(crate) fn overlap_penalty_targets_score_weighted_logslope_span() {
1306 let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1307 Array2::from_shape_vec((4, 1), vec![0.0, 1.0, 2.0, 3.0]).unwrap(),
1308 ));
1309 let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1310 Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1311 ));
1312 let z = Array1::from_vec(vec![0.0, 1.0, 2.0, 3.0]);
1313 let row_metric = Array1::ones(4);
1314 let offsets = Array1::zeros(4);
1315
1316 let penalty = marginal_logslope_overlap_penalty(
1317 &marginal,
1318 &logslope,
1319 &z,
1320 &row_metric,
1321 &offsets,
1322 &offsets,
1323 0.0,
1324 0.0,
1325 1.0,
1326 )
1327 .expect("overlap penalty should build")
1328 .expect("marginal signal lies in the pilot logslope Jacobian span");
1329
1330 assert_eq!(penalty.dim(), (1, 1));
1331 assert!((penalty[[0, 0]] - 14.0).abs() < 1.0e-6);
1332 }
1333
1334 #[test]
1335 pub(crate) fn overlap_penalty_skips_weight_orthogonal_channels() {
1336 let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1337 Array2::from_shape_vec((4, 1), vec![-1.0, 1.0, -1.0, 1.0]).unwrap(),
1338 ));
1339 let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1340 Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1341 ));
1342 let z = Array1::ones(4);
1343 let row_metric = Array1::ones(4);
1344 let offsets = Array1::zeros(4);
1345
1346 let penalty = marginal_logslope_overlap_penalty(
1347 &marginal,
1348 &logslope,
1349 &z,
1350 &row_metric,
1351 &offsets,
1352 &offsets,
1353 0.0,
1354 0.0,
1355 1.0,
1356 )
1357 .expect("overlap penalty should build");
1358
1359 assert!(penalty.is_none());
1360 }
1361
1362 fn dense_marginal_design(
1370 x: Array2<f64>,
1371 intercept_range: std::ops::Range<usize>,
1372 linear_ranges: Vec<(String, std::ops::Range<usize>)>,
1373 ) -> TermCollectionDesign {
1374 TermCollectionDesign {
1375 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
1376 penalties: Vec::new(),
1377 nullspace_dims: Vec::new(),
1378 penaltyinfo: Vec::new(),
1379 dropped_penaltyinfo: Vec::new(),
1380 coefficient_lower_bounds: None,
1381 linear_constraints: None,
1382 intercept_range,
1383 linear_ranges,
1384 random_effect_ranges: Vec::new(),
1385 random_effect_levels: Vec::new(),
1386 smooth: gam_terms::smooth::SmoothDesign {
1387 term_designs: Vec::new(),
1388 penalties: Vec::new(),
1389 nullspace_dims: Vec::new(),
1390 penaltyinfo: Vec::new(),
1391 dropped_penaltyinfo: Vec::new(),
1392 terms: Vec::new(),
1393 coefficient_lower_bounds: None,
1394 linear_constraints: None,
1395 },
1396 }
1397 }
1398
1399 fn linear_term(name: &str, feature_col: usize) -> LinearTermSpec {
1400 LinearTermSpec {
1401 name: name.to_string(),
1402 feature_col,
1403 feature_cols: vec![feature_col],
1404 categorical_levels: vec![],
1405 double_penalty: false,
1406 coefficient_geometry: LinearCoefficientGeometry::default(),
1407 coefficient_min: None,
1408 coefficient_max: None,
1409 }
1410 }
1411
1412 fn empty_spec() -> TermCollectionSpec {
1413 TermCollectionSpec {
1414 linear_terms: Vec::new(),
1415 random_effect_terms: Vec::new(),
1416 smooth_terms: Vec::new(),
1417 }
1418 }
1419
1420 #[test]
1427 pub(crate) fn runaway_guard_silent_when_huge_beta_cancels_to_bounded_eta() {
1428 let x = Array2::<f64>::from_shape_vec((4, 2), vec![1.0; 8]).unwrap();
1430 let design = dense_marginal_design(x, 0..0, Vec::new());
1431 let beta = Array1::from_vec(vec![60.0, -60.0]);
1432
1433 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1434 beta.view(),
1435 &design,
1436 &empty_spec(),
1437 true,
1438 "regression-fixture",
1439 );
1440 assert!(
1441 msg.is_none(),
1442 "huge cancelling β with bounded fitted η must NOT trip the runaway guard; got {msg:?}"
1443 );
1444 }
1445
1446 #[test]
1450 pub(crate) fn runaway_guard_fires_when_fitted_eta_exceeds_threshold() {
1451 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1452 let design = dense_marginal_design(x, 0..0, Vec::new());
1453 let beta = Array1::from_vec(vec![40.0]);
1454
1455 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1456 beta.view(),
1457 &design,
1458 &empty_spec(),
1459 true,
1460 "separation-fixture",
1461 )
1462 .expect("fitted |η|∞=40 ≥ 35 must trip the runaway guard");
1463
1464 assert!(msg.contains("marginal/logslope runaway"));
1465 assert!(msg.contains("|η|∞"));
1466 assert!(msg.contains("4.000e1"));
1467 assert!(msg.contains("score is correlated with the shared surface covariates"));
1468 assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1469 assert!(msg.contains("KKT certificate"));
1470 }
1471
1472 #[test]
1476 pub(crate) fn runaway_guard_names_unpenalized_parametric_direction_via_fitted_eta() {
1477 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1478 let design = dense_marginal_design(x, 0..0, vec![("sex".to_string(), 0..1)]);
1479 let mut spec = empty_spec();
1480 spec.linear_terms.push(linear_term("sex", 0));
1481 let beta = Array1::from_vec(vec![41.0]);
1482
1483 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1484 beta.view(),
1485 &design,
1486 &spec,
1487 true,
1488 "parametric-fixture",
1489 )
1490 .expect("parametric fitted |η|∞=41 ≥ 35 must trip the runaway guard");
1491
1492 assert!(msg.contains("unpenalized parametric marginal direction"));
1493 assert!(msg.contains("|η|∞"));
1494 assert!(msg.contains("robust Jeffreys curvature path is already installed"));
1495 assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1496 }
1497
1498 #[test]
1502 pub(crate) fn runaway_guard_silent_for_nonconverged_but_bounded_eta() {
1503 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1504 let design = dense_marginal_design(x, 0..0, Vec::new());
1505 let beta = Array1::from_vec(vec![5.0]);
1506
1507 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1508 beta.view(),
1509 &design,
1510 &empty_spec(),
1511 false,
1512 "nonconverged-fixture",
1513 );
1514 assert!(
1515 msg.is_none(),
1516 "bounded fitted η must not raise the separation error even when the inner solve did not converge; got {msg:?}"
1517 );
1518 }
1519
1520 #[test]
1523 pub(crate) fn runaway_guard_fires_for_nonconverged_separating_eta() {
1524 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1525 let design = dense_marginal_design(x, 0..0, Vec::new());
1526 let beta = Array1::from_vec(vec![50.0]);
1527
1528 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1529 beta.view(),
1530 &design,
1531 &empty_spec(),
1532 false,
1533 "nonconverged-separating-fixture",
1534 )
1535 .expect("separating |η|∞ at non-convergence must still trip the guard");
1536
1537 assert!(msg.contains(
1538 "the inner solve failed while already carrying a separation-scale predictor"
1539 ));
1540 }
1541
1542 #[test]
1555 pub(crate) fn bms_block_jacobians_self_compute_at_audit_empty_beta_nonzero_logslope_baseline() {
1556 use std::sync::Arc;
1557 let n = 4usize;
1558 let marginal =
1559 Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1560 let logslope =
1561 Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1562 let offset_m = Array1::<f64>::zeros(n);
1563 let g_baseline = 0.3_f64;
1566 let offset_s = Array1::<f64>::from_elem(n, g_baseline);
1567 let z = Arc::new(Array1::from_vec(vec![-0.7, 0.2, 0.9, 1.4]));
1568 let s = 1.0_f64;
1569
1570 let beta: Vec<f64> = Vec::new();
1572 let state = FamilyLinearizationState {
1573 beta: &beta,
1574 family_scalars: None,
1575 channel_hessian: None,
1576 probit_frailty_scale: s,
1577 };
1578
1579 let marginal_jac = BmsMarginalJacobian::new(
1580 Arc::clone(&marginal),
1581 Arc::clone(&logslope),
1582 offset_m.clone(),
1583 offset_s.clone(),
1584 1,
1585 );
1586 let j_m = marginal_jac
1587 .effective_jacobian_rows(&state, 0..n)
1588 .expect("BMS marginal Jacobian must self-compute at audit empty β (gam#370)");
1589 let c_expected = (1.0 + (s * g_baseline).powi(2)).sqrt();
1591 assert_eq!(j_m.dim(), (n, 1));
1592 for i in 0..n {
1593 assert!(
1594 (j_m[[i, 0]] - c_expected).abs() < 1e-12,
1595 "marginal J[{i}] = {} != closed-form c_i = {c_expected}",
1596 j_m[[i, 0]]
1597 );
1598 }
1599
1600 let logslope_jac = BmsLogslopeJacobian::new(
1601 Arc::clone(&marginal),
1602 Arc::clone(&logslope),
1603 offset_m,
1604 offset_s,
1605 Arc::clone(&z),
1606 1,
1607 );
1608 let j_s = logslope_jac
1609 .effective_jacobian_rows(&state, 0..n)
1610 .expect("BMS logslope Jacobian must self-compute at audit empty β (gam#370)");
1611 assert_eq!(j_s.dim(), (n, 1));
1614 for i in 0..n {
1615 let expected = s * z[i];
1616 assert!(
1617 (j_s[[i, 0]] - expected).abs() < 1e-12,
1618 "logslope J[{i}] = {} != closed-form factor {expected}",
1619 j_s[[i, 0]]
1620 );
1621 assert!(j_s[[i, 0]].is_finite());
1622 }
1623 }
1624}
1625
1626pub(crate) fn build_marginal_blockspec_bms(
1627 design: &TermCollectionDesign,
1628 baseline: f64,
1629 offset: &Array1<f64>,
1630 rho: Array1<f64>,
1631 beta_hint: Option<Array1<f64>>,
1632 logslope_design: &TermCollectionDesign,
1633 logslope_offset: &Array1<f64>,
1634 logslope_baseline: f64,
1635 p_marginal: usize,
1636 influence_columns: Option<&Array2<f64>>,
1637) -> Result<ParameterBlockSpec, String> {
1638 let offset_m = offset + baseline;
1639 let offset_s = logslope_offset + logslope_baseline;
1640 let raw_marginal_dense = design
1641 .design
1642 .try_to_dense_arc("build_marginal_blockspec_bms::marginal")?;
1643 let marginal_dense =
1644 widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1645 let logslope_dense = logslope_design
1646 .design
1647 .try_to_dense_arc("build_marginal_blockspec_bms::logslope")?;
1648 let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsMarginalJacobian {
1649 marginal_dense: Arc::clone(&marginal_dense),
1650 logslope_dense,
1651 offset_m: offset_m.clone(),
1652 offset_s,
1653 p_marginal,
1654 });
1655 let (penalties, nullspace_dims, initial_log_lambdas) =
1656 marginal_penalties_with_influence_ridge(design, &rho, influence_columns)?;
1657 Ok(ParameterBlockSpec {
1658 name: "marginal_surface".to_string(),
1659 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1660 (*marginal_dense).clone(),
1661 )),
1662 offset: offset_m,
1663 penalties,
1664 nullspace_dims,
1665 initial_log_lambdas,
1666 initial_beta: widen_marginal_beta_hint(beta_hint, p_marginal),
1667 gauge_priority: GAUGE_PRIORITY_MARGINAL,
1680 jacobian_callback: Some(callback),
1681 stacked_design: None,
1682 stacked_offset: None,
1683 })
1684}
1685
1686pub(crate) fn build_logslope_blockspec_bms(
1687 design: &TermCollectionDesign,
1688 baseline: f64,
1689 offset: &Array1<f64>,
1690 rho: Array1<f64>,
1691 beta_hint: Option<Array1<f64>>,
1692 marginal_design: &TermCollectionDesign,
1693 marginal_offset: &Array1<f64>,
1694 marginal_baseline: f64,
1695 z: Arc<Array1<f64>>,
1696 p_marginal: usize,
1697 influence_columns: Option<&Array2<f64>>,
1698) -> Result<ParameterBlockSpec, String> {
1699 let offset_s = offset + baseline;
1700 let offset_m = marginal_offset + marginal_baseline;
1701 let raw_marginal_dense = marginal_design
1702 .design
1703 .try_to_dense_arc("build_logslope_blockspec_bms::marginal")?;
1704 let marginal_dense =
1709 widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1710 let logslope_dense = design
1711 .design
1712 .try_to_dense_arc("build_logslope_blockspec_bms::logslope")?;
1713 let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsLogslopeJacobian {
1714 marginal_dense,
1715 logslope_dense: Arc::clone(&logslope_dense),
1716 offset_m,
1717 offset_s: offset_s.clone(),
1718 z,
1719 p_marginal,
1720 });
1721 Ok(ParameterBlockSpec {
1722 name: "logslope_surface".to_string(),
1723 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1724 (*logslope_dense).clone(),
1725 )),
1726 offset: offset_s,
1727 penalties: design.penalties_as_penalty_matrix(),
1728 nullspace_dims: design.nullspace_dims.clone(),
1729 initial_log_lambdas: rho,
1730 initial_beta: beta_hint,
1731 gauge_priority: GAUGE_PRIORITY_LOGSLOPE,
1739 jacobian_callback: Some(callback),
1740 stacked_design: None,
1741 stacked_offset: None,
1742 })
1743}
1744
1745pub(crate) fn build_deviation_aux_blockspec(
1746 name: &str,
1747 prepared: &DeviationPrepared,
1748 rho: Array1<f64>,
1749 beta_hint: Option<Array1<f64>>,
1750) -> Result<ParameterBlockSpec, String> {
1751 let mut block = prepared.block.clone();
1752 block.initial_log_lambdas = Some(rho);
1753 let candidate_beta = beta_hint.or_else(|| Some(Array1::<f64>::zeros(block.design.ncols())));
1754 block.initial_beta = candidate_beta
1755 .map(|beta| {
1756 let zero = Array1::<f64>::zeros(beta.len());
1757 project_monotone_feasible_beta(&prepared.runtime, &zero, &beta, name)
1758 })
1759 .transpose()?;
1760 let mut spec = block.intospec(name)?;
1761 spec.gauge_priority = match name {
1771 "link_dev" => GAUGE_PRIORITY_LINK_DEV,
1772 "score_warp_dev" => GAUGE_PRIORITY_SCORE_WARP_DEV,
1779 _ => GAUGE_PRIORITY_DEVIATION_DEFAULT,
1780 };
1781 Ok(spec)
1782}
1783
1784pub(crate) fn push_deviation_aux_blockspecs(
1785 blocks: &mut Vec<ParameterBlockSpec>,
1786 rho: &Array1<f64>,
1787 cursor: &mut usize,
1788 score_warp_prepared: Option<&DeviationPrepared>,
1789 link_dev_prepared: Option<&DeviationPrepared>,
1790 score_warp_beta_hint: Option<Array1<f64>>,
1791 link_dev_beta_hint: Option<Array1<f64>>,
1792) -> Result<(), String> {
1793 if let Some(prepared) = score_warp_prepared {
1794 let rho_h = rho
1795 .slice(s![*cursor..*cursor + prepared.block.penalties.len()])
1796 .to_owned();
1797 *cursor += prepared.block.penalties.len();
1798 blocks.push(build_deviation_aux_blockspec(
1799 "score_warp_dev",
1800 prepared,
1801 rho_h,
1802 score_warp_beta_hint,
1803 )?);
1804 }
1805 if let Some(prepared) = link_dev_prepared {
1806 let rho_w = rho
1807 .slice(s![*cursor..*cursor + prepared.block.penalties.len()])
1808 .to_owned();
1809 blocks.push(build_deviation_aux_blockspec(
1810 "link_dev",
1811 prepared,
1812 rho_w,
1813 link_dev_beta_hint,
1814 )?);
1815 }
1816 Ok(())
1817}
1818
1819fn inner_fit(
1820 family: &BernoulliMarginalSlopeFamily,
1821 blocks: &[ParameterBlockSpec],
1822 options: &BlockwiseFitOptions,
1823) -> Result<UnifiedFitResult, String> {
1824 let mut options = options.clone();
1825 options.use_outer_hessian = false;
1830 options.outer_tol = options.outer_tol.max(2.0e-5);
1831 fit_custom_family(family, blocks, &options).map_err(|e| e.to_string())
1832}
1833
1834pub fn fit_bernoulli_marginal_slope_terms(
1835 data: ArrayView2<'_, f64>,
1836 spec: BernoulliMarginalSlopeTermSpec,
1837 options: &BlockwiseFitOptions,
1838 kappa_options: &SpatialLengthScaleOptimizationOptions,
1839 policy: &gam_runtime::resource::ResourcePolicy,
1840) -> Result<BernoulliMarginalSlopeFitResult, String> {
1841 let mut spec = spec;
1842 let data_view = data;
1843 validate_spec(data_view, &spec)?;
1844 let mjs_frozen_marginal =
1853 gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.marginalspec);
1854 let mjs_frozen_logslope =
1855 gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.logslopespec);
1856 if mjs_frozen_marginal + mjs_frozen_logslope > 0 {
1857 log::info!(
1858 "[BMS spatial] froze measure-jet length-scale learning on {} marginal + {} log-slope \
1859 term(s): the coupled surface keeps ℓ at its conditioned auto value (#1116)",
1860 mjs_frozen_marginal,
1861 mjs_frozen_logslope
1862 );
1863 }
1864 let mut effective_kappa_options = kappa_options.clone();
1865 let kappa_locked_marginal =
1875 gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.marginalspec);
1876 let kappa_locked_logslope =
1877 gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.logslopespec);
1878 if effective_kappa_options.enabled && kappa_locked_marginal && kappa_locked_logslope {
1879 log::info!(
1880 "[BMS spatial] disabling κ/ψ optimization: every spatial term has an \
1881 explicit length_scale and no anisotropy; user-supplied kernel scale is fixed"
1882 );
1883 effective_kappa_options.enabled = false;
1884 }
1885 let flex_spatial_pilot_path = (spec.score_warp.is_some() || spec.link_dev.is_some())
1886 && spec.y.len() >= BMS_FLEX_SPATIAL_OUTER_PILOT_ROW_THRESHOLD
1887 && effective_kappa_options.enabled;
1888 if flex_spatial_pilot_path {
1889 let marginal_terms = spatial_length_scale_term_indices(&spec.marginalspec);
1890 let logslope_terms = spatial_length_scale_term_indices(&spec.logslopespec);
1891 let marginal_updates = apply_spatial_anisotropy_pilot_initializer(
1892 data_view,
1893 &mut spec.marginalspec,
1894 &marginal_terms,
1895 effective_kappa_options.pilot_subsample_threshold,
1896 &effective_kappa_options,
1897 );
1898 let logslope_updates = apply_spatial_anisotropy_pilot_initializer(
1899 data_view,
1900 &mut spec.logslopespec,
1901 &logslope_terms,
1902 effective_kappa_options.pilot_subsample_threshold,
1903 &effective_kappa_options,
1904 );
1905 effective_kappa_options.enabled = false;
1906 log::info!(
1907 "[BMS spatial] n={} flex=true pilot_geometry_updates={} iterative_spatial_outer=false reason=large-flex-spatial-pilot",
1908 spec.y.len(),
1909 marginal_updates + logslope_updates,
1910 );
1911 }
1912 let (z_standardized, z_normalization) = standardize_latent_z_with_policy(
1913 &spec.z,
1914 &spec.weights,
1915 "bernoulli-marginal-slope",
1916 &spec.latent_z_policy,
1917 )?;
1918 spec.z = z_standardized;
1919 let sigma_learnable = matches!(
1920 &spec.frailty,
1921 FrailtySpec::GaussianShift { sigma_fixed: None }
1922 );
1923 let initial_sigma = match &spec.frailty {
1924 FrailtySpec::GaussianShift {
1925 sigma_fixed: Some(s),
1926 } => Some(*s),
1927 FrailtySpec::GaussianShift { sigma_fixed: None } => Some(0.5),
1928 FrailtySpec::None => None,
1929 FrailtySpec::HazardMultiplier { .. } => {
1930 return Err(
1931 "internal: validate_spec should have rejected unsupported marginal-slope frailty"
1932 .to_string(),
1933 );
1934 }
1935 };
1936 let probit_scale = probit_frailty_scale(initial_sigma);
1937 let (_raw_joint_designs, mut joint_specs) = build_term_collection_designs_and_freeze_joint(
1938 data_view,
1939 &[spec.marginalspec.clone(), spec.logslopespec.clone()],
1940 )
1941 .map_err(|e| e.to_string())?;
1942 let marginalspec_boot = joint_specs.remove(0);
1943 let logslopespec_boot = joint_specs.remove(0);
1944 let (mut joint_designs, _) = build_term_collection_designs_and_freeze_joint(
1961 data_view,
1962 &[marginalspec_boot.clone(), logslopespec_boot.clone()],
1963 )
1964 .map_err(|e| format!("failed to rebuild frozen probe BMS joint designs: {e}"))?;
1965 let marginal_design = joint_designs.remove(0);
1966 let logslope_design = joint_designs.remove(0);
1967 let absorber_active = spec
1974 .score_influence_jacobian
1975 .as_ref()
1976 .is_some_and(|j| j.ncols() > 0);
1977 let conditioning_dense = if absorber_active {
1978 None
1979 } else {
1980 Some(
1981 marginal_design
1982 .design
1983 .try_to_dense_arc("bernoulli marginal-slope conditional latent-z gate")?,
1984 )
1985 };
1986 let (latent_measure, latent_z_calibration) = build_latent_measure_with_geometry(
1987 &spec.z,
1988 &spec.weights,
1989 &spec.latent_z_policy,
1990 conditioning_dense.as_ref().map(|d| d.view()),
1991 )?;
1992 if latent_measure.is_empirical() && sigma_learnable {
1993 return Err("empirical latent-measure marginal-slope calibration requires fixed GaussianShift sigma; learnable sigma derivatives must be fit under the standard-normal latent measure"
1994 .to_string());
1995 }
1996
1997 let y = Arc::new(spec.y.clone());
1998 let weights = Arc::new(spec.weights.clone());
1999 let z = match &latent_z_calibration {
2004 LatentMeasureCalibration::None => Arc::new(spec.z.clone()),
2005 LatentMeasureCalibration::RankInverseNormal(cal) => {
2006 Arc::new(cal.apply_to_training(&spec.z)?)
2007 }
2008 LatentMeasureCalibration::ConditionalLocationScale(cal) => {
2009 let a_block = conditioning_dense.as_ref().ok_or_else(|| {
2012 "conditional latent calibration requires the marginal conditioning block"
2013 .to_string()
2014 })?;
2015 Arc::new(cal.apply(spec.z.view(), a_block.view())?)
2016 }
2017 };
2018 let z_train = z.as_ref();
2019 let pilot_baseline = pooled_probit_baseline(&spec.y, z_train, &spec.weights)?;
2020 let baseline = (
2021 bernoulli_marginal_slope_eta_from_probability(
2022 &spec.base_link,
2023 normal_cdf(pilot_baseline.0),
2024 "bernoulli marginal-slope baseline link inversion",
2025 )?,
2026 pilot_baseline.1 / probit_scale,
2027 );
2028
2029 let rigid_pilot_eta = rigid_pooled_probit_pilot_eta(
2072 &spec.base_link,
2073 z_train,
2074 &spec.marginal_offset,
2075 &spec.logslope_offset,
2076 baseline.0,
2077 baseline.1,
2078 probit_scale,
2079 )?;
2080 let cross_block_pilot_w_score_warp =
2081 pilot_irls_hessian_row_metric_at_eta(&rigid_pilot_eta, &spec.weights);
2082
2083 let influence_columns = if let Some(jac) = spec
2095 .score_influence_jacobian
2096 .as_ref()
2097 .filter(|j| j.ncols() > 0)
2098 {
2099 let protected_design = DesignMatrix::hstack(vec![
2100 marginal_design.design.clone(),
2101 logslope_design.design.clone(),
2102 ])
2103 .map_err(|e| {
2104 format!(
2105 "bernoulli marginal-slope influence-block protected projection stack failed to concatenate marginal + logslope design: {e}"
2106 )
2107 })?;
2108 let protected_dense_for_proj = protected_design
2109 .try_to_dense_arc("bernoulli marginal-slope influence-block protected projection")?;
2110 let protected_dense = protected_dense_for_proj.as_ref();
2111 if jac.nrows() != protected_dense.nrows() {
2112 return Err(format!(
2113 "influence block: Jacobian has {} rows, protected design has {}",
2114 jac.nrows(),
2115 protected_dense.nrows()
2116 ));
2117 }
2118 let rigid_logslope_at_rows = &spec.logslope_offset + baseline.1;
2129 let residualized = crate::marginal_slope_orthogonal::residualized_influence_block(
2130 jac,
2131 z_train,
2132 &rigid_logslope_at_rows,
2133 probit_scale,
2134 protected_dense.view(),
2135 &cross_block_pilot_w_score_warp,
2136 )?;
2137 Some(residualized)
2138 } else {
2139 None
2140 };
2141 let mut cross_block_warnings: Vec<CrossBlockIdentifiabilityWarning> = Vec::new();
2142 let score_warp_prepared = if let Some(cfg) = spec.score_warp.as_ref() {
2143 use super::deviation_runtime::ParametricAnchorBlock;
2144 let mut prepared = build_score_warp_deviation_block_from_seed(z_train, cfg)?;
2145 let outcome = install_compiled_flex_block_into_runtime(
2150 &mut prepared,
2151 z_train,
2152 cfg,
2153 &[
2154 (&marginal_design.design, ParametricAnchorBlock::Marginal),
2155 (&logslope_design.design, ParametricAnchorBlock::Logslope),
2156 ],
2157 &[],
2158 &cross_block_pilot_w_score_warp,
2159 )?;
2160 match outcome {
2161 FlexCompileOutcome::Reparameterised => Some(prepared),
2162 FlexCompileOutcome::FullyAliased { reason } => {
2163 cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2169 candidate_label: "score_warp",
2170 anchor_summary: "marginal+logslope".to_string(),
2171 reason,
2172 });
2173 Some(prepared)
2174 }
2175 }
2176 } else {
2177 None
2178 };
2179 let link_dev_prepared = if let Some(cfg) = spec.link_dev.as_ref() {
2205 let eta_pilot = pilot_eta_for_link_dev_orthogonalisation(
2206 &spec.base_link,
2207 &spec.y,
2208 z_train,
2209 &spec.weights,
2210 &marginal_design.design,
2211 &spec.marginal_offset,
2212 &spec.logslope_offset,
2213 baseline.0,
2214 baseline.1,
2215 probit_scale,
2216 )?;
2217 let link_dev_seed = padded_deviation_seed(&eta_pilot, 1.0, 0.5);
2218 let mut prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
2219 &link_dev_seed,
2220 &eta_pilot,
2221 cfg,
2222 )?;
2223 let score_warp_anchor_design = score_warp_prepared
2260 .as_ref()
2261 .map(|sw| sw.runtime.design_at_training_with_residual(z_train))
2262 .transpose()?;
2263 use super::deviation_runtime::ParametricAnchorBlock;
2264 let parametric_anchors: [(&DesignMatrix, ParametricAnchorBlock); 2] = [
2265 (&marginal_design.design, ParametricAnchorBlock::Marginal),
2266 (&logslope_design.design, ParametricAnchorBlock::Logslope),
2267 ];
2268 let flex_anchor_slot: Option<&Array2<f64>> = score_warp_anchor_design.as_ref();
2269 let flex_anchors: Vec<&Array2<f64>> = flex_anchor_slot.into_iter().collect();
2270 let cross_block_pilot_w_link_dev =
2275 pilot_irls_hessian_row_metric_at_eta(&eta_pilot, &spec.weights);
2276 let outcome = install_compiled_flex_block_into_runtime(
2277 &mut prepared,
2278 &eta_pilot,
2279 cfg,
2280 ¶metric_anchors,
2281 &flex_anchors,
2282 &cross_block_pilot_w_link_dev,
2283 )?;
2284 match outcome {
2285 FlexCompileOutcome::Reparameterised => Some(prepared),
2286 FlexCompileOutcome::FullyAliased { reason } => {
2287 cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2293 candidate_label: "link_deviation",
2294 anchor_summary: "marginal+logslope+score_warp".to_string(),
2295 reason,
2296 });
2297 Some(prepared)
2298 }
2299 }
2300 } else {
2301 None
2302 };
2303 let extra_rho0 = {
2304 let mut out = Vec::new();
2305 if let Some(ref prepared) = score_warp_prepared {
2306 out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2307 }
2308 if let Some(ref prepared) = link_dev_prepared {
2309 out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2310 }
2311 out
2312 };
2313 let logslope_reduced_reparam: Option<ReducedLogslopeReparam> = build_reduced_logslope_reparam(
2326 &marginal_design,
2327 &logslope_design,
2328 z.as_ref(),
2329 &cross_block_pilot_w_score_warp,
2330 &spec.marginal_offset,
2331 &spec.logslope_offset,
2332 baseline.0,
2333 baseline.1,
2334 probit_scale,
2335 )?;
2336 let reduce_logslope_design =
2342 |logslope_design: &TermCollectionDesign| -> Result<TermCollectionDesign, String> {
2343 match logslope_reduced_reparam.as_ref() {
2344 Some(reparam) => reparameterize_logslope_design_reduced(logslope_design, reparam),
2345 None => Ok(logslope_design.clone()),
2346 }
2347 };
2348
2349 let absorber_slots = usize::from(influence_columns.is_some());
2353 let absorber_rho0 = influence_columns
2354 .as_ref()
2355 .map(|_| influence_absorber_log_lambda(spec.z.len()).clamp(-12.0, 12.0));
2356 let marginal_penalty_count = marginal_design.penalties.len() + absorber_slots;
2357 let setup = joint_setup(
2358 data_view,
2359 &marginalspec_boot,
2360 &logslopespec_boot,
2361 marginal_penalty_count,
2362 logslope_design.penalties.len(),
2363 absorber_rho0,
2364 &extra_rho0,
2365 &effective_kappa_options,
2366 );
2367 let setup = if sigma_learnable {
2368 setup.with_auxiliary(
2369 Array1::from_vec(vec![initial_sigma.expect("learnable sigma seed").ln()]),
2370 Array1::from_vec(vec![0.01_f64.ln()]),
2371 Array1::from_vec(vec![5.0_f64.ln()]),
2372 )
2373 } else {
2374 setup
2375 };
2376 let final_sigma_cell = std::cell::Cell::new(initial_sigma);
2377 let exact_warm_start = RefCell::new(None::<CustomFamilyWarmStart>);
2378 let runaway_error = RefCell::new(None::<String>);
2379 let pending_beta_seed = RefCell::new(None::<Array1<f64>>);
2386 let hints = RefCell::new(ThetaHints::default());
2387 let score_warp_runtime = score_warp_prepared.as_ref().map(|p| p.runtime.clone());
2388 let link_dev_runtime = link_dev_prepared.as_ref().map(|p| p.runtime.clone());
2389
2390 let build_blocks = |rho: &Array1<f64>,
2391 marginal_design: &TermCollectionDesign,
2392 logslope_design: &TermCollectionDesign|
2393 -> Result<Vec<ParameterBlockSpec>, String> {
2394 let hints = hints.borrow();
2395 let mut cursor = 0usize;
2396 let logslope_design_reduced = reduce_logslope_design(logslope_design)?;
2403 let logslope_design = &logslope_design_reduced;
2404 let marginal_rho_len = marginal_design.penalties.len() + absorber_slots;
2409 let rho_marginal = rho.slice(s![cursor..cursor + marginal_rho_len]).to_owned();
2410 cursor += marginal_rho_len;
2411 let rho_logslope = rho
2412 .slice(s![cursor..cursor + logslope_design.penalties.len()])
2413 .to_owned();
2414 cursor += logslope_design.penalties.len();
2415 let p_m = marginal_design.design.ncols()
2416 + influence_columns.as_ref().map(|z| z.ncols()).unwrap_or(0);
2417 let mut blocks = vec![
2418 build_marginal_blockspec_bms(
2419 marginal_design,
2420 baseline.0,
2421 &spec.marginal_offset,
2422 rho_marginal,
2423 hints.marginal_beta.clone(),
2424 logslope_design,
2425 &spec.logslope_offset,
2426 baseline.1,
2427 p_m,
2428 influence_columns.as_ref(),
2429 )?,
2430 build_logslope_blockspec_bms(
2431 logslope_design,
2432 baseline.1,
2433 &spec.logslope_offset,
2434 rho_logslope,
2435 hints.logslope_beta.clone(),
2436 marginal_design,
2437 &spec.marginal_offset,
2438 baseline.0,
2439 Arc::clone(&z),
2440 p_m,
2441 influence_columns.as_ref(),
2442 )?,
2443 ];
2444 push_deviation_aux_blockspecs(
2445 &mut blocks,
2446 rho,
2447 &mut cursor,
2448 score_warp_prepared.as_ref(),
2449 link_dev_prepared.as_ref(),
2450 hints.score_warp_beta.clone(),
2451 hints.link_dev_beta.clone(),
2452 )?;
2453 Ok(blocks)
2454 };
2455
2456 let intercept_warm_starts = new_intercept_warm_start_cache(y.len());
2457 let cell_moment_lru = new_cell_moment_lru_cache(policy);
2458 let cell_moment_cache_stats = new_cell_moment_cache_stats();
2459 let make_family = |marginal_design: &TermCollectionDesign,
2460 logslope_design: &TermCollectionDesign,
2461 sigma: Option<f64>|
2462 -> BernoulliMarginalSlopeFamily {
2463 let kernel_marginal_design = match influence_columns.as_ref() {
2469 Some(z_infl) => {
2470 let raw = marginal_design
2471 .design
2472 .try_to_dense_arc("make_family::widened-marginal")
2473 .expect("dense marginal design for influence widening");
2474 let widened = widen_marginal_dense_with_influence(&raw, Some(z_infl))
2475 .expect("widen marginal design with influence columns");
2476 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
2477 (*widened).clone(),
2478 ))
2479 }
2480 None => marginal_design.design.clone(),
2481 };
2482 let kernel_logslope_design = reduce_logslope_design(logslope_design)
2488 .expect("reduce logslope design for family construction")
2489 .design;
2490 BernoulliMarginalSlopeFamily {
2491 y: Arc::clone(&y),
2492 weights: Arc::clone(&weights),
2493 z: Arc::clone(&z),
2494 latent_measure: latent_measure.clone(),
2495 gaussian_frailty_sd: sigma,
2496 base_link: spec.base_link.clone(),
2497 marginal_design: kernel_marginal_design,
2498 logslope_design: kernel_logslope_design,
2499 score_warp: score_warp_runtime.clone(),
2500 link_dev: link_dev_runtime.clone(),
2501 policy: policy.clone(),
2502 cell_moment_lru: Arc::clone(&cell_moment_lru),
2503 cell_moment_cache_stats: Arc::clone(&cell_moment_cache_stats),
2504 intercept_warm_starts: Some(Arc::clone(&intercept_warm_starts)),
2505 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2506 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
2507 }
2508 };
2509
2510 let marginal_terms = spatial_length_scale_term_indices(&marginalspec_boot);
2511 let logslope_terms = spatial_length_scale_term_indices(&logslopespec_boot);
2512 let marginal_has_spatial = !marginal_terms.is_empty();
2513 let logslope_has_spatial = !logslope_terms.is_empty();
2514 let analytic_joint_derivatives_available =
2515 marginal_has_spatial || logslope_has_spatial || setup.log_kappa_dim() == 0;
2516 if setup.log_kappa_dim() > 0 && !analytic_joint_derivatives_available {
2517 return Err("exact bernoulli marginal-slope spatial optimization requires analytic joint psi derivatives"
2518 .to_string());
2519 }
2520 let initial_rho = setup.theta0().slice(s![..setup.rho_dim()]).to_owned();
2521 let initial_blocks = build_blocks(&initial_rho, &marginal_design, &logslope_design)?;
2522 let initial_family = make_family(&marginal_design, &logslope_design, initial_sigma);
2523 let (joint_gradient, joint_hessian) =
2524 custom_family_outer_derivatives(&initial_family, &initial_blocks, options);
2525 let analytic_joint_gradient_available = analytic_joint_derivatives_available
2526 && matches!(joint_gradient, gam_problem::Derivative::Analytic);
2527 let analytic_joint_hessian_available =
2533 analytic_joint_derivatives_available && joint_hessian.is_analytic();
2534 let kappa_options_ref: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
2535 let sigma_from_theta = |theta: &Array1<f64>| -> Option<f64> {
2536 if sigma_learnable {
2537 Some(theta[setup.rho_dim() + setup.log_kappa_dim()].exp())
2538 } else {
2539 initial_sigma
2540 }
2541 };
2542 let derivative_block_cache = RefCell::new(
2543 None::<(
2544 Array1<f64>,
2545 Arc<Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>>>,
2546 )>,
2547 );
2548 let theta_matches = |left: &Array1<f64>, right: &Array1<f64>| -> bool {
2549 left.len() == right.len()
2550 && left
2551 .iter()
2552 .zip(right.iter())
2553 .all(|(lhs, rhs)| (*lhs - *rhs).abs() <= 1e-12 * (1.0 + lhs.abs().max(rhs.abs())))
2554 };
2555 let get_derivative_blocks = |theta: &Array1<f64>,
2556 specs: &[TermCollectionSpec],
2557 designs: &[TermCollectionDesign]|
2558 -> Result<
2559 Arc<Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>>>,
2560 String,
2561 > {
2562 if let Some((cached_theta, cached_blocks)) = derivative_block_cache.borrow().as_ref()
2563 && theta_matches(cached_theta, theta)
2564 {
2565 return Ok(Arc::clone(cached_blocks));
2566 }
2567
2568 let built = |specs: &[TermCollectionSpec],
2569 designs: &[TermCollectionDesign]|
2570 -> Result<
2571 Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>>,
2572 String,
2573 > {
2574 let marginal_psi_derivs = if marginal_has_spatial {
2575 build_block_spatial_psi_derivatives(data_view, &specs[0], &designs[0])?.ok_or_else(
2576 || {
2577 "bernoulli marginal-slope: marginal block has spatial terms \
2578 but spatial psi derivatives are unavailable"
2579 .to_string()
2580 },
2581 )?
2582 } else {
2583 Vec::new()
2584 };
2585 let logslope_psi_derivs = if logslope_has_spatial {
2586 build_block_spatial_psi_derivatives(data_view, &specs[1], &designs[1])?.ok_or_else(
2587 || {
2588 "bernoulli marginal-slope: logslope block has spatial terms \
2589 but spatial psi derivatives are unavailable"
2590 .to_string()
2591 },
2592 )?
2593 } else {
2594 Vec::new()
2595 };
2596 let mut derivative_blocks = vec![marginal_psi_derivs, logslope_psi_derivs];
2597 if score_warp_runtime.is_some() {
2598 derivative_blocks.push(Vec::new());
2599 }
2600 if link_dev_runtime.is_some() {
2601 derivative_blocks.push(Vec::new());
2602 }
2603 if sigma_learnable {
2604 derivative_blocks
2605 .last_mut()
2606 .expect("bernoulli derivative block list is non-empty")
2607 .push(crate::custom_family::CustomFamilyBlockPsiDerivative::new(
2608 None,
2609 Array2::zeros((0, 0)),
2610 Array2::zeros((0, 0)),
2611 None,
2612 None,
2613 None,
2614 None,
2615 ));
2616 }
2617 Ok(derivative_blocks)
2618 }(specs, designs)?;
2619 let built = Arc::new(built);
2620 derivative_block_cache.replace(Some((theta.clone(), Arc::clone(&built))));
2621 Ok(built)
2622 };
2623
2624 let outer_policy = {
2629 let psi_dim = setup.theta0().len() - setup.rho_dim();
2630 initial_family.outer_derivative_policy(&initial_blocks, psi_dim, options)
2631 };
2632 let exact_spatial_outer_tol = kappa_options_ref.rel_tol.max(EXACT_SPATIAL_OUTER_TOL_FLOOR);
2633 let solved = optimize_spatial_length_scale_exact_joint(
2634 data_view,
2635 &[marginalspec_boot.clone(), logslopespec_boot.clone()],
2636 &[marginal_terms.clone(), logslope_terms.clone()],
2637 kappa_options_ref,
2638 &setup,
2639 gam_solve::seeding::SeedRiskProfile::GeneralizedLinear,
2640 analytic_joint_gradient_available,
2641 analytic_joint_hessian_available,
2642 true,
2643 None,
2644 outer_policy,
2645 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
2646 if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2647 return Err(err);
2648 }
2649 assert_eq!(
2650 specs.len(),
2651 designs.len(),
2652 "spatial joint optimizer must supply one spec per design",
2653 );
2654 let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2655 let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2656 let sigma = sigma_from_theta(theta);
2657 final_sigma_cell.set(sigma);
2658 let family = make_family(&designs[0], &designs[1], sigma);
2659 let fit = inner_fit(&family, &blocks, options)?;
2660 if let Some(block) = fit.block_states.first()
2661 && let Some(err) = bernoulli_marginal_slope_runaway_error_from_beta(
2662 block.beta.view(),
2663 &designs[0],
2664 &specs[0],
2665 true,
2666 "final fit",
2667 )
2668 {
2669 runaway_error.replace(Some(err.clone()));
2670 return Err(err);
2671 }
2672 let mut hints_mut = hints.borrow_mut();
2673 let mut bidx = 0usize;
2674 if let Some(block) = fit.block_states.get(bidx) {
2675 hints_mut.marginal_beta = Some(block.beta.clone());
2676 }
2677 bidx += 1;
2678 if let Some(block) = fit.block_states.get(bidx) {
2679 hints_mut.logslope_beta = Some(block.beta.clone());
2680 }
2681 bidx += 1;
2682 if score_warp_prepared.is_some() {
2683 if let Some(block) = fit.block_states.get(bidx) {
2684 hints_mut.score_warp_beta = Some(block.beta.clone());
2685 }
2686 bidx += 1;
2687 }
2688 if link_dev_prepared.is_some()
2689 && let Some(block) = fit.block_states.get(bidx)
2690 {
2691 hints_mut.link_dev_beta = Some(block.beta.clone());
2692 }
2693 Ok(fit)
2694 },
2695 |theta,
2696 specs: &[TermCollectionSpec],
2697 designs: &[TermCollectionDesign],
2698 eval_mode,
2699 row_set: &crate::row_kernel::RowSet| {
2700 if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2701 return Err(err);
2702 }
2703 use gam_problem::EvalMode;
2704 static BMS_OUTER_EVAL_ROWSET_LOGGED: std::sync::Once = std::sync::Once::new();
2711 BMS_OUTER_EVAL_ROWSET_LOGGED.call_once(|| {
2712 let row_set_rows = match row_set {
2713 crate::row_kernel::RowSet::All => spec.y.len(),
2714 crate::row_kernel::RowSet::Subsample { rows, .. } => rows.len(),
2715 };
2716 log::debug!(
2717 "[BMS exact outer eval] mode={eval_mode:?} row_set_rows={row_set_rows}"
2718 );
2719 });
2720 let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2721 let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2722 if let Some(beta_seed) = pending_beta_seed.borrow_mut().take() {
2726 let widths: Vec<usize> = blocks.iter().map(|b| b.design.ncols()).collect();
2727 match CustomFamilyWarmStart::from_cached_beta(&widths, &beta_seed) {
2728 Ok(ws) => {
2729 exact_warm_start.replace(Some(ws));
2730 }
2731 Err(e) => {
2732 log::warn!(
2733 "[BMS] outer ρ-cache β-warm-start rejected: {e}; falling back to cold β"
2734 );
2735 }
2736 }
2737 }
2738 let sigma = sigma_from_theta(theta);
2739 final_sigma_cell.set(sigma);
2740 let family = make_family(&designs[0], &designs[1], sigma);
2741 let derivative_blocks = get_derivative_blocks(theta, specs, designs)?;
2742 let effective_mode = match eval_mode {
2746 EvalMode::ValueGradientHessian if !analytic_joint_hessian_available => {
2747 EvalMode::ValueAndGradient
2748 }
2749 other => other,
2750 };
2751 let mut eval_options =
2752 joint_hyper_options_for_outer_tolerance(options, exact_spatial_outer_tol);
2753 if let crate::row_kernel::RowSet::Subsample { rows, n_full } = row_set {
2754 let subsample = crate::outer_subsample::OuterScoreSubsample::from_weighted_rows(
2755 rows.as_ref().clone(),
2756 *n_full,
2757 0,
2758 );
2759 eval_options.outer_score_subsample = Some(Arc::new(subsample));
2760 eval_options.auto_outer_subsample = false;
2761 }
2762 let eval = evaluate_custom_family_joint_hyper_shared(
2763 &family,
2764 &blocks,
2765 &eval_options,
2766 &rho,
2767 derivative_blocks,
2768 exact_warm_start.borrow().as_ref(),
2769 effective_mode,
2770 )?;
2771 if let Some(err) = bernoulli_marginal_slope_runaway_error(
2772 &eval.warm_start,
2773 &designs[0],
2774 &specs[0],
2775 eval.inner_converged,
2776 "exact outer evaluation",
2777 ) {
2778 runaway_error.replace(Some(err.clone()));
2779 return Err(err);
2780 }
2781 exact_warm_start.replace(Some(eval.warm_start.clone()));
2782 if !eval.inner_converged {
2783 return Err(
2784 "exact bernoulli marginal-slope inner solve did not converge".to_string(),
2785 );
2786 }
2787 if matches!(eval_mode, EvalMode::ValueGradientHessian)
2788 && analytic_joint_hessian_available
2789 && !eval.outer_hessian.is_analytic()
2790 {
2791 return Err("exact bernoulli marginal-slope joint [rho, psi] objective did not return an outer Hessian"
2792 .to_string());
2793 }
2794 Ok((eval.objective, eval.gradient, eval.outer_hessian))
2795 },
2796 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign]| {
2797 if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2798 return Err(err);
2799 }
2800 let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2801 let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2802 if let Some(beta_seed) = pending_beta_seed.borrow_mut().take() {
2803 let widths: Vec<usize> = blocks.iter().map(|b| b.design.ncols()).collect();
2804 match CustomFamilyWarmStart::from_cached_beta(&widths, &beta_seed) {
2805 Ok(ws) => {
2806 exact_warm_start.replace(Some(ws));
2807 }
2808 Err(e) => {
2809 log::warn!(
2810 "[BMS] outer ρ-cache β-warm-start rejected (efs): {e}; falling back to cold β"
2811 );
2812 }
2813 }
2814 }
2815 let sigma = sigma_from_theta(theta);
2816 final_sigma_cell.set(sigma);
2817 let family = make_family(&designs[0], &designs[1], sigma);
2818 let derivative_blocks = get_derivative_blocks(theta, specs, designs)?;
2819 let eval = evaluate_custom_family_joint_hyper_efs_shared(
2820 &family,
2821 &blocks,
2822 &joint_hyper_options_for_outer_tolerance(options, exact_spatial_outer_tol),
2823 &rho,
2824 derivative_blocks,
2825 exact_warm_start.borrow().as_ref(),
2826 )?;
2827 if let Some(err) = bernoulli_marginal_slope_runaway_error(
2828 &eval.warm_start,
2829 &designs[0],
2830 &specs[0],
2831 eval.inner_converged,
2832 "EFS outer evaluation",
2833 ) {
2834 runaway_error.replace(Some(err.clone()));
2835 return Err(err);
2836 }
2837 exact_warm_start.replace(Some(eval.warm_start.clone()));
2838 if !eval.inner_converged {
2839 return Err(
2840 "exact bernoulli marginal-slope EFS inner solve did not converge".to_string(),
2841 );
2842 }
2843 Ok(eval.efs_eval)
2844 },
2845 crate::marginal_slope_shared::make_beta_seed_validator(&pending_beta_seed),
2846 )?;
2847
2848 let mut resolved_specs = solved.resolved_specs;
2849 let mut designs = solved.designs;
2850 let mut solved_fit = solved.fit;
2862 if let Some(reparam) = logslope_reduced_reparam.as_ref() {
2863 let r = reparam.reduced_cols();
2864 if let Some(block) = solved_fit.blocks.get_mut(1)
2865 && block.beta.len() == r
2866 {
2867 block.beta = reparam.recover_original_logslope_beta(&block.beta)?;
2868 }
2869 if let Some(state) = solved_fit.block_states.get_mut(1)
2870 && state.beta.len() == r
2871 {
2872 state.beta = reparam.recover_original_logslope_beta(&state.beta)?;
2873 }
2874 }
2875 let (latent_z_rank_int_calibration, latent_z_conditional_calibration) =
2917 match latent_z_calibration {
2918 LatentMeasureCalibration::None => (None, None),
2919 LatentMeasureCalibration::RankInverseNormal(cal) => (Some(cal), None),
2920 LatentMeasureCalibration::ConditionalLocationScale(cal) => (None, Some(cal)),
2921 };
2922 if let Some(cal) = latent_z_conditional_calibration.as_ref()
2936 && matches!(latent_measure, LatentMeasureKind::StandardNormal)
2937 && let Some(vb) = solved_fit.covariance_conditional.clone()
2938 {
2939 let p_beta = vb.nrows();
2940 let marginal_dense = marginal_design
2941 .design
2942 .try_to_dense_arc("bms generated-regressor marginal design")?;
2943 let logslope_reduced = reduce_logslope_design(&logslope_design)?;
2944 let logslope_reduced_dense = logslope_reduced
2945 .design
2946 .try_to_dense_arc("bms generated-regressor reduced logslope design")?;
2947 let p_m = marginal_dense.ncols();
2948 let r = logslope_reduced_dense.ncols();
2949 if p_beta != vb.ncols() {
2950 return Err(format!(
2951 "bms generated-regressor: covariance_conditional must be square, got {}×{}",
2952 vb.nrows(),
2953 vb.ncols()
2954 ));
2955 }
2956 if p_beta == p_m + r {
2960 let marginal_eta = &solved_fit.block_states[0].eta;
2961 let slope_eta = &solved_fit.block_states[1].eta;
2962 let probit_scale = probit_frailty_scale(final_sigma_cell.get());
2963 let s = rigid_standard_normal_score_zeta_sensitivity(
2964 &spec.base_link,
2965 marginal_eta,
2966 slope_eta,
2967 z.as_ref(),
2968 y.as_ref(),
2969 weights.as_ref(),
2970 probit_scale,
2971 marginal_dense.view(),
2972 logslope_reduced_dense.view(),
2973 p_beta,
2974 )?;
2975 let correction = cal.generated_regressor_correction(
2983 s.view(),
2984 spec.z.view(),
2985 marginal_dense.view(),
2986 vb.view(),
2987 )?;
2988 if let Some(cov) = solved_fit.covariance_conditional.as_mut() {
2989 *cov = &*cov + &correction;
2990 }
2991 if let Some(cov) = solved_fit.covariance_corrected.as_mut() {
2992 *cov = &*cov + &correction;
2993 }
2994 log::info!(
2995 "[BMS latent-z] Murphy–Topel generated-regressor SE correction applied: \
2996 p_beta={p_beta} theta1_dim={} max_diag_inflation={:.3e}",
2997 cal.theta1_dim(),
2998 (0..p_beta)
2999 .map(|i| correction[[i, i]])
3000 .fold(0.0_f64, f64::max),
3001 );
3002 } else {
3003 log::info!(
3004 "[BMS latent-z] Murphy–Topel generated-regressor SE correction skipped: \
3005 aux deviation blocks present (p_beta={p_beta} > marginal({p_m})+logslope({r})); \
3006 rigid-kernel z-channel does not yet cover score_warp/link_dev deviations"
3007 );
3008 }
3009 }
3010 Ok(BernoulliMarginalSlopeFitResult {
3023 fit: solved_fit,
3024 marginalspec_resolved: resolved_specs.remove(0),
3025 logslopespec_resolved: resolved_specs.remove(0),
3026 marginal_design: designs.remove(0),
3027 logslope_design: designs.remove(0),
3028 baseline_marginal: baseline.0,
3029 baseline_logslope: baseline.1,
3030 z_normalization,
3031 latent_measure,
3032 score_warp_runtime,
3033 link_dev_runtime,
3034 gaussian_frailty_sd: final_sigma_cell.get(),
3035 cross_block_warnings,
3036 latent_z_rank_int_calibration,
3037 latent_z_conditional_calibration,
3038 })
3039}