1use super::family::clamp_bernoulli_link_probability;
2use super::*;
3use gam_linalg::faer_ndarray::FaerEigh;
4use gam_linalg::matrix::{FiniteSignedWeightsView, LinearOperator};
5use gam_math::jet_scalar::SymmetricQuadraticCoefficients;
6use gam_math::probability::normal_logcdf_derivatives;
7use gam_row_macros::row_program;
8use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, escalate_ridge};
9
10pub(crate) fn standardize_latent_z_with_policy(
11 z: &Array1<f64>,
12 weights: &Array1<f64>,
13 context: &str,
14 policy: &LatentZPolicy,
15) -> Result<(Array1<f64>, LatentZNormalization), String> {
16 if z.len() != weights.len() {
17 return Err(format!(
18 "{context} latent-score normalization length mismatch: z={}, weights={}",
19 z.len(),
20 weights.len()
21 ));
22 }
23 let weight_sum = weights.iter().copied().sum::<f64>();
24 let weight_sq_sum = weights.iter().map(|&w| w * w).sum::<f64>();
25 if !(weight_sum.is_finite()
26 && weight_sum > 0.0
27 && weight_sq_sum.is_finite()
28 && weight_sq_sum > 0.0)
29 {
30 return Err(format!("{context} requires positive finite total weight"));
31 }
32 let effective_n = weight_sum * weight_sum / weight_sq_sum;
33 if !(effective_n.is_finite() && effective_n > 1.0) {
34 return Err(format!(
35 "{context} requires at least two effective observations for latent-score normalization"
36 ));
37 }
38 let mean = z
39 .iter()
40 .zip(weights.iter())
41 .map(|(&zi, &wi)| wi * zi)
42 .sum::<f64>()
43 / weight_sum;
44 let var = z
45 .iter()
46 .zip(weights.iter())
47 .map(|(&zi, &wi)| wi * (zi - mean) * (zi - mean))
48 .sum::<f64>()
49 / weight_sum;
50 let sd = var.sqrt();
51 if !(sd.is_finite() && sd > BMS_VARIANCE_FLOOR) {
52 return Err(format!(
53 "{context} requires z with positive finite weighted standard deviation"
54 ));
55 }
56 let target_norm = match policy.normalization {
57 LatentZNormalizationMode::None => LatentZNormalization { mean: 0.0, sd: 1.0 },
58 LatentZNormalizationMode::FitWeighted => LatentZNormalization { mean, sd },
59 LatentZNormalizationMode::Frozen {
60 mean: frozen_mean,
61 sd: frozen_sd,
62 } => LatentZNormalization {
63 mean: frozen_mean,
64 sd: frozen_sd,
65 },
66 };
67 let mean_tol = policy.mean_tol_multiplier / effective_n.sqrt();
68 let sd_tol = policy.sd_tol_multiplier / (2.0 * (effective_n - 1.0).max(1.0)).sqrt();
69 let check_msg = || {
70 format!(
71 "{context} requires z to already be approximately latent N(0,1) before identification normalization; got mean={mean:.6e}, sd={sd:.6e}, effective_n={effective_n:.1}, allowed_mean={mean_tol:.3e}, allowed_sd={sd_tol:.3e}"
72 )
73 };
74 if mean.abs() > mean_tol || (sd - 1.0).abs() > sd_tol {
75 match policy.check_mode {
76 LatentZCheckMode::Strict => return Err(check_msg()),
77 LatentZCheckMode::WarnOnly => log::warn!("{}", check_msg()),
78 LatentZCheckMode::Off => {}
79 }
80 }
81
82 let normalization = target_norm;
83 let z_std = normalization.apply(z, context)?;
84 let std_mean = z_std
90 .iter()
91 .zip(weights.iter())
92 .map(|(&zi, &wi)| wi * zi)
93 .sum::<f64>()
94 / weight_sum;
95 let std_var = (z_std
96 .iter()
97 .zip(weights.iter())
98 .map(|(&zi, &wi)| wi * (zi - std_mean) * (zi - std_mean))
99 .sum::<f64>()
100 / weight_sum)
101 .max(f64::MIN_POSITIVE);
102 let skew = z_std
103 .iter()
104 .zip(weights.iter())
105 .map(|(&zi, &wi)| wi * (zi - std_mean).powi(3))
106 .sum::<f64>()
107 / weight_sum
108 / std_var.powf(1.5);
109 let kurt = z_std
110 .iter()
111 .zip(weights.iter())
112 .map(|(&zi, &wi)| wi * (zi - std_mean).powi(4))
113 .sum::<f64>()
114 / weight_sum
115 / (std_var * std_var)
116 - 3.0;
117 if skew.abs() > policy.max_abs_skew || kurt.abs() > policy.max_abs_excess_kurtosis {
118 let msg = format!(
119 "{context} requires z to be approximately Gaussian after identification normalization; got skewness={skew:.3}, excess_kurtosis={kurt:.3}"
120 );
121 match policy.check_mode {
122 LatentZCheckMode::Strict => return Err(msg),
123 LatentZCheckMode::WarnOnly => log::warn!("{}", msg),
124 LatentZCheckMode::Off => {}
125 }
126 }
127 if skew.abs() > 0.75 || kurt.abs() > 2.0 {
128 log::warn!(
129 "{context}: z has skewness={skew:.3} and excess kurtosis={kurt:.3}; latent-measure auto-selection will use empirical calibration unless stricter diagnostics pass"
130 );
131 }
132 Ok((z_std, normalization))
133}
134
135pub fn padded_deviation_seed(seed: &Array1<f64>, min_iqr: f64, pad_fraction: f64) -> Array1<f64> {
136 let mut sorted = seed.to_vec();
137 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
138
139 if sorted.len() < 4 {
140 return seed.clone();
141 }
142
143 let n = sorted.len();
144 let q1 = sorted[n / 4];
145 let q3 = sorted[3 * n / 4];
146 let iqr = (q3 - q1).max(min_iqr);
147 let pad = pad_fraction * iqr;
148
149 let mut out = seed.to_vec();
150 out.push(sorted[0] - pad);
151 out.push(sorted[n - 1] + pad);
152 Array1::from_vec(out)
153}
154
155const POOLED_PILOT_MAX_NEWTON_ITERS: usize = 50;
166pub(crate) const POOLED_PILOT_RIDGE_INIT: f64 = 1e-8;
168pub(crate) const POOLED_PILOT_DET_FLOOR: f64 = 1e-18;
171pub(crate) const POOLED_PILOT_RIDGE_GROWTH: f64 = 10.0;
173pub(crate) const POOLED_PILOT_RIDGE_MAX: f64 = 1e6;
176const POOLED_PILOT_MAX_BACKTRACKS: usize = 25;
178pub(crate) const POOLED_PILOT_BACKTRACK_SHRINK: f64 = 0.5;
180pub(crate) const POOLED_PILOT_STALL_TOL: f64 = 1e-10;
183pub(crate) const POOLED_PILOT_MIN_ABS_SLOPE: f64 = 1e-6;
186
187pub(super) fn pooled_probit_baseline(
188 y: &Array1<f64>,
189 z: &Array1<f64>,
190 weights: &Array1<f64>,
191) -> Result<(f64, f64), String> {
192 if y.len() != z.len() || y.len() != weights.len() {
193 return Err(format!(
194 "pooled bernoulli-marginal-slope pilot length mismatch: y={}, z={}, weights={}",
195 y.len(),
196 z.len(),
197 weights.len()
198 ));
199 }
200 let weight_sum = weights.iter().copied().sum::<f64>();
201 if !weight_sum.is_finite() || weight_sum <= 0.0 {
202 return Err(
203 "pooled bernoulli-marginal-slope pilot requires positive finite total weight"
204 .to_string(),
205 );
206 }
207 let prevalence = y
208 .iter()
209 .zip(weights.iter())
210 .map(|(&yi, &wi)| yi * wi)
211 .sum::<f64>()
212 / weight_sum;
213 let prevalence = prevalence.clamp(1e-6, 1.0 - 1e-6);
214 let z_mean = z
215 .iter()
216 .zip(weights.iter())
217 .map(|(&zi, &wi)| zi * wi)
218 .sum::<f64>()
219 / weight_sum;
220 let z_var = z
221 .iter()
222 .zip(weights.iter())
223 .map(|(&zi, &wi)| wi * (zi - z_mean) * (zi - z_mean))
224 .sum::<f64>()
225 / weight_sum;
226 let yz_cov = y
227 .iter()
228 .zip(z.iter())
229 .zip(weights.iter())
230 .map(|((&yi, &zi), &wi)| wi * (yi - prevalence) * (zi - z_mean))
231 .sum::<f64>()
232 / weight_sum;
233 let mut beta0 = standard_normal_quantile(prevalence).map_err(|e| {
234 format!("failed to initialize pooled bernoulli-marginal-slope pilot intercept: {e}")
235 })?;
236 let mut beta1 = if z_var > BMS_VARIANCE_FLOOR {
237 yz_cov / z_var
238 } else {
239 0.0
240 };
241
242 let objective_grad_hess =
243 |intercept: f64, slope: f64| -> Result<(f64, f64, f64, f64, f64, f64), String> {
244 let mut obj = 0.0;
245 let mut g0 = 0.0;
246 let mut g1 = 0.0;
247 let mut h00 = 0.0;
248 let mut h01 = 0.0;
249 let mut h11 = 0.0;
250 for ((&yi, &zi), &wi) in y.iter().zip(z.iter()).zip(weights.iter()) {
251 if wi == 0.0 {
252 continue;
253 }
254 let eta = intercept + slope * zi;
255 let s = 2.0 * yi - 1.0;
256 let margin = s * eta;
257 let probit = normal_logcdf_derivatives(margin);
258 let logcdf = probit[0];
259 let lambda = probit[1];
260 let g_eta = -wi * s * lambda;
261 let h_eta = -wi * probit[2];
262 obj -= wi * logcdf;
263 g0 += g_eta;
264 g1 += g_eta * zi;
265 h00 += h_eta;
266 h01 += h_eta * zi;
267 h11 += h_eta * zi * zi;
268 }
269 Ok((obj, g0, g1, h00, h01, h11))
270 };
271
272 let mut obj_prev = f64::INFINITY;
273 for _ in 0..POOLED_PILOT_MAX_NEWTON_ITERS {
274 let (obj, g0, g1, h00, h01, h11) = objective_grad_hess(beta0, beta1)?;
275 if !obj.is_finite() || !g0.is_finite() || !g1.is_finite() {
276 return Err(
277 "pooled bernoulli-marginal-slope pilot produced non-finite objective or gradient"
278 .to_string(),
279 );
280 }
281 let grad_max = g0.abs().max(g1.abs());
282 if grad_max < BMS_DERIV_TOL {
283 break;
284 }
285 let ridge_trials = (POOLED_PILOT_RIDGE_MAX / POOLED_PILOT_RIDGE_INIT)
289 .log10()
290 .ceil() as usize
291 + 1;
292 let (step0, step1) = escalate_ridge(
293 RidgeSchedule {
294 initial: POOLED_PILOT_RIDGE_INIT,
295 growth: POOLED_PILOT_RIDGE_GROWTH,
296 max_escalations: ridge_trials,
297 },
298 |ridge| {
299 let h00_r = h00 + ridge;
300 let h11_r = h11 + ridge;
301 let det = h00_r * h11_r - h01 * h01;
302 if !(det.is_finite() && det.abs() > POOLED_PILOT_DET_FLOOR) {
303 return None;
304 }
305 let s0 = (h11_r * g0 - h01 * g1) / det;
306 let s1 = (-h01 * g0 + h00_r * g1) / det;
307 (s0.is_finite() && s1.is_finite()).then_some((s0, s1))
308 },
309 )
310 .map(|success| success.value)
311 .map_err(|_| "pooled bernoulli-marginal-slope pilot Hessian solve failed".to_string())?;
312 let accepted = backtracking_line_search::<_, String>(
313 BacktrackConfig {
314 contraction: POOLED_PILOT_BACKTRACK_SHRINK,
315 max_steps: POOLED_PILOT_MAX_BACKTRACKS,
316 ..BacktrackConfig::default()
317 },
318 |step_scale| {
319 let cand0 = beta0 - step_scale * step0;
320 let cand1 = beta1 - step_scale * step1;
321 let (cand_obj, _, _, _, _, _) = objective_grad_hess(cand0, cand1)?;
322 Ok(Some((cand_obj, (cand0, cand1))))
323 },
324 |_, cand_obj| cand_obj.is_finite() && cand_obj <= obj,
325 )?;
326 match accepted {
327 Some(step) => {
328 (beta0, beta1) = step.payload;
329 obj_prev = step.value;
330 }
331 None => {
332 if (obj_prev - obj).abs() < POOLED_PILOT_STALL_TOL {
333 break;
334 }
335 return Err("pooled bernoulli-marginal-slope pilot line search failed".to_string());
336 }
337 }
338 }
339 let a = beta0;
340 let b = if beta1.abs() < POOLED_PILOT_MIN_ABS_SLOPE {
342 if beta1.is_sign_negative() {
343 -POOLED_PILOT_MIN_ABS_SLOPE
344 } else {
345 POOLED_PILOT_MIN_ABS_SLOPE
346 }
347 } else {
348 beta1
349 };
350 Ok((a / (1.0 + b * b).sqrt(), b))
351}
352
353pub(super) fn pilot_irls_hessian_row_metric_at_eta(
393 eta_pilot: &Array1<f64>,
394 sample_weights: &Array1<f64>,
395) -> Array1<f64> {
396 let n = eta_pilot.len();
397 let mut w = Array1::<f64>::zeros(n);
398 for i in 0..n {
399 let eta = eta_pilot[i];
400 let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
401 let phi = normal_pdf(eta).max(1e-300);
402 let var = (mu * (1.0 - mu)).max(1e-300);
403 w[i] = sample_weights[i] * (phi * phi) / var;
404 }
405 w
406}
407
408pub(super) fn rigid_pooled_probit_pilot_eta(
415 base_link: &InverseLink,
416 z: &Array1<f64>,
417 marginal_offset: &Array1<f64>,
418 logslope_offset: &Array1<f64>,
419 baseline_marginal: f64,
420 baseline_logslope: f64,
421 probit_scale: f64,
422) -> Result<Array1<f64>, String> {
423 let n = z.len();
424 let mut out = Array1::<f64>::zeros(n);
425 for i in 0..n {
426 let a_pre = baseline_marginal + marginal_offset[i];
427 let b_pre = baseline_logslope + logslope_offset[i];
428 let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
429 .map_err(|e| format!("rigid_pooled_probit_pilot_eta marginal link map: {e}"))?
430 .q;
431 out[i] = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
432 }
433 Ok(out)
434}
435
436pub(crate) const PILOT_RIDGE_DIAG_FRACTION: f64 = 1e-6;
442pub(crate) const PILOT_RIDGE_DIAG_FLOOR: f64 = 1e-12;
445
446pub(super) fn pilot_eta_for_link_dev_orthogonalisation(
447 base_link: &InverseLink,
448 y: &Array1<f64>,
449 z: &Array1<f64>,
450 weights: &Array1<f64>,
451 marginal_design: &DesignMatrix,
452 marginal_offset: &Array1<f64>,
453 logslope_offset: &Array1<f64>,
454 baseline_marginal: f64,
455 baseline_logslope: f64,
456 probit_scale: f64,
457) -> Result<Array1<f64>, String> {
458 use gam_linalg::faer_ndarray::FaerCholesky;
459
460 let n = y.len();
461 if marginal_design.nrows() != n {
462 return Err(format!(
463 "pilot_eta_for_link_dev_orthogonalisation: marginal design has {} rows, expected {}",
464 marginal_design.nrows(),
465 n,
466 ));
467 }
468 let mut working_eta = Array1::<f64>::zeros(n);
469 let mut w_irls = Array1::<f64>::zeros(n);
470 let mut residual = Array1::<f64>::zeros(n);
471 for i in 0..n {
472 let a_pre = baseline_marginal + marginal_offset[i];
473 let b_pre = baseline_logslope + logslope_offset[i];
474 let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
475 .map_err(|e| {
476 format!("pilot_eta_for_link_dev_orthogonalisation marginal link map: {e}")
477 })?
478 .q;
479 let eta = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
480 working_eta[i] = eta;
481 let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
482 let phi = normal_pdf(eta).max(1e-300);
483 let var = (mu * (1.0 - mu)).max(1e-300);
484 w_irls[i] = weights[i] * (phi * phi) / var;
485 residual[i] = (y[i] - mu) / phi;
486 }
487 let p_marg = marginal_design.ncols();
488 if p_marg == 0 {
489 return Ok(working_eta);
490 }
491 let xtwr = marginal_design.compute_xtwy(&w_irls, &residual)?;
492 let mut xtwx =
493 marginal_design.xt_diag_x_signed_op(FiniteSignedWeightsView::try_from_array(&w_irls)?)?;
494 let trace_diag: f64 = (0..p_marg).map(|i| xtwx[[i, i]]).sum();
495 let ridge =
496 (trace_diag / p_marg as f64).max(PILOT_RIDGE_DIAG_FLOOR) * PILOT_RIDGE_DIAG_FRACTION;
497 for i in 0..p_marg {
498 xtwx[[i, i]] += ridge;
499 }
500 let factor = xtwx
501 .cholesky(faer::Side::Lower)
502 .map_err(|e| format!("pilot_eta_for_link_dev_orthogonalisation Cholesky failed: {e}"))?;
503 let delta_beta_marg = factor.solvevec(&xtwr);
504 let marg_contrib = marginal_design.dot(&delta_beta_marg);
505 Ok(&working_eta + &marg_contrib)
506}
507
508pub(super) fn joint_setup(
509 data: ArrayView2<'_, f64>,
510 marginalspec: &TermCollectionSpec,
511 logslopespec: &TermCollectionSpec,
512 marginal_penalties: usize,
513 logslope_penalties: usize,
514 absorber_rho0: Option<f64>,
515 extra_rho0: &[f64],
516 kappa_options: &SpatialLengthScaleOptimizationOptions,
517) -> Result<ExactJointHyperSetup, gam_terms::basis::BasisError> {
518 let marginal_terms = spatial_length_scale_term_indices(marginalspec);
519 let logslope_terms = spatial_length_scale_term_indices(logslopespec);
520 let rho_dim = marginal_penalties + logslope_penalties + extra_rho0.len();
521 let mut rho0vec = Array1::<f64>::zeros(rho_dim);
522 if let Some(seed) = absorber_rho0 {
526 assert!(
527 marginal_penalties > 0,
528 "an absorber rho0 seed requires at least one marginal penalty to land in"
529 );
530 rho0vec[marginal_penalties - 1] = seed;
531 }
532 for (idx, &value) in extra_rho0.iter().enumerate() {
533 rho0vec[marginal_penalties + logslope_penalties + idx] = value;
534 }
535 let rho_lower = Array1::<f64>::from_elem(rho_dim, -12.0);
536 let rho_upper = Array1::<f64>::from_elem(rho_dim, 12.0);
537 let marginal_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
538 marginalspec,
539 &marginal_terms,
540 kappa_options,
541 )
542 .reseed_from_data(data, marginalspec, &marginal_terms, kappa_options)?;
543 let logslope_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
544 logslopespec,
545 &logslope_terms,
546 kappa_options,
547 )
548 .reseed_from_data(data, logslopespec, &logslope_terms, kappa_options)?;
549 let mut values = marginal_kappa.as_array().to_vec();
550 values.extend(logslope_kappa.as_array().iter());
551 let marginal_dims = marginal_kappa.dims_per_term().to_vec();
552 let logslope_dims = logslope_kappa.dims_per_term().to_vec();
553 let mut dims = marginal_dims.clone();
554 dims.extend(logslope_dims.iter().copied());
555 let log_kappa0 = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(values), dims.clone());
556 let marginal_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
558 data,
559 marginalspec,
560 &marginal_terms,
561 &marginal_dims,
562 kappa_options,
563 )?;
564 let logslope_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
565 data,
566 logslopespec,
567 &logslope_terms,
568 &logslope_dims,
569 kappa_options,
570 )?;
571 let mut lower_vals = marginal_lower.as_array().to_vec();
572 lower_vals.extend(logslope_lower.as_array().iter());
573 let log_kappa_lower =
574 SpatialLogKappaCoords::new_with_dims(Array1::from_vec(lower_vals), dims.clone());
575 let marginal_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
576 data,
577 marginalspec,
578 &marginal_terms,
579 &marginal_dims,
580 kappa_options,
581 )?;
582 let logslope_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
583 data,
584 logslopespec,
585 &logslope_terms,
586 &logslope_dims,
587 kappa_options,
588 )?;
589 let mut upper_vals = marginal_upper.as_array().to_vec();
590 upper_vals.extend(logslope_upper.as_array().iter());
591 let log_kappa_upper = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(upper_vals), dims);
592 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
595 Ok(ExactJointHyperSetup::new(
596 rho0vec,
597 rho_lower,
598 rho_upper,
599 log_kappa0,
600 log_kappa_lower,
601 log_kappa_upper,
602 ))
603}
604
605#[inline]
606pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth_numeric(
607 signed_margin: f64,
608 weight: f64,
609) -> (f64, f64, f64, f64) {
610 if weight == 0.0 || signed_margin == f64::INFINITY {
611 return (0.0, 0.0, 0.0, 0.0);
612 }
613 if signed_margin.is_nan() {
614 return (f64::NAN, f64::NAN, f64::NAN, f64::NAN);
615 }
616 let d = normal_logcdf_derivatives(signed_margin);
617 (
618 -weight * d[1],
619 -weight * d[2],
620 -weight * d[3],
621 -weight * d[4],
622 )
623}
624
625pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth(
633 signed_margin: f64,
634 weight: f64,
635) -> Result<(f64, f64, f64, f64), String> {
636 if weight == 0.0 || signed_margin == f64::INFINITY {
637 return Ok((0.0, 0.0, 0.0, 0.0));
638 }
639 if !signed_margin.is_finite() {
640 return Err(format!(
641 "non-finite signed margin in exact probit derivative helper: {signed_margin}"
642 ));
643 }
644 Ok(signed_probit_neglog_derivatives_up_to_fourth_numeric(
645 signed_margin,
646 weight,
647 ))
648}
649
650#[inline]
672pub(crate) fn signed_probit_neglog_unary_stack(signed_margin: f64, weight: f64) -> [f64; 5] {
673 if weight == 0.0 || signed_margin == f64::INFINITY {
674 return [0.0; 5];
675 }
676 if signed_margin.is_nan() {
677 return [f64::NAN; 5];
678 }
679 let d = normal_logcdf_derivatives(signed_margin);
680 [
681 -weight * d[0],
682 -weight * d[1],
683 -weight * d[2],
684 -weight * d[3],
685 -weight * d[4],
686 ]
687}
688
689#[inline]
701pub(super) fn rigid_observed_slope(slope: f64, probit_scale: f64) -> f64 {
702 probit_scale * slope
703}
704
705#[inline]
706pub(super) fn rigid_observed_scale(logslope: f64, probit_scale: f64) -> f64 {
707 let observed_slope = rigid_observed_slope(logslope, probit_scale);
708 (1.0 + observed_slope * observed_slope).sqrt()
709}
710
711#[inline]
712pub(super) fn rigid_intercept_from_marginal(
713 marginal_eta: f64,
714 logslope: f64,
715 probit_scale: f64,
716) -> f64 {
717 marginal_eta * rigid_observed_scale(logslope, probit_scale)
718}
719
720#[inline]
721pub(super) fn rigid_prescale_intercept_from_marginal(
722 marginal_eta: f64,
723 logslope: f64,
724 probit_scale: f64,
725) -> f64 {
726 rigid_intercept_from_marginal(marginal_eta, logslope, probit_scale) / probit_scale
727}
728
729#[inline]
730pub(super) fn rigid_prescale_intercept_derivative_abs(
731 marginal_eta: f64,
732 logslope: f64,
733 probit_scale: f64,
734) -> f64 {
735 let c = rigid_observed_scale(logslope, probit_scale);
736 probit_scale * normal_pdf(marginal_eta) / c
737}
738
739#[inline]
740pub(super) fn rigid_observed_eta(
741 marginal_eta: f64,
742 logslope: f64,
743 z: f64,
744 probit_scale: f64,
745) -> f64 {
746 marginal_slope_standard_normal_scalar_eta(marginal_eta, logslope, z, probit_scale)
747}
748
749#[inline]
750pub(super) fn marginal_slope_standard_normal_scalar_eta(
751 q: f64,
752 slope: f64,
753 z: f64,
754 probit_scale: f64,
755) -> f64 {
756 let observed_slope = rigid_observed_slope(slope, probit_scale);
757 q * (1.0 + observed_slope * observed_slope).sqrt() + observed_slope * z
758}
759
760pub(super) fn unary_derivatives_normal_cdf(x: f64) -> [f64; 5] {
761 let pdf = normal_pdf(x);
762 [
763 normal_cdf(x),
764 pdf,
765 -x * pdf,
766 (x * x - 1.0) * pdf,
767 (-x.powi(3) + 3.0 * x) * pdf,
768 ]
769}
770
771#[inline]
778pub(super) fn lse_accumulate(log_max: &mut f64, sum: &mut f64, log_term: f64) {
779 if !log_term.is_finite() {
780 return;
781 }
782 if log_term > *log_max {
783 if log_max.is_finite() {
784 *sum = *sum * (*log_max - log_term).exp() + 1.0;
785 } else {
786 *sum = 1.0;
787 }
788 *log_max = log_term;
789 } else {
790 *sum += (log_term - *log_max).exp();
791 }
792}
793
794#[derive(Clone, Copy, Debug, PartialEq, Eq)]
795pub enum MarginalSlopeCovarianceShape {
796 Diagonal,
797 Full,
798 LowRank,
799}
800
801#[derive(Clone, Debug, PartialEq)]
802enum MarginalSlopeCovarianceStorage {
803 Diagonal {
804 covariance: Array1<f64>,
805 },
806 Full {
807 covariance: Array2<f64>,
808 square_root_factor: Array2<f64>,
810 },
811 LowRank {
813 factor: Array2<f64>,
814 },
815}
816
817#[derive(Clone, Debug)]
826pub struct MarginalSlopeCovariance {
827 storage: MarginalSlopeCovarianceStorage,
828 ones_quadratic_form: f64,
829}
830
831impl PartialEq for MarginalSlopeCovariance {
832 fn eq(&self, other: &Self) -> bool {
833 self.storage == other.storage
834 }
835}
836
837#[derive(Clone, Copy, Debug)]
838pub(crate) enum MarginalSlopeCovarianceRef<'a> {
839 Diagonal(&'a Array1<f64>),
840 Full(&'a Array2<f64>),
841 LowRank(&'a Array2<f64>),
842}
843
844impl MarginalSlopeCovariance {
845 pub fn diagonal(covariance: Array1<f64>) -> Result<Self, String> {
846 if covariance.is_empty() {
847 return Err("marginal-slope diagonal covariance is empty".to_string());
848 }
849 let mut ones_quadratic_form = 0.0;
850 for (axis, &value) in covariance.iter().enumerate() {
851 if !(value.is_finite() && value >= 0.0) {
852 return Err(format!(
853 "marginal-slope diagonal covariance entry {axis} must be finite and non-negative, got {value}"
854 ));
855 }
856 ones_quadratic_form += value;
857 }
858 if !ones_quadratic_form.is_finite() {
859 return Err("marginal-slope diagonal covariance geometry overflowed".to_string());
860 }
861 Ok(Self {
862 storage: MarginalSlopeCovarianceStorage::Diagonal { covariance },
863 ones_quadratic_form,
864 })
865 }
866
867 pub fn full(covariance: Array2<f64>) -> Result<Self, String> {
868 if covariance.nrows() == 0 || covariance.nrows() != covariance.ncols() {
869 return Err(format!(
870 "marginal-slope full covariance must be non-empty and square, got {}x{}",
871 covariance.nrows(),
872 covariance.ncols(),
873 ));
874 }
875 for ((row, column), &value) in covariance.indexed_iter() {
876 if !value.is_finite() {
877 return Err(format!(
878 "marginal-slope full covariance entry ({row},{column}) is non-finite"
879 ));
880 }
881 }
882 for row in 0..covariance.nrows() {
883 for column in (row + 1)..covariance.ncols() {
884 if covariance[[row, column]] != covariance[[column, row]] {
885 return Err(format!(
886 "marginal-slope full covariance must be exactly symmetric at ({row},{column}): upper={}, lower={}",
887 covariance[[row, column]],
888 covariance[[column, row]],
889 ));
890 }
891 }
892 }
893 let (eigenvalues, eigenvectors) = covariance.eigh(faer::Side::Lower).map_err(|error| {
894 format!("marginal-slope covariance eigendecomposition failed: {error}")
895 })?;
896 let dimension = covariance.nrows();
897 let spectral_magnitude = eigenvalues
911 .iter()
912 .fold(0.0_f64, |magnitude, &value| magnitude.max(value.abs()));
913 let psd_roundoff = 128.0 * dimension as f64 * f64::EPSILON * spectral_magnitude;
914 let mut square_root_factor = Array2::<f64>::zeros((dimension, dimension));
915 for (eigen_axis, &eigenvalue) in eigenvalues.iter().enumerate() {
916 if !eigenvalue.is_finite() || eigenvalue < -psd_roundoff {
917 return Err(format!(
918 "marginal-slope full covariance must be positive semidefinite; eigenvalue {eigen_axis} is {eigenvalue} (admissible eigensolver band -{psd_roundoff:.3e})"
919 ));
920 }
921 let scale = eigenvalue.max(0.0).sqrt();
926 for axis in 0..dimension {
927 square_root_factor[[eigen_axis, axis]] = scale * eigenvectors[[axis, eigen_axis]];
928 }
929 }
930 let mut ones_quadratic_form = 0.0;
931 for factor_row in square_root_factor.rows() {
932 let projection = factor_row.sum();
933 ones_quadratic_form += projection * projection;
934 }
935 if !ones_quadratic_form.is_finite() {
936 return Err("marginal-slope full covariance geometry overflowed".to_string());
937 }
938 Ok(Self {
939 storage: MarginalSlopeCovarianceStorage::Full {
940 covariance,
941 square_root_factor,
942 },
943 ones_quadratic_form,
944 })
945 }
946
947 pub fn low_rank(factor: Array2<f64>) -> Result<Self, String> {
948 if factor.nrows() == 0 {
949 return Err("marginal-slope low-rank covariance factor has zero rows".to_string());
950 }
951 for ((row, column), &value) in factor.indexed_iter() {
952 if !value.is_finite() {
953 return Err(format!(
954 "marginal-slope low-rank covariance factor entry ({row},{column}) is non-finite"
955 ));
956 }
957 }
958 let mut ones_quadratic_form = 0.0;
959 for factor_column in factor.columns() {
960 let projection = factor_column.sum();
961 ones_quadratic_form += projection * projection;
962 }
963 if !ones_quadratic_form.is_finite() {
964 return Err("marginal-slope low-rank covariance geometry overflowed".to_string());
965 }
966 Ok(Self {
967 storage: MarginalSlopeCovarianceStorage::LowRank { factor },
968 ones_quadratic_form,
969 })
970 }
971
972 pub fn to_dense(&self) -> Array2<f64> {
973 match &self.storage {
974 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
975 Array2::from_diag(covariance)
976 }
977 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.clone(),
978 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.dot(&factor.t()),
979 }
980 }
981
982 pub fn shape(&self) -> MarginalSlopeCovarianceShape {
983 match &self.storage {
984 MarginalSlopeCovarianceStorage::Diagonal { .. } => {
985 MarginalSlopeCovarianceShape::Diagonal
986 }
987 MarginalSlopeCovarianceStorage::Full { .. } => MarginalSlopeCovarianceShape::Full,
988 MarginalSlopeCovarianceStorage::LowRank { .. } => MarginalSlopeCovarianceShape::LowRank,
989 }
990 }
991
992 pub fn dim(&self) -> usize {
993 match &self.storage {
994 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => covariance.len(),
995 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.nrows(),
996 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.nrows(),
997 }
998 }
999
1000 pub fn ones_quadratic_form(&self) -> f64 {
1001 self.ones_quadratic_form
1002 }
1003
1004 pub(crate) fn representation(&self) -> MarginalSlopeCovarianceRef<'_> {
1005 match &self.storage {
1006 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
1007 MarginalSlopeCovarianceRef::Diagonal(covariance)
1008 }
1009 MarginalSlopeCovarianceStorage::Full { covariance, .. } => {
1010 MarginalSlopeCovarianceRef::Full(covariance)
1011 }
1012 MarginalSlopeCovarianceStorage::LowRank { factor } => {
1013 MarginalSlopeCovarianceRef::LowRank(factor)
1014 }
1015 }
1016 }
1017
1018 #[inline(always)]
1019 pub(crate) fn quadratic_form_unchecked(&self, vector: &[f64]) -> f64 {
1020 <Self as SymmetricQuadraticCoefficients>::quadratic_value(self, vector, |value| *value)
1021 }
1022
1023 pub fn quadratic_form(&self, vector: &[f64]) -> Result<f64, String> {
1024 if vector.len() != self.dim() {
1025 return Err(format!(
1026 "marginal-slope covariance dimension mismatch: vector={}, covariance={}",
1027 vector.len(),
1028 self.dim()
1029 ));
1030 }
1031 if vector.iter().any(|value| !value.is_finite()) {
1032 return Err("marginal-slope covariance vector contains non-finite values".to_string());
1033 }
1034 let value = self.quadratic_form_unchecked(vector);
1035 if !value.is_finite() {
1036 return Err(format!(
1037 "marginal-slope covariance quadratic form is non-finite: {value}"
1038 ));
1039 }
1040 Ok(value)
1041 }
1042}
1043
1044enum VectorSupport {
1045 Zero,
1046 Singleton { axis: usize, value: f64 },
1047 Multiple,
1048}
1049
1050#[inline(always)]
1051fn vector_support(input: &[f64]) -> VectorSupport {
1052 let mut singleton = None;
1053 for (axis, &value) in input.iter().enumerate() {
1054 if value == 0.0 {
1055 continue;
1056 }
1057 if singleton.is_some() {
1058 return VectorSupport::Multiple;
1059 }
1060 singleton = Some((axis, value));
1061 }
1062 match singleton {
1063 None => VectorSupport::Zero,
1064 Some((axis, value)) => VectorSupport::Singleton { axis, value },
1065 }
1066}
1067
1068impl SymmetricQuadraticCoefficients for MarginalSlopeCovariance {
1069 fn dimension(&self) -> usize {
1070 self.dim()
1071 }
1072
1073 fn multiply(&self, input: &[f64], output: &mut [f64]) {
1074 assert_eq!(input.len(), self.dim());
1075 assert_eq!(output.len(), self.dim());
1076 match self.representation() {
1077 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1078 for axis in 0..input.len() {
1079 output[axis] = diagonal[axis] * input[axis];
1080 }
1081 }
1082 MarginalSlopeCovarianceRef::Full(matrix) => {
1083 match vector_support(input) {
1084 VectorSupport::Zero => {
1085 output.fill(0.0);
1086 return;
1087 }
1088 VectorSupport::Singleton { axis, value } => {
1089 for row in 0..input.len() {
1090 output[row] = matrix[[row, axis]] * value;
1091 }
1092 return;
1093 }
1094 VectorSupport::Multiple => {}
1095 }
1096 for row in 0..input.len() {
1097 let mut value = 0.0;
1098 for column in 0..input.len() {
1099 value += matrix[[row, column]] * input[column];
1100 }
1101 output[row] = value;
1102 }
1103 }
1104 MarginalSlopeCovarianceRef::LowRank(factor) => {
1105 output.fill(0.0);
1106 match vector_support(input) {
1107 VectorSupport::Zero => return,
1108 VectorSupport::Singleton { axis, value } => {
1109 for rank in 0..factor.ncols() {
1110 let projection = factor[[axis, rank]] * value;
1111 for row in 0..input.len() {
1112 output[row] += factor[[row, rank]] * projection;
1113 }
1114 }
1115 return;
1116 }
1117 VectorSupport::Multiple => {}
1118 }
1119 for rank in 0..factor.ncols() {
1120 let mut projection = 0.0;
1121 for row in 0..input.len() {
1122 projection += factor[[row, rank]] * input[row];
1123 }
1124 for row in 0..input.len() {
1125 output[row] += factor[[row, rank]] * projection;
1126 }
1127 }
1128 }
1129 }
1130 }
1131
1132 fn coefficient(&self, row: usize, column: usize) -> f64 {
1133 match self.representation() {
1134 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1135 if row == column {
1136 diagonal[row]
1137 } else {
1138 0.0
1139 }
1140 }
1141 MarginalSlopeCovarianceRef::Full(matrix) => matrix[[row, column]],
1142 MarginalSlopeCovarianceRef::LowRank(factor) => {
1143 let mut value = 0.0;
1144 for rank in 0..factor.ncols() {
1145 value += factor[[row, rank]] * factor[[column, rank]];
1146 }
1147 value
1148 }
1149 }
1150 }
1151
1152 fn visit_upper_triangle(
1153 &self,
1154 direction: &mut [f64],
1155 projected: &mut [f64],
1156 mut visit: impl FnMut(usize, usize, f64),
1157 ) {
1158 let dimension = self.dim();
1159 assert_eq!(direction.len(), dimension);
1160 assert_eq!(projected.len(), dimension);
1161 match self.representation() {
1162 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1163 for column in 0..dimension {
1164 for row in 0..=column {
1165 visit(row, column, if row == column { diagonal[row] } else { 0.0 });
1166 }
1167 }
1168 }
1169 MarginalSlopeCovarianceRef::Full(matrix) => {
1170 for column in 0..dimension {
1171 for row in 0..=column {
1172 visit(row, column, matrix[[row, column]]);
1173 }
1174 }
1175 }
1176 MarginalSlopeCovarianceRef::LowRank(factor) => {
1177 for column in 0..dimension {
1178 for row in 0..=column {
1179 let mut value = 0.0;
1180 for rank in 0..factor.ncols() {
1181 value += factor[[row, rank]] * factor[[column, rank]];
1182 }
1183 visit(row, column, value);
1184 }
1185 }
1186 }
1187 }
1188 }
1189
1190 fn quadratic_value<T, F>(&self, input: &[T], value: F) -> f64
1191 where
1192 F: Fn(&T) -> f64,
1193 {
1194 assert_eq!(input.len(), self.dim());
1195 match &self.storage {
1196 MarginalSlopeCovarianceStorage::Diagonal { covariance } => input
1197 .iter()
1198 .zip(covariance)
1199 .map(|(input, &covariance)| {
1200 let input = value(input);
1201 covariance * input * input
1202 })
1203 .sum(),
1204 MarginalSlopeCovarianceStorage::Full {
1205 square_root_factor, ..
1206 } => {
1207 let mut total = 0.0;
1208 for factor_row in square_root_factor.rows() {
1209 let mut projection = 0.0;
1210 for axis in 0..input.len() {
1211 projection += factor_row[axis] * value(&input[axis]);
1212 }
1213 total += projection * projection;
1214 }
1215 total
1216 }
1217 MarginalSlopeCovarianceStorage::LowRank { factor } => {
1218 let mut total = 0.0;
1221 for rank in 0..factor.ncols() {
1222 let mut projection = 0.0;
1223 for row in 0..input.len() {
1224 projection += factor[[row, rank]] * value(&input[row]);
1225 }
1226 total += projection * projection;
1227 }
1228 total
1229 }
1230 }
1231 }
1232}
1233
1234pub fn marginal_slope_covariance_from_scores(
1261 scores: ArrayView2<'_, f64>,
1262 weights: &Array1<f64>,
1263) -> Result<MarginalSlopeCovariance, String> {
1264 let (n, k) = scores.dim();
1265 if k == 0 {
1266 return Err("marginal-slope score matrix must have at least one column".to_string());
1267 }
1268 if weights.len() != n {
1269 return Err(format!(
1270 "marginal-slope covariance weight length mismatch: weights={}, rows={n}",
1271 weights.len()
1272 ));
1273 }
1274 let total_weight = weights.iter().copied().sum::<f64>();
1275 if !(total_weight.is_finite() && total_weight > 0.0) {
1276 return Err("marginal-slope covariance needs positive finite total weight".to_string());
1277 }
1278 let mut mean = Array1::<f64>::zeros(k);
1279 for i in 0..n {
1280 let weight = weights[i];
1281 if !(weight.is_finite() && weight >= 0.0) {
1282 return Err(format!(
1283 "marginal-slope covariance weight {i} must be finite and non-negative, got {weight}"
1284 ));
1285 }
1286 for j in 0..k {
1287 let score = scores[[i, j]];
1288 if !score.is_finite() {
1289 return Err(format!(
1290 "marginal-slope covariance score ({i},{j}) is non-finite"
1291 ));
1292 }
1293 mean[j] += weight * score;
1294 }
1295 }
1296 mean.mapv_inplace(|value| value / total_weight);
1297
1298 let mut cov = Array2::<f64>::zeros((k, k));
1299 for i in 0..n {
1300 let weight = weights[i];
1301 for a in 0..k {
1302 let da = scores[[i, a]] - mean[a];
1303 for b in 0..=a {
1304 let value = weight * da * (scores[[i, b]] - mean[b]) / total_weight;
1305 cov[[a, b]] += value;
1306 if a != b {
1307 cov[[b, a]] += value;
1308 }
1309 }
1310 }
1311 }
1312
1313 let is_diagonal = (0..k).all(|row| ((row + 1)..k).all(|column| cov[[row, column]] == 0.0));
1317 if is_diagonal {
1318 MarginalSlopeCovariance::diagonal(cov.diag().to_owned())
1319 } else {
1320 MarginalSlopeCovariance::full(cov)
1321 }
1322}
1323
1324pub fn marginal_slope_preserving_scale(
1325 slopes: &[f64],
1326 covariance: &MarginalSlopeCovariance,
1327 probit_scale: f64,
1328) -> Result<f64, String> {
1329 if !probit_scale.is_finite() {
1330 return Err(format!(
1331 "marginal-slope probit scale must be finite, got {probit_scale}"
1332 ));
1333 }
1334 let variance = probit_scale * probit_scale * covariance.quadratic_form(slopes)?;
1335 if !variance.is_finite() {
1336 return Err("marginal-slope preserving variance is non-finite".to_string());
1337 }
1338 Ok((1.0 + variance).sqrt())
1339}
1340
1341pub fn marginal_slope_probit_eta(
1342 q: f64,
1343 z: &[f64],
1344 slopes: &[f64],
1345 covariance: &MarginalSlopeCovariance,
1346 probit_scale: f64,
1347) -> Result<f64, String> {
1348 if z.len() != slopes.len() {
1349 return Err(format!(
1350 "marginal-slope score/slope dimension mismatch: z={}, slopes={}",
1351 z.len(),
1352 slopes.len()
1353 ));
1354 }
1355 if slopes.len() != covariance.dim() {
1356 return Err(format!(
1357 "marginal-slope covariance dimension mismatch: slopes={}, covariance={}",
1358 slopes.len(),
1359 covariance.dim()
1360 ));
1361 }
1362 if !q.is_finite() || z.iter().any(|value| !value.is_finite()) {
1363 return Err("marginal-slope probit eta inputs must be finite".to_string());
1364 }
1365 let scale = marginal_slope_preserving_scale(slopes, covariance, probit_scale)?;
1366 let linear = z
1367 .iter()
1368 .zip(slopes.iter())
1369 .map(|(&score, &slope)| probit_scale * slope * score)
1370 .sum::<f64>();
1371 Ok(q * scale + linear)
1372}
1373
1374pub(super) fn empirical_rigid_calibration_eval(
1404 intercept: f64,
1405 log_target_mu: f64,
1406 slope: f64,
1407 probit_scale: f64,
1408 nodes: &[f64],
1409 weights: &[f64],
1410) -> Result<(f64, f64, f64), String> {
1411 if !intercept.is_finite() {
1412 return Err(format!(
1413 "empirical latent calibration: non-finite intercept {intercept}"
1414 ));
1415 }
1416 let observed_slope = rigid_observed_slope(slope, probit_scale);
1417 const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_8; let mut log_max_phi = f64::NEG_INFINITY;
1421 let mut sum_phi = 0.0_f64;
1422 let mut log_max_cdf = f64::NEG_INFINITY;
1423 let mut sum_cdf = 0.0_f64;
1424
1425 let mut log_max_pos = f64::NEG_INFINITY;
1429 let mut sum_pos = 0.0_f64;
1430 let mut log_max_neg = f64::NEG_INFINITY;
1431 let mut sum_neg = 0.0_f64;
1432
1433 for (&node, &weight) in nodes.iter().zip(weights.iter()) {
1434 if !(weight.is_finite() && weight > 0.0) {
1435 continue;
1436 }
1437 let eta = intercept + observed_slope * node;
1438 if !eta.is_finite() {
1439 return Err(format!(
1440 "empirical latent calibration: non-finite η at intercept={intercept}, slope={slope}, node={node}"
1441 ));
1442 }
1443 let log_w = weight.ln();
1444 let log_phi = -0.5 * eta * eta - HALF_LOG_2PI;
1445 let log_term_phi = log_w + log_phi;
1446 let log_term_cdf = log_w + normal_logcdf(eta);
1447
1448 lse_accumulate(&mut log_max_phi, &mut sum_phi, log_term_phi);
1449 lse_accumulate(&mut log_max_cdf, &mut sum_cdf, log_term_cdf);
1450
1451 if eta != 0.0 {
1452 let log_term_eta_phi = log_term_phi + eta.abs().ln();
1453 if eta > 0.0 {
1454 lse_accumulate(&mut log_max_pos, &mut sum_pos, log_term_eta_phi);
1455 } else {
1456 lse_accumulate(&mut log_max_neg, &mut sum_neg, log_term_eta_phi);
1457 }
1458 }
1459 }
1460
1461 if !(sum_phi.is_finite() && sum_cdf.is_finite() && sum_phi > 0.0 && sum_cdf > 0.0) {
1462 return Err(format!(
1463 "empirical latent calibration: log-space accumulation failed (sum_phi={sum_phi}, sum_cdf={sum_cdf}, intercept={intercept})"
1464 ));
1465 }
1466
1467 let log_s_phi = log_max_phi + sum_phi.ln();
1468 let log_s_cdf = log_max_cdf + sum_cdf.ln();
1469
1470 let f = log_s_cdf - log_target_mu;
1472 let log_f_prime = log_s_phi - log_s_cdf;
1484 let f_prime = if log_f_prime > -740.0 {
1485 log_f_prime.exp()
1486 } else {
1487 f64::MIN_POSITIVE
1488 };
1489
1490 let exp_safe = |log_x: f64| -> f64 { if log_x > -740.0 { log_x.exp() } else { 0.0 } };
1498 let pos_over_cdf = if sum_pos > 0.0 {
1499 exp_safe(log_max_pos + sum_pos.ln() - log_s_cdf)
1500 } else {
1501 0.0
1502 };
1503 let neg_over_cdf = if sum_neg > 0.0 {
1504 exp_safe(log_max_neg + sum_neg.ln() - log_s_cdf)
1505 } else {
1506 0.0
1507 };
1508 let s_etaphi_over_s_cdf = pos_over_cdf - neg_over_cdf;
1509 let f_double_prime = -s_etaphi_over_s_cdf - f_prime * f_prime;
1510
1511 if !(f.is_finite() && f_prime.is_finite() && f_prime > 0.0 && f_double_prime.is_finite()) {
1512 return Err(format!(
1513 "empirical latent calibration: non-finite log-space state f={f}, f'={f_prime}, f''={f_double_prime} at intercept={intercept}"
1514 ));
1515 }
1516 Ok((f, f_prime, f_double_prime))
1517}
1518
1519pub(crate) fn empirical_intercept_from_marginal(
1520 target_mu: f64,
1521 target_q: f64,
1522 slope: f64,
1523 probit_scale: f64,
1524 nodes: &[f64],
1525 weights: &[f64],
1526 initial: Option<f64>,
1527) -> Result<f64, String> {
1528 if !(target_mu.is_finite() && target_mu > 0.0 && target_mu < 1.0) {
1529 return Err(format!(
1530 "empirical latent calibration requires target mu in (0,1), got {target_mu}"
1531 ));
1532 }
1533 let log_target_mu = target_mu.ln();
1534 let closed_form_seed = rigid_intercept_from_marginal(target_q, slope, probit_scale);
1535 let seed = initial.unwrap_or(closed_form_seed);
1536 let eval = |a: f64| {
1537 empirical_rigid_calibration_eval(a, log_target_mu, slope, probit_scale, nodes, weights)
1538 };
1539 let abs_tol = 1e-13_f64.max(4.0 * f64::EPSILON);
1546 let solve_from = |s: f64| {
1547 crate::monotone_root::solve_monotone_root(
1548 eval,
1549 s,
1550 "empirical latent intercept",
1551 abs_tol,
1552 64,
1553 48,
1554 )
1555 .map_err(|e| e.to_string())
1558 };
1559 let (root, _, f_best) = match solve_from(seed) {
1570 Ok(v) => v,
1571 Err(first_err) => {
1572 if seed == closed_form_seed {
1573 return Err(first_err);
1574 }
1575 solve_from(closed_form_seed).map_err(|retry_err| {
1576 format!("{first_err}; closed-form retry from a={closed_form_seed:.6}: {retry_err}")
1577 })?
1578 }
1579 };
1580 if f_best.abs() > abs_tol {
1581 return Err(format!(
1582 "empirical latent intercept solve failed: log-residual={f_best:.3e} at a={root:.6}, target mu={target_mu:.6}"
1583 ));
1584 }
1585 Ok(root)
1586}
1587
1588#[inline]
1589pub(super) fn rigid_standard_normal_neglog_only(
1590 q: f64,
1591 g: f64,
1592 z: f64,
1593 y: f64,
1594 w: f64,
1595 probit_scale: f64,
1596) -> Result<f64, String> {
1597 let s = 2.0 * y - 1.0;
1598 let eta = marginal_slope_standard_normal_scalar_eta(q, g, z, probit_scale);
1599 let m = s * eta;
1600 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
1601 if !logcdf.is_finite() {
1602 return Err(format!(
1603 "rigid probit neglog_only: non-finite log Φ at q={q}, g={g}, z={z}, y={y}"
1604 ));
1605 }
1606 Ok(-w * logcdf)
1607}
1608
1609#[inline(always)]
1610fn rigid_supplied_link_stack(
1611 composition_point: f64,
1612 value: f64,
1613 first: f64,
1614 second: f64,
1615 third: f64,
1616 fourth: f64,
1617) -> [f64; 5] {
1618 if composition_point.is_nan() {
1619 [f64::NAN; 5]
1620 } else {
1621 [value, first, second, third, fourth]
1622 }
1623}
1624
1625#[inline(always)]
1626fn rigid_observed_scale_stack(observed_slope: f64) -> [f64; 5] {
1627 let scale = (1.0 + observed_slope * observed_slope).sqrt();
1628 let inverse = scale.recip();
1629 let inverse_squared = inverse * inverse;
1630 let inverse_cubed = inverse_squared * inverse;
1631 let inverse_fifth = inverse_cubed * inverse_squared;
1632 let inverse_seventh = inverse_fifth * inverse_squared;
1633 [
1634 scale,
1635 observed_slope * inverse,
1636 inverse_cubed,
1637 -3.0 * observed_slope * inverse_fifth,
1638 (12.0 * observed_slope * observed_slope - 3.0) * inverse_seventh,
1639 ]
1640}
1641
1642row_program! {
1643 fn rigid_standard_normal_program(
1644 marginal_eta,
1645 logslope;
1646 marginal_q,
1647 marginal_q1,
1648 marginal_q2,
1649 marginal_q3,
1650 marginal_q4,
1651 probit_scale,
1652 latent_score,
1653 outcome_sign,
1654 weight
1655 )
1656 emit [generic, order2, third, fourth, full];
1657 leaves {
1658 supplied_link => rigid_supplied_link_stack => rigid_supplied_link_stack_cuda,
1659 observed_scale => rigid_observed_scale_stack => rigid_observed_scale_stack_cuda,
1660 signed_probit => signed_probit_neglog_unary_stack => signed_probit_neglog_unary_stack_cuda,
1661 }
1662 witnesses [];
1663 {
1664 let q = compose(
1665 supplied_link,
1666 marginal_eta,
1667 marginal_q,
1668 marginal_q1,
1669 marginal_q2,
1670 marginal_q3,
1671 marginal_q4
1672 );
1673 let observed_logslope = scale(logslope, probit_scale);
1674 let observed_scale_value = compose(observed_scale, observed_logslope);
1675 let latent_index = add(
1676 mul(q, observed_scale_value),
1677 scale(observed_logslope, latent_score)
1678 );
1679 let signed_margin = scale(latent_index, outcome_sign);
1680 return compose(signed_probit, signed_margin, weight);
1681 }
1682}
1683
1684#[inline]
1714pub(crate) fn rigid_standard_normal_row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1715 p: &[S; 2],
1716 marginal: BernoulliMarginalLinkMap,
1717 z: f64,
1718 y: f64,
1719 w: f64,
1720 probit_scale: f64,
1721) -> Result<S, String> {
1722 let outcome_sign = 2.0 * y - 1.0;
1723 let m = outcome_sign
1724 * marginal_slope_standard_normal_scalar_eta(marginal.q, p[1].value(), z, probit_scale);
1725 if !(m.is_finite() || m == f64::INFINITY) {
1726 return Err(format!(
1727 "non-finite signed margin in rigid probit row NLL: {m}"
1728 ));
1729 }
1730 let (nll, []) = rigid_standard_normal_program(
1731 &p[0],
1732 &p[1],
1733 marginal.q,
1734 marginal.q1,
1735 marginal.q2,
1736 marginal.q3,
1737 marginal.q4,
1738 probit_scale,
1739 z,
1740 outcome_sign,
1741 w,
1742 );
1743 Ok(nll)
1744}
1745
1746#[inline]
1747pub(super) fn rigid_standard_normal_row_kernel(
1748 marginal: BernoulliMarginalLinkMap,
1749 g: f64,
1750 z: f64,
1751 y: f64,
1752 w: f64,
1753 probit_scale: f64,
1754) -> Result<(f64, [f64; 2], [[f64; 2]; 2]), String> {
1755 let outcome_sign = 2.0 * y - 1.0;
1756 let signed_margin =
1757 outcome_sign * marginal_slope_standard_normal_scalar_eta(marginal.q, g, z, probit_scale);
1758 if !(signed_margin.is_finite() || signed_margin == f64::INFINITY) {
1759 return Err(format!(
1760 "non-finite signed margin in rigid probit row NLL: {signed_margin}"
1761 ));
1762 }
1763 let (value, gradient, hessian, []) = rigid_standard_normal_program_order2(
1764 marginal.eta_value(),
1765 g,
1766 marginal.q,
1767 marginal.q1,
1768 marginal.q2,
1769 marginal.q3,
1770 marginal.q4,
1771 probit_scale,
1772 z,
1773 outcome_sign,
1774 w,
1775 );
1776 Ok((value, gradient, hessian))
1777}
1778
1779#[inline]
1817pub(super) fn rigid_standard_normal_mixed_z_sensitivity(
1818 marginal: BernoulliMarginalLinkMap,
1819 g: f64,
1820 z: f64,
1821 y: f64,
1822 w: f64,
1823 probit_scale: f64,
1824) -> Result<[f64; 2], String> {
1825 use gam_math::jet_tower::Tower2;
1836 let mut q = Tower2::<3>::constant(marginal.q);
1837 q.g[0] = marginal.q1;
1838 q.h[0][0] = marginal.q2;
1839 let slope = Tower2::<3>::variable(g, 1);
1840 let z_var = Tower2::<3>::variable(z, 2);
1841 let observed_logslope = slope * probit_scale;
1842 let c = (observed_logslope * observed_logslope + 1.0).sqrt();
1843 let eta = q * c + slope * (z_var * probit_scale);
1847 let signed = eta * (2.0 * y - 1.0);
1848 if !(signed.v.is_finite() || signed.v == f64::INFINITY) {
1850 return Err(format!(
1851 "rigid probit mixed-z sensitivity: non-finite signed margin {} at q={}, g={g}, z={z}, y={y}",
1852 signed.v, marginal.q
1853 ));
1854 }
1855 let stack = signed_probit_neglog_unary_stack(signed.v, w);
1856 if !stack[0].is_finite() {
1857 return Err(format!(
1858 "rigid probit mixed-z sensitivity: non-finite log Φ at q={}, g={g}, z={z}, y={y}",
1859 marginal.q
1860 ));
1861 }
1862 let tower = signed.compose_unary([stack[0], stack[1], stack[2]]);
1865 let s_q = -tower.h[0][2];
1872 let s_g = -tower.h[1][2];
1873 if !(s_q.is_finite() && s_g.is_finite()) {
1874 return Err(format!(
1875 "rigid probit mixed-z sensitivity: non-finite ∂²(log L)/∂(q,g)∂z = [{s_q}, {s_g}] at q={}, g={g}, z={z}",
1876 marginal.q
1877 ));
1878 }
1879 Ok([s_q, s_g])
1880}
1881
1882pub(super) fn rigid_standard_normal_score_zeta_sensitivity(
1906 base_link: &InverseLink,
1907 marginal_eta: &Array1<f64>,
1908 slope_eta: &Array1<f64>,
1909 z: &Array1<f64>,
1910 y: &Array1<f64>,
1911 weights: &Array1<f64>,
1912 probit_scale: f64,
1913 marginal_design: ArrayView2<'_, f64>,
1914 logslope_design: ArrayView2<'_, f64>,
1915 p_beta: usize,
1916) -> Result<Array2<f64>, String> {
1917 let n = marginal_eta.len();
1918 let p_m = marginal_design.ncols();
1919 let r = logslope_design.ncols();
1920 if slope_eta.len() != n
1921 || z.len() != n
1922 || y.len() != n
1923 || weights.len() != n
1924 || marginal_design.nrows() != n
1925 || logslope_design.nrows() != n
1926 {
1927 return Err(format!(
1928 "score_zeta_sensitivity row mismatch: marginal_eta={n}, slope_eta={}, z={}, y={}, \
1929 weights={}, marginal_design rows={}, logslope_design rows={}",
1930 slope_eta.len(),
1931 z.len(),
1932 y.len(),
1933 weights.len(),
1934 marginal_design.nrows(),
1935 logslope_design.nrows()
1936 ));
1937 }
1938 if p_m + r != p_beta {
1939 return Err(format!(
1940 "rigid score_zeta_sensitivity width mismatch: marginal({p_m}) + logslope({r}) != p_beta({p_beta})"
1941 ));
1942 }
1943 let mut s = Array2::<f64>::zeros((n, p_beta));
1944 for i in 0..n {
1945 let marginal = bernoulli_marginal_link_map(base_link, marginal_eta[i])?;
1946 let [s_q, s_g] = rigid_standard_normal_mixed_z_sensitivity(
1947 marginal,
1948 slope_eta[i],
1949 z[i],
1950 y[i],
1951 weights[i],
1952 probit_scale,
1953 )?;
1954 if s_q != 0.0 {
1957 let m_row = marginal_design.row(i);
1958 for (j, &mij) in m_row.iter().enumerate() {
1959 s[[i, j]] = s_q * mij;
1960 }
1961 }
1962 if s_g != 0.0 {
1963 let g_row = logslope_design.row(i);
1964 for (j, &gij) in g_row.iter().enumerate() {
1965 s[[i, p_m + j]] = s_g * gij;
1966 }
1967 }
1968 }
1969 Ok(s)
1970}
1971
1972#[inline]
1978pub(super) fn rigid_standard_normal_third_full(
1979 marginal: BernoulliMarginalLinkMap,
1980 g: f64,
1981 z: f64,
1982 y: f64,
1983 w: f64,
1984 probit_scale: f64,
1985) -> Result<[[[f64; 2]; 2]; 2], String> {
1986 let outcome_sign = 2.0 * y - 1.0;
1987 let signed_margin =
1988 outcome_sign * marginal_slope_standard_normal_scalar_eta(marginal.q, g, z, probit_scale);
1989 if !(signed_margin.is_finite() || signed_margin == f64::INFINITY) {
1990 return Err(format!(
1991 "non-finite signed margin in rigid probit row NLL: {signed_margin}"
1992 ));
1993 }
1994 Ok(rigid_standard_normal_program_third_full(
1995 marginal.eta_value(),
1996 g,
1997 marginal.q,
1998 marginal.q1,
1999 marginal.q2,
2000 marginal.q3,
2001 marginal.q4,
2002 probit_scale,
2003 z,
2004 outcome_sign,
2005 w,
2006 ))
2007}
2008
2009#[inline]
2010pub(super) fn rigid_standard_normal_third_contracted_generated(
2011 marginal: BernoulliMarginalLinkMap,
2012 g: f64,
2013 z: f64,
2014 y: f64,
2015 w: f64,
2016 probit_scale: f64,
2017 direction: &[f64; 2],
2018) -> Result<[[f64; 2]; 2], String> {
2019 let outcome_sign = 2.0 * y - 1.0;
2020 let signed_margin =
2021 outcome_sign * marginal_slope_standard_normal_scalar_eta(marginal.q, g, z, probit_scale);
2022 if !(signed_margin.is_finite() || signed_margin == f64::INFINITY) {
2023 return Err(format!(
2024 "non-finite signed margin in rigid probit row NLL: {signed_margin}"
2025 ));
2026 }
2027 Ok(rigid_standard_normal_program_third_contracted(
2028 marginal.eta_value(),
2029 g,
2030 marginal.q,
2031 marginal.q1,
2032 marginal.q2,
2033 marginal.q3,
2034 marginal.q4,
2035 probit_scale,
2036 z,
2037 outcome_sign,
2038 w,
2039 direction,
2040 ))
2041}
2042
2043#[inline]
2048pub(super) fn contract_third_full(t: &[[[f64; 2]; 2]; 2], d_eta: f64, d_g: f64) -> [[f64; 2]; 2] {
2049 [
2050 [
2051 t[0][0][0] * d_eta + t[0][0][1] * d_g,
2052 t[0][1][0] * d_eta + t[0][1][1] * d_g,
2053 ],
2054 [
2055 t[1][0][0] * d_eta + t[1][0][1] * d_g,
2056 t[1][1][0] * d_eta + t[1][1][1] * d_g,
2057 ],
2058 ]
2059}
2060
2061#[inline]
2066pub(super) fn rigid_standard_normal_fourth_full(
2067 marginal: BernoulliMarginalLinkMap,
2068 g: f64,
2069 z: f64,
2070 y: f64,
2071 w: f64,
2072 probit_scale: f64,
2073) -> Result<[[[[f64; 2]; 2]; 2]; 2], String> {
2074 let outcome_sign = 2.0 * y - 1.0;
2078 let signed_margin =
2079 outcome_sign * marginal_slope_standard_normal_scalar_eta(marginal.q, g, z, probit_scale);
2080 if !(signed_margin.is_finite() || signed_margin == f64::INFINITY) {
2081 return Err(format!(
2082 "non-finite signed margin in rigid probit row NLL: {signed_margin}"
2083 ));
2084 }
2085 Ok(rigid_standard_normal_program_fourth_full(
2086 marginal.eta_value(),
2087 g,
2088 marginal.q,
2089 marginal.q1,
2090 marginal.q2,
2091 marginal.q3,
2092 marginal.q4,
2093 probit_scale,
2094 z,
2095 outcome_sign,
2096 w,
2097 ))
2098}
2099
2100#[inline]
2119pub(super) fn contract_fourth_full(
2120 t: &[[[[f64; 2]; 2]; 2]; 2],
2121 u_eta: f64,
2122 u_g: f64,
2123 v_eta: f64,
2124 v_g: f64,
2125) -> [[f64; 2]; 2] {
2126 let mut out = [[0.0; 2]; 2];
2127 for a in 0..2 {
2128 for b in 0..2 {
2129 let mut sum = 0.0;
2130 sum += t[a][b][0][0] * u_eta * v_eta;
2131 sum += t[a][b][0][1] * u_eta * v_g;
2132 sum += t[a][b][1][0] * u_g * v_eta;
2133 sum += t[a][b][1][1] * u_g * v_g;
2134 out[a][b] = sum;
2135 }
2136 }
2137 out
2138}
2139
2140pub(super) fn ensure_finite_third_full_cache_row(
2141 t: &[[[f64; 2]; 2]; 2],
2142 context: &str,
2143) -> Result<(), String> {
2144 if t.iter().flatten().flatten().all(|value| value.is_finite()) {
2145 Ok(())
2146 } else {
2147 Err(format!(
2148 "{context}: warmed third-derivative cache row contains a non-finite value"
2149 ))
2150 }
2151}
2152
2153pub(super) fn ensure_finite_fourth_full_cache_row(
2154 t: &[[[[f64; 2]; 2]; 2]; 2],
2155 context: &str,
2156) -> Result<(), String> {
2157 if t.iter()
2158 .flatten()
2159 .flatten()
2160 .flatten()
2161 .all(|value| value.is_finite())
2162 {
2163 Ok(())
2164 } else {
2165 Err(format!(
2166 "{context}: warmed fourth-derivative cache row contains a non-finite value"
2167 ))
2168 }
2169}
2170
2171pub(crate) fn unary_derivatives_sqrt(x: f64) -> [f64; 5] {
2172 let s = x.max(1e-300).sqrt();
2173 let x1 = x.max(1e-300);
2174 let x2 = x1 * x1;
2175 let x3 = x2 * x1;
2176 [
2177 s,
2178 0.5 / s,
2179 -0.25 / (x1 * s),
2180 3.0 / (8.0 * x2 * s),
2181 -15.0 / (16.0 * x3 * s),
2182 ]
2183}
2184pub(crate) fn unary_derivatives_inverse_sqrt(x: f64) -> [f64; 5] {
2200 let x1 = x.max(1e-300);
2201 let s = x1.sqrt();
2202 let r = 1.0 / s;
2203 let x2 = x1 * x1;
2204 let x3 = x2 * x1;
2205 let x4 = x3 * x1;
2206 [
2207 r,
2208 -0.5 * r / x1,
2209 0.75 * r / x2,
2210 -1.875 * r / x3,
2211 6.5625 * r / x4,
2212 ]
2213}
2214
2215pub(crate) fn unary_derivatives_neglog_phi(x: f64, weight: f64) -> [f64; 5] {
2216 signed_probit_neglog_unary_stack(x, weight)
2221}
2222
2223pub(crate) fn unary_derivatives_log(x: f64) -> [f64; 5] {
2241 let x2 = x * x;
2242 let x3 = x2 * x;
2243 let x4 = x3 * x;
2244 [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
2245}
2246
2247pub(crate) fn unary_derivatives_log_normal_pdf(x: f64) -> [f64; 5] {
2249 let c = 0.5 * (2.0 * std::f64::consts::PI).ln();
2250 [-0.5 * x * x - c, -x, -1.0, 0.0, 0.0]
2251}
2252
2253#[cfg(test)]
2254mod covariance_admission_tests {
2255 use super::*;
2256 use ndarray::array;
2257
2258 #[test]
2259 fn full_covariance_admission_rejects_one_ulp_asymmetry_932() {
2260 let upper = 0.25_f64;
2261 let lower = f64::from_bits(upper.to_bits() + 1);
2262 let error = MarginalSlopeCovariance::full(array![[1.0, upper], [lower, 1.0]])
2263 .expect_err("any asymmetric full operator must be rejected");
2264 assert!(error.contains("must be exactly symmetric"), "{error}");
2265 }
2266
2267 #[test]
2268 fn full_covariance_admission_rejects_indefinite_matrix_before_row_use_932() {
2269 let error = MarginalSlopeCovariance::full(array![[1.0, 2.0], [2.0, 1.0]])
2270 .expect_err("an indefinite full operator is not a covariance");
2271 assert!(error.contains("must be positive semidefinite"), "{error}");
2272 }
2273
2274 #[test]
2275 fn full_covariance_admission_accepts_exact_singular_psd_932() {
2276 MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, 0.0]])
2277 .expect("an exact singular PSD covariance is admissible");
2278 }
2279
2280 #[test]
2288 fn full_covariance_admission_decides_psd_against_the_eigensolver_band() {
2289 let admitted = MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, -1.0e-17]])
2290 .expect("a negative eigenvalue inside the eigensolver band is a zero one");
2291 assert!(
2292 admitted.ones_quadratic_form().is_finite(),
2293 "clamping inside the band must keep the square-root factor real"
2294 );
2295 let error = MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, -1.0e-12]])
2296 .expect_err("a negative eigenvalue outside the band is material indefiniteness");
2297 assert!(error.contains("must be positive semidefinite"), "{error}");
2298 }
2299
2300 #[test]
2301 fn full_covariance_admission_accepts_coupled_singular_psd_932() {
2302 let covariance = MarginalSlopeCovariance::full(array![[1.0, 1.0], [1.0, 1.0]])
2303 .expect("a coupled singular PSD covariance is admissible");
2304 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2305 assert_eq!(covariance.to_dense(), array![[1.0, 1.0], [1.0, 1.0]]);
2306 }
2307
2308 #[test]
2309 fn exact_nonzero_offdiagonal_classifier_retains_full_geometry_932() {
2310 let epsilon = 1.0e-14;
2311 let scores = array![[-1.0, -epsilon], [1.0, epsilon], [0.0, -1.0], [0.0, 1.0]];
2312 let covariance =
2313 marginal_slope_covariance_from_scores(scores.view(), &Array1::ones(4)).unwrap();
2314 let dense = covariance.to_dense();
2315 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2316 assert_ne!(dense[[0, 1]], 0.0);
2317 assert_eq!(dense[[0, 1]], dense[[1, 0]]);
2318 let direction = [0.75, -1.25];
2319 let expected = direction[0] * (dense[[0, 0]] * direction[0] + dense[[0, 1]] * direction[1])
2320 + direction[1] * (dense[[1, 0]] * direction[0] + dense[[1, 1]] * direction[1]);
2321 let actual = covariance.quadratic_form(&direction).unwrap();
2322 assert!((actual - expected).abs() <= 2.0e-15);
2323 }
2324
2325 #[test]
2326 fn diagonal_covariance_entries_are_the_exact_geometry_authority_932() {
2327 let covariance = MarginalSlopeCovariance::diagonal(array![3.75]).unwrap();
2328 assert_eq!(covariance.ones_quadratic_form(), 3.75);
2329 assert_eq!(covariance.quadratic_form(&[1.0]).unwrap(), 3.75);
2330 }
2331
2332 #[test]
2333 fn equal_dense_covariance_quadratic_forms_match_all_representations_932() {
2334 let diagonal = MarginalSlopeCovariance::diagonal(array![1.2, 0.7]).unwrap();
2335 let full = MarginalSlopeCovariance::full(array![[1.2, 0.0], [0.0, 0.7]]).unwrap();
2336 let low_rank =
2337 MarginalSlopeCovariance::low_rank(array![[1.2_f64.sqrt(), 0.0], [0.0, 0.7_f64.sqrt()]])
2338 .unwrap();
2339 let direction = [0.35, -0.8];
2340 let expected = diagonal.quadratic_form(&direction).unwrap();
2341 for covariance in [&full, &low_rank] {
2342 let actual = covariance.quadratic_form(&direction).unwrap();
2343 assert!((actual - expected).abs() <= 2.0e-15);
2344 assert!(
2345 (covariance.ones_quadratic_form() - diagonal.ones_quadratic_form()).abs()
2346 <= 2.0e-15
2347 );
2348 }
2349 }
2350}
2351
2352#[cfg(test)]
2353mod jet_tower_oracle_tests {
2354 use super::*;
2376
2377 use crate::bms::test_support::{
2378 RigidStandardNormalRow, rigid_standard_normal_tower,
2379 };
2380
2381 #[test]
2382 fn signed_probit_stack_preserves_extreme_tail_derivatives_and_weight_sign() {
2383 let positive = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, 2.0);
2384 assert_eq!(
2385 positive,
2386 [f64::INFINITY, f64::NEG_INFINITY, 2.0, -0.0, -0.0]
2387 );
2388 let negative = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, -2.0);
2389 assert_eq!(negative, [f64::NEG_INFINITY, f64::INFINITY, -2.0, 0.0, 0.0]);
2390
2391 let right = signed_probit_neglog_unary_stack(38.6, 1.0);
2392 assert_eq!(right[1], -0.0);
2393 assert!(right[2] > 0.0 && right[2].is_subnormal());
2394 assert!(right[3] < 0.0 && right[3].is_subnormal());
2395 assert!(right[4] > 0.0 && right[4].is_subnormal());
2396
2397 let left = signed_probit_neglog_unary_stack(-1.0e100, 1.0);
2398 assert_eq!(left[1], -1.0e100);
2399 assert_eq!(left[2], 1.0);
2400 assert!(left[3] < 0.0 && left[3].is_finite());
2401 assert_eq!(left[4], -0.0);
2402 }
2403
2404 fn rigid_standard_normal_third_and_fourth_full(
2412 marginal: BernoulliMarginalLinkMap,
2413 g: f64,
2414 z: f64,
2415 y: f64,
2416 w: f64,
2417 probit_scale: f64,
2418 ) -> Result<([[[f64; 2]; 2]; 2], [[[[f64; 2]; 2]; 2]; 2]), String> {
2419 let tower = rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?;
2420 Ok((tower.t3, tower.t4))
2421 }
2422 use gam_math::jet_tower::{
2423 KernelChannels, RowProgram, program_full_tower, verify_kernel_channels,
2424 };
2425
2426 struct BernoulliRigidStandardNormalNllProgram {
2429 primaries: Vec<[f64; 2]>,
2431 z: Vec<f64>,
2433 y: Vec<f64>,
2434 w: Vec<f64>,
2435 probit_scale: f64,
2436 }
2437
2438 impl RowProgram<2> for BernoulliRigidStandardNormalNllProgram {
2439 fn n_rows(&self) -> usize {
2440 self.primaries.len()
2441 }
2442
2443 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2444 self.primaries
2445 .get(row)
2446 .copied()
2447 .ok_or_else(|| format!("bernoulli rigid nll program: row {row} out of range"))
2448 }
2449
2450 fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
2451 &self,
2452 row: usize,
2453 p: &[S; 2],
2454 ) -> Result<S, String> {
2455 let z = self.z[row];
2456 let y = self.y[row];
2457 let w = self.w[row];
2458 let s = self.probit_scale;
2459 let eta_marginal = p[0];
2463 let link = bernoulli_marginal_link_map(
2464 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2465 eta_marginal.value(),
2466 )?;
2467 let q = eta_marginal.compose_unary([link.q, link.q1, link.q2, link.q3, link.q4]);
2468 let g = p[1];
2469 let observed_slope = g.scale(s);
2471 let one_plus_slope_squared = observed_slope.mul(&observed_slope).add(&S::constant(1.0));
2472 let c = one_plus_slope_squared
2473 .compose_unary(unary_derivatives_sqrt(one_plus_slope_squared.value()));
2474 let eta = q.mul(&c).add(&observed_slope.scale(z));
2476 let signed = eta.scale(2.0 * y - 1.0);
2477 Ok(signed.compose_unary(unary_derivatives_neglog_phi(signed.value(), w)))
2479 }
2480 }
2481
2482 fn scalar_nll(eta_marginal: f64, g: f64, z: f64, y: f64, w: f64, s: f64) -> f64 {
2485 let link = bernoulli_marginal_link_map(
2486 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2487 eta_marginal,
2488 )
2489 .unwrap();
2490 let observed_slope = g * s;
2491 let c = (observed_slope * observed_slope + 1.0).sqrt();
2492 let eta = link.q * c + observed_slope * z;
2493 let signed = (2.0 * y - 1.0) * eta;
2494 let cdf = 0.5 * libm::erfc(-signed / std::f64::consts::SQRT_2);
2495 -w * cdf.max(1e-300).ln()
2496 }
2497
2498 #[test]
2499 fn rigid_bernoulli_row_kernel_agrees_with_jet_tower_program_all_channels() {
2500 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2504 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2505 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2506 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2507 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2508 let n = eta.len();
2509
2510 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2512
2513 for &probit_scale in &[1.0_f64, 0.8] {
2514 let program = BernoulliRigidStandardNormalNllProgram {
2515 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2516 z: z.to_vec(),
2517 y: y.to_vec(),
2518 w: w.to_vec(),
2519 probit_scale,
2520 };
2521
2522 for row in 0..n {
2523 let tower = program_full_tower(&program, row).expect("tower evaluation");
2524
2525 let marginal = bernoulli_marginal_link_map(
2527 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2528 eta[row],
2529 )
2530 .expect("link map");
2531 let (value, gradient, hessian) = rigid_standard_normal_row_kernel(
2532 marginal,
2533 g[row],
2534 z[row],
2535 y[row],
2536 w[row],
2537 probit_scale,
2538 )
2539 .expect("production row kernel");
2540
2541 let (third_full, fourth_full) = rigid_standard_normal_third_and_fourth_full(
2548 marginal,
2549 g[row],
2550 z[row],
2551 y[row],
2552 w[row],
2553 probit_scale,
2554 )
2555 .expect("production third+fourth");
2556 let third: Vec<([f64; 2], [[f64; 2]; 2])> = dirs
2557 .iter()
2558 .map(|d| (*d, contract_third_full(&third_full, d[0], d[1])))
2559 .collect();
2560
2561 let fourth: Vec<([f64; 2], [f64; 2], [[f64; 2]; 2])> = dirs
2562 .iter()
2563 .enumerate()
2564 .map(|(i, u)| {
2565 let v = dirs[(i + 1) % dirs.len()];
2566 (
2567 *u,
2568 v,
2569 contract_fourth_full(&fourth_full, u[0], u[1], v[0], v[1]),
2570 )
2571 })
2572 .collect();
2573
2574 let claims = KernelChannels {
2575 value,
2576 gradient,
2577 hessian,
2578 third,
2579 fourth,
2580 };
2581
2582 verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
2583 panic!(
2584 "probit_scale {probit_scale} row {row}: production rigid Bernoulli \
2585 RowKernel disagrees with #932 jet-tower truth: {e}"
2586 )
2587 });
2588
2589 let h = 1e-3;
2593 let f = |de: f64, dg: f64| {
2594 scalar_nll(
2595 eta[row] + de,
2596 g[row] + dg,
2597 z[row],
2598 y[row],
2599 w[row],
2600 probit_scale,
2601 )
2602 };
2603 let f0 = f(0.0, 0.0);
2604 assert!(
2605 (f0 - tower.v).abs() <= 1e-9 * f0.abs().max(1.0),
2606 "row {row}: independent scalar NLL {f0:+.12e} != tower value {:+.12e}",
2607 tower.v
2608 );
2609 let g_eta = (f(-2.0 * h, 0.0) - 8.0 * f(-h, 0.0) + 8.0 * f(h, 0.0)
2611 - f(2.0 * h, 0.0))
2612 / (12.0 * h);
2613 let g_g = (f(0.0, -2.0 * h) - 8.0 * f(0.0, -h) + 8.0 * f(0.0, h) - f(0.0, 2.0 * h))
2614 / (12.0 * h);
2615 for (label, fd, ad) in [("∂η", g_eta, tower.g[0]), ("∂g", g_g, tower.g[1])] {
2616 assert!(
2617 (fd - ad).abs() <= 1e-5 * ad.abs().max(1.0),
2618 "row {row} {label}: FD witness {fd:+.6e} != tower grad {ad:+.6e}"
2619 );
2620 }
2621 }
2622 }
2623 }
2624
2625 #[test]
2635 fn rigid_third_and_fourth_full_shares_one_tower_bit_identical() {
2636 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2637 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2638 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2639 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2640 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2641 for &probit_scale in &[1.0_f64, 0.8] {
2642 for r in 0..eta.len() {
2643 let marginal = bernoulli_marginal_link_map(
2644 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2645 eta[r],
2646 )
2647 .expect("link map");
2648 let t3_sep = rigid_standard_normal_third_full(
2649 marginal,
2650 g[r],
2651 z[r],
2652 y[r],
2653 w[r],
2654 probit_scale,
2655 )
2656 .expect("separate third");
2657 let t4_sep = rigid_standard_normal_fourth_full(
2658 marginal,
2659 g[r],
2660 z[r],
2661 y[r],
2662 w[r],
2663 probit_scale,
2664 )
2665 .expect("separate fourth");
2666 let (t3_comb, t4_comb) = rigid_standard_normal_third_and_fourth_full(
2667 marginal,
2668 g[r],
2669 z[r],
2670 y[r],
2671 w[r],
2672 probit_scale,
2673 )
2674 .expect("combined third+fourth");
2675 for a in 0..2 {
2677 for b in 0..2 {
2678 for c in 0..2 {
2679 assert_eq!(
2680 t3_comb[a][b][c], t3_sep[a][b][c],
2681 "t3[{a}][{b}][{c}] row {r} scale {probit_scale} not bit-identical"
2682 );
2683 for d in 0..2 {
2684 assert_eq!(
2685 t4_comb[a][b][c][d], t4_sep[a][b][c][d],
2686 "t4[{a}][{b}][{c}][{d}] row {r} scale {probit_scale} not bit-identical"
2687 );
2688 }
2689 }
2690 }
2691 }
2692 }
2693 }
2694 }
2695
2696 #[test]
2706 fn rigid_bernoulli_generic_program_matches_independent_program_all_channels() {
2707 use gam_math::jet_tower::{
2708 program_fourth_contracted, program_row_kernel, program_third_contracted,
2709 };
2710
2711 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2712 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2713 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2714 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2715 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2716 let n = eta.len();
2717 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2718
2719 let close = |a: f64, b: f64, label: &str| {
2720 let band = 1e-12 + 1e-12 * a.abs().max(b.abs());
2721 assert!(
2722 (a - b).abs() <= band,
2723 "{label}: generic {a:+.15e} vs Tower4-program {b:+.15e} (band {band:.3e})"
2724 );
2725 };
2726
2727 for &probit_scale in &[1.0_f64, 0.8] {
2728 let tower_program = BernoulliRigidStandardNormalNllProgram {
2730 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2731 z: z.to_vec(),
2732 y: y.to_vec(),
2733 w: w.to_vec(),
2734 probit_scale,
2735 };
2736
2737 for row in 0..n {
2738 let truth = program_full_tower(&tower_program, row).expect("program tower");
2739
2740 let marginal = bernoulli_marginal_link_map(
2741 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2742 eta[row],
2743 )
2744 .expect("link map");
2745 let program = RigidStandardNormalRow {
2746 marginal,
2747 g: g[row],
2748 z: z[row],
2749 y: y[row],
2750 w: w[row],
2751 probit_scale,
2752 };
2753
2754 let full = program_full_tower(&program, 0).expect("generic full tower");
2757 close(full.v, truth.v, "full value");
2758 for a in 0..2 {
2759 close(full.g[a], truth.g[a], "full grad");
2760 for b in 0..2 {
2761 close(full.h[a][b], truth.h[a][b], "full hess");
2762 for c in 0..2 {
2763 close(full.t3[a][b][c], truth.t3[a][b][c], "full t3");
2764 for d in 0..2 {
2765 close(full.t4[a][b][c][d], truth.t4[a][b][c][d], "full t4");
2766 }
2767 }
2768 }
2769 }
2770
2771 let (val, grad, hess) =
2773 program_row_kernel(&program, 0).expect("generic row kernel");
2774 close(val, truth.v, "order2 value");
2775 for a in 0..2 {
2776 close(grad[a], truth.g[a], "order2 grad");
2777 for b in 0..2 {
2778 close(hess[a][b], truth.h[a][b], "order2 hess");
2779 }
2780 }
2781
2782 for dir in &dirs {
2785 let third = program_third_contracted(&program, 0, dir)
2786 .expect("generic third contracted");
2787 let truth3 = truth.third_contracted(dir);
2788 for a in 0..2 {
2789 for b in 0..2 {
2790 close(third[a][b], truth3[a][b], "third contracted");
2791 }
2792 }
2793 }
2794
2795 for (i, u) in dirs.iter().enumerate() {
2798 let v = dirs[(i + 1) % dirs.len()];
2799 let fourth = program_fourth_contracted(&program, 0, u, &v)
2800 .expect("generic fourth contracted");
2801 let truth4 = truth.fourth_contracted(u, &v);
2802 for a in 0..2 {
2803 for b in 0..2 {
2804 close(fourth[a][b], truth4[a][b], "fourth contracted");
2805 }
2806 }
2807 }
2808 }
2809 }
2810 }
2811
2812 #[inline(always)]
2818 fn hand_rigid_vgh(
2819 marginal: BernoulliMarginalLinkMap,
2820 g: f64,
2821 z: f64,
2822 y: f64,
2823 w: f64,
2824 probit_scale: f64,
2825 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2826 let s = 2.0 * y - 1.0;
2827 let observed_logslope = probit_scale * g;
2828 let g2 = observed_logslope * observed_logslope;
2829 let c = (1.0 + g2).sqrt();
2830 let c1 = probit_scale * observed_logslope / c;
2831 let c_inv3 = 1.0 / (c * c * c);
2832 let c2 = probit_scale * probit_scale * c_inv3;
2833 let q = marginal.q;
2834 let eta = q * c + observed_logslope * z;
2836 let m = s * eta;
2837 let stack = signed_probit_neglog_unary_stack(m, w);
2838 let (k1, k2) = (stack[1], stack[2]);
2839 let u1 = s * k1;
2840 let u2 = k2;
2841 let eta_q = c;
2842 let eta_g = q * c1 + probit_scale * z;
2843 let value = stack[0];
2844 let gradient = [u1 * eta_q * marginal.q1, u1 * eta_g];
2846 let h00 = u2 * eta_q * eta_q;
2848 let h01 = u2 * eta_q * eta_g + u1 * c1;
2849 let h11 = u2 * eta_g * eta_g + u1 * q * c2;
2850 let grad_q = u1 * eta_q;
2852 let hessian = [
2853 [
2854 h00 * marginal.q1 * marginal.q1 + grad_q * marginal.q2,
2855 h01 * marginal.q1,
2856 ],
2857 [h01 * marginal.q1, h11],
2858 ];
2859 (value, gradient, hessian)
2860 }
2861
2862 #[inline(never)]
2863 fn measured_production_rigid_vgh(
2864 marginal: BernoulliMarginalLinkMap,
2865 g: f64,
2866 z: f64,
2867 y: f64,
2868 w: f64,
2869 probit_scale: f64,
2870 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2871 rigid_standard_normal_row_kernel(marginal, g, z, y, w, probit_scale)
2872 .expect("generated rigid row")
2873 }
2874
2875 #[inline(never)]
2876 fn measured_hand_rigid_vgh(
2877 marginal: BernoulliMarginalLinkMap,
2878 g: f64,
2879 z: f64,
2880 y: f64,
2881 w: f64,
2882 probit_scale: f64,
2883 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2884 hand_rigid_vgh(marginal, g, z, y, w, probit_scale)
2885 }
2886
2887 #[test]
2893 fn rigid_bernoulli_row_kernel_matches_hand_chain_witness() {
2894 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2895 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2896 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2897 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2898 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2899 let close = |a: f64, b: f64, label: &str| {
2900 let band = 1e-12 + 1e-9 * a.abs().max(b.abs());
2901 assert!(
2902 (a - b).abs() <= band,
2903 "{label}: jet {a:+.15e} vs hand {b:+.15e} (band {band:.3e})"
2904 );
2905 };
2906 for &probit_scale in &[1.0_f64, 0.8] {
2907 for r in 0..eta.len() {
2908 let marginal = bernoulli_marginal_link_map(
2909 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2910 eta[r],
2911 )
2912 .expect("link map");
2913 let (jv, jg, jh) = rigid_standard_normal_row_kernel(
2914 marginal,
2915 g[r],
2916 z[r],
2917 y[r],
2918 w[r],
2919 probit_scale,
2920 )
2921 .expect("jet kernel");
2922 let (hv, hg, hh) = hand_rigid_vgh(marginal, g[r], z[r], y[r], w[r], probit_scale);
2923 close(jv, hv, "value");
2924 for a in 0..2 {
2925 close(jg[a], hg[a], "grad");
2926 for b in 0..2 {
2927 close(jh[a][b], hh[a][b], "hess");
2928 }
2929 }
2930 }
2931 }
2932 }
2933
2934 #[test]
2943 fn release_measure_rigid_bernoulli_vgh_vs_hand_chain_932() {
2944 use gam_math::paired_timing::paired_interleaved;
2945
2946 let cases = [
2950 (0.3_f64, 0.2_f64, 0.4_f64, 1.0_f64, 1.0_f64),
2951 (-0.7, -0.5, -1.1, 0.0, 0.8),
2952 ];
2953 let probit_scale = 0.8;
2954
2955 let reps = 15usize;
2962 let iterations = 300_000usize;
2963
2964 for (case_idx, &(eta, g, z, y, w)) in cases.iter().enumerate() {
2965 let marginal = bernoulli_marginal_link_map(
2966 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2967 eta,
2968 )
2969 .expect("link map");
2970
2971 let (jet_value, ..) =
2976 rigid_standard_normal_row_kernel(marginal, g, z, y, w, probit_scale)
2977 .expect("jet kernel");
2978 let (hand_value, ..) = hand_rigid_vgh(marginal, g, z, y, w, probit_scale);
2979 let band = 1e-12 + 1e-9 * jet_value.abs().max(hand_value.abs());
2980 assert!(
2981 (jet_value - hand_value).abs() <= band,
2982 "y={y:.0} value: jet {jet_value:+.15e} vs hand {hand_value:+.15e}"
2983 );
2984
2985 let timing = paired_interleaved(
2990 reps,
2991 iterations,
2992 0x9320_0BAD ^ case_idx as u64,
2993 |nudge| {
2994 let (value, gradient, hessian) = measured_production_rigid_vgh(
2995 marginal,
2996 g + nudge,
2997 z,
2998 y,
2999 w,
3000 probit_scale,
3001 );
3002 value + gradient[0] + hessian[0][0]
3003 },
3004 |nudge| {
3005 let (value, gradient, hessian) =
3006 measured_hand_rigid_vgh(marginal, g + nudge, z, y, w, probit_scale);
3007 value + gradient[0] + hessian[0][0]
3008 },
3009 );
3010
3011 eprintln!(
3014 "RIGID-BERNOULLI-VGH-932 y={y:.0} {}",
3015 timing.summary("production", "hand"),
3016 );
3017
3018 assert!(
3039 cfg!(debug_assertions)
3040 || (timing.median_ratio() > 1.0 && timing.wins_fraction() >= 0.75),
3041 "generated rigid BMS y={y:.0} must beat strongest hand: {}",
3042 timing.summary("production", "hand"),
3043 );
3044 }
3045 }
3046}
3047
3048#[cfg(test)]
3049mod flex_primary_hessian_oracle_tests {
3050 use super::*;
3074 use super::family::*;
3081 use gam_linalg::matrix::DenseDesignMatrix;
3082 use ndarray::Array1;
3083 use ndarray::Array2;
3084 use std::sync::Arc;
3085 use std::sync::Mutex;
3086
3087 fn make_flex_oracle_family(
3093 n: usize,
3094 ) -> (BernoulliMarginalSlopeFamily, Vec<ParameterBlockState>) {
3095 let score_seed = Array1::linspace(-2.0, 2.0, n.max(6));
3096 let link_seed = Array1::linspace(-1.8, 1.8, n.max(6));
3097 let cfg = DeviationBlockConfig {
3098 num_internal_knots: 3,
3099 ..DeviationBlockConfig::default()
3100 };
3101 let score_prepared = build_score_warp_deviation_block_from_seed(&score_seed, &cfg)
3102 .expect("build score warp block");
3103 let link_prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
3104 &link_seed, &link_seed, &cfg,
3105 )
3106 .expect("build link deviation block");
3107
3108 let y: Array1<f64> =
3109 Array1::from_iter((0..n).map(|i| if (i * 17 + 3) % 7 >= 4 { 1.0 } else { 0.0 }));
3110 let weights: Array1<f64> =
3111 Array1::from_iter((0..n).map(|i| 0.75 + ((i * 11 + 5) % 5) as f64 * 0.05));
3112 let z: Array1<f64> =
3113 Array1::from_iter((0..n).map(|i| -1.7 + 3.4 * (i as f64 + 0.5) / n as f64));
3114 let marginal_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3115 if j == 0 {
3116 1.0
3117 } else {
3118 -0.4 + 0.8 * ((i * 19 + 7) % n) as f64 / n as f64
3119 }
3120 });
3121 let logslope_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3122 if j == 0 {
3123 1.0
3124 } else {
3125 0.3 - 0.6 * ((i * 23 + 11) % n) as f64 / n as f64
3126 }
3127 });
3128
3129 let family = BernoulliMarginalSlopeFamily {
3130 y: Arc::new(y),
3131 weights: Arc::new(weights),
3132 z: Arc::new(z.clone()),
3133 latent_measure: LatentMeasureKind::StandardNormal,
3134 gaussian_frailty_sd: Some(0.15),
3135 base_link: InverseLink::Standard(gam_problem::StandardLink::Probit),
3136 marginal_design: DesignMatrix::Dense(DenseDesignMatrix::from(marginal_x.clone())),
3137 logslope_design: DesignMatrix::Dense(DenseDesignMatrix::from(logslope_x.clone())),
3138 score_warp: Some(score_prepared.runtime.clone()),
3139 link_dev: Some(link_prepared.runtime.clone()),
3140 policy: gam_runtime::resource::ResourcePolicy::default_library(),
3141 cell_moment_lru: Arc::new(exact_kernel::CellMomentLruCache::new(1024)),
3142 cell_moment_cache_stats: Arc::new(exact_kernel::CellMomentCacheStats::default()),
3143 intercept_warm_starts: None,
3144 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
3145 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
3146 };
3147
3148 let beta_m = Array1::from_vec(vec![0.12, -0.04]);
3149 let beta_g = Array1::from_vec(vec![0.35, 0.03]);
3150 let beta_h = Array1::from_iter(
3151 (0..score_prepared.runtime.basis_dim()).map(|idx| 0.0015 * (idx as f64 + 1.0)),
3152 );
3153 let beta_w = Array1::from_iter(
3154 (0..link_prepared.runtime.basis_dim()).map(|idx| -0.001 * (idx as f64 + 1.0)),
3155 );
3156 let states = vec![
3157 ParameterBlockState {
3158 eta: marginal_x.dot(&beta_m),
3159 beta: beta_m,
3160 },
3161 ParameterBlockState {
3162 eta: logslope_x.dot(&beta_g),
3163 beta: beta_g,
3164 },
3165 ParameterBlockState {
3166 beta: beta_h,
3167 eta: Array1::zeros(z.len()),
3168 },
3169 ParameterBlockState {
3170 beta: beta_w,
3171 eta: Array1::zeros(z.len()),
3172 },
3173 ];
3174 (family, states)
3175 }
3176
3177 fn flex_gradient_at_perturbed(
3185 family: &BernoulliMarginalSlopeFamily,
3186 states: &[ParameterBlockState],
3187 primary: &super::super::hessian_paths::PrimarySlices,
3188 row: usize,
3189 u: usize,
3190 delta: f64,
3191 ) -> Array1<f64> {
3192 let mut states = states.to_vec();
3193 if u == primary.q {
3199 states[0].eta[row] += delta;
3200 } else if u == primary.logslope {
3201 states[1].eta[row] += delta;
3202 } else if let Some(h_range) = primary.h.as_ref()
3203 && h_range.contains(&u)
3204 {
3205 states[2].beta[u - h_range.start] += delta;
3206 } else if let Some(w_range) = primary.w.as_ref()
3207 && w_range.contains(&u)
3208 {
3209 states[3].beta[u - w_range.start] += delta;
3210 } else {
3211 panic!("primary coordinate {u} out of range for flex oracle");
3212 }
3213 let row_ctx = family
3214 .build_row_exact_context_with_stats_and_cell_cache(row, &states, None, false)
3215 .expect("perturbed row context");
3216 let (_neglog, grad, _hess) = family
3217 .compute_row_primary_gradient_hessian(row, &states, primary, &row_ctx)
3218 .expect("perturbed gradient");
3219 grad
3220 }
3221
3222 fn flex_nll_gradient_at_perturbed_z(
3228 family: &BernoulliMarginalSlopeFamily,
3229 states: &[ParameterBlockState],
3230 primary: &super::super::hessian_paths::PrimarySlices,
3231 row: usize,
3232 delta: f64,
3233 ) -> Array1<f64> {
3234 let mut perturbed = family.clone();
3235 let mut z = family.z.as_ref().clone();
3236 z[row] += delta;
3237 perturbed.z = Arc::new(z);
3238 let row_ctx = perturbed
3239 .build_row_exact_context_with_stats_and_cell_cache(row, states, None, false)
3240 .expect("z-perturbed row context");
3241 let mut scratch =
3242 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3243 perturbed
3244 .lower_bms_flex_row_order2(row, states, primary, &row_ctx, None, false, &mut scratch)
3245 .expect("z-perturbed flex gradient");
3246 scratch.grad
3247 }
3248
3249 #[test]
3255 fn flex_score_zeta_sensitivity_covers_all_active_deviation_blocks_2303() {
3256 let n = 12usize;
3257 let (family, states) = make_flex_oracle_family(n);
3258 let cache = family
3259 .build_exact_eval_cache(&states)
3260 .expect("flex exact eval cache");
3261 let primary = &cache.primary;
3262 let row = 5usize;
3263 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3264 let mut scratch =
3265 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3266 family
3267 .lower_bms_flex_row_order2(row, &states, primary, row_ctx, None, false, &mut scratch)
3268 .expect("analytic flex z-sensitivity");
3269 let analytic = scratch.score_zeta.clone();
3270
3271 let h = 1e-5_f64;
3272 let nll_plus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, h);
3273 let nll_minus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, -h);
3274 for u in 0..primary.total {
3275 let finite_difference = -(nll_plus[u] - nll_minus[u]) / (2.0 * h);
3277 let scale = 1.0 + analytic[u].abs().max(finite_difference.abs());
3278 let relative_error = (analytic[u] - finite_difference).abs() / scale;
3279 assert!(
3280 relative_error <= 2e-6,
3281 "flex score-zeta primary {u}: analytic={} FD={} relative_error={relative_error}",
3282 analytic[u],
3283 finite_difference
3284 );
3285 }
3286 let h_range = primary.h.as_ref().expect("score-warp primary range");
3287 let w_range = primary.w.as_ref().expect("link-deviation primary range");
3288 assert!(
3289 analytic
3290 .slice(s![h_range.start..h_range.end])
3291 .iter()
3292 .any(|value| value.abs() > 1e-10),
3293 "score-warp z-sensitivity must not be zero-filled"
3294 );
3295 assert!(
3296 analytic
3297 .slice(s![w_range.start..w_range.end])
3298 .iter()
3299 .any(|value| value.abs() > 1e-10),
3300 "link-deviation z-sensitivity must not be zero-filled"
3301 );
3302
3303 let coefficient = family
3304 .flex_score_zeta_sensitivity(
3305 &states,
3306 &BlockwiseFitOptions::default(),
3307 cache.slices.total,
3308 )
3309 .expect("full coefficient score-zeta sensitivity");
3310 assert_eq!(coefficient.dim(), (n, cache.slices.total));
3311 for range in [
3312 cache.slices.h.as_ref().expect("score-warp beta range"),
3313 cache.slices.w.as_ref().expect("link-deviation beta range"),
3314 ] {
3315 assert!(
3316 coefficient
3317 .slice(s![.., range.start..range.end])
3318 .iter()
3319 .any(|value| value.abs() > 1e-10),
3320 "active deviation coefficient range {range:?} must carry Murphy-Topel sensitivity"
3321 );
3322 }
3323
3324 let wrong_width = cache
3325 .slices
3326 .total
3327 .checked_sub(1)
3328 .expect("nonempty coefficient frame");
3329 let error = family
3330 .flex_score_zeta_sensitivity(&states, &BlockwiseFitOptions::default(), wrong_width)
3331 .expect_err("partial Murphy-Topel covariance frame must be rejected");
3332 assert!(
3333 error.contains("covariance/frame mismatch"),
3334 "unexpected partial-frame error: {error}"
3335 );
3336 }
3337
3338 #[test]
3341 fn flex_primary_hessian_matches_central_fd_of_gradient() {
3342 let n = 12usize;
3343 let (family, states) = make_flex_oracle_family(n);
3344 let cache = family
3345 .build_exact_eval_cache(&states)
3346 .expect("flex exact eval cache");
3347 let primary = &cache.primary;
3348 let r = primary.total;
3349 assert!(
3350 r >= 4,
3351 "flex fixture must carry q + logslope + deviation blocks"
3352 );
3353
3354 let h = 1e-4;
3358 let mut max_rel = 0.0_f64;
3359
3360 for &row in &[2usize, 5, 8] {
3363 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3364 let (_neglog, _grad, analytic_hess) = family
3365 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3366 .expect("analytic flex gradient + hessian");
3367
3368 for u in 0..r {
3369 let grad_plus = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3370 let grad_minus = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3371 for v in 0..r {
3372 let fd = (grad_plus[v] - grad_minus[v]) / (2.0 * h);
3373 let analytic = analytic_hess[[v, u]];
3374 let denom = 1.0 + analytic.abs().max(fd.abs());
3375 let rel = (analytic - fd).abs() / denom;
3376 max_rel = max_rel.max(rel);
3377 assert!(
3378 rel <= 1e-6,
3379 "flex hand Hessian H[{v}][{u}] = {analytic:.6e} disagrees with central \
3380 FD of the gradient {fd:.6e} at row {row} (rel {rel:.3e}); a product-rule \
3381 term is dropped or mis-signed"
3382 );
3383 }
3384 }
3385 }
3386 assert!(
3388 max_rel <= 1e-6,
3389 "flex Hessian FD oracle max rel {max_rel:.3e}"
3390 );
3391 }
3392
3393 #[test]
3402 fn arbiter_flex_hessian_h00_fd_step_scaling() {
3403 let n = 12usize;
3404 let (family, states) = make_flex_oracle_family(n);
3405 let cache = family
3406 .build_exact_eval_cache(&states)
3407 .expect("flex exact eval cache");
3408 let primary = &cache.primary;
3409 let row = 2usize;
3410 let u = primary.q; let v = primary.q;
3412
3413 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3414 let (_neglog, _grad, analytic_hess) = family
3415 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3416 .expect("analytic flex gradient + hessian");
3417 let analytic = analytic_hess[[v, u]];
3418
3419 let fd_at = |h: f64| -> f64 {
3420 let gp = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3421 let gm = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3422 (gp[v] - gm[v]) / (2.0 * h)
3423 };
3424
3425 let h = 1e-3_f64;
3432 let fd_h = fd_at(h);
3433 let fd_half = fd_at(h * 0.5);
3434 let fd_quarter = fd_at(h * 0.25);
3435 let gap_h = (analytic - fd_h).abs();
3436 let gap_half = (analytic - fd_half).abs();
3437 let gap_quarter = (analytic - fd_quarter).abs();
3438 let rich = (4.0 * fd_half - fd_h) / 3.0;
3439 let rich_gap = (analytic - rich).abs();
3440 let denom = analytic.abs().max(1.0);
3441
3442 let record = format!(
3444 "FLEX H[0][0] ARBITER row 2: analytic={analytic:+.12e} \
3445 fd(h)={fd_h:+.12e} fd(h/2)={fd_half:+.12e} fd(h/4)={fd_quarter:+.12e} \
3446 gap(h)={gap_h:.3e} gap(h/2)={gap_half:.3e} gap(h/4)={gap_quarter:.3e} \
3447 ratio_h_over_half={:.3} ratio_half_over_quarter={:.3} \
3448 richardson={rich:+.12e} richardson_gap={rich_gap:.3e} (rich_rel={:.3e})",
3449 gap_h / gap_half.max(f64::MIN_POSITIVE),
3450 gap_half / gap_quarter.max(f64::MIN_POSITIVE),
3451 rich_gap / denom,
3452 );
3453
3454 assert!(
3460 rich_gap / denom <= 1e-7,
3461 "{record}\nVERDICT: Richardson residual exceeds the FD-truncation floor — \
3462 the hand H[0][0] genuinely diverges (real dropped/mis-signed term), NOT FD noise"
3463 );
3464 }
3465}