1use super::empirical_measure_sensitivity::{
2 EmpiricalGeneratedRegressorChannel, classify_empirical_generated_regressor_channel,
3 rigid_empirical_score_zeta_channels,
4};
5use super::family::*;
6use super::gradient_paths::*;
7use super::hessian_paths::{new_cell_moment_cache_stats, new_cell_moment_lru_cache};
8use super::install_flex::validate_spec;
9use super::*;
10use crate::marginal_slope_orthogonal::influence_absorber_log_lambda;
11use faer::Side;
12use gam_linalg::faer_ndarray::{FaerEigh, fast_ab, fast_atb, fast_xt_diag_x};
13
14pub(crate) const BMS_PROBIT_SEPARATION_ETA_INF: f64 = 35.0;
29
30pub(super) const GAUGE_PRIORITY_ANCHOR: u8 = 200;
45pub(super) const GAUGE_PRIORITY_MARGINAL: u8 = 150;
48pub(super) const GAUGE_PRIORITY_LOGSLOPE: u8 = 120;
50pub(super) const GAUGE_PRIORITY_CANDIDATE_FLEX: u8 = 100;
53pub(super) const GAUGE_PRIORITY_SCORE_WARP_DEV: u8 = 80;
56pub(super) const GAUGE_PRIORITY_DEVIATION_DEFAULT: u8 = 70;
60pub(super) const GAUGE_PRIORITY_LINK_DEV: u8 = 60;
62
63pub(crate) const EXACT_SPATIAL_OUTER_TOL_FLOOR: f64 = 1e-6;
69
70pub struct BmsMarginalJacobian {
109 pub marginal_dense: Arc<Array2<f64>>,
111 pub logslope_dense: Arc<Array2<f64>>,
113 pub offset_m: Array1<f64>,
114 pub offset_s: Array1<f64>,
115 pub p_marginal: usize,
117}
118
119impl BmsMarginalJacobian {
120 pub fn new(
121 marginal_dense: Arc<Array2<f64>>,
122 logslope_dense: Arc<Array2<f64>>,
123 offset_m: Array1<f64>,
124 offset_s: Array1<f64>,
125 p_marginal: usize,
126 ) -> Self {
127 Self {
128 marginal_dense,
129 logslope_dense,
130 offset_m,
131 offset_s,
132 p_marginal,
133 }
134 }
135}
136
137impl BlockEffectiveJacobian for BmsMarginalJacobian {
138 fn effective_jacobian_rows(
139 &self,
140 state: &FamilyLinearizationState<'_>,
141 rows: std::ops::Range<usize>,
142 ) -> Result<Array2<f64>, String> {
143 let beta = state.beta;
144 let s = state.probit_frailty_scale;
145 let p_m = self.p_marginal;
146 let p_s_block = self.logslope_dense.ncols();
147 let beta_s_raw = if beta.len() > p_m {
148 &beta[p_m..]
149 } else {
150 &[][..]
151 };
152 let p_s_use = p_s_block.min(beta_s_raw.len());
153 let beta_s = &beta_s_raw[..p_s_use];
154 let n = self.marginal_dense.nrows();
155 let rows = rows.start.min(n)..rows.end.min(n);
156 let p_block = self.marginal_dense.ncols();
157
158 let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_block));
168 for i in rows.clone() {
169 let g_i = self.offset_s[i]
170 + self
171 .logslope_dense
172 .row(i)
173 .slice(ndarray::s![..p_s_use])
174 .dot(&ArrayView1::from(beta_s));
175 let sg = s * g_i;
176 let c_i = (1.0 + sg * sg).sqrt();
177 let m_row = self.marginal_dense.row(i);
179 out.row_mut(i - rows.start).assign(&m_row.mapv(|x| c_i * x));
180 }
181 Ok(out)
182 }
183
184 fn n_outputs(&self) -> usize {
185 1
186 }
187
188 fn locks_raw_width_reduction(&self) -> bool {
189 true
199 }
200}
201
202pub struct BmsLogslopeJacobian {
214 pub marginal_dense: Arc<Array2<f64>>,
216 pub logslope_dense: Arc<Array2<f64>>,
218 pub offset_m: Array1<f64>,
219 pub offset_s: Array1<f64>,
220 pub z: Arc<Array1<f64>>,
221 pub p_marginal: usize,
223}
224
225impl BmsLogslopeJacobian {
226 pub fn new(
227 marginal_dense: Arc<Array2<f64>>,
228 logslope_dense: Arc<Array2<f64>>,
229 offset_m: Array1<f64>,
230 offset_s: Array1<f64>,
231 z: Arc<Array1<f64>>,
232 p_marginal: usize,
233 ) -> Self {
234 Self {
235 marginal_dense,
236 logslope_dense,
237 offset_m,
238 offset_s,
239 z,
240 p_marginal,
241 }
242 }
243}
244
245impl BlockEffectiveJacobian for BmsLogslopeJacobian {
246 fn effective_jacobian_rows(
247 &self,
248 state: &FamilyLinearizationState<'_>,
249 rows: std::ops::Range<usize>,
250 ) -> Result<Array2<f64>, String> {
251 let beta = state.beta;
252 let s = state.probit_frailty_scale;
253 let p_m = self.p_marginal;
254 let p_m_use = p_m.min(beta.len());
255 let beta_m = &beta[..p_m_use];
256 let beta_s_raw = if beta.len() > p_m {
257 &beta[p_m..]
258 } else {
259 &[][..]
260 };
261 let p_s_block = self.logslope_dense.ncols();
262 let p_s_use = p_s_block.min(beta_s_raw.len());
263 let beta_s = &beta_s_raw[..p_s_use];
264 let n = self.logslope_dense.nrows();
265 let rows = rows.start.min(n)..rows.end.min(n);
266
267 let mut out = Array2::<f64>::zeros((rows.end - rows.start, p_s_block));
280 for i in rows.clone() {
281 let q_i = self.offset_m[i]
282 + self
283 .marginal_dense
284 .row(i)
285 .slice(ndarray::s![..p_m_use])
286 .dot(&ArrayView1::from(beta_m));
287 let g_i = self.offset_s[i]
288 + self
289 .logslope_dense
290 .row(i)
291 .slice(ndarray::s![..p_s_use])
292 .dot(&ArrayView1::from(beta_s));
293 let sg = s * g_i;
294 let c_i = (1.0 + sg * sg).sqrt();
295 let z_i = self.z[i];
296 let factor = q_i * s * s * g_i / c_i + s * z_i;
298 let g_row = self.logslope_dense.row(i);
300 out.row_mut(i - rows.start)
301 .assign(&g_row.mapv(|x| factor * x));
302 }
303 Ok(out)
304 }
305
306 fn n_outputs(&self) -> usize {
307 1
308 }
309
310 fn locks_raw_width_reduction(&self) -> bool {
311 true
322 }
323}
324
325pub(crate) fn widen_marginal_dense_with_influence(
335 marginal_dense: &Arc<Array2<f64>>,
336 influence_columns: Option<&Array2<f64>>,
337) -> Result<Arc<Array2<f64>>, String> {
338 let Some(z_infl) = influence_columns else {
339 return Ok(Arc::clone(marginal_dense));
340 };
341 let n = marginal_dense.nrows();
342 if z_infl.nrows() != n {
343 return Err(format!(
344 "influence block: residualised columns have {} rows, marginal design has {n}",
345 z_infl.nrows()
346 ));
347 }
348 let p_m = marginal_dense.ncols();
349 let p1 = z_infl.ncols();
350 let mut widened = Array2::<f64>::zeros((n, p_m + p1));
351 widened
352 .slice_mut(s![.., ..p_m])
353 .assign(marginal_dense.as_ref());
354 widened.slice_mut(s![.., p_m..]).assign(z_infl);
355 Ok(Arc::new(widened))
356}
357
358pub(crate) const LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL: f64 = 1.0e-6;
366
367#[derive(Debug, Clone)]
419pub(super) struct ReducedLogslopeReparam {
420 transform: Array2<f64>,
423}
424
425impl ReducedLogslopeReparam {
426 #[inline]
428 pub(super) fn original_cols(&self) -> usize {
429 self.transform.nrows()
430 }
431
432 #[inline]
434 pub(super) fn reduced_cols(&self) -> usize {
435 self.transform.ncols()
436 }
437
438 pub(super) fn recover_original_logslope_beta(
442 &self,
443 beta_reduced: &Array1<f64>,
444 ) -> Result<Array1<f64>, String> {
445 if beta_reduced.len() != self.reduced_cols() {
446 return Err(format!(
447 "reduced logslope reparam: β' length ({}) != reduced width ({})",
448 beta_reduced.len(),
449 self.reduced_cols()
450 ));
451 }
452 Ok(self.transform.dot(beta_reduced))
453 }
454}
455
456fn build_reduced_logslope_reparam(
467 marginal_design: &TermCollectionDesign,
468 logslope_design: &TermCollectionDesign,
469 z: &Array1<f64>,
470 row_metric: &Array1<f64>,
471 marginal_offset: &Array1<f64>,
472 logslope_offset: &Array1<f64>,
473 marginal_baseline: f64,
474 logslope_baseline: f64,
475 probit_scale: f64,
476) -> Result<Option<ReducedLogslopeReparam>, String> {
477 let marginal = marginal_design
478 .design
479 .try_to_dense_arc("build_reduced_logslope_reparam::marginal")?;
480 let logslope = logslope_design
481 .design
482 .try_to_dense_arc("build_reduced_logslope_reparam::logslope")?;
483 let n = marginal.nrows();
484 if logslope.nrows() != n
485 || z.len() != n
486 || row_metric.len() != n
487 || marginal_offset.len() != n
488 || logslope_offset.len() != n
489 {
490 return Err(format!(
491 "reduced logslope reparam row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
492 marginal.nrows(),
493 logslope.nrows(),
494 z.len(),
495 row_metric.len(),
496 marginal_offset.len(),
497 logslope_offset.len(),
498 ));
499 }
500 let p_m = marginal.ncols();
501 let p_g = logslope.ncols();
502 if p_m == 0 || p_g == 0 {
503 return Ok(None);
504 }
505 if !marginal_baseline.is_finite()
506 || !logslope_baseline.is_finite()
507 || !probit_scale.is_finite()
508 || probit_scale <= 0.0
509 || z.iter().any(|v| !v.is_finite())
510 || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
511 || marginal_offset.iter().any(|v| !v.is_finite())
512 || logslope_offset.iter().any(|v| !v.is_finite())
513 {
514 return Err(
515 "reduced logslope reparam requires finite pilot geometry and finite non-negative row metric"
516 .to_string(),
517 );
518 }
519
520 match reduced_logslope_transform_effective(
531 marginal.view(),
532 logslope.view(),
533 z,
534 row_metric,
535 marginal_offset,
536 logslope_offset,
537 marginal_baseline,
538 logslope_baseline,
539 probit_scale,
540 )? {
541 ReducedLogslopeOutcome::Reduced(transform) => {
542 Ok(Some(ReducedLogslopeReparam { transform }))
543 }
544 ReducedLogslopeOutcome::FullRank => Ok(None),
545 ReducedLogslopeOutcome::FullyConfounded => Err(
546 "BMS score-slope block is fully confounded with the marginal index: every \
547 effective logslope direction diag(f)·G·v is W-explained by the effective \
548 marginal span at the rigid pilot, so the data identify only the sum of the \
549 marginal and score-slope surfaces and the smoothing penalty would select an \
550 arbitrary decomposition between them. Refusing to fit; remove the score-slope \
551 terms or supply covariates that separate them from the marginal index."
552 .to_string(),
553 ),
554 }
555}
556
557#[derive(Debug)]
564pub(crate) enum ReducedLogslopeOutcome {
565 FullRank,
568 Reduced(Array2<f64>),
571 FullyConfounded,
574}
575
576pub(crate) fn reduced_logslope_transform_effective(
596 marginal: ArrayView2<'_, f64>,
597 logslope: ArrayView2<'_, f64>,
598 z: &Array1<f64>,
599 row_metric: &Array1<f64>,
600 marginal_offset: &Array1<f64>,
601 logslope_offset: &Array1<f64>,
602 marginal_baseline: f64,
603 logslope_baseline: f64,
604 probit_scale: f64,
605) -> Result<ReducedLogslopeOutcome, String> {
606 let n = marginal.nrows();
607 let p_m = marginal.ncols();
608 let p_g = logslope.ncols();
609 if p_m == 0 || p_g == 0 {
610 return Ok(ReducedLogslopeOutcome::FullRank);
611 }
612
613 let mut m_eff = Array2::<f64>::zeros((n, p_m));
615 let mut g_eff = Array2::<f64>::zeros((n, p_g));
616 for i in 0..n {
617 let q_i = marginal_offset[i] + marginal_baseline;
618 let g_i = logslope_offset[i] + logslope_baseline;
619 let sg = probit_scale * g_i;
620 let c_i = (1.0 + sg * sg).sqrt();
621 let f_i = q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
622 for j in 0..p_m {
623 m_eff[[i, j]] = c_i * marginal[[i, j]];
624 }
625 for j in 0..p_g {
626 g_eff[[i, j]] = f_i * logslope[[i, j]];
627 }
628 }
629
630 let c_gram = fast_xt_diag_x(&g_eff, row_metric);
633 let energy_scale = (0..p_g).map(|i| c_gram[[i, i]]).fold(0.0_f64, f64::max);
634 if !energy_scale.is_finite() {
635 return Err(
636 "reduced logslope reparam: effective logslope Gram produced non-finite energy"
637 .to_string(),
638 );
639 }
640 if energy_scale <= 0.0 {
641 return Ok(ReducedLogslopeOutcome::FullyConfounded);
644 }
645
646 let mut a_gram = fast_xt_diag_x(&m_eff, row_metric);
650 let a_scale = (0..p_m).map(|i| a_gram[[i, i]]).fold(0.0_f64, f64::max);
651 let a_ridge = (a_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL).max(f64::EPSILON);
652 for i in 0..p_m {
653 a_gram[[i, i]] += a_ridge;
654 }
655
656 let b_cross = gam_linalg::faer_ndarray::fast_xt_diag_y(&m_eff, row_metric, &g_eff);
658 let a_view = gam_linalg::faer_ndarray::FaerArrayView::new(&a_gram);
659 let a_factor =
660 gam_linalg::faer_ndarray::factorize_symmetricwith_fallback(a_view.as_ref(), Side::Lower)
661 .map_err(|e| {
662 format!(
663 "reduced logslope reparam: effective marginal Gram factorization failed: {e}"
664 )
665 })?;
666 let b_view = gam_linalg::faer_ndarray::FaerArrayView::new(&b_cross);
667 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)]);
669 let schur = fast_atb(&b_cross, &a_inv_b); let mut stt = &c_gram - &schur;
671 stt = (&stt + &stt.t()) * 0.5;
672 if stt.iter().any(|v| !v.is_finite()) {
673 return Err(
674 "reduced logslope reparam: effective Schur Gram produced non-finite entries"
675 .to_string(),
676 );
677 }
678
679 let (evals, evecs) = stt
680 .eigh(Side::Lower)
681 .map_err(|e| format!("reduced logslope reparam: eigendecomposition failed: {e:?}"))?;
682 let tol = energy_scale * LOGSLOPE_REDUCED_BASIS_RELATIVE_TOL;
686 let mut kept: Vec<usize> = (0..evals.len()).filter(|&i| evals[i] > tol).collect();
687 kept.sort_by(|&a, &b| {
688 evals[b]
689 .partial_cmp(&evals[a])
690 .unwrap_or(std::cmp::Ordering::Equal)
691 });
692 let r = kept.len();
693 if r == p_g {
699 return Ok(ReducedLogslopeOutcome::FullRank);
700 }
701 if r == 0 {
702 return Ok(ReducedLogslopeOutcome::FullyConfounded);
703 }
704 let mut transform = Array2::<f64>::zeros((p_g, r));
705 for (out_col, &src) in kept.iter().enumerate() {
706 transform.column_mut(out_col).assign(&evecs.column(src));
707 }
708 if transform.iter().any(|v| !v.is_finite()) {
709 return Err(
710 "reduced logslope reparam: reduced transform produced non-finite entries".to_string(),
711 );
712 }
713 Ok(ReducedLogslopeOutcome::Reduced(transform))
714}
715
716fn reparameterize_logslope_design_reduced(
723 logslope_design: &TermCollectionDesign,
724 reparam: &ReducedLogslopeReparam,
725) -> Result<TermCollectionDesign, String> {
726 let g = logslope_design
727 .design
728 .try_to_dense_arc("reparameterize_logslope_design_reduced::logslope")?;
729 let p_g = g.ncols();
730 if p_g != reparam.original_cols() {
731 return Err(format!(
732 "reduced logslope reparam width mismatch: design has {p_g} cols, transform expects {}",
733 reparam.original_cols()
734 ));
735 }
736 let t = &reparam.transform;
737 let r = reparam.reduced_cols();
738 let g_reduced = fast_ab(&g, t);
740
741 let mut new_penalties: Vec<gam_terms::smooth::BlockwisePenalty> =
744 Vec::with_capacity(logslope_design.penalties.len());
745 let mut new_nullspace_dims: Vec<usize> = Vec::with_capacity(logslope_design.penalties.len());
746 for bp in &logslope_design.penalties {
747 let mut full = Array2::<f64>::zeros((p_g, p_g));
748 full.slice_mut(s![bp.col_range.clone(), bp.col_range.clone()])
749 .assign(&bp.local);
750 let st = fast_ab(&full, t); let mut s_reduced = fast_atb(t, &st); s_reduced = (&s_reduced + &s_reduced.t()) * 0.5;
754 let (evals, _) = s_reduced
756 .eigh(Side::Lower)
757 .map_err(|e| format!("reduced logslope penalty eigendecomposition failed: {e:?}"))?;
758 let max_eval = evals.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
759 let pen_tol = (max_eval * 1.0e-12).max(f64::EPSILON);
760 let rank = evals.iter().filter(|&&v| v.abs() > pen_tol).count();
761 let nullspace_dim = r.saturating_sub(rank);
762 new_penalties.push(gam_terms::smooth::BlockwisePenalty::new(0..r, s_reduced));
763 new_nullspace_dims.push(nullspace_dim);
764 }
765
766 let new_design = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(g_reduced));
767 Ok(TermCollectionDesign {
773 design: new_design,
774 affine_offset: logslope_design.affine_offset.clone(),
778 penalties: new_penalties,
779 nullspace_dims: new_nullspace_dims,
780 penaltyinfo: Vec::new(),
781 dropped_penaltyinfo: Vec::new(),
782 coefficient_lower_bounds: None,
783 linear_constraints: None,
784 intercept_range: 0..0,
785 linear_ranges: Vec::new(),
786 linear_function_masses: Vec::new(),
787 random_effect_ranges: Vec::new(),
788 random_effect_levels: Vec::new(),
789 smooth: gam_terms::smooth::SmoothDesign {
790 term_designs: Vec::new(),
791 penalties: Vec::new(),
792 nullspace_dims: Vec::new(),
793 penaltyinfo: Vec::new(),
794 dropped_penaltyinfo: Vec::new(),
795 terms: Vec::new(),
796 coefficient_lower_bounds: None,
797 linear_constraints: None,
798 },
799 })
800}
801
802pub(crate) fn marginal_penalties_with_influence_ridge(
825 design: &TermCollectionDesign,
826 rho_marginal: &Array1<f64>,
827 influence_columns: Option<&Array2<f64>>,
828) -> Result<(Vec<PenaltyMatrix>, Vec<usize>, Array1<f64>), String> {
829 let p_m = design.design.ncols();
830 let p1 = influence_columns.map(|z| z.ncols()).unwrap_or(0);
831 let total_dim = p_m + p1;
832 let expected_rho = design.penalties.len() + usize::from(p1 > 0);
833 if rho_marginal.len() != expected_rho {
834 return Err(format!(
835 "marginal rho width {} != smooth penalties {} + absorber slot {}",
836 rho_marginal.len(),
837 design.penalties.len(),
838 usize::from(p1 > 0),
839 ));
840 }
841 let mut penalties: Vec<PenaltyMatrix> = design
844 .penalties
845 .iter()
846 .map(|bp| bp.to_penalty_matrix(total_dim))
847 .collect();
848 let mut nullspace_dims = design.nullspace_dims.clone();
849 let log_lambdas = rho_marginal.to_vec();
850
851 if p1 > 0 {
857 penalties.push(PenaltyMatrix::Blockwise {
858 local: Array2::<f64>::eye(p1),
859 col_range: p_m..total_dim,
860 total_dim,
861 });
862 nullspace_dims.push(0);
863 }
864
865 Ok((penalties, nullspace_dims, Array1::from_vec(log_lambdas)))
866}
867
868pub(crate) fn widen_marginal_beta_hint(
871 beta_hint: Option<Array1<f64>>,
872 p_marginal_widened: usize,
873) -> Option<Array1<f64>> {
874 beta_hint.map(|hint| {
875 if hint.len() == p_marginal_widened {
876 hint
877 } else {
878 let mut widened = Array1::<f64>::zeros(p_marginal_widened);
879 let copy = hint.len().min(p_marginal_widened);
880 widened
881 .slice_mut(s![..copy])
882 .assign(&hint.slice(s![..copy]));
883 widened
884 }
885 })
886}
887
888fn marginal_fitted_eta_sup_norm(design: &TermCollectionDesign, masked_beta: &Array1<f64>) -> f64 {
898 let x = &design.design;
899 let n = x.nrows();
900 if n == 0 || x.ncols() == 0 {
901 return 0.0;
902 }
903 let mut sup = 0.0_f64;
904 for row in 0..n {
905 let eta = x.dot_row_view(row, masked_beta.view());
906 if eta.is_finite() {
907 sup = sup.max(eta.abs());
908 }
909 }
910 sup
911}
912
913fn marginal_design_beta(
916 design: &TermCollectionDesign,
917 block_beta: ArrayView1<'_, f64>,
918) -> Array1<f64> {
919 let ncols = design.design.ncols();
920 let mut masked = Array1::<f64>::zeros(ncols);
921 let copy = ncols.min(block_beta.len());
922 masked
923 .slice_mut(s![..copy])
924 .assign(&block_beta.slice(s![..copy]));
925 masked
926}
927
928fn mask_parametric_columns(
934 design: &TermCollectionDesign,
935 spec: &TermCollectionSpec,
936 full: &Array1<f64>,
937) -> Array1<f64> {
938 let ncols = design.design.ncols();
939 let mut masked = Array1::<f64>::zeros(ncols);
940 if design.intercept_range.len() == 1 {
941 let idx = design.intercept_range.start;
942 if idx < ncols {
943 masked[idx] = full[idx];
944 }
945 }
946 for (linear, (_, range)) in spec.linear_terms.iter().zip(design.linear_ranges.iter()) {
947 if linear.double_penalty {
948 continue;
949 }
950 for col in range.clone() {
951 if col < ncols {
952 masked[col] = full[col];
953 }
954 }
955 }
956 masked
957}
958
959pub(crate) fn bernoulli_marginal_slope_runaway_error_from_beta(
970 block_beta: ArrayView1<'_, f64>,
971 design: &TermCollectionDesign,
972 spec: &TermCollectionSpec,
973 inner_converged: bool,
974 eval_label: &str,
975) -> Option<String> {
976 let full_beta = marginal_design_beta(design, block_beta);
977 let parametric_beta = mask_parametric_columns(design, spec, &full_beta);
978
979 let eta_parametric = marginal_fitted_eta_sup_norm(design, ¶metric_beta);
980 let eta_full = marginal_fitted_eta_sup_norm(design, &full_beta);
981
982 let (eta_inf, explanation) = if eta_parametric >= BMS_PROBIT_SEPARATION_ETA_INF {
983 (
984 eta_parametric,
985 "an unpenalized parametric marginal direction has no stable finite probit optimum and its fitted predictor has run to the probit underflow scale",
986 )
987 } else if eta_full >= BMS_PROBIT_SEPARATION_ETA_INF {
988 (
989 eta_full,
990 "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",
991 )
992 } else {
993 return None;
997 };
998
999 let inner_status = if inner_converged {
1000 "the inner solve reached a KKT certificate at this separation-scale predictor"
1001 } else {
1002 "the inner solve failed while already carrying a separation-scale predictor"
1003 };
1004 let beta_abs = full_beta
1006 .iter()
1007 .copied()
1008 .filter(|v| v.is_finite())
1009 .fold(0.0_f64, |acc, v| acc.max(v.abs()));
1010
1011 Some(format!(
1012 "bernoulli marginal-slope probit marginal/logslope runaway detected in block \
1013 'marginal_surface' during {eval_label}: the fitted marginal predictor has \
1014 |η|∞={eta_inf:.3e} (numerical-degeneracy threshold \
1015 {BMS_PROBIT_SEPARATION_ETA_INF:.1}; raw |β|∞={beta_abs:.3e} is reported for \
1016 context only and does not gate this diagnostic). The joint design is \
1017 identifiable; {explanation}. {inner_status}. The robust Jeffreys curvature \
1018 path is already installed for this fit, so this diagnostic means the current \
1019 coupled surface still drives the linear predictor to the probit underflow \
1020 scale rather than a request for an external bias-reduction prior. Reduce or \
1021 reparameterize the coupled marginal/logslope surface, or use a \
1022 lower-dimensional logslope interaction. This is not a \
1023 Matérn/Duchon polynomial-nullspace or cross-block gauge-priority \
1024 failure."
1025 ))
1026}
1027
1028pub(crate) fn bernoulli_marginal_slope_runaway_error(
1029 warm_start: &CustomFamilyWarmStart,
1030 design: &TermCollectionDesign,
1031 spec: &TermCollectionSpec,
1032 inner_converged: bool,
1033 eval_label: &str,
1034) -> Option<String> {
1035 let block_beta = warm_start.block_beta_view(0)?;
1036 bernoulli_marginal_slope_runaway_error_from_beta(
1037 block_beta,
1038 design,
1039 spec,
1040 inner_converged,
1041 eval_label,
1042 )
1043}
1044
1045#[cfg(test)]
1046mod runaway_tests {
1047 use super::*;
1048 use gam_linalg::faer_ndarray::{
1049 FaerArrayView, factorize_symmetricwith_fallback, fast_xt_diag_y,
1050 };
1051 use gam_terms::smooth::{LinearCoefficientGeometry, LinearTermSpec};
1052
1053 pub(crate) fn marginal_logslope_overlap_penalty(
1059 marginal_design: &DesignMatrix,
1060 logslope_design: &DesignMatrix,
1061 z: &Array1<f64>,
1062 row_metric: &Array1<f64>,
1063 marginal_offset: &Array1<f64>,
1064 logslope_offset: &Array1<f64>,
1065 marginal_baseline: f64,
1066 logslope_baseline: f64,
1067 probit_scale: f64,
1068 ) -> Result<Option<Array2<f64>>, String> {
1069 let marginal =
1070 marginal_design.try_to_dense_arc("marginal_logslope_overlap_penalty::marginal")?;
1071 let logslope =
1072 logslope_design.try_to_dense_arc("marginal_logslope_overlap_penalty::logslope")?;
1073 let n = marginal.nrows();
1074 if logslope.nrows() != n
1075 || z.len() != n
1076 || row_metric.len() != n
1077 || marginal_offset.len() != n
1078 || logslope_offset.len() != n
1079 {
1080 return Err(format!(
1081 "marginal/logslope overlap penalty row mismatch: marginal={}, logslope={}, z={}, row_metric={}, marginal_offset={}, logslope_offset={}",
1082 marginal.nrows(),
1083 logslope.nrows(),
1084 z.len(),
1085 row_metric.len(),
1086 marginal_offset.len(),
1087 logslope_offset.len(),
1088 ));
1089 }
1090 let p_m = marginal.ncols();
1091 let p_g = logslope.ncols();
1092 if p_m == 0 || p_g == 0 {
1093 return Ok(None);
1094 }
1095 if !marginal_baseline.is_finite()
1096 || !logslope_baseline.is_finite()
1097 || !probit_scale.is_finite()
1098 || probit_scale <= 0.0
1099 || z.iter().any(|v| !v.is_finite())
1100 || row_metric.iter().any(|v| !v.is_finite() || *v < 0.0)
1101 || marginal_offset.iter().any(|v| !v.is_finite())
1102 || logslope_offset.iter().any(|v| !v.is_finite())
1103 {
1104 return Err(
1105 "marginal/logslope overlap penalty requires finite pilot geometry and finite non-negative row metric"
1106 .to_string(),
1107 );
1108 }
1109
1110 let mut marginal_effective = Array2::<f64>::zeros((n, p_m));
1111 let mut effective_logslope = Array2::<f64>::zeros((n, p_g));
1112 for i in 0..n {
1113 let q_i = marginal_offset[i] + marginal_baseline;
1114 let g_i = logslope_offset[i] + logslope_baseline;
1115 let sg = probit_scale * g_i;
1116 let c_i = (1.0 + sg * sg).sqrt();
1117 let logslope_factor =
1118 q_i * probit_scale * probit_scale * g_i / c_i + probit_scale * z[i];
1119 for j in 0..p_m {
1120 marginal_effective[[i, j]] = c_i * marginal[[i, j]];
1121 }
1122 for j in 0..p_g {
1123 effective_logslope[[i, j]] = logslope_factor * logslope[[i, j]];
1124 }
1125 }
1126 if effective_logslope.iter().all(|v| v.abs() <= f64::EPSILON) {
1127 return Ok(None);
1128 }
1129
1130 let mut gram = fast_xt_diag_x(&effective_logslope, row_metric);
1131 let gram_scale = gram.diag().iter().copied().fold(0.0_f64, f64::max);
1132 if !gram_scale.is_finite() || gram_scale <= 0.0 {
1133 return Ok(None);
1134 }
1135 let projection_ridge = (gram_scale * 1.0e-10).max(f64::EPSILON);
1136 for i in 0..p_g {
1137 gram[[i, i]] += projection_ridge;
1138 }
1139 let cross = fast_xt_diag_y(&effective_logslope, row_metric, &marginal_effective);
1140 let gram_view = FaerArrayView::new(&gram);
1141 let factor = factorize_symmetricwith_fallback(gram_view.as_ref(), Side::Lower)
1142 .map_err(|e| format!("marginal/logslope overlap Gram factorization failed: {e}"))?;
1143 let rhsview = FaerArrayView::new(&cross);
1144 let coeffs_mat = factor.solve(rhsview.as_ref());
1145 let coeffs = Array2::from_shape_fn((p_g, p_m), |(i, j)| coeffs_mat[(i, j)]);
1146 let projected_marginal = fast_ab(&effective_logslope, &coeffs);
1147 let mut penalty = fast_xt_diag_y(&marginal_effective, row_metric, &projected_marginal);
1148 penalty = (&penalty + &penalty.t()) * 0.5;
1149 let max_abs = penalty.iter().map(|v| v.abs()).fold(0.0_f64, f64::max);
1150 if !max_abs.is_finite() || max_abs <= 1.0e-12 {
1151 return Ok(None);
1152 }
1153 Ok(Some(penalty))
1154 }
1155
1156 #[test]
1164 pub(crate) fn effective_reduction_drops_score_weighted_confound_raw_audit_misses() {
1165 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1167 let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 1.0, 2.0, 2.0, 3.0, 9.0]).unwrap();
1168 let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1169 let w = Array1::<f64>::ones(3);
1170 let zero = Array1::<f64>::zeros(3);
1171
1172 let reparam = match reduced_logslope_transform_effective(
1175 m.view(),
1176 g.view(),
1177 &z,
1178 &w,
1179 &zero,
1180 &zero,
1181 0.0,
1182 0.0,
1183 1.0,
1184 )
1185 .expect("effective reduction must succeed")
1186 {
1187 ReducedLogslopeOutcome::Reduced(t) => t,
1188 other => panic!(
1189 "effective audit must reduce the score-weighted confound (raw audit would not), got {}",
1190 match other {
1191 ReducedLogslopeOutcome::FullRank => "FullRank",
1192 ReducedLogslopeOutcome::FullyConfounded => "FullyConfounded",
1193 ReducedLogslopeOutcome::Reduced(_) => unreachable!(),
1194 }
1195 ),
1196 };
1197 assert_eq!(
1198 reparam.ncols(),
1199 1,
1200 "exactly one effective-identifiable logslope direction should survive"
1201 );
1202
1203 let g_eff = {
1207 let mut e = Array2::<f64>::zeros((3, 2));
1208 for i in 0..3 {
1209 for j in 0..2 {
1210 e[[i, j]] = z[i] * g[[i, j]];
1211 }
1212 }
1213 e
1214 };
1215 let img = g_eff.dot(&reparam.column(0));
1216 let mean = img.iter().sum::<f64>() / 3.0;
1217 let var = img.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / 3.0;
1218 assert!(
1219 var > 1.0e-6,
1220 "kept direction must be the identifiable (non-constant) effective column, var={var}"
1221 );
1222 }
1223
1224 #[test]
1231 pub(crate) fn effective_reduction_fully_confounded_single_column_is_distinct_outcome() {
1232 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1233 let g = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 2.0, 3.0]).unwrap();
1234 let z = Array1::from_vec(vec![1.0, 0.5, 1.0 / 3.0]);
1235 let w = Array1::<f64>::ones(3);
1236 let zero = Array1::<f64>::zeros(3);
1237 let outcome = reduced_logslope_transform_effective(
1238 m.view(),
1239 g.view(),
1240 &z,
1241 &w,
1242 &zero,
1243 &zero,
1244 0.0,
1245 0.0,
1246 1.0,
1247 )
1248 .expect("effective reduction must succeed");
1249 assert!(
1250 matches!(outcome, ReducedLogslopeOutcome::FullyConfounded),
1251 "fully effective-confounded logslope must surface the distinct FullyConfounded outcome"
1252 );
1253 }
1254
1255 #[test]
1258 pub(crate) fn effective_reduction_no_confound_returns_none() {
1259 let m = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1260 let g = Array2::<f64>::from_shape_vec((3, 2), vec![1.0, 0.0, 0.0, 1.0, 0.0, 0.0]).unwrap();
1262 let z = Array1::from_vec(vec![1.0, 2.0, 3.0]);
1263 let w = Array1::<f64>::ones(3);
1264 let zero = Array1::<f64>::zeros(3);
1265 let outcome = reduced_logslope_transform_effective(
1266 m.view(),
1267 g.view(),
1268 &z,
1269 &w,
1270 &zero,
1271 &zero,
1272 0.0,
1273 0.0,
1274 1.0,
1275 )
1276 .expect("effective reduction must succeed");
1277 assert!(
1278 matches!(outcome, ReducedLogslopeOutcome::FullRank),
1279 "no effective confound ⇒ FullRank (raw design kept unchanged)"
1280 );
1281 }
1282
1283 #[test]
1284 pub(crate) fn spatial_joint_setup_counts_only_learned_penalties_in_rho() {
1285 let data = Array2::<f64>::zeros((3, 1));
1286 let empty_terms = TermCollectionSpec {
1287 linear_terms: Vec::new(),
1288 random_effect_terms: Vec::new(),
1289 smooth_terms: Vec::new(),
1290 };
1291 let setup = joint_setup(
1292 data.view(),
1293 &empty_terms,
1294 &empty_terms,
1295 2,
1296 3,
1297 Some(2.5),
1298 &[0.4],
1299 &SpatialLengthScaleOptimizationOptions::default(),
1300 )
1301 .expect("empty spatial geometry is valid");
1302
1303 assert_eq!(
1304 setup.rho_dim(),
1305 6,
1306 "BMS spatial setup rho holds every learned marginal/logslope/auxiliary penalty; the #461 absorber ridge occupies the trailing marginal slot"
1307 );
1308 assert_eq!(
1309 setup.theta0()[1],
1310 2.5,
1311 "absorber ridge seeds the trailing marginal rho coordinate at the ln(n) leakage scale"
1312 );
1313 }
1314
1315 #[test]
1316 pub(crate) fn overlap_penalty_targets_score_weighted_logslope_span() {
1317 let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1318 Array2::from_shape_vec((4, 1), vec![0.0, 1.0, 2.0, 3.0]).unwrap(),
1319 ));
1320 let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1321 Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1322 ));
1323 let z = Array1::from_vec(vec![0.0, 1.0, 2.0, 3.0]);
1324 let row_metric = Array1::ones(4);
1325 let offsets = Array1::zeros(4);
1326
1327 let penalty = marginal_logslope_overlap_penalty(
1328 &marginal,
1329 &logslope,
1330 &z,
1331 &row_metric,
1332 &offsets,
1333 &offsets,
1334 0.0,
1335 0.0,
1336 1.0,
1337 )
1338 .expect("overlap penalty should build")
1339 .expect("marginal signal lies in the pilot logslope Jacobian span");
1340
1341 assert_eq!(penalty.dim(), (1, 1));
1342 assert!((penalty[[0, 0]] - 14.0).abs() < 1.0e-6);
1343 }
1344
1345 #[test]
1346 pub(crate) fn overlap_penalty_skips_weight_orthogonal_channels() {
1347 let marginal = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1348 Array2::from_shape_vec((4, 1), vec![-1.0, 1.0, -1.0, 1.0]).unwrap(),
1349 ));
1350 let logslope = DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1351 Array2::from_shape_vec((4, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap(),
1352 ));
1353 let z = Array1::ones(4);
1354 let row_metric = Array1::ones(4);
1355 let offsets = Array1::zeros(4);
1356
1357 let penalty = marginal_logslope_overlap_penalty(
1358 &marginal,
1359 &logslope,
1360 &z,
1361 &row_metric,
1362 &offsets,
1363 &offsets,
1364 0.0,
1365 0.0,
1366 1.0,
1367 )
1368 .expect("overlap penalty should build");
1369
1370 assert!(penalty.is_none());
1371 }
1372
1373 fn dense_marginal_design(
1381 x: Array2<f64>,
1382 intercept_range: std::ops::Range<usize>,
1383 linear_ranges: Vec<(String, std::ops::Range<usize>)>,
1384 ) -> TermCollectionDesign {
1385 let nrows = x.nrows();
1386 TermCollectionDesign {
1387 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(x)),
1388 affine_offset: Array1::zeros(nrows),
1389 penalties: Vec::new(),
1390 nullspace_dims: Vec::new(),
1391 penaltyinfo: Vec::new(),
1392 dropped_penaltyinfo: Vec::new(),
1393 coefficient_lower_bounds: None,
1394 linear_constraints: None,
1395 intercept_range,
1396 linear_ranges,
1397 linear_function_masses: Vec::new(),
1398 random_effect_ranges: Vec::new(),
1399 random_effect_levels: Vec::new(),
1400 smooth: gam_terms::smooth::SmoothDesign {
1401 term_designs: Vec::new(),
1402 penalties: Vec::new(),
1403 nullspace_dims: Vec::new(),
1404 penaltyinfo: Vec::new(),
1405 dropped_penaltyinfo: Vec::new(),
1406 terms: Vec::new(),
1407 coefficient_lower_bounds: None,
1408 linear_constraints: None,
1409 },
1410 }
1411 }
1412
1413 fn linear_term(name: &str, feature_col: usize) -> LinearTermSpec {
1414 LinearTermSpec {
1415 name: name.to_string(),
1416 feature_col,
1417 feature_cols: vec![feature_col],
1418 categorical_levels: vec![],
1419 double_penalty: false,
1420 coefficient_geometry: LinearCoefficientGeometry::default(),
1421 coefficient_min: None,
1422 coefficient_max: None,
1423 frozen_function_mass: None,
1424 }
1425 }
1426
1427 fn empty_spec() -> TermCollectionSpec {
1428 TermCollectionSpec {
1429 linear_terms: Vec::new(),
1430 random_effect_terms: Vec::new(),
1431 smooth_terms: Vec::new(),
1432 }
1433 }
1434
1435 #[test]
1442 pub(crate) fn runaway_guard_silent_when_huge_beta_cancels_to_bounded_eta() {
1443 let x = Array2::<f64>::from_shape_vec((4, 2), vec![1.0; 8]).unwrap();
1445 let design = dense_marginal_design(x, 0..0, Vec::new());
1446 let beta = Array1::from_vec(vec![60.0, -60.0]);
1447
1448 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1449 beta.view(),
1450 &design,
1451 &empty_spec(),
1452 true,
1453 "regression-fixture",
1454 );
1455 assert!(
1456 msg.is_none(),
1457 "huge cancelling β with bounded fitted η must NOT trip the runaway guard; got {msg:?}"
1458 );
1459 }
1460
1461 #[test]
1465 pub(crate) fn runaway_guard_fires_when_fitted_eta_exceeds_threshold() {
1466 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1467 let design = dense_marginal_design(x, 0..0, Vec::new());
1468 let beta = Array1::from_vec(vec![40.0]);
1469
1470 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1471 beta.view(),
1472 &design,
1473 &empty_spec(),
1474 true,
1475 "separation-fixture",
1476 )
1477 .expect("fitted |η|∞=40 ≥ 35 must trip the runaway guard");
1478
1479 assert!(msg.contains("marginal/logslope runaway"));
1480 assert!(msg.contains("|η|∞"));
1481 assert!(msg.contains("4.000e1"));
1482 assert!(msg.contains("score is correlated with the shared surface covariates"));
1483 assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1484 assert!(msg.contains("KKT certificate"));
1485 }
1486
1487 #[test]
1491 pub(crate) fn runaway_guard_names_unpenalized_parametric_direction_via_fitted_eta() {
1492 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1493 let design = dense_marginal_design(x, 0..0, vec![("sex".to_string(), 0..1)]);
1494 let mut spec = empty_spec();
1495 spec.linear_terms.push(linear_term("sex", 0));
1496 let beta = Array1::from_vec(vec![41.0]);
1497
1498 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1499 beta.view(),
1500 &design,
1501 &spec,
1502 true,
1503 "parametric-fixture",
1504 )
1505 .expect("parametric fitted |η|∞=41 ≥ 35 must trip the runaway guard");
1506
1507 assert!(msg.contains("unpenalized parametric marginal direction"));
1508 assert!(msg.contains("|η|∞"));
1509 assert!(msg.contains("robust Jeffreys curvature path is already installed"));
1510 assert!(msg.contains("not a Matérn/Duchon polynomial-nullspace"));
1511 }
1512
1513 #[test]
1517 pub(crate) fn runaway_guard_silent_for_nonconverged_but_bounded_eta() {
1518 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1519 let design = dense_marginal_design(x, 0..0, Vec::new());
1520 let beta = Array1::from_vec(vec![5.0]);
1521
1522 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1523 beta.view(),
1524 &design,
1525 &empty_spec(),
1526 false,
1527 "nonconverged-fixture",
1528 );
1529 assert!(
1530 msg.is_none(),
1531 "bounded fitted η must not raise the separation error even when the inner solve did not converge; got {msg:?}"
1532 );
1533 }
1534
1535 #[test]
1538 pub(crate) fn runaway_guard_fires_for_nonconverged_separating_eta() {
1539 let x = Array2::<f64>::from_shape_vec((3, 1), vec![1.0, 1.0, 1.0]).unwrap();
1540 let design = dense_marginal_design(x, 0..0, Vec::new());
1541 let beta = Array1::from_vec(vec![50.0]);
1542
1543 let msg = bernoulli_marginal_slope_runaway_error_from_beta(
1544 beta.view(),
1545 &design,
1546 &empty_spec(),
1547 false,
1548 "nonconverged-separating-fixture",
1549 )
1550 .expect("separating |η|∞ at non-convergence must still trip the guard");
1551
1552 assert!(msg.contains(
1553 "the inner solve failed while already carrying a separation-scale predictor"
1554 ));
1555 }
1556
1557 #[test]
1570 pub(crate) fn bms_block_jacobians_self_compute_at_audit_empty_beta_nonzero_logslope_baseline() {
1571 use std::sync::Arc;
1572 let n = 4usize;
1573 let marginal =
1574 Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1575 let logslope =
1576 Arc::new(Array2::<f64>::from_shape_vec((n, 1), vec![1.0, 1.0, 1.0, 1.0]).unwrap());
1577 let offset_m = Array1::<f64>::zeros(n);
1578 let g_baseline = 0.3_f64;
1581 let offset_s = Array1::<f64>::from_elem(n, g_baseline);
1582 let z = Arc::new(Array1::from_vec(vec![-0.7, 0.2, 0.9, 1.4]));
1583 let s = 1.0_f64;
1584
1585 let beta: Vec<f64> = Vec::new();
1587 let state = FamilyLinearizationState {
1588 beta: &beta,
1589 family_scalars: None,
1590 channel_hessian: None,
1591 probit_frailty_scale: s,
1592 };
1593
1594 let marginal_jac = BmsMarginalJacobian::new(
1595 Arc::clone(&marginal),
1596 Arc::clone(&logslope),
1597 offset_m.clone(),
1598 offset_s.clone(),
1599 1,
1600 );
1601 let j_m = marginal_jac
1602 .effective_jacobian_rows(&state, 0..n)
1603 .expect("BMS marginal Jacobian must self-compute at audit empty β (gam#370)");
1604 let c_expected = (1.0 + (s * g_baseline).powi(2)).sqrt();
1606 assert_eq!(j_m.dim(), (n, 1));
1607 for i in 0..n {
1608 assert!(
1609 (j_m[[i, 0]] - c_expected).abs() < 1e-12,
1610 "marginal J[{i}] = {} != closed-form c_i = {c_expected}",
1611 j_m[[i, 0]]
1612 );
1613 }
1614
1615 let logslope_jac = BmsLogslopeJacobian::new(
1616 Arc::clone(&marginal),
1617 Arc::clone(&logslope),
1618 offset_m,
1619 offset_s,
1620 Arc::clone(&z),
1621 1,
1622 );
1623 let j_s = logslope_jac
1624 .effective_jacobian_rows(&state, 0..n)
1625 .expect("BMS logslope Jacobian must self-compute at audit empty β (gam#370)");
1626 assert_eq!(j_s.dim(), (n, 1));
1629 for i in 0..n {
1630 let expected = s * z[i];
1631 assert!(
1632 (j_s[[i, 0]] - expected).abs() < 1e-12,
1633 "logslope J[{i}] = {} != closed-form factor {expected}",
1634 j_s[[i, 0]]
1635 );
1636 assert!(j_s[[i, 0]].is_finite());
1637 }
1638 }
1639}
1640
1641pub(crate) fn build_marginal_blockspec_bms(
1642 design: &TermCollectionDesign,
1643 baseline: f64,
1644 offset: &Array1<f64>,
1645 rho: Array1<f64>,
1646 beta_hint: Option<Array1<f64>>,
1647 logslope_design: &TermCollectionDesign,
1648 logslope_offset: &Array1<f64>,
1649 logslope_baseline: f64,
1650 p_marginal: usize,
1651 influence_columns: Option<&Array2<f64>>,
1652) -> Result<ParameterBlockSpec, String> {
1653 let offset_m = offset + baseline;
1654 let offset_s = logslope_offset + logslope_baseline;
1655 let raw_marginal_dense = design
1656 .design
1657 .try_to_dense_arc("build_marginal_blockspec_bms::marginal")?;
1658 let marginal_dense =
1659 widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1660 let logslope_dense = logslope_design
1661 .design
1662 .try_to_dense_arc("build_marginal_blockspec_bms::logslope")?;
1663 let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsMarginalJacobian {
1664 marginal_dense: Arc::clone(&marginal_dense),
1665 logslope_dense,
1666 offset_m: offset_m.clone(),
1667 offset_s,
1668 p_marginal,
1669 });
1670 let (penalties, nullspace_dims, initial_log_lambdas) =
1671 marginal_penalties_with_influence_ridge(design, &rho, influence_columns)?;
1672 Ok(ParameterBlockSpec {
1673 name: "marginal_surface".to_string(),
1674 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1675 (*marginal_dense).clone(),
1676 )),
1677 offset: offset_m,
1678 penalties,
1679 nullspace_dims,
1680 initial_log_lambdas,
1681 initial_beta: widen_marginal_beta_hint(beta_hint, p_marginal),
1682 gauge_priority: GAUGE_PRIORITY_MARGINAL,
1695 jacobian_callback: Some(callback),
1696 stacked_design: None,
1697 stacked_offset: None,
1698 })
1699}
1700
1701pub(crate) fn build_logslope_blockspec_bms(
1702 design: &TermCollectionDesign,
1703 baseline: f64,
1704 offset: &Array1<f64>,
1705 rho: Array1<f64>,
1706 beta_hint: Option<Array1<f64>>,
1707 marginal_design: &TermCollectionDesign,
1708 marginal_offset: &Array1<f64>,
1709 marginal_baseline: f64,
1710 z: Arc<Array1<f64>>,
1711 p_marginal: usize,
1712 influence_columns: Option<&Array2<f64>>,
1713) -> Result<ParameterBlockSpec, String> {
1714 let offset_s = offset + baseline;
1715 let offset_m = marginal_offset + marginal_baseline;
1716 let raw_marginal_dense = marginal_design
1717 .design
1718 .try_to_dense_arc("build_logslope_blockspec_bms::marginal")?;
1719 let marginal_dense =
1724 widen_marginal_dense_with_influence(&raw_marginal_dense, influence_columns)?;
1725 let logslope_dense = design
1726 .design
1727 .try_to_dense_arc("build_logslope_blockspec_bms::logslope")?;
1728 let callback: Arc<dyn BlockEffectiveJacobian> = Arc::new(BmsLogslopeJacobian {
1729 marginal_dense,
1730 logslope_dense: Arc::clone(&logslope_dense),
1731 offset_m,
1732 offset_s: offset_s.clone(),
1733 z,
1734 p_marginal,
1735 });
1736 Ok(ParameterBlockSpec {
1737 name: "logslope_surface".to_string(),
1738 design: DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
1739 (*logslope_dense).clone(),
1740 )),
1741 offset: offset_s,
1742 penalties: design.penalties_as_penalty_matrix(),
1743 nullspace_dims: design.nullspace_dims.clone(),
1744 initial_log_lambdas: rho,
1745 initial_beta: beta_hint,
1746 gauge_priority: GAUGE_PRIORITY_LOGSLOPE,
1754 jacobian_callback: Some(callback),
1755 stacked_design: None,
1756 stacked_offset: None,
1757 })
1758}
1759
1760pub(crate) fn build_deviation_aux_blockspec(
1761 name: &str,
1762 prepared: &DeviationPrepared,
1763 rho: Array1<f64>,
1764 beta_hint: Option<Array1<f64>>,
1765) -> Result<ParameterBlockSpec, String> {
1766 let mut block = prepared.block.clone();
1767 block.initial_log_lambdas = Some(rho);
1768 let candidate_beta = beta_hint.or_else(|| Some(Array1::<f64>::zeros(block.design.ncols())));
1769 block.initial_beta = candidate_beta
1770 .map(|beta| {
1771 let zero = Array1::<f64>::zeros(beta.len());
1772 project_monotone_feasible_beta(&prepared.runtime, &zero, &beta, name)
1773 })
1774 .transpose()?;
1775 let mut spec = block.intospec(name)?;
1776 spec.gauge_priority = match name {
1786 "link_dev" => GAUGE_PRIORITY_LINK_DEV,
1787 "score_warp_dev" => GAUGE_PRIORITY_SCORE_WARP_DEV,
1794 _ => GAUGE_PRIORITY_DEVIATION_DEFAULT,
1795 };
1796 Ok(spec)
1797}
1798
1799pub(crate) fn push_deviation_aux_blockspecs(
1800 blocks: &mut Vec<ParameterBlockSpec>,
1801 rho: &Array1<f64>,
1802 cursor: &mut usize,
1803 score_warp_prepared: Option<&DeviationPrepared>,
1804 link_dev_prepared: Option<&DeviationPrepared>,
1805 score_warp_beta_hint: Option<Array1<f64>>,
1806 link_dev_beta_hint: Option<Array1<f64>>,
1807) -> Result<(), String> {
1808 fn take_rho_slice(
1809 rho: &Array1<f64>,
1810 cursor: &mut usize,
1811 count: usize,
1812 block_name: &str,
1813 ) -> Result<Array1<f64>, String> {
1814 let start = *cursor;
1815 let end = start.checked_add(count).ok_or_else(|| {
1816 format!("{block_name} penalty-rho range overflow: start={start}, count={count}")
1817 })?;
1818 if end > rho.len() {
1819 return Err(format!(
1820 "{block_name} penalty-rho range {start}..{end} exceeds rho length {}",
1821 rho.len()
1822 ));
1823 }
1824 let slice = rho.slice(s![start..end]).to_owned();
1825 *cursor = end;
1826 Ok(slice)
1827 }
1828
1829 if let Some(prepared) = score_warp_prepared {
1830 let rho_h = take_rho_slice(
1831 rho,
1832 cursor,
1833 prepared.block.penalties.len(),
1834 "score_warp_dev",
1835 )?;
1836 blocks.push(build_deviation_aux_blockspec(
1837 "score_warp_dev",
1838 prepared,
1839 rho_h,
1840 score_warp_beta_hint,
1841 )?);
1842 }
1843 if let Some(prepared) = link_dev_prepared {
1844 let rho_w = take_rho_slice(rho, cursor, prepared.block.penalties.len(), "link_dev")?;
1845 blocks.push(build_deviation_aux_blockspec(
1846 "link_dev",
1847 prepared,
1848 rho_w,
1849 link_dev_beta_hint,
1850 )?);
1851 }
1852 Ok(())
1853}
1854
1855#[cfg(test)]
1856mod deviation_penalty_layout_tests {
1857 use super::*;
1858
1859 fn prepared_with_penalty_orders(orders: Vec<usize>) -> DeviationPrepared {
1860 let seed = Array1::linspace(-1.0, 1.0, 48);
1861 let config = DeviationBlockConfig {
1862 degree: 3,
1863 num_internal_knots: 6,
1864 penalty_order: *orders.first().expect("test requires a penalty order"),
1865 penalty_orders: orders,
1866 double_penalty: false,
1867 monotonicity_eps: 0.0,
1868 };
1869 build_score_warp_deviation_block_from_seed(&seed, &config)
1870 .expect("test deviation block must build")
1871 }
1872
1873 #[test]
1874 fn composed_score_link_influence_rho_layout_advances_by_emitted_counts_2315() {
1875 let score_warp = prepared_with_penalty_orders(vec![1, 2]);
1880 let link_dev = prepared_with_penalty_orders(vec![1, 2, 3]);
1881 assert_eq!(score_warp.block.penalties.len(), 2);
1882 assert_eq!(link_dev.block.penalties.len(), 3);
1883
1884 let rho = Array1::from_vec(vec![-101.0, 11.0, 12.0, 21.0, 22.0, 23.0, 31.0]);
1887 let mut cursor = 1usize;
1888 let mut blocks = Vec::new();
1889 push_deviation_aux_blockspecs(
1890 &mut blocks,
1891 &rho,
1892 &mut cursor,
1893 Some(&score_warp),
1894 Some(&link_dev),
1895 None,
1896 None,
1897 )
1898 .expect("composed deviation layout must be realized");
1899
1900 assert_eq!(blocks.len(), 2);
1901 assert_eq!(blocks[0].name, "score_warp_dev");
1902 assert_eq!(
1903 blocks[0].initial_log_lambdas.as_slice(),
1904 Some(&[11.0, 12.0][..])
1905 );
1906 assert_eq!(blocks[1].name, "link_dev");
1907 assert_eq!(
1908 blocks[1].initial_log_lambdas.as_slice(),
1909 Some(&[21.0, 22.0, 23.0][..])
1910 );
1911 assert_eq!(
1912 cursor, 6,
1913 "the next consumer must start after every emitted deviation penalty"
1914 );
1915
1916 let influence_rho = rho.slice(s![cursor..cursor + 1]).to_owned();
1917 assert_eq!(influence_rho.as_slice(), Some(&[31.0][..]));
1918 assert_ne!(
1919 influence_rho[0], blocks[1].initial_log_lambdas[0],
1920 "the influence absorber must not reuse link_dev's first rho coordinate"
1921 );
1922 }
1923}
1924
1925fn inner_fit(
1926 family: &BernoulliMarginalSlopeFamily,
1927 blocks: &[ParameterBlockSpec],
1928 options: &BlockwiseFitOptions,
1929) -> Result<UnifiedFitResult, String> {
1930 let mut options = options.clone();
1931 options.use_outer_hessian = false;
1936 options.outer_tol = options.outer_tol.max(2.0e-5);
1937 fit_custom_family(family, blocks, &options).map_err(|e| e.to_string())
1938}
1939
1940fn inner_fit_from_certified_outer(
1941 family: &BernoulliMarginalSlopeFamily,
1942 blocks: &[ParameterBlockSpec],
1943 options: &BlockwiseFitOptions,
1944 mode: CustomFamilyJointHyperModeSelection,
1945 theta: &Array1<f64>,
1946 outer: &gam_solve::rho_optimizer::CertifiedOuterResult,
1947) -> Result<UnifiedFitResult, String> {
1948 let mut options = crate::outer_subsample::exact_outer_options_for_row_set(
1949 options,
1950 &crate::row_kernel::RowSet::All,
1951 );
1952 options.use_outer_hessian = false;
1953 options.outer_tol = options.outer_tol.max(2.0e-5);
1954 fit_custom_family_fixed_log_lambdas_from_mode_selection(
1955 family, blocks, &options, mode, theta, outer,
1956 )
1957 .map_err(|error| error.to_string())
1958}
1959
1960pub fn fit_bernoulli_marginal_slope_terms(
1961 data: ArrayView2<'_, f64>,
1962 spec: BernoulliMarginalSlopeTermSpec,
1963 options: &BlockwiseFitOptions,
1964 kappa_options: &SpatialLengthScaleOptimizationOptions,
1965 policy: &gam_runtime::resource::ResourcePolicy,
1966) -> Result<BernoulliMarginalSlopeFitResult, String> {
1967 let mut spec = spec;
1968 let data_view = data;
1969 validate_spec(data_view, &spec)?;
1970 let mjs_frozen_marginal =
1979 gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.marginalspec);
1980 let mjs_frozen_logslope =
1981 gam_terms::smooth::freeze_measure_jet_length_scale_learning(&mut spec.logslopespec);
1982 if mjs_frozen_marginal + mjs_frozen_logslope > 0 {
1983 log::info!(
1984 "[BMS spatial] froze measure-jet length-scale learning on {} marginal + {} log-slope \
1985 term(s): the coupled surface keeps ℓ at its conditioned auto value (#1116)",
1986 mjs_frozen_marginal,
1987 mjs_frozen_logslope
1988 );
1989 }
1990 let mut effective_kappa_options = kappa_options.clone();
1991 let kappa_locked_marginal =
2001 gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.marginalspec);
2002 let kappa_locked_logslope =
2003 gam_terms::smooth::all_spatial_terms_kappa_fixed(&spec.logslopespec);
2004 if effective_kappa_options.enabled && kappa_locked_marginal && kappa_locked_logslope {
2005 log::info!(
2006 "[BMS spatial] disabling κ/ψ optimization: every spatial term has an \
2007 explicit length_scale and no anisotropy; user-supplied kernel scale is fixed"
2008 );
2009 effective_kappa_options.enabled = false;
2010 }
2011 let flex_spatial_pilot_path = (spec.score_warp.is_some() || spec.link_dev.is_some())
2012 && spec.y.len() >= BMS_FLEX_SPATIAL_OUTER_PILOT_ROW_THRESHOLD
2013 && effective_kappa_options.enabled;
2014 if flex_spatial_pilot_path {
2015 let marginal_terms = spatial_length_scale_term_indices(&spec.marginalspec);
2016 let logslope_terms = spatial_length_scale_term_indices(&spec.logslopespec);
2017 let marginal_updates = apply_spatial_anisotropy_pilot_initializer(
2018 data_view,
2019 &mut spec.marginalspec,
2020 &marginal_terms,
2021 effective_kappa_options.pilot_subsample_threshold,
2022 &effective_kappa_options,
2023 )
2024 .map_err(|error| error.to_string())?;
2025 let logslope_updates = apply_spatial_anisotropy_pilot_initializer(
2026 data_view,
2027 &mut spec.logslopespec,
2028 &logslope_terms,
2029 effective_kappa_options.pilot_subsample_threshold,
2030 &effective_kappa_options,
2031 )
2032 .map_err(|error| error.to_string())?;
2033 effective_kappa_options.enabled = false;
2034 log::info!(
2035 "[BMS spatial] n={} flex=true pilot_geometry_updates={} iterative_spatial_outer=false reason=large-flex-spatial-pilot",
2036 spec.y.len(),
2037 marginal_updates + logslope_updates,
2038 );
2039 }
2040 let (z_standardized, z_normalization) = standardize_latent_z_with_policy(
2041 &spec.z,
2042 &spec.weights,
2043 "bernoulli-marginal-slope",
2044 &spec.latent_z_policy,
2045 )?;
2046 spec.z = z_standardized;
2047 {
2084 let marginal_seeded = crate::fit_orchestration::drivers::seed_measure_jet_auto_ranges(
2085 data_view,
2086 spec.y.view(),
2087 spec.weights.view(),
2088 &mut spec.marginalspec,
2089 );
2090 let logslope_seeded = match
2091 crate::fit_orchestration::drivers::marginal_slope_logslope_screen_response(
2092 spec.y.view(),
2093 spec.z.view(),
2094 spec.weights.view(),
2095 ) {
2096 Some(surrogate) => crate::fit_orchestration::drivers::seed_measure_jet_auto_ranges(
2097 data_view,
2098 surrogate.view(),
2099 spec.weights.view(),
2100 &mut spec.logslopespec,
2101 ),
2102 None => 0,
2103 };
2104 if marginal_seeded + logslope_seeded > 0 {
2105 log::info!(
2106 "[BMS spatial] #2750 screened the representer range of {marginal_seeded} marginal \
2107 + {logslope_seeded} log-slope auto measure-jet term(s) against the response \
2108 before the BMS design build"
2109 );
2110 }
2111 }
2112 let sigma_learnable = matches!(
2113 &spec.frailty,
2114 FrailtySpec::GaussianShift {
2115 scale: FrailtyScale::Learned { .. }
2116 }
2117 );
2118 let initial_sigma = match &spec.frailty {
2119 FrailtySpec::GaussianShift {
2120 scale: FrailtyScale::Fixed { sigma },
2121 } => Some(*sigma),
2122 FrailtySpec::GaussianShift {
2123 scale: FrailtyScale::Learned { initial_sigma },
2124 } => Some(*initial_sigma),
2125 FrailtySpec::None => None,
2126 FrailtySpec::HazardMultiplier { .. } => {
2127 return Err(
2128 "internal: validate_spec should have rejected unsupported marginal-slope frailty"
2129 .to_string(),
2130 );
2131 }
2132 };
2133 let probit_scale = probit_frailty_scale(initial_sigma);
2134 let (_raw_joint_designs, mut joint_specs) = build_term_collection_designs_and_freeze_joint(
2135 data_view,
2136 &[spec.marginalspec.clone(), spec.logslopespec.clone()],
2137 )
2138 .map_err(|e| e.to_string())?;
2139 let marginalspec_boot = joint_specs.remove(0);
2140 let logslopespec_boot = joint_specs.remove(0);
2141 let (mut joint_designs, _) = build_term_collection_designs_and_freeze_joint(
2156 data_view,
2157 &[marginalspec_boot.clone(), logslopespec_boot.clone()],
2158 )
2159 .map_err(|e| format!("failed to rebuild frozen probe BMS joint designs: {e}"))?;
2160 let marginal_design = joint_designs.remove(0);
2161 let logslope_design = joint_designs.remove(0);
2162 spec.marginal_offset = marginal_design
2163 .compose_offset(spec.marginal_offset.view(), "BMS marginal block")
2164 .map_err(|error| error.to_string())?;
2165 spec.logslope_offset = logslope_design
2166 .compose_offset(spec.logslope_offset.view(), "BMS logslope block")
2167 .map_err(|error| error.to_string())?;
2168 let absorber_active = spec
2175 .score_influence_jacobian
2176 .as_ref()
2177 .is_some_and(|j| j.ncols() > 0);
2178 let conditioning_dense = if absorber_active {
2179 None
2180 } else {
2181 Some(
2182 marginal_design
2183 .design
2184 .try_to_dense_arc("bernoulli marginal-slope conditional latent-z gate")?,
2185 )
2186 };
2187 let (latent_measure, latent_z_calibration, latent_measure_build) =
2188 build_latent_measure_with_geometry(
2189 &spec.z,
2190 &spec.weights,
2191 &spec.latent_z_policy,
2192 conditioning_dense.as_ref().map(|d| d.view()),
2193 )?;
2194 if latent_measure.is_empirical() && sigma_learnable {
2195 return Err("empirical latent-measure marginal-slope calibration requires fixed GaussianShift sigma; learnable sigma derivatives must be fit under the standard-normal latent measure"
2196 .to_string());
2197 }
2198
2199 let y = Arc::new(spec.y.clone());
2200 let weights = Arc::new(spec.weights.clone());
2201 let z = match &latent_z_calibration {
2206 LatentMeasureCalibration::None => Arc::new(spec.z.clone()),
2207 LatentMeasureCalibration::RankInverseNormal(cal) => {
2208 Arc::new(cal.apply_to_training(&spec.z)?)
2209 }
2210 LatentMeasureCalibration::ConditionalLocationScale(cal) => {
2211 let a_block = conditioning_dense.as_ref().ok_or_else(|| {
2214 "conditional latent calibration requires the marginal conditioning block"
2215 .to_string()
2216 })?;
2217 Arc::new(cal.apply(spec.z.view(), a_block.view())?)
2218 }
2219 };
2220 let z_train = z.as_ref();
2221 let pilot_baseline = pooled_probit_baseline(&spec.y, z_train, &spec.weights)?;
2222 let baseline = (
2223 bernoulli_marginal_slope_eta_from_probability(
2224 &spec.base_link,
2225 normal_cdf(pilot_baseline.0),
2226 "bernoulli marginal-slope baseline link inversion",
2227 )?,
2228 pilot_baseline.1 / probit_scale,
2229 );
2230
2231 let rigid_pilot_eta = rigid_pooled_probit_pilot_eta(
2274 &spec.base_link,
2275 z_train,
2276 &spec.marginal_offset,
2277 &spec.logslope_offset,
2278 baseline.0,
2279 baseline.1,
2280 probit_scale,
2281 )?;
2282 let cross_block_pilot_w_score_warp =
2283 pilot_irls_hessian_row_metric_at_eta(&rigid_pilot_eta, &spec.weights);
2284
2285 let influence_columns = if let Some(jac) = spec
2297 .score_influence_jacobian
2298 .as_ref()
2299 .filter(|j| j.ncols() > 0)
2300 {
2301 let protected_design = DesignMatrix::hstack(vec![
2302 marginal_design.design.clone(),
2303 logslope_design.design.clone(),
2304 ])
2305 .map_err(|e| {
2306 format!(
2307 "bernoulli marginal-slope influence-block protected projection stack failed to concatenate marginal + logslope design: {e}"
2308 )
2309 })?;
2310 let protected_dense_for_proj = protected_design
2311 .try_to_dense_arc("bernoulli marginal-slope influence-block protected projection")?;
2312 let protected_dense = protected_dense_for_proj.as_ref();
2313 if jac.nrows() != protected_dense.nrows() {
2314 return Err(format!(
2315 "influence block: Jacobian has {} rows, protected design has {}",
2316 jac.nrows(),
2317 protected_dense.nrows()
2318 ));
2319 }
2320 let rigid_logslope_at_rows = &spec.logslope_offset + baseline.1;
2331 let residualized = crate::marginal_slope_orthogonal::residualized_influence_block(
2332 jac,
2333 z_train,
2334 &rigid_logslope_at_rows,
2335 probit_scale,
2336 protected_dense.view(),
2337 &cross_block_pilot_w_score_warp,
2338 )?;
2339 Some(residualized)
2340 } else {
2341 None
2342 };
2343 let mut cross_block_warnings: Vec<CrossBlockIdentifiabilityWarning> = Vec::new();
2344 let score_warp_prepared = if let Some(cfg) = spec.score_warp.as_ref() {
2345 use super::deviation_runtime::ParametricAnchorBlock;
2346 let mut prepared = build_score_warp_deviation_block_from_seed(z_train, cfg)?;
2347 let outcome = install_compiled_flex_block_into_runtime(
2352 &mut prepared,
2353 z_train,
2354 cfg,
2355 &[
2356 (&marginal_design.design, ParametricAnchorBlock::Marginal),
2357 (&logslope_design.design, ParametricAnchorBlock::Logslope),
2358 ],
2359 &[],
2360 &cross_block_pilot_w_score_warp,
2361 )?;
2362 match outcome {
2363 FlexCompileOutcome::Reparameterised => Some(prepared),
2364 FlexCompileOutcome::FullyAliased { reason } => {
2365 cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2371 candidate_label: "score_warp",
2372 anchor_summary: "marginal+logslope".to_string(),
2373 reason,
2374 });
2375 Some(prepared)
2376 }
2377 }
2378 } else {
2379 None
2380 };
2381 let link_dev_prepared = if let Some(cfg) = spec.link_dev.as_ref() {
2407 let eta_pilot = pilot_eta_for_link_dev_orthogonalisation(
2408 &spec.base_link,
2409 &spec.y,
2410 z_train,
2411 &spec.weights,
2412 &marginal_design.design,
2413 &spec.marginal_offset,
2414 &spec.logslope_offset,
2415 baseline.0,
2416 baseline.1,
2417 probit_scale,
2418 )?;
2419 let link_dev_seed = padded_deviation_seed(&eta_pilot, 1.0, 0.5);
2420 let mut prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
2421 &link_dev_seed,
2422 &eta_pilot,
2423 cfg,
2424 )?;
2425 let score_warp_anchor_design = score_warp_prepared
2462 .as_ref()
2463 .map(|sw| sw.runtime.design_at_training_with_residual(z_train))
2464 .transpose()?;
2465 use super::deviation_runtime::ParametricAnchorBlock;
2466 let parametric_anchors: [(&DesignMatrix, ParametricAnchorBlock); 2] = [
2467 (&marginal_design.design, ParametricAnchorBlock::Marginal),
2468 (&logslope_design.design, ParametricAnchorBlock::Logslope),
2469 ];
2470 let flex_anchor_slot: Option<&Array2<f64>> = score_warp_anchor_design.as_ref();
2471 let flex_anchors: Vec<&Array2<f64>> = flex_anchor_slot.into_iter().collect();
2472 let cross_block_pilot_w_link_dev =
2477 pilot_irls_hessian_row_metric_at_eta(&eta_pilot, &spec.weights);
2478 let outcome = install_compiled_flex_block_into_runtime(
2479 &mut prepared,
2480 &eta_pilot,
2481 cfg,
2482 ¶metric_anchors,
2483 &flex_anchors,
2484 &cross_block_pilot_w_link_dev,
2485 )?;
2486 match outcome {
2487 FlexCompileOutcome::Reparameterised => Some(prepared),
2488 FlexCompileOutcome::FullyAliased { reason } => {
2489 cross_block_warnings.push(CrossBlockIdentifiabilityWarning {
2495 candidate_label: "link_deviation",
2496 anchor_summary: "marginal+logslope+score_warp".to_string(),
2497 reason,
2498 });
2499 Some(prepared)
2500 }
2501 }
2502 } else {
2503 None
2504 };
2505 let extra_rho0 = {
2506 let mut out = Vec::new();
2507 if let Some(ref prepared) = score_warp_prepared {
2508 out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2509 }
2510 if let Some(ref prepared) = link_dev_prepared {
2511 out.extend(std::iter::repeat_n(0.0, prepared.block.penalties.len()));
2512 }
2513 out
2514 };
2515 let logslope_reduced_reparam: Option<ReducedLogslopeReparam> = build_reduced_logslope_reparam(
2528 &marginal_design,
2529 &logslope_design,
2530 z.as_ref(),
2531 &cross_block_pilot_w_score_warp,
2532 &spec.marginal_offset,
2533 &spec.logslope_offset,
2534 baseline.0,
2535 baseline.1,
2536 probit_scale,
2537 )?;
2538 let reduce_logslope_design =
2544 |logslope_design: &TermCollectionDesign| -> Result<TermCollectionDesign, String> {
2545 match logslope_reduced_reparam.as_ref() {
2546 Some(reparam) => reparameterize_logslope_design_reduced(logslope_design, reparam),
2547 None => Ok(logslope_design.clone()),
2548 }
2549 };
2550
2551 let absorber_slots = usize::from(influence_columns.is_some());
2555 let absorber_rho0 = influence_columns
2556 .as_ref()
2557 .map(|_| influence_absorber_log_lambda(spec.z.len()).clamp(-12.0, 12.0));
2558 let marginal_penalty_count = marginal_design.penalties.len() + absorber_slots;
2559 let setup = joint_setup(
2560 data_view,
2561 &marginalspec_boot,
2562 &logslopespec_boot,
2563 marginal_penalty_count,
2564 logslope_design.penalties.len(),
2565 absorber_rho0,
2566 &extra_rho0,
2567 &effective_kappa_options,
2568 )
2569 .map_err(|error| error.to_string())?;
2570 let setup = if sigma_learnable {
2571 setup.with_auxiliary(
2572 Array1::from_vec(vec![initial_sigma.expect("learnable sigma seed").ln()]),
2573 Array1::from_vec(vec![0.01_f64.ln()]),
2574 Array1::from_vec(vec![5.0_f64.ln()]),
2575 )
2576 } else {
2577 setup
2578 };
2579 let final_sigma_cell = std::cell::Cell::new(initial_sigma);
2580 let exact_mode_branch = RefCell::new(ExactCoefficientModeBranch::default());
2581 let runaway_error = RefCell::new(None::<String>);
2582 let pending_beta_seed = RefCell::new(None::<Array1<f64>>);
2589 let hints = RefCell::new(ThetaHints::default());
2590 let score_warp_runtime = score_warp_prepared.as_ref().map(|p| p.runtime.clone());
2591 let link_dev_runtime = link_dev_prepared.as_ref().map(|p| p.runtime.clone());
2592
2593 let build_blocks = |rho: &Array1<f64>,
2594 marginal_design: &TermCollectionDesign,
2595 logslope_design: &TermCollectionDesign|
2596 -> Result<Vec<ParameterBlockSpec>, String> {
2597 let hints = hints.borrow();
2598 let mut cursor = 0usize;
2599 let logslope_design_reduced = reduce_logslope_design(logslope_design)?;
2606 let logslope_design = &logslope_design_reduced;
2607 let marginal_rho_len = marginal_design.penalties.len() + absorber_slots;
2612 let rho_marginal = rho.slice(s![cursor..cursor + marginal_rho_len]).to_owned();
2613 cursor += marginal_rho_len;
2614 let rho_logslope = rho
2615 .slice(s![cursor..cursor + logslope_design.penalties.len()])
2616 .to_owned();
2617 cursor += logslope_design.penalties.len();
2618 let p_m = marginal_design.design.ncols()
2619 + influence_columns.as_ref().map(|z| z.ncols()).unwrap_or(0);
2620 let mut blocks = vec![
2621 build_marginal_blockspec_bms(
2622 marginal_design,
2623 baseline.0,
2624 &spec.marginal_offset,
2625 rho_marginal,
2626 hints.marginal_beta.clone(),
2627 logslope_design,
2628 &spec.logslope_offset,
2629 baseline.1,
2630 p_m,
2631 influence_columns.as_ref(),
2632 )?,
2633 build_logslope_blockspec_bms(
2634 logslope_design,
2635 baseline.1,
2636 &spec.logslope_offset,
2637 rho_logslope,
2638 hints.logslope_beta.clone(),
2639 marginal_design,
2640 &spec.marginal_offset,
2641 baseline.0,
2642 Arc::clone(&z),
2643 p_m,
2644 influence_columns.as_ref(),
2645 )?,
2646 ];
2647 push_deviation_aux_blockspecs(
2648 &mut blocks,
2649 rho,
2650 &mut cursor,
2651 score_warp_prepared.as_ref(),
2652 link_dev_prepared.as_ref(),
2653 hints.score_warp_beta.clone(),
2654 hints.link_dev_beta.clone(),
2655 )?;
2656 Ok(blocks)
2657 };
2658
2659 let intercept_warm_starts = new_intercept_warm_start_cache(y.len());
2660 let cell_moment_lru = new_cell_moment_lru_cache(policy);
2661 let cell_moment_cache_stats = new_cell_moment_cache_stats();
2662 let make_family = |marginal_design: &TermCollectionDesign,
2663 logslope_design: &TermCollectionDesign,
2664 sigma: Option<f64>|
2665 -> BernoulliMarginalSlopeFamily {
2666 let kernel_marginal_design = match influence_columns.as_ref() {
2672 Some(z_infl) => {
2673 let raw = marginal_design
2674 .design
2675 .try_to_dense_arc("make_family::widened-marginal")
2676 .expect("dense marginal design for influence widening");
2677 let widened = widen_marginal_dense_with_influence(&raw, Some(z_infl))
2678 .expect("widen marginal design with influence columns");
2679 DesignMatrix::Dense(gam_linalg::matrix::DenseDesignMatrix::from(
2680 (*widened).clone(),
2681 ))
2682 }
2683 None => marginal_design.design.clone(),
2684 };
2685 let kernel_logslope_design = reduce_logslope_design(logslope_design)
2691 .expect("reduce logslope design for family construction")
2692 .design;
2693 BernoulliMarginalSlopeFamily {
2694 y: Arc::clone(&y),
2695 weights: Arc::clone(&weights),
2696 z: Arc::clone(&z),
2697 latent_measure: latent_measure.clone(),
2698 gaussian_frailty_sd: sigma,
2699 base_link: spec.base_link.clone(),
2700 marginal_design: kernel_marginal_design,
2701 logslope_design: kernel_logslope_design,
2702 score_warp: score_warp_runtime.clone(),
2703 link_dev: link_dev_runtime.clone(),
2704 policy: policy.clone(),
2705 cell_moment_lru: Arc::clone(&cell_moment_lru),
2706 cell_moment_cache_stats: Arc::clone(&cell_moment_cache_stats),
2707 intercept_warm_starts: Some(Arc::clone(&intercept_warm_starts)),
2708 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2709 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
2710 }
2711 };
2712
2713 let marginal_terms = spatial_length_scale_term_indices(&marginalspec_boot);
2714 let logslope_terms = spatial_length_scale_term_indices(&logslopespec_boot);
2715 let marginal_has_spatial = !marginal_terms.is_empty();
2716 let logslope_has_spatial = !logslope_terms.is_empty();
2717 let analytic_joint_derivatives_available =
2718 marginal_has_spatial || logslope_has_spatial || setup.log_kappa_dim() == 0;
2719 if setup.log_kappa_dim() > 0 && !analytic_joint_derivatives_available {
2720 return Err("exact bernoulli marginal-slope spatial optimization requires analytic joint psi derivatives"
2721 .to_string());
2722 }
2723 let initial_rho = setup.theta0().slice(s![..setup.rho_dim()]).to_owned();
2724 let initial_blocks = build_blocks(&initial_rho, &marginal_design, &logslope_design)?;
2725 let initial_family = make_family(&marginal_design, &logslope_design, initial_sigma);
2726 let (joint_gradient, joint_hessian) =
2727 custom_family_outer_derivatives(&initial_family, &initial_blocks, options);
2728 let analytic_joint_gradient_available = analytic_joint_derivatives_available
2729 && matches!(joint_gradient, gam_problem::Derivative::Analytic);
2730 let analytic_joint_hessian_available =
2736 analytic_joint_derivatives_available && joint_hessian.is_analytic();
2737 let kappa_options_ref: &SpatialLengthScaleOptimizationOptions = &effective_kappa_options;
2738 let sigma_from_theta = |theta: &Array1<f64>| -> Option<f64> {
2739 if sigma_learnable {
2740 Some(theta[setup.rho_dim() + setup.log_kappa_dim()].exp())
2741 } else {
2742 initial_sigma
2743 }
2744 };
2745 let hyper_layout_cache = RefCell::new(
2746 None::<(
2747 Array1<f64>,
2748 crate::custom_family::SharedCustomFamilyHyperLayout,
2749 )>,
2750 );
2751 let theta_matches = |left: &Array1<f64>, right: &Array1<f64>| -> bool {
2752 left.len() == right.len()
2753 && left
2754 .iter()
2755 .zip(right.iter())
2756 .all(|(lhs, rhs)| lhs.to_bits() == rhs.to_bits())
2757 };
2758 let get_hyper_layout =
2759 |theta: &Array1<f64>,
2760 specs: &[TermCollectionSpec],
2761 designs: &[TermCollectionDesign]|
2762 -> Result<crate::custom_family::SharedCustomFamilyHyperLayout, String> {
2763 if let Some((cached_theta, cached_layout)) = hyper_layout_cache.borrow().as_ref()
2764 && theta_matches(cached_theta, theta)
2765 {
2766 return Ok(Arc::clone(cached_layout));
2767 }
2768
2769 let built = |specs: &[TermCollectionSpec],
2770 designs: &[TermCollectionDesign]|
2771 -> Result<
2772 Vec<Vec<crate::custom_family::CustomFamilyBlockPsiDerivative>>,
2773 String,
2774 > {
2775 let marginal_psi_derivs = if marginal_has_spatial {
2776 build_block_spatial_psi_derivatives(data_view, &specs[0], &designs[0])?
2777 .ok_or_else(|| {
2778 "bernoulli marginal-slope: marginal block has spatial terms \
2779 but spatial psi derivatives are unavailable"
2780 .to_string()
2781 })?
2782 } else {
2783 Vec::new()
2784 };
2785 let logslope_psi_derivs = if logslope_has_spatial {
2786 let built = if let Some(reparam) = logslope_reduced_reparam.as_ref() {
2787 let transform =
2788 CoefficientSpatialPsiBlockTransform::new(&reparam.transform)?;
2789 build_block_spatial_psi_derivatives_with_transform(
2790 data_view,
2791 &specs[1],
2792 &designs[1],
2793 &transform,
2794 )?
2795 } else {
2796 build_block_spatial_psi_derivatives(data_view, &specs[1], &designs[1])?
2797 };
2798 built.ok_or_else(|| {
2799 "bernoulli marginal-slope: logslope block has spatial terms \
2800 but spatial psi derivatives are unavailable"
2801 .to_string()
2802 })?
2803 } else {
2804 Vec::new()
2805 };
2806 let mut derivative_blocks = vec![marginal_psi_derivs, logslope_psi_derivs];
2807 if score_warp_runtime.is_some() {
2808 derivative_blocks.push(Vec::new());
2809 }
2810 if link_dev_runtime.is_some() {
2811 derivative_blocks.push(Vec::new());
2812 }
2813 Ok(derivative_blocks)
2814 }(specs, designs)?;
2815 let family_axes = if sigma_learnable { vec![0] } else { Vec::new() };
2816 let hyper_values = theta.slice(s![setup.rho_dim()..]).to_owned();
2817 let layout = Arc::new(crate::custom_family::CustomFamilyHyperLayout::new(
2818 built,
2819 family_axes,
2820 hyper_values,
2821 )?);
2822 hyper_layout_cache.replace(Some((theta.clone(), Arc::clone(&layout))));
2823 Ok(layout)
2824 };
2825
2826 let outer_policy = {
2831 let psi_dim = setup.theta0().len() - setup.rho_dim();
2832 initial_family.outer_derivative_policy(&initial_blocks, psi_dim, options)
2833 };
2834 let exact_spatial_outer_tol = kappa_options_ref.rel_tol.max(EXACT_SPATIAL_OUTER_TOL_FLOOR);
2835 let solved = optimize_spatial_length_scale_exact_joint(
2836 data_view,
2837 &[marginalspec_boot.clone(), logslopespec_boot.clone()],
2838 &[marginal_terms.clone(), logslope_terms.clone()],
2839 kappa_options_ref,
2840 &setup,
2841 gam_solve::seeding::SeedRiskProfile::GeneralizedLinear,
2842 analytic_joint_gradient_available,
2843 analytic_joint_hessian_available,
2844 true,
2845 None,
2846 outer_policy,
2847 |theta, specs: &[TermCollectionSpec], designs: &[TermCollectionDesign], provenance| {
2848 if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2849 return Err(err);
2850 }
2851 assert_eq!(
2852 specs.len(),
2853 designs.len(),
2854 "spatial joint optimizer must supply one spec per design",
2855 );
2856 let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2857 let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2858 let sigma = sigma_from_theta(theta);
2859 final_sigma_cell.set(sigma);
2860 let family = make_family(&designs[0], &designs[1], sigma);
2861 let fit = match provenance {
2862 SpatialFitProvenance::NoOuterOptimization => inner_fit(&family, &blocks, options)?,
2863 SpatialFitProvenance::Certified { outer, mode } => {
2864 inner_fit_from_certified_outer(&family, &blocks, options, mode, theta, outer)?
2865 }
2866 };
2867 if let Some(block) = fit.block_states.first()
2868 && let Some(err) = bernoulli_marginal_slope_runaway_error_from_beta(
2869 block.beta.view(),
2870 &designs[0],
2871 &specs[0],
2872 true,
2873 "final fit",
2874 )
2875 {
2876 runaway_error.replace(Some(err.clone()));
2877 return Err(err);
2878 }
2879 let mut hints_mut = hints.borrow_mut();
2880 let mut bidx = 0usize;
2881 if let Some(block) = fit.block_states.get(bidx) {
2882 hints_mut.marginal_beta = Some(block.beta.clone());
2883 }
2884 bidx += 1;
2885 if let Some(block) = fit.block_states.get(bidx) {
2886 hints_mut.logslope_beta = Some(block.beta.clone());
2887 }
2888 bidx += 1;
2889 if score_warp_prepared.is_some() {
2890 if let Some(block) = fit.block_states.get(bidx) {
2891 hints_mut.score_warp_beta = Some(block.beta.clone());
2892 }
2893 bidx += 1;
2894 }
2895 if link_dev_prepared.is_some()
2896 && let Some(block) = fit.block_states.get(bidx)
2897 {
2898 hints_mut.link_dev_beta = Some(block.beta.clone());
2899 }
2900 Ok(fit)
2901 },
2902 |theta,
2903 specs: &[TermCollectionSpec],
2904 designs: &[TermCollectionDesign],
2905 eval_mode,
2906 row_set: &crate::row_kernel::RowSet,
2907 _| {
2908 if let Some(err) = runaway_error.borrow().as_ref().cloned() {
2909 return Err(err);
2910 }
2911 use gam_problem::EvalMode;
2912 static BMS_OUTER_EVAL_ROWSET_LOGGED: std::sync::Once = std::sync::Once::new();
2919 BMS_OUTER_EVAL_ROWSET_LOGGED.call_once(|| {
2920 let row_set_rows = match row_set {
2921 crate::row_kernel::RowSet::All => spec.y.len(),
2922 crate::row_kernel::RowSet::Subsample { rows, .. } => rows.len(),
2923 };
2924 log::debug!(
2925 "[BMS exact outer eval] mode={eval_mode:?} row_set_rows={row_set_rows}"
2926 );
2927 });
2928 let rho = theta.slice(s![..setup.rho_dim()]).to_owned();
2929 let blocks = build_blocks(&rho, &designs[0], &designs[1])?;
2930 if let Some(beta_seed) = pending_beta_seed.borrow_mut().take() {
2934 let widths: Vec<usize> = blocks.iter().map(|b| b.design.ncols()).collect();
2935 match CustomFamilyWarmStart::from_cached_beta(&widths, &beta_seed) {
2936 Ok(ws) => {
2937 if !exact_mode_branch.borrow_mut().install_seed(ws) {
2938 log::debug!(
2939 "[BMS] ignored a late outer-cache coefficient seed after the exact mode branch froze"
2940 );
2941 }
2942 }
2943 Err(e) => {
2944 log::warn!(
2945 "[BMS] outer ρ-cache β-warm-start rejected: {e}; falling back to cold β"
2946 );
2947 }
2948 }
2949 }
2950 let sigma = sigma_from_theta(theta);
2951 final_sigma_cell.set(sigma);
2952 let family = make_family(&designs[0], &designs[1], sigma);
2953 let hyper_layout = get_hyper_layout(theta, specs, designs)?;
2954 let effective_mode = match eval_mode {
2958 EvalMode::ValueGradientHessian if !analytic_joint_hessian_available => {
2959 EvalMode::ValueAndGradient
2960 }
2961 other => other,
2962 };
2963 let tolerance_options =
2964 joint_hyper_options_for_outer_tolerance(options, exact_spatial_outer_tol);
2965 let eval_options = crate::outer_subsample::exact_outer_options_for_row_set(
2966 &tolerance_options,
2967 row_set,
2968 );
2969 let (froze, candidates) = exact_mode_branch
2970 .borrow_mut()
2971 .candidates(effective_mode, &rho);
2972 if froze {
2973 log::info!(
2974 "[BMS] froze deterministic exact coefficient-mode branch at the first derivative-bearing outer evaluation"
2975 );
2976 }
2977 let selection = evaluate_custom_family_joint_hyper_best_mode_shared(
2978 &family,
2979 &blocks,
2980 &eval_options,
2981 &rho,
2982 hyper_layout,
2983 &candidates,
2984 effective_mode,
2985 )
2986 .map_err(|error| error.to_string())?;
2987 if let Some(err) = bernoulli_marginal_slope_runaway_error(
2988 &selection.result.warm_start,
2989 &designs[0],
2990 &specs[0],
2991 selection.result.inner_converged,
2992 "exact outer evaluation",
2993 ) {
2994 runaway_error.replace(Some(err.clone()));
2995 return Err(err);
2996 }
2997 exact_mode_branch
2998 .borrow_mut()
2999 .record_value(eval_mode, selection.result.warm_start.clone());
3000 if !selection.result.inner_converged {
3001 return Err(
3002 "exact bernoulli marginal-slope inner solve did not converge".to_string(),
3003 );
3004 }
3005 if matches!(eval_mode, EvalMode::ValueGradientHessian)
3006 && analytic_joint_hessian_available
3007 && !selection.result.outer_hessian.is_analytic()
3008 {
3009 return Err("exact bernoulli marginal-slope joint [rho, psi] objective did not return an outer Hessian"
3010 .to_string());
3011 }
3012 Ok(ExactJointEvaluation {
3013 objective: selection.result.objective,
3014 gradient: selection.result.gradient.clone(),
3015 hessian: selection.result.outer_hessian.clone(),
3016 mode: selection,
3017 })
3018 },
3019 |_, _, _, _| {
3020 Err::<ExactJointEfsEvaluation<CustomFamilyJointHyperModeSelection>, String>(
3021 "bernoulli marginal-slope EFS callback invoked even though fixed-point optimization is disabled for beta-dependent exact curvature".to_string(),
3022 )
3023 },
3024 crate::marginal_slope_shared::make_beta_seed_validator(&pending_beta_seed),
3025 )?;
3026
3027 let mut resolved_specs = solved.resolved_specs;
3028 let mut designs = solved.designs;
3029 let mut solved_fit = solved.fit;
3030 let (latent_z_rank_int_calibration, latent_z_conditional_calibration) =
3066 match latent_z_calibration {
3067 LatentMeasureCalibration::None => (None, None),
3068 LatentMeasureCalibration::RankInverseNormal(cal) => (Some(cal), None),
3069 LatentMeasureCalibration::ConditionalLocationScale(cal) => (None, Some(cal)),
3070 };
3071 let flex_active = score_warp_runtime.is_some() || link_dev_runtime.is_some();
3115 let empirical_channel = classify_empirical_generated_regressor_channel(
3116 &latent_measure,
3117 latent_measure_build.as_ref(),
3118 flex_active,
3119 );
3120 if solved_fit.covariance_conditional.is_some()
3126 && latent_z_conditional_calibration.is_some()
3127 && let EmpiricalGeneratedRegressorChannel::Unavailable {
3128 latent_measure: latent_measure_label,
3129 unavailable_channel,
3130 } = &empirical_channel
3131 {
3132 solved_fit.covariance_conditional = None;
3133 solved_fit.covariance_corrected = None;
3134 if let Some(inference) = solved_fit.inference.as_mut() {
3135 inference.beta_covariance = None;
3136 inference.beta_standard_errors = None;
3137 inference.beta_covariance_corrected = None;
3138 inference.beta_standard_errors_corrected = None;
3139 }
3140 let declined = gam_solve::estimate::CovarianceDeclined::
3141 BmsGeneratedRegressorLatentMeasureNotStandardNormal {
3142 latent_measure: latent_measure_label.clone(),
3143 unavailable_channel: unavailable_channel.clone(),
3144 };
3145 log::warn!("[BMS latent-z] {}", declined.explain());
3146 solved_fit.artifacts.covariance_declined = Some(declined);
3147 }
3148 if let Some(cal) = latent_z_conditional_calibration.as_ref()
3149 && let Some(vb) = solved_fit.covariance_conditional.clone()
3150 {
3151 let p_beta = vb.nrows();
3152 let calibration_marginal_dense = marginal_design
3153 .design
3154 .try_to_dense_arc("bms generated-regressor marginal design")?;
3155 if p_beta != vb.ncols() {
3156 return Err(format!(
3157 "bms generated-regressor: covariance_conditional must be square, got {}×{}",
3158 vb.nrows(),
3159 vb.ncols()
3160 ));
3161 }
3162 let correction_family =
3163 make_family(&marginal_design, &logslope_design, final_sigma_cell.get());
3164 let s = if flex_active {
3165 correction_family.flex_score_zeta_sensitivity(
3166 &solved_fit.block_states,
3167 options,
3168 p_beta,
3169 )?
3170 } else {
3171 let score_marginal_dense = correction_family
3175 .marginal_design
3176 .try_to_dense_arc("bms generated-regressor fitted marginal design")?;
3177 let score_logslope_dense = correction_family
3178 .logslope_design
3179 .try_to_dense_arc("bms generated-regressor fitted logslope design")?;
3180 let p_score = score_marginal_dense.ncols() + score_logslope_dense.ncols();
3181 if p_beta != p_score {
3182 return Err(format!(
3183 "bms generated-regressor rigid covariance/frame mismatch: covariance width {p_beta} != marginal({}) + logslope({})",
3184 score_marginal_dense.ncols(),
3185 score_logslope_dense.ncols()
3186 ));
3187 }
3188 let marginal_eta = &solved_fit.block_states[0].eta;
3189 let slope_eta = &solved_fit.block_states[1].eta;
3190 let probit_scale = probit_frailty_scale(final_sigma_cell.get());
3191 match &empirical_channel {
3192 EmpiricalGeneratedRegressorChannel::Empirical(build) => {
3200 let grid = match &latent_measure {
3201 LatentMeasureKind::GlobalEmpirical { grid } => grid,
3202 _ => {
3203 return Err(
3204 "bms generated-regressor: an empirical build record without a \
3205 global-empirical measure"
3206 .to_string(),
3207 );
3208 }
3209 };
3210 let channels = rigid_empirical_score_zeta_channels(
3211 &spec.base_link,
3212 marginal_eta,
3213 slope_eta,
3214 z.as_ref(),
3215 y.as_ref(),
3216 weights.as_ref(),
3217 probit_scale,
3218 grid,
3219 score_marginal_dense.view(),
3220 score_logslope_dense.view(),
3221 p_beta,
3222 )?;
3223 let cross_row = build.node_zeta_vjp(channels.node.view())?;
3224 if cross_row.nrows() != channels.direct.nrows() {
3225 return Err(format!(
3226 "bms generated-regressor: the empirical cross-row channel has {} \
3227 rows against the fit's {}",
3228 cross_row.nrows(),
3229 channels.direct.nrows()
3230 ));
3231 }
3232 log::info!(
3233 "[BMS latent-z] empirical generated-regressor channels: nodes={} \
3234 |S_direct|_max={:.6e} |D^T U_Q^T|_max={:.6e}",
3235 grid.nodes.len(),
3236 channels
3237 .direct
3238 .iter()
3239 .fold(0.0_f64, |acc, v| acc.max(v.abs())),
3240 cross_row.iter().fold(0.0_f64, |acc, v| acc.max(v.abs())),
3241 );
3242 channels.direct + cross_row
3243 }
3244 EmpiricalGeneratedRegressorChannel::ClosedForm => {
3245 rigid_standard_normal_score_zeta_sensitivity(
3246 &spec.base_link,
3247 marginal_eta,
3248 slope_eta,
3249 z.as_ref(),
3250 y.as_ref(),
3251 weights.as_ref(),
3252 probit_scale,
3253 score_marginal_dense.view(),
3254 score_logslope_dense.view(),
3255 p_beta,
3256 )?
3257 }
3258 EmpiricalGeneratedRegressorChannel::Unavailable {
3261 latent_measure: measure,
3262 unavailable_channel: channel,
3263 } => {
3264 return Err(format!(
3265 "bms generated-regressor: reached the correction with a {measure} \
3266 measure whose channel is unavailable ({channel})"
3267 ));
3268 }
3269 }
3270 };
3271 let correction = cal.generated_regressor_correction(
3275 s.view(),
3276 spec.z.view(),
3277 calibration_marginal_dense.view(),
3278 vb.view(),
3279 )?;
3280 if let Some(cov) = solved_fit.covariance_conditional.as_mut() {
3281 *cov = &*cov + &correction;
3282 }
3283 if let Some(cov) = solved_fit.covariance_corrected.as_mut() {
3284 *cov = &*cov + &correction;
3285 }
3286 log::info!(
3287 "[BMS latent-z] Murphy–Topel generated-regressor SE correction applied: \
3288 p_beta={p_beta} flex_active={flex_active} theta1_dim={} max_diag_inflation={:.3e}",
3289 cal.theta1_dim(),
3290 (0..p_beta)
3291 .map(|i| correction[[i, i]])
3292 .fold(0.0_f64, f64::max),
3293 );
3294 }
3295 if let Some(reparam) = logslope_reduced_reparam.as_ref() {
3299 let r = reparam.reduced_cols();
3300 if let Some(block) = solved_fit.blocks.get_mut(1)
3301 && block.beta.len() == r
3302 {
3303 block.beta = reparam.recover_original_logslope_beta(&block.beta)?;
3304 }
3305 if let Some(state) = solved_fit.block_states.get_mut(1)
3306 && state.beta.len() == r
3307 {
3308 state.beta = reparam.recover_original_logslope_beta(&state.beta)?;
3309 }
3310 }
3311 Ok(BernoulliMarginalSlopeFitResult {
3324 fit: solved_fit,
3325 marginalspec_resolved: resolved_specs.remove(0),
3326 logslopespec_resolved: resolved_specs.remove(0),
3327 marginal_design: designs.remove(0),
3328 logslope_design: designs.remove(0),
3329 baseline_marginal: baseline.0,
3330 baseline_logslope: baseline.1,
3331 z_normalization,
3332 latent_measure,
3333 score_warp_runtime,
3334 link_dev_runtime,
3335 gaussian_frailty_sd: final_sigma_cell.get(),
3336 cross_block_warnings,
3337 latent_z_rank_int_calibration,
3338 latent_z_conditional_calibration,
3339 })
3340}