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::jet_tower::Tower4;
7use gam_math::probability::normal_logcdf_derivatives;
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 |_scale, 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]
690pub(super) fn rigid_observed_logslope(logslope: f64, probit_scale: f64) -> f64 {
691 probit_scale * logslope
692}
693
694#[inline]
695pub(super) fn rigid_observed_scale(logslope: f64, probit_scale: f64) -> f64 {
696 let observed_logslope = rigid_observed_logslope(logslope, probit_scale);
697 (1.0 + observed_logslope * observed_logslope).sqrt()
698}
699
700#[inline]
701pub(super) fn rigid_intercept_from_marginal(
702 marginal_eta: f64,
703 logslope: f64,
704 probit_scale: f64,
705) -> f64 {
706 marginal_eta * rigid_observed_scale(logslope, probit_scale)
707}
708
709#[inline]
710pub(super) fn rigid_prescale_intercept_from_marginal(
711 marginal_eta: f64,
712 logslope: f64,
713 probit_scale: f64,
714) -> f64 {
715 rigid_intercept_from_marginal(marginal_eta, logslope, probit_scale) / probit_scale
716}
717
718#[inline]
719pub(super) fn rigid_prescale_intercept_derivative_abs(
720 marginal_eta: f64,
721 logslope: f64,
722 probit_scale: f64,
723) -> f64 {
724 let c = rigid_observed_scale(logslope, probit_scale);
725 probit_scale * normal_pdf(marginal_eta) / c
726}
727
728#[inline]
729pub(super) fn rigid_observed_eta(
730 marginal_eta: f64,
731 logslope: f64,
732 z: f64,
733 probit_scale: f64,
734) -> f64 {
735 marginal_slope_standard_normal_scalar_eta(marginal_eta, logslope, z, probit_scale)
736}
737
738#[inline]
739pub(super) fn marginal_slope_standard_normal_scalar_eta(
740 q: f64,
741 slope: f64,
742 z: f64,
743 probit_scale: f64,
744) -> f64 {
745 let observed_slope = rigid_observed_logslope(slope, probit_scale);
746 q * (1.0 + observed_slope * observed_slope).sqrt() + observed_slope * z
747}
748
749pub(super) fn unary_derivatives_normal_cdf(x: f64) -> [f64; 5] {
750 let pdf = normal_pdf(x);
751 [
752 normal_cdf(x),
753 pdf,
754 -x * pdf,
755 (x * x - 1.0) * pdf,
756 (-x.powi(3) + 3.0 * x) * pdf,
757 ]
758}
759
760#[inline]
767pub(super) fn lse_accumulate(log_max: &mut f64, sum: &mut f64, log_term: f64) {
768 if !log_term.is_finite() {
769 return;
770 }
771 if log_term > *log_max {
772 if log_max.is_finite() {
773 *sum = *sum * (*log_max - log_term).exp() + 1.0;
774 } else {
775 *sum = 1.0;
776 }
777 *log_max = log_term;
778 } else {
779 *sum += (log_term - *log_max).exp();
780 }
781}
782
783#[derive(Clone, Copy, Debug, PartialEq, Eq)]
784pub enum MarginalSlopeCovarianceShape {
785 Diagonal,
786 Full,
787 LowRank,
788}
789
790#[derive(Clone, Debug, PartialEq)]
791enum MarginalSlopeCovarianceStorage {
792 Diagonal {
793 covariance: Array1<f64>,
794 },
795 Full {
796 covariance: Array2<f64>,
797 square_root_factor: Array2<f64>,
799 },
800 LowRank {
802 factor: Array2<f64>,
803 },
804}
805
806#[derive(Clone, Debug)]
815pub struct MarginalSlopeCovariance {
816 storage: MarginalSlopeCovarianceStorage,
817 ones_quadratic_form: f64,
818}
819
820impl PartialEq for MarginalSlopeCovariance {
821 fn eq(&self, other: &Self) -> bool {
822 self.storage == other.storage
823 }
824}
825
826#[derive(Clone, Copy, Debug)]
827pub(crate) enum MarginalSlopeCovarianceRef<'a> {
828 Diagonal(&'a Array1<f64>),
829 Full(&'a Array2<f64>),
830 LowRank(&'a Array2<f64>),
831}
832
833impl MarginalSlopeCovariance {
834 pub fn diagonal(covariance: Array1<f64>) -> Result<Self, String> {
835 if covariance.is_empty() {
836 return Err("marginal-slope diagonal covariance is empty".to_string());
837 }
838 let mut ones_quadratic_form = 0.0;
839 for (axis, &value) in covariance.iter().enumerate() {
840 if !(value.is_finite() && value >= 0.0) {
841 return Err(format!(
842 "marginal-slope diagonal covariance entry {axis} must be finite and non-negative, got {value}"
843 ));
844 }
845 ones_quadratic_form += value;
846 }
847 if !ones_quadratic_form.is_finite() {
848 return Err("marginal-slope diagonal covariance geometry overflowed".to_string());
849 }
850 Ok(Self {
851 storage: MarginalSlopeCovarianceStorage::Diagonal { covariance },
852 ones_quadratic_form,
853 })
854 }
855
856 pub fn full(covariance: Array2<f64>) -> Result<Self, String> {
857 if covariance.nrows() == 0 || covariance.nrows() != covariance.ncols() {
858 return Err(format!(
859 "marginal-slope full covariance must be non-empty and square, got {}x{}",
860 covariance.nrows(),
861 covariance.ncols(),
862 ));
863 }
864 for ((row, column), &value) in covariance.indexed_iter() {
865 if !value.is_finite() {
866 return Err(format!(
867 "marginal-slope full covariance entry ({row},{column}) is non-finite"
868 ));
869 }
870 }
871 for row in 0..covariance.nrows() {
872 for column in (row + 1)..covariance.ncols() {
873 if covariance[[row, column]] != covariance[[column, row]] {
874 return Err(format!(
875 "marginal-slope full covariance must be exactly symmetric at ({row},{column}): upper={}, lower={}",
876 covariance[[row, column]],
877 covariance[[column, row]],
878 ));
879 }
880 }
881 }
882 let (eigenvalues, eigenvectors) = covariance.eigh(faer::Side::Lower).map_err(|error| {
883 format!("marginal-slope covariance eigendecomposition failed: {error}")
884 })?;
885 let dimension = covariance.nrows();
886 let spectral_magnitude = eigenvalues
900 .iter()
901 .fold(0.0_f64, |magnitude, &value| magnitude.max(value.abs()));
902 let psd_roundoff = 128.0 * dimension as f64 * f64::EPSILON * spectral_magnitude;
903 let mut square_root_factor = Array2::<f64>::zeros((dimension, dimension));
904 for (eigen_axis, &eigenvalue) in eigenvalues.iter().enumerate() {
905 if !eigenvalue.is_finite() || eigenvalue < -psd_roundoff {
906 return Err(format!(
907 "marginal-slope full covariance must be positive semidefinite; eigenvalue {eigen_axis} is {eigenvalue} (admissible eigensolver band -{psd_roundoff:.3e})"
908 ));
909 }
910 let scale = eigenvalue.max(0.0).sqrt();
915 for axis in 0..dimension {
916 square_root_factor[[eigen_axis, axis]] = scale * eigenvectors[[axis, eigen_axis]];
917 }
918 }
919 let mut ones_quadratic_form = 0.0;
920 for factor_row in square_root_factor.rows() {
921 let projection = factor_row.sum();
922 ones_quadratic_form += projection * projection;
923 }
924 if !ones_quadratic_form.is_finite() {
925 return Err("marginal-slope full covariance geometry overflowed".to_string());
926 }
927 Ok(Self {
928 storage: MarginalSlopeCovarianceStorage::Full {
929 covariance,
930 square_root_factor,
931 },
932 ones_quadratic_form,
933 })
934 }
935
936 pub fn low_rank(factor: Array2<f64>) -> Result<Self, String> {
937 if factor.nrows() == 0 {
938 return Err("marginal-slope low-rank covariance factor has zero rows".to_string());
939 }
940 for ((row, column), &value) in factor.indexed_iter() {
941 if !value.is_finite() {
942 return Err(format!(
943 "marginal-slope low-rank covariance factor entry ({row},{column}) is non-finite"
944 ));
945 }
946 }
947 let mut ones_quadratic_form = 0.0;
948 for factor_column in factor.columns() {
949 let projection = factor_column.sum();
950 ones_quadratic_form += projection * projection;
951 }
952 if !ones_quadratic_form.is_finite() {
953 return Err("marginal-slope low-rank covariance geometry overflowed".to_string());
954 }
955 Ok(Self {
956 storage: MarginalSlopeCovarianceStorage::LowRank { factor },
957 ones_quadratic_form,
958 })
959 }
960
961 pub fn to_dense(&self) -> Array2<f64> {
962 match &self.storage {
963 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
964 Array2::from_diag(covariance)
965 }
966 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.clone(),
967 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.dot(&factor.t()),
968 }
969 }
970
971 pub fn shape(&self) -> MarginalSlopeCovarianceShape {
972 match &self.storage {
973 MarginalSlopeCovarianceStorage::Diagonal { .. } => {
974 MarginalSlopeCovarianceShape::Diagonal
975 }
976 MarginalSlopeCovarianceStorage::Full { .. } => MarginalSlopeCovarianceShape::Full,
977 MarginalSlopeCovarianceStorage::LowRank { .. } => MarginalSlopeCovarianceShape::LowRank,
978 }
979 }
980
981 pub fn dim(&self) -> usize {
982 match &self.storage {
983 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => covariance.len(),
984 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.nrows(),
985 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.nrows(),
986 }
987 }
988
989 pub fn ones_quadratic_form(&self) -> f64 {
990 self.ones_quadratic_form
991 }
992
993 pub(crate) fn representation(&self) -> MarginalSlopeCovarianceRef<'_> {
994 match &self.storage {
995 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
996 MarginalSlopeCovarianceRef::Diagonal(covariance)
997 }
998 MarginalSlopeCovarianceStorage::Full { covariance, .. } => {
999 MarginalSlopeCovarianceRef::Full(covariance)
1000 }
1001 MarginalSlopeCovarianceStorage::LowRank { factor } => {
1002 MarginalSlopeCovarianceRef::LowRank(factor)
1003 }
1004 }
1005 }
1006
1007 #[inline(always)]
1008 pub(crate) fn quadratic_form_unchecked(&self, vector: &[f64]) -> f64 {
1009 <Self as SymmetricQuadraticCoefficients>::quadratic_value(self, vector, |value| *value)
1010 }
1011
1012 pub fn quadratic_form(&self, vector: &[f64]) -> Result<f64, String> {
1013 if vector.len() != self.dim() {
1014 return Err(format!(
1015 "marginal-slope covariance dimension mismatch: vector={}, covariance={}",
1016 vector.len(),
1017 self.dim()
1018 ));
1019 }
1020 if vector.iter().any(|value| !value.is_finite()) {
1021 return Err("marginal-slope covariance vector contains non-finite values".to_string());
1022 }
1023 let value = self.quadratic_form_unchecked(vector);
1024 if !value.is_finite() {
1025 return Err(format!(
1026 "marginal-slope covariance quadratic form is non-finite: {value}"
1027 ));
1028 }
1029 Ok(value)
1030 }
1031}
1032
1033enum VectorSupport {
1034 Zero,
1035 Singleton { axis: usize, value: f64 },
1036 Multiple,
1037}
1038
1039#[inline(always)]
1040fn vector_support(input: &[f64]) -> VectorSupport {
1041 let mut singleton = None;
1042 for (axis, &value) in input.iter().enumerate() {
1043 if value == 0.0 {
1044 continue;
1045 }
1046 if singleton.is_some() {
1047 return VectorSupport::Multiple;
1048 }
1049 singleton = Some((axis, value));
1050 }
1051 match singleton {
1052 None => VectorSupport::Zero,
1053 Some((axis, value)) => VectorSupport::Singleton { axis, value },
1054 }
1055}
1056
1057impl SymmetricQuadraticCoefficients for MarginalSlopeCovariance {
1058 fn dimension(&self) -> usize {
1059 self.dim()
1060 }
1061
1062 fn multiply(&self, input: &[f64], output: &mut [f64]) {
1063 assert_eq!(input.len(), self.dim());
1064 assert_eq!(output.len(), self.dim());
1065 match self.representation() {
1066 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1067 for axis in 0..input.len() {
1068 output[axis] = diagonal[axis] * input[axis];
1069 }
1070 }
1071 MarginalSlopeCovarianceRef::Full(matrix) => {
1072 match vector_support(input) {
1073 VectorSupport::Zero => {
1074 output.fill(0.0);
1075 return;
1076 }
1077 VectorSupport::Singleton { axis, value } => {
1078 for row in 0..input.len() {
1079 output[row] = matrix[[row, axis]] * value;
1080 }
1081 return;
1082 }
1083 VectorSupport::Multiple => {}
1084 }
1085 for row in 0..input.len() {
1086 let mut value = 0.0;
1087 for column in 0..input.len() {
1088 value += matrix[[row, column]] * input[column];
1089 }
1090 output[row] = value;
1091 }
1092 }
1093 MarginalSlopeCovarianceRef::LowRank(factor) => {
1094 output.fill(0.0);
1095 match vector_support(input) {
1096 VectorSupport::Zero => return,
1097 VectorSupport::Singleton { axis, value } => {
1098 for rank in 0..factor.ncols() {
1099 let projection = factor[[axis, rank]] * value;
1100 for row in 0..input.len() {
1101 output[row] += factor[[row, rank]] * projection;
1102 }
1103 }
1104 return;
1105 }
1106 VectorSupport::Multiple => {}
1107 }
1108 for rank in 0..factor.ncols() {
1109 let mut projection = 0.0;
1110 for row in 0..input.len() {
1111 projection += factor[[row, rank]] * input[row];
1112 }
1113 for row in 0..input.len() {
1114 output[row] += factor[[row, rank]] * projection;
1115 }
1116 }
1117 }
1118 }
1119 }
1120
1121 fn coefficient(&self, row: usize, column: usize) -> f64 {
1122 match self.representation() {
1123 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1124 if row == column {
1125 diagonal[row]
1126 } else {
1127 0.0
1128 }
1129 }
1130 MarginalSlopeCovarianceRef::Full(matrix) => matrix[[row, column]],
1131 MarginalSlopeCovarianceRef::LowRank(factor) => {
1132 let mut value = 0.0;
1133 for rank in 0..factor.ncols() {
1134 value += factor[[row, rank]] * factor[[column, rank]];
1135 }
1136 value
1137 }
1138 }
1139 }
1140
1141 fn visit_upper_triangle(
1142 &self,
1143 direction: &mut [f64],
1144 projected: &mut [f64],
1145 mut visit: impl FnMut(usize, usize, f64),
1146 ) {
1147 let dimension = self.dim();
1148 assert_eq!(direction.len(), dimension);
1149 assert_eq!(projected.len(), dimension);
1150 match self.representation() {
1151 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1152 for column in 0..dimension {
1153 for row in 0..=column {
1154 visit(row, column, if row == column { diagonal[row] } else { 0.0 });
1155 }
1156 }
1157 }
1158 MarginalSlopeCovarianceRef::Full(matrix) => {
1159 for column in 0..dimension {
1160 for row in 0..=column {
1161 visit(row, column, matrix[[row, column]]);
1162 }
1163 }
1164 }
1165 MarginalSlopeCovarianceRef::LowRank(factor) => {
1166 for column in 0..dimension {
1167 for row in 0..=column {
1168 let mut value = 0.0;
1169 for rank in 0..factor.ncols() {
1170 value += factor[[row, rank]] * factor[[column, rank]];
1171 }
1172 visit(row, column, value);
1173 }
1174 }
1175 }
1176 }
1177 }
1178
1179 fn quadratic_value<T, F>(&self, input: &[T], value: F) -> f64
1180 where
1181 F: Fn(&T) -> f64,
1182 {
1183 assert_eq!(input.len(), self.dim());
1184 match &self.storage {
1185 MarginalSlopeCovarianceStorage::Diagonal { covariance } => input
1186 .iter()
1187 .zip(covariance)
1188 .map(|(input, &covariance)| {
1189 let input = value(input);
1190 covariance * input * input
1191 })
1192 .sum(),
1193 MarginalSlopeCovarianceStorage::Full {
1194 square_root_factor, ..
1195 } => {
1196 let mut total = 0.0;
1197 for factor_row in square_root_factor.rows() {
1198 let mut projection = 0.0;
1199 for axis in 0..input.len() {
1200 projection += factor_row[axis] * value(&input[axis]);
1201 }
1202 total += projection * projection;
1203 }
1204 total
1205 }
1206 MarginalSlopeCovarianceStorage::LowRank { factor } => {
1207 let mut total = 0.0;
1210 for rank in 0..factor.ncols() {
1211 let mut projection = 0.0;
1212 for row in 0..input.len() {
1213 projection += factor[[row, rank]] * value(&input[row]);
1214 }
1215 total += projection * projection;
1216 }
1217 total
1218 }
1219 }
1220 }
1221}
1222
1223pub fn marginal_slope_covariance_from_scores(
1250 scores: ArrayView2<'_, f64>,
1251 weights: &Array1<f64>,
1252) -> Result<MarginalSlopeCovariance, String> {
1253 let (n, k) = scores.dim();
1254 if k == 0 {
1255 return Err("marginal-slope score matrix must have at least one column".to_string());
1256 }
1257 if weights.len() != n {
1258 return Err(format!(
1259 "marginal-slope covariance weight length mismatch: weights={}, rows={n}",
1260 weights.len()
1261 ));
1262 }
1263 let total_weight = weights.iter().copied().sum::<f64>();
1264 if !(total_weight.is_finite() && total_weight > 0.0) {
1265 return Err("marginal-slope covariance needs positive finite total weight".to_string());
1266 }
1267 let mut mean = Array1::<f64>::zeros(k);
1268 for i in 0..n {
1269 let weight = weights[i];
1270 if !(weight.is_finite() && weight >= 0.0) {
1271 return Err(format!(
1272 "marginal-slope covariance weight {i} must be finite and non-negative, got {weight}"
1273 ));
1274 }
1275 for j in 0..k {
1276 let score = scores[[i, j]];
1277 if !score.is_finite() {
1278 return Err(format!(
1279 "marginal-slope covariance score ({i},{j}) is non-finite"
1280 ));
1281 }
1282 mean[j] += weight * score;
1283 }
1284 }
1285 mean.mapv_inplace(|value| value / total_weight);
1286
1287 let mut cov = Array2::<f64>::zeros((k, k));
1288 for i in 0..n {
1289 let weight = weights[i];
1290 for a in 0..k {
1291 let da = scores[[i, a]] - mean[a];
1292 for b in 0..=a {
1293 let value = weight * da * (scores[[i, b]] - mean[b]) / total_weight;
1294 cov[[a, b]] += value;
1295 if a != b {
1296 cov[[b, a]] += value;
1297 }
1298 }
1299 }
1300 }
1301
1302 let is_diagonal = (0..k).all(|row| ((row + 1)..k).all(|column| cov[[row, column]] == 0.0));
1306 if is_diagonal {
1307 MarginalSlopeCovariance::diagonal(cov.diag().to_owned())
1308 } else {
1309 MarginalSlopeCovariance::full(cov)
1310 }
1311}
1312
1313pub fn marginal_slope_preserving_scale(
1314 slopes: &[f64],
1315 covariance: &MarginalSlopeCovariance,
1316 probit_scale: f64,
1317) -> Result<f64, String> {
1318 if !probit_scale.is_finite() {
1319 return Err(format!(
1320 "marginal-slope probit scale must be finite, got {probit_scale}"
1321 ));
1322 }
1323 let variance = probit_scale * probit_scale * covariance.quadratic_form(slopes)?;
1324 if !variance.is_finite() {
1325 return Err("marginal-slope preserving variance is non-finite".to_string());
1326 }
1327 Ok((1.0 + variance).sqrt())
1328}
1329
1330pub fn marginal_slope_probit_eta(
1331 q: f64,
1332 z: &[f64],
1333 slopes: &[f64],
1334 covariance: &MarginalSlopeCovariance,
1335 probit_scale: f64,
1336) -> Result<f64, String> {
1337 if z.len() != slopes.len() {
1338 return Err(format!(
1339 "marginal-slope score/slope dimension mismatch: z={}, slopes={}",
1340 z.len(),
1341 slopes.len()
1342 ));
1343 }
1344 if slopes.len() != covariance.dim() {
1345 return Err(format!(
1346 "marginal-slope covariance dimension mismatch: slopes={}, covariance={}",
1347 slopes.len(),
1348 covariance.dim()
1349 ));
1350 }
1351 if !q.is_finite() || z.iter().any(|value| !value.is_finite()) {
1352 return Err("marginal-slope probit eta inputs must be finite".to_string());
1353 }
1354 let scale = marginal_slope_preserving_scale(slopes, covariance, probit_scale)?;
1355 let linear = z
1356 .iter()
1357 .zip(slopes.iter())
1358 .map(|(&score, &slope)| probit_scale * slope * score)
1359 .sum::<f64>();
1360 Ok(q * scale + linear)
1361}
1362
1363pub(super) fn empirical_rigid_calibration_eval(
1393 intercept: f64,
1394 log_target_mu: f64,
1395 slope: f64,
1396 probit_scale: f64,
1397 nodes: &[f64],
1398 weights: &[f64],
1399) -> Result<(f64, f64, f64), String> {
1400 if !intercept.is_finite() {
1401 return Err(format!(
1402 "empirical latent calibration: non-finite intercept {intercept}"
1403 ));
1404 }
1405 let observed_slope = rigid_observed_logslope(slope, probit_scale);
1406 const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_8; let mut log_max_phi = f64::NEG_INFINITY;
1410 let mut sum_phi = 0.0_f64;
1411 let mut log_max_cdf = f64::NEG_INFINITY;
1412 let mut sum_cdf = 0.0_f64;
1413
1414 let mut log_max_pos = f64::NEG_INFINITY;
1418 let mut sum_pos = 0.0_f64;
1419 let mut log_max_neg = f64::NEG_INFINITY;
1420 let mut sum_neg = 0.0_f64;
1421
1422 for (&node, &weight) in nodes.iter().zip(weights.iter()) {
1423 if !(weight.is_finite() && weight > 0.0) {
1424 continue;
1425 }
1426 let eta = intercept + observed_slope * node;
1427 if !eta.is_finite() {
1428 return Err(format!(
1429 "empirical latent calibration: non-finite η at intercept={intercept}, slope={slope}, node={node}"
1430 ));
1431 }
1432 let log_w = weight.ln();
1433 let log_phi = -0.5 * eta * eta - HALF_LOG_2PI;
1434 let log_term_phi = log_w + log_phi;
1435 let log_term_cdf = log_w + normal_logcdf(eta);
1436
1437 lse_accumulate(&mut log_max_phi, &mut sum_phi, log_term_phi);
1438 lse_accumulate(&mut log_max_cdf, &mut sum_cdf, log_term_cdf);
1439
1440 if eta != 0.0 {
1441 let log_term_eta_phi = log_term_phi + eta.abs().ln();
1442 if eta > 0.0 {
1443 lse_accumulate(&mut log_max_pos, &mut sum_pos, log_term_eta_phi);
1444 } else {
1445 lse_accumulate(&mut log_max_neg, &mut sum_neg, log_term_eta_phi);
1446 }
1447 }
1448 }
1449
1450 if !(sum_phi.is_finite() && sum_cdf.is_finite() && sum_phi > 0.0 && sum_cdf > 0.0) {
1451 return Err(format!(
1452 "empirical latent calibration: log-space accumulation failed (sum_phi={sum_phi}, sum_cdf={sum_cdf}, intercept={intercept})"
1453 ));
1454 }
1455
1456 let log_s_phi = log_max_phi + sum_phi.ln();
1457 let log_s_cdf = log_max_cdf + sum_cdf.ln();
1458
1459 let f = log_s_cdf - log_target_mu;
1461 let log_f_prime = log_s_phi - log_s_cdf;
1473 let f_prime = if log_f_prime > -740.0 {
1474 log_f_prime.exp()
1475 } else {
1476 f64::MIN_POSITIVE
1477 };
1478
1479 let exp_safe = |log_x: f64| -> f64 { if log_x > -740.0 { log_x.exp() } else { 0.0 } };
1487 let pos_over_cdf = if sum_pos > 0.0 {
1488 exp_safe(log_max_pos + sum_pos.ln() - log_s_cdf)
1489 } else {
1490 0.0
1491 };
1492 let neg_over_cdf = if sum_neg > 0.0 {
1493 exp_safe(log_max_neg + sum_neg.ln() - log_s_cdf)
1494 } else {
1495 0.0
1496 };
1497 let s_etaphi_over_s_cdf = pos_over_cdf - neg_over_cdf;
1498 let f_double_prime = -s_etaphi_over_s_cdf - f_prime * f_prime;
1499
1500 if !(f.is_finite() && f_prime.is_finite() && f_prime > 0.0 && f_double_prime.is_finite()) {
1501 return Err(format!(
1502 "empirical latent calibration: non-finite log-space state f={f}, f'={f_prime}, f''={f_double_prime} at intercept={intercept}"
1503 ));
1504 }
1505 Ok((f, f_prime, f_double_prime))
1506}
1507
1508pub(crate) fn empirical_intercept_from_marginal(
1509 target_mu: f64,
1510 target_q: f64,
1511 slope: f64,
1512 probit_scale: f64,
1513 nodes: &[f64],
1514 weights: &[f64],
1515 initial: Option<f64>,
1516) -> Result<f64, String> {
1517 if !(target_mu.is_finite() && target_mu > 0.0 && target_mu < 1.0) {
1518 return Err(format!(
1519 "empirical latent calibration requires target mu in (0,1), got {target_mu}"
1520 ));
1521 }
1522 let log_target_mu = target_mu.ln();
1523 let closed_form_seed = rigid_intercept_from_marginal(target_q, slope, probit_scale);
1524 let seed = initial.unwrap_or(closed_form_seed);
1525 let eval = |a: f64| {
1526 empirical_rigid_calibration_eval(a, log_target_mu, slope, probit_scale, nodes, weights)
1527 };
1528 let abs_tol = 1e-13_f64.max(4.0 * f64::EPSILON);
1535 let solve_from = |s: f64| {
1536 crate::monotone_root::solve_monotone_root(
1537 eval,
1538 s,
1539 "empirical latent intercept",
1540 abs_tol,
1541 64,
1542 48,
1543 )
1544 .map_err(|e| e.to_string())
1547 };
1548 let (root, _, f_best) = match solve_from(seed) {
1559 Ok(v) => v,
1560 Err(first_err) => {
1561 if seed == closed_form_seed {
1562 return Err(first_err);
1563 }
1564 solve_from(closed_form_seed).map_err(|retry_err| {
1565 format!("{first_err}; closed-form retry from a={closed_form_seed:.6}: {retry_err}")
1566 })?
1567 }
1568 };
1569 if f_best.abs() > abs_tol {
1570 return Err(format!(
1571 "empirical latent intercept solve failed: log-residual={f_best:.3e} at a={root:.6}, target mu={target_mu:.6}"
1572 ));
1573 }
1574 Ok(root)
1575}
1576
1577#[inline]
1578pub(super) fn rigid_standard_normal_neglog_only(
1579 q: f64,
1580 g: f64,
1581 z: f64,
1582 y: f64,
1583 w: f64,
1584 probit_scale: f64,
1585) -> Result<f64, String> {
1586 let s = 2.0 * y - 1.0;
1587 let eta = marginal_slope_standard_normal_scalar_eta(q, g, z, probit_scale);
1588 let m = s * eta;
1589 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
1590 if !logcdf.is_finite() {
1591 return Err(format!(
1592 "rigid probit neglog_only: non-finite log Φ at q={q}, g={g}, z={z}, y={y}"
1593 ));
1594 }
1595 Ok(-w * logcdf)
1596}
1597
1598#[inline]
1628pub(crate) fn rigid_standard_normal_row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1629 p: &[S; 2],
1630 marginal: BernoulliMarginalLinkMap,
1631 z: f64,
1632 y: f64,
1633 w: f64,
1634 probit_scale: f64,
1635) -> Result<S, String> {
1636 let signed = rigid_standard_normal_signed_margin(p, marginal, z, y, probit_scale);
1640 let m = signed.value();
1643 if !(m.is_finite() || m == f64::INFINITY) {
1644 return Err(format!(
1645 "non-finite signed margin in rigid probit row NLL: {m}"
1646 ));
1647 }
1648 Ok(signed.compose_unary(signed_probit_neglog_unary_stack(m, w)))
1650}
1651
1652#[inline]
1662pub(crate) fn rigid_standard_normal_signed_margin<S: gam_math::jet_scalar::JetScalar<2>>(
1663 p: &[S; 2],
1664 marginal: BernoulliMarginalLinkMap,
1665 z: f64,
1666 y: f64,
1667 probit_scale: f64,
1668) -> S {
1669 let q = p[0].compose_unary([
1671 marginal.q,
1672 marginal.q1,
1673 marginal.q2,
1674 marginal.q3,
1675 marginal.q4,
1676 ]);
1677 let slope = p[1];
1678 let observed_slope = slope.scale(probit_scale);
1680 let b2 = observed_slope.mul(&observed_slope);
1681 let c = b2.add(&S::constant(1.0)).sqrt();
1682 let eta = q.mul(&c).add(&observed_slope.scale(z));
1684 eta.scale(2.0 * y - 1.0)
1685}
1686
1687pub(crate) struct RigidStandardNormalRow {
1700 pub(crate) marginal: BernoulliMarginalLinkMap,
1701 pub(crate) g: f64,
1702 pub(crate) z: f64,
1703 pub(crate) y: f64,
1704 pub(crate) w: f64,
1705 pub(crate) probit_scale: f64,
1706}
1707
1708impl gam_math::jet_tower::RowProgram<2> for RigidStandardNormalRow {
1709 fn n_rows(&self) -> usize {
1710 1
1711 }
1712
1713 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
1714 if row != 0 {
1715 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1716 }
1717 Ok([self.marginal.eta_value(), self.g])
1718 }
1719
1720 fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
1721 &self,
1722 row: usize,
1723 p: &[S; 2],
1724 ) -> Result<S, String> {
1725 if row != 0 {
1726 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1727 }
1728 rigid_standard_normal_row_nll_generic(
1729 p,
1730 self.marginal,
1731 self.z,
1732 self.y,
1733 self.w,
1734 self.probit_scale,
1735 )
1736 }
1737}
1738
1739#[inline]
1740pub(crate) fn rigid_standard_normal_tower(
1741 marginal: BernoulliMarginalLinkMap,
1742 g: f64,
1743 z: f64,
1744 y: f64,
1745 w: f64,
1746 probit_scale: f64,
1747) -> Result<Tower4<2>, String> {
1748 let program = RigidStandardNormalRow {
1755 marginal,
1756 g,
1757 z,
1758 y,
1759 w,
1760 probit_scale,
1761 };
1762 gam_math::jet_tower::program_full_tower(&program, 0).map(|tower| *tower)
1763}
1764
1765#[inline]
1778fn rigid_standard_normal_signed_jet(
1779 marginal: BernoulliMarginalLinkMap,
1780 g: f64,
1781 z: f64,
1782 y: f64,
1783 probit_scale: f64,
1784) -> Tower4<2> {
1785 let p = [
1788 Tower4::<2>::variable(marginal.eta_value(), 0),
1789 Tower4::<2>::variable(g, 1),
1790 ];
1791 rigid_standard_normal_signed_margin(&p, marginal, z, y, probit_scale)
1792}
1793
1794#[inline]
1822pub(super) fn rigid_standard_normal_towers_batch<T>(
1823 marginals: &[BernoulliMarginalLinkMap],
1824 slopes: &[f64],
1825 zs: &[f64],
1826 ys: &[f64],
1827 weights: &[f64],
1828 probit_scale: f64,
1829 out: &mut [T],
1830 mut fill: impl FnMut(&Tower4<2>) -> Result<T, String>,
1831) -> Result<(), String> {
1832 let chunk = marginals.len();
1833 if slopes.len() != chunk
1834 || zs.len() != chunk
1835 || ys.len() != chunk
1836 || weights.len() != chunk
1837 || out.len() != chunk
1838 {
1839 return Err(format!(
1840 "rigid_standard_normal_towers_batch length mismatch: marginals={chunk}, \
1841 slopes={}, zs={}, ys={}, weights={}, out={}",
1842 slopes.len(),
1843 zs.len(),
1844 ys.len(),
1845 weights.len(),
1846 out.len()
1847 ));
1848 }
1849
1850 let mut signed: Vec<Tower4<2>> = Vec::with_capacity(chunk);
1852 let mut margins: Vec<f64> = Vec::with_capacity(chunk);
1853 for i in 0..chunk {
1854 let jet =
1855 rigid_standard_normal_signed_jet(marginals[i], slopes[i], zs[i], ys[i], probit_scale);
1856 margins.push(jet.v);
1857 signed.push(jet);
1858 }
1859
1860 let mut stacks: Vec<[f64; 5]> = Vec::with_capacity(chunk);
1864 for i in 0..chunk {
1865 let m = margins[i];
1866 if !(m.is_finite() || m == f64::INFINITY) {
1867 return Err(format!(
1868 "non-finite signed margin in rigid probit tower batch: {m}"
1869 ));
1870 }
1871 stacks.push(signed_probit_neglog_unary_stack(m, weights[i]));
1872 }
1873
1874 for i in 0..chunk {
1876 let tower = signed[i].compose_unary(stacks[i]);
1877 out[i] = fill(&tower)?;
1878 }
1879 Ok(())
1880}
1881
1882#[inline]
1883pub(super) fn rigid_standard_normal_row_kernel(
1884 marginal: BernoulliMarginalLinkMap,
1885 g: f64,
1886 z: f64,
1887 y: f64,
1888 w: f64,
1889 probit_scale: f64,
1890) -> Result<(f64, [f64; 2], [[f64; 2]; 2]), String> {
1891 let program = RigidStandardNormalRow {
1900 marginal,
1901 g,
1902 z,
1903 y,
1904 w,
1905 probit_scale,
1906 };
1907 gam_math::jet_tower::program_row_kernel(&program, 0)
1908}
1909
1910#[inline]
1948pub(super) fn rigid_standard_normal_mixed_z_sensitivity(
1949 marginal: BernoulliMarginalLinkMap,
1950 g: f64,
1951 z: f64,
1952 y: f64,
1953 w: f64,
1954 probit_scale: f64,
1955) -> Result<[f64; 2], String> {
1956 use gam_math::jet_tower::Tower2;
1967 let mut q = Tower2::<3>::constant(marginal.q);
1968 q.g[0] = marginal.q1;
1969 q.h[0][0] = marginal.q2;
1970 let slope = Tower2::<3>::variable(g, 1);
1971 let z_var = Tower2::<3>::variable(z, 2);
1972 let observed_logslope = slope * probit_scale;
1973 let c = (observed_logslope * observed_logslope + 1.0).sqrt();
1974 let eta = q * c + slope * (z_var * probit_scale);
1978 let signed = eta * (2.0 * y - 1.0);
1979 if !(signed.v.is_finite() || signed.v == f64::INFINITY) {
1981 return Err(format!(
1982 "rigid probit mixed-z sensitivity: non-finite signed margin {} at q={}, g={g}, z={z}, y={y}",
1983 signed.v, marginal.q
1984 ));
1985 }
1986 let stack = signed_probit_neglog_unary_stack(signed.v, w);
1987 if !stack[0].is_finite() {
1988 return Err(format!(
1989 "rigid probit mixed-z sensitivity: non-finite log Φ at q={}, g={g}, z={z}, y={y}",
1990 marginal.q
1991 ));
1992 }
1993 let tower = signed.compose_unary([stack[0], stack[1], stack[2]]);
1996 let s_q = -tower.h[0][2];
2003 let s_g = -tower.h[1][2];
2004 if !(s_q.is_finite() && s_g.is_finite()) {
2005 return Err(format!(
2006 "rigid probit mixed-z sensitivity: non-finite ∂²(log L)/∂(q,g)∂z = [{s_q}, {s_g}] at q={}, g={g}, z={z}",
2007 marginal.q
2008 ));
2009 }
2010 Ok([s_q, s_g])
2011}
2012
2013pub(super) fn rigid_standard_normal_score_zeta_sensitivity(
2037 base_link: &InverseLink,
2038 marginal_eta: &Array1<f64>,
2039 slope_eta: &Array1<f64>,
2040 z: &Array1<f64>,
2041 y: &Array1<f64>,
2042 weights: &Array1<f64>,
2043 probit_scale: f64,
2044 marginal_design: ArrayView2<'_, f64>,
2045 logslope_design: ArrayView2<'_, f64>,
2046 p_beta: usize,
2047) -> Result<Array2<f64>, String> {
2048 let n = marginal_eta.len();
2049 let p_m = marginal_design.ncols();
2050 let r = logslope_design.ncols();
2051 if slope_eta.len() != n
2052 || z.len() != n
2053 || y.len() != n
2054 || weights.len() != n
2055 || marginal_design.nrows() != n
2056 || logslope_design.nrows() != n
2057 {
2058 return Err(format!(
2059 "score_zeta_sensitivity row mismatch: marginal_eta={n}, slope_eta={}, z={}, y={}, \
2060 weights={}, marginal_design rows={}, logslope_design rows={}",
2061 slope_eta.len(),
2062 z.len(),
2063 y.len(),
2064 weights.len(),
2065 marginal_design.nrows(),
2066 logslope_design.nrows()
2067 ));
2068 }
2069 if p_m + r != p_beta {
2070 return Err(format!(
2071 "rigid score_zeta_sensitivity width mismatch: marginal({p_m}) + logslope({r}) != p_beta({p_beta})"
2072 ));
2073 }
2074 let mut s = Array2::<f64>::zeros((n, p_beta));
2075 for i in 0..n {
2076 let marginal = bernoulli_marginal_link_map(base_link, marginal_eta[i])?;
2077 let [s_q, s_g] = rigid_standard_normal_mixed_z_sensitivity(
2078 marginal,
2079 slope_eta[i],
2080 z[i],
2081 y[i],
2082 weights[i],
2083 probit_scale,
2084 )?;
2085 if s_q != 0.0 {
2088 let m_row = marginal_design.row(i);
2089 for (j, &mij) in m_row.iter().enumerate() {
2090 s[[i, j]] = s_q * mij;
2091 }
2092 }
2093 if s_g != 0.0 {
2094 let g_row = logslope_design.row(i);
2095 for (j, &gij) in g_row.iter().enumerate() {
2096 s[[i, p_m + j]] = s_g * gij;
2097 }
2098 }
2099 }
2100 Ok(s)
2101}
2102
2103#[inline]
2104pub(super) fn rigid_standard_normal_third_full(
2105 marginal: BernoulliMarginalLinkMap,
2106 g: f64,
2107 z: f64,
2108 y: f64,
2109 w: f64,
2110 probit_scale: f64,
2111) -> Result<[[[f64; 2]; 2]; 2], String> {
2112 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t3)
2113}
2114
2115#[inline]
2120pub(super) fn contract_third_full(t: &[[[f64; 2]; 2]; 2], d_eta: f64, d_g: f64) -> [[f64; 2]; 2] {
2121 [
2122 [
2123 t[0][0][0] * d_eta + t[0][0][1] * d_g,
2124 t[0][1][0] * d_eta + t[0][1][1] * d_g,
2125 ],
2126 [
2127 t[1][0][0] * d_eta + t[1][0][1] * d_g,
2128 t[1][1][0] * d_eta + t[1][1][1] * d_g,
2129 ],
2130 ]
2131}
2132
2133#[inline]
2134pub(super) fn rigid_standard_normal_fourth_full(
2135 marginal: BernoulliMarginalLinkMap,
2136 g: f64,
2137 z: f64,
2138 y: f64,
2139 w: f64,
2140 probit_scale: f64,
2141) -> Result<[[[[f64; 2]; 2]; 2]; 2], String> {
2142 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t4)
2156}
2157
2158#[inline]
2177pub(super) fn contract_fourth_full(
2178 t: &[[[[f64; 2]; 2]; 2]; 2],
2179 u_eta: f64,
2180 u_g: f64,
2181 v_eta: f64,
2182 v_g: f64,
2183) -> [[f64; 2]; 2] {
2184 let mut out = [[0.0; 2]; 2];
2185 for a in 0..2 {
2186 for b in 0..2 {
2187 let mut sum = 0.0;
2188 sum += t[a][b][0][0] * u_eta * v_eta;
2189 sum += t[a][b][0][1] * u_eta * v_g;
2190 sum += t[a][b][1][0] * u_g * v_eta;
2191 sum += t[a][b][1][1] * u_g * v_g;
2192 out[a][b] = sum;
2193 }
2194 }
2195 out
2196}
2197
2198pub(super) fn ensure_finite_third_full_cache_row(
2199 t: &[[[f64; 2]; 2]; 2],
2200 context: &str,
2201) -> Result<(), String> {
2202 if t.iter().flatten().flatten().all(|value| value.is_finite()) {
2203 Ok(())
2204 } else {
2205 Err(format!(
2206 "{context}: warmed third-derivative cache row contains a non-finite value"
2207 ))
2208 }
2209}
2210
2211pub(super) fn ensure_finite_fourth_full_cache_row(
2212 t: &[[[[f64; 2]; 2]; 2]; 2],
2213 context: &str,
2214) -> Result<(), String> {
2215 if t.iter()
2216 .flatten()
2217 .flatten()
2218 .flatten()
2219 .all(|value| value.is_finite())
2220 {
2221 Ok(())
2222 } else {
2223 Err(format!(
2224 "{context}: warmed fourth-derivative cache row contains a non-finite value"
2225 ))
2226 }
2227}
2228
2229pub(crate) fn unary_derivatives_sqrt(x: f64) -> [f64; 5] {
2230 let s = x.max(1e-300).sqrt();
2231 let x1 = x.max(1e-300);
2232 let x2 = x1 * x1;
2233 let x3 = x2 * x1;
2234 [
2235 s,
2236 0.5 / s,
2237 -0.25 / (x1 * s),
2238 3.0 / (8.0 * x2 * s),
2239 -15.0 / (16.0 * x3 * s),
2240 ]
2241}
2242pub(crate) fn unary_derivatives_neglog_phi(x: f64, weight: f64) -> [f64; 5] {
2243 signed_probit_neglog_unary_stack(x, weight)
2248}
2249
2250pub(crate) fn unary_derivatives_log(x: f64) -> [f64; 5] {
2268 let x2 = x * x;
2269 let x3 = x2 * x;
2270 let x4 = x3 * x;
2271 [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
2272}
2273
2274pub(crate) fn unary_derivatives_log_normal_pdf(x: f64) -> [f64; 5] {
2276 let c = 0.5 * (2.0 * std::f64::consts::PI).ln();
2277 [-0.5 * x * x - c, -x, -1.0, 0.0, 0.0]
2278}
2279
2280#[cfg(test)]
2281mod covariance_admission_tests {
2282 use super::*;
2283 use ndarray::array;
2284
2285 #[test]
2286 fn full_covariance_admission_rejects_one_ulp_asymmetry_932() {
2287 let upper = 0.25_f64;
2288 let lower = f64::from_bits(upper.to_bits() + 1);
2289 let error = MarginalSlopeCovariance::full(array![[1.0, upper], [lower, 1.0]])
2290 .expect_err("any asymmetric full operator must be rejected");
2291 assert!(error.contains("must be exactly symmetric"), "{error}");
2292 }
2293
2294 #[test]
2295 fn full_covariance_admission_rejects_indefinite_matrix_before_row_use_932() {
2296 let error = MarginalSlopeCovariance::full(array![[1.0, 2.0], [2.0, 1.0]])
2297 .expect_err("an indefinite full operator is not a covariance");
2298 assert!(error.contains("must be positive semidefinite"), "{error}");
2299 }
2300
2301 #[test]
2302 fn full_covariance_admission_accepts_exact_singular_psd_932() {
2303 MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, 0.0]])
2304 .expect("an exact singular PSD covariance is admissible");
2305 }
2306
2307 #[test]
2315 fn full_covariance_admission_decides_psd_against_the_eigensolver_band() {
2316 let admitted = MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, -1.0e-17]])
2317 .expect("a negative eigenvalue inside the eigensolver band is a zero one");
2318 assert!(
2319 admitted.ones_quadratic_form().is_finite(),
2320 "clamping inside the band must keep the square-root factor real"
2321 );
2322 let error = MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, -1.0e-12]])
2323 .expect_err("a negative eigenvalue outside the band is material indefiniteness");
2324 assert!(error.contains("must be positive semidefinite"), "{error}");
2325 }
2326
2327 #[test]
2328 fn full_covariance_admission_accepts_coupled_singular_psd_932() {
2329 let covariance = MarginalSlopeCovariance::full(array![[1.0, 1.0], [1.0, 1.0]])
2330 .expect("a coupled singular PSD covariance is admissible");
2331 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2332 assert_eq!(covariance.to_dense(), array![[1.0, 1.0], [1.0, 1.0]]);
2333 }
2334
2335 #[test]
2336 fn exact_nonzero_offdiagonal_classifier_retains_full_geometry_932() {
2337 let epsilon = 1.0e-14;
2338 let scores = array![[-1.0, -epsilon], [1.0, epsilon], [0.0, -1.0], [0.0, 1.0]];
2339 let covariance =
2340 marginal_slope_covariance_from_scores(scores.view(), &Array1::ones(4)).unwrap();
2341 let dense = covariance.to_dense();
2342 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2343 assert_ne!(dense[[0, 1]], 0.0);
2344 assert_eq!(dense[[0, 1]], dense[[1, 0]]);
2345 let direction = [0.75, -1.25];
2346 let expected = direction[0] * (dense[[0, 0]] * direction[0] + dense[[0, 1]] * direction[1])
2347 + direction[1] * (dense[[1, 0]] * direction[0] + dense[[1, 1]] * direction[1]);
2348 let actual = covariance.quadratic_form(&direction).unwrap();
2349 assert!((actual - expected).abs() <= 2.0e-15);
2350 }
2351
2352 #[test]
2353 fn diagonal_covariance_entries_are_the_exact_geometry_authority_932() {
2354 let covariance = MarginalSlopeCovariance::diagonal(array![3.75]).unwrap();
2355 assert_eq!(covariance.ones_quadratic_form(), 3.75);
2356 assert_eq!(covariance.quadratic_form(&[1.0]).unwrap(), 3.75);
2357 }
2358
2359 #[test]
2360 fn equal_dense_covariance_quadratic_forms_match_all_representations_932() {
2361 let diagonal = MarginalSlopeCovariance::diagonal(array![1.2, 0.7]).unwrap();
2362 let full = MarginalSlopeCovariance::full(array![[1.2, 0.0], [0.0, 0.7]]).unwrap();
2363 let low_rank =
2364 MarginalSlopeCovariance::low_rank(array![[1.2_f64.sqrt(), 0.0], [0.0, 0.7_f64.sqrt()]])
2365 .unwrap();
2366 let direction = [0.35, -0.8];
2367 let expected = diagonal.quadratic_form(&direction).unwrap();
2368 for covariance in [&full, &low_rank] {
2369 let actual = covariance.quadratic_form(&direction).unwrap();
2370 assert!((actual - expected).abs() <= 2.0e-15);
2371 assert!(
2372 (covariance.ones_quadratic_form() - diagonal.ones_quadratic_form()).abs()
2373 <= 2.0e-15
2374 );
2375 }
2376 }
2377}
2378
2379#[cfg(test)]
2380mod jet_tower_oracle_tests {
2381 use super::*;
2403
2404 #[test]
2405 fn signed_probit_stack_preserves_extreme_tail_derivatives_and_weight_sign() {
2406 let positive = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, 2.0);
2407 assert_eq!(
2408 positive,
2409 [f64::INFINITY, f64::NEG_INFINITY, 2.0, -0.0, -0.0]
2410 );
2411 let negative = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, -2.0);
2412 assert_eq!(negative, [f64::NEG_INFINITY, f64::INFINITY, -2.0, 0.0, 0.0]);
2413
2414 let right = signed_probit_neglog_unary_stack(38.6, 1.0);
2415 assert_eq!(right[1], -0.0);
2416 assert!(right[2] > 0.0 && right[2].is_subnormal());
2417 assert!(right[3] < 0.0 && right[3].is_subnormal());
2418 assert!(right[4] > 0.0 && right[4].is_subnormal());
2419
2420 let left = signed_probit_neglog_unary_stack(-1.0e100, 1.0);
2421 assert_eq!(left[1], -1.0e100);
2422 assert_eq!(left[2], 1.0);
2423 assert!(left[3] < 0.0 && left[3].is_finite());
2424 assert_eq!(left[4], -0.0);
2425 }
2426
2427 fn rigid_standard_normal_third_and_fourth_full(
2435 marginal: BernoulliMarginalLinkMap,
2436 g: f64,
2437 z: f64,
2438 y: f64,
2439 w: f64,
2440 probit_scale: f64,
2441 ) -> Result<([[[f64; 2]; 2]; 2], [[[[f64; 2]; 2]; 2]; 2]), String> {
2442 let tower = rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?;
2443 Ok((tower.t3, tower.t4))
2444 }
2445 use gam_math::jet_tower::{
2446 KernelChannels, RowProgram, program_full_tower, verify_kernel_channels,
2447 };
2448
2449 struct BernoulliRigidStandardNormalNllProgram {
2452 primaries: Vec<[f64; 2]>,
2454 z: Vec<f64>,
2456 y: Vec<f64>,
2457 w: Vec<f64>,
2458 probit_scale: f64,
2459 }
2460
2461 impl RowProgram<2> for BernoulliRigidStandardNormalNllProgram {
2462 fn n_rows(&self) -> usize {
2463 self.primaries.len()
2464 }
2465
2466 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2467 self.primaries
2468 .get(row)
2469 .copied()
2470 .ok_or_else(|| format!("bernoulli rigid nll program: row {row} out of range"))
2471 }
2472
2473 fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
2474 &self,
2475 row: usize,
2476 p: &[S; 2],
2477 ) -> Result<S, String> {
2478 let z = self.z[row];
2479 let y = self.y[row];
2480 let w = self.w[row];
2481 let s = self.probit_scale;
2482 let eta_marginal = p[0];
2486 let link = bernoulli_marginal_link_map(
2487 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2488 eta_marginal.value(),
2489 )?;
2490 let q = eta_marginal.compose_unary([link.q, link.q1, link.q2, link.q3, link.q4]);
2491 let g = p[1];
2492 let observed_slope = g.scale(s);
2494 let one_plus_slope_squared = observed_slope.mul(&observed_slope).add(&S::constant(1.0));
2495 let c = one_plus_slope_squared
2496 .compose_unary(unary_derivatives_sqrt(one_plus_slope_squared.value()));
2497 let eta = q.mul(&c).add(&observed_slope.scale(z));
2499 let signed = eta.scale(2.0 * y - 1.0);
2500 Ok(signed.compose_unary(unary_derivatives_neglog_phi(signed.value(), w)))
2502 }
2503 }
2504
2505 fn scalar_nll(eta_marginal: f64, g: f64, z: f64, y: f64, w: f64, s: f64) -> f64 {
2508 let link = bernoulli_marginal_link_map(
2509 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2510 eta_marginal,
2511 )
2512 .unwrap();
2513 let observed_slope = g * s;
2514 let c = (observed_slope * observed_slope + 1.0).sqrt();
2515 let eta = link.q * c + observed_slope * z;
2516 let signed = (2.0 * y - 1.0) * eta;
2517 let cdf = 0.5 * libm::erfc(-signed / std::f64::consts::SQRT_2);
2518 -w * cdf.max(1e-300).ln()
2519 }
2520
2521 #[test]
2522 fn rigid_bernoulli_row_kernel_agrees_with_jet_tower_program_all_channels() {
2523 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2527 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2528 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2529 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2530 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2531 let n = eta.len();
2532
2533 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2535
2536 for &probit_scale in &[1.0_f64, 0.8] {
2537 let program = BernoulliRigidStandardNormalNllProgram {
2538 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2539 z: z.to_vec(),
2540 y: y.to_vec(),
2541 w: w.to_vec(),
2542 probit_scale,
2543 };
2544
2545 for row in 0..n {
2546 let tower = program_full_tower(&program, row).expect("tower evaluation");
2547
2548 let marginal = bernoulli_marginal_link_map(
2550 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2551 eta[row],
2552 )
2553 .expect("link map");
2554 let (value, gradient, hessian) = rigid_standard_normal_row_kernel(
2555 marginal,
2556 g[row],
2557 z[row],
2558 y[row],
2559 w[row],
2560 probit_scale,
2561 )
2562 .expect("production row kernel");
2563
2564 let (third_full, fourth_full) = rigid_standard_normal_third_and_fourth_full(
2571 marginal,
2572 g[row],
2573 z[row],
2574 y[row],
2575 w[row],
2576 probit_scale,
2577 )
2578 .expect("production third+fourth");
2579 let third: Vec<([f64; 2], [[f64; 2]; 2])> = dirs
2580 .iter()
2581 .map(|d| (*d, contract_third_full(&third_full, d[0], d[1])))
2582 .collect();
2583
2584 let fourth: Vec<([f64; 2], [f64; 2], [[f64; 2]; 2])> = dirs
2585 .iter()
2586 .enumerate()
2587 .map(|(i, u)| {
2588 let v = dirs[(i + 1) % dirs.len()];
2589 (
2590 *u,
2591 v,
2592 contract_fourth_full(&fourth_full, u[0], u[1], v[0], v[1]),
2593 )
2594 })
2595 .collect();
2596
2597 let claims = KernelChannels {
2598 value,
2599 gradient,
2600 hessian,
2601 third,
2602 fourth,
2603 };
2604
2605 verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
2606 panic!(
2607 "probit_scale {probit_scale} row {row}: production rigid Bernoulli \
2608 RowKernel disagrees with #932 jet-tower truth: {e}"
2609 )
2610 });
2611
2612 let h = 1e-3;
2616 let f = |de: f64, dg: f64| {
2617 scalar_nll(
2618 eta[row] + de,
2619 g[row] + dg,
2620 z[row],
2621 y[row],
2622 w[row],
2623 probit_scale,
2624 )
2625 };
2626 let f0 = f(0.0, 0.0);
2627 assert!(
2628 (f0 - tower.v).abs() <= 1e-9 * f0.abs().max(1.0),
2629 "row {row}: independent scalar NLL {f0:+.12e} != tower value {:+.12e}",
2630 tower.v
2631 );
2632 let g_eta = (f(-2.0 * h, 0.0) - 8.0 * f(-h, 0.0) + 8.0 * f(h, 0.0)
2634 - f(2.0 * h, 0.0))
2635 / (12.0 * h);
2636 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))
2637 / (12.0 * h);
2638 for (label, fd, ad) in [("∂η", g_eta, tower.g[0]), ("∂g", g_g, tower.g[1])] {
2639 assert!(
2640 (fd - ad).abs() <= 1e-5 * ad.abs().max(1.0),
2641 "row {row} {label}: FD witness {fd:+.6e} != tower grad {ad:+.6e}"
2642 );
2643 }
2644 }
2645 }
2646 }
2647
2648 #[test]
2658 fn rigid_third_and_fourth_full_shares_one_tower_bit_identical() {
2659 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2660 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2661 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2662 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2663 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2664 for &probit_scale in &[1.0_f64, 0.8] {
2665 for r in 0..eta.len() {
2666 let marginal = bernoulli_marginal_link_map(
2667 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2668 eta[r],
2669 )
2670 .expect("link map");
2671 let t3_sep = rigid_standard_normal_third_full(
2672 marginal,
2673 g[r],
2674 z[r],
2675 y[r],
2676 w[r],
2677 probit_scale,
2678 )
2679 .expect("separate third");
2680 let t4_sep = rigid_standard_normal_fourth_full(
2681 marginal,
2682 g[r],
2683 z[r],
2684 y[r],
2685 w[r],
2686 probit_scale,
2687 )
2688 .expect("separate fourth");
2689 let (t3_comb, t4_comb) = rigid_standard_normal_third_and_fourth_full(
2690 marginal,
2691 g[r],
2692 z[r],
2693 y[r],
2694 w[r],
2695 probit_scale,
2696 )
2697 .expect("combined third+fourth");
2698 for a in 0..2 {
2700 for b in 0..2 {
2701 for c in 0..2 {
2702 assert_eq!(
2703 t3_comb[a][b][c], t3_sep[a][b][c],
2704 "t3[{a}][{b}][{c}] row {r} scale {probit_scale} not bit-identical"
2705 );
2706 for d in 0..2 {
2707 assert_eq!(
2708 t4_comb[a][b][c][d], t4_sep[a][b][c][d],
2709 "t4[{a}][{b}][{c}][{d}] row {r} scale {probit_scale} not bit-identical"
2710 );
2711 }
2712 }
2713 }
2714 }
2715 }
2716 }
2717 }
2718
2719 #[test]
2729 fn rigid_bernoulli_generic_program_matches_independent_program_all_channels() {
2730 use gam_math::jet_tower::{
2731 program_fourth_contracted, program_row_kernel, program_third_contracted,
2732 };
2733
2734 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2735 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2736 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2737 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2738 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2739 let n = eta.len();
2740 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2741
2742 let close = |a: f64, b: f64, label: &str| {
2743 let band = 1e-12 + 1e-12 * a.abs().max(b.abs());
2744 assert!(
2745 (a - b).abs() <= band,
2746 "{label}: generic {a:+.15e} vs Tower4-program {b:+.15e} (band {band:.3e})"
2747 );
2748 };
2749
2750 for &probit_scale in &[1.0_f64, 0.8] {
2751 let tower_program = BernoulliRigidStandardNormalNllProgram {
2753 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2754 z: z.to_vec(),
2755 y: y.to_vec(),
2756 w: w.to_vec(),
2757 probit_scale,
2758 };
2759
2760 for row in 0..n {
2761 let truth = program_full_tower(&tower_program, row).expect("program tower");
2762
2763 let marginal = bernoulli_marginal_link_map(
2764 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2765 eta[row],
2766 )
2767 .expect("link map");
2768 let program = RigidStandardNormalRow {
2769 marginal,
2770 g: g[row],
2771 z: z[row],
2772 y: y[row],
2773 w: w[row],
2774 probit_scale,
2775 };
2776
2777 let full = program_full_tower(&program, 0).expect("generic full tower");
2780 close(full.v, truth.v, "full value");
2781 for a in 0..2 {
2782 close(full.g[a], truth.g[a], "full grad");
2783 for b in 0..2 {
2784 close(full.h[a][b], truth.h[a][b], "full hess");
2785 for c in 0..2 {
2786 close(full.t3[a][b][c], truth.t3[a][b][c], "full t3");
2787 for d in 0..2 {
2788 close(full.t4[a][b][c][d], truth.t4[a][b][c][d], "full t4");
2789 }
2790 }
2791 }
2792 }
2793
2794 let (val, grad, hess) =
2796 program_row_kernel(&program, 0).expect("generic row kernel");
2797 close(val, truth.v, "order2 value");
2798 for a in 0..2 {
2799 close(grad[a], truth.g[a], "order2 grad");
2800 for b in 0..2 {
2801 close(hess[a][b], truth.h[a][b], "order2 hess");
2802 }
2803 }
2804
2805 for dir in &dirs {
2808 let third = program_third_contracted(&program, 0, dir)
2809 .expect("generic third contracted");
2810 let truth3 = truth.third_contracted(dir);
2811 for a in 0..2 {
2812 for b in 0..2 {
2813 close(third[a][b], truth3[a][b], "third contracted");
2814 }
2815 }
2816 }
2817
2818 for (i, u) in dirs.iter().enumerate() {
2821 let v = dirs[(i + 1) % dirs.len()];
2822 let fourth = program_fourth_contracted(&program, 0, u, &v)
2823 .expect("generic fourth contracted");
2824 let truth4 = truth.fourth_contracted(u, &v);
2825 for a in 0..2 {
2826 for b in 0..2 {
2827 close(fourth[a][b], truth4[a][b], "fourth contracted");
2828 }
2829 }
2830 }
2831 }
2832 }
2833 }
2834
2835 fn hand_rigid_vgh(
2844 marginal: BernoulliMarginalLinkMap,
2845 g: f64,
2846 z: f64,
2847 y: f64,
2848 w: f64,
2849 probit_scale: f64,
2850 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2851 let s = 2.0 * y - 1.0;
2852 let observed_logslope = probit_scale * g;
2853 let g2 = observed_logslope * observed_logslope;
2854 let c = (1.0 + g2).sqrt();
2855 let c1 = probit_scale * observed_logslope / c;
2856 let c_inv3 = 1.0 / (c * c * c);
2857 let c2 = probit_scale * probit_scale * c_inv3;
2858 let q = marginal.q;
2859 let eta = q * c + observed_logslope * z;
2861 let m = s * eta;
2862 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
2863 let (k1, k2, _k3, _k4) =
2867 signed_probit_neglog_derivatives_up_to_fourth(m, w).expect("hand kernel");
2868 let u1 = s * k1;
2869 let u2 = k2;
2870 let eta_q = c;
2871 let eta_g = q * c1 + probit_scale * z;
2872 let value = -w * logcdf;
2874 let gradient = [u1 * eta_q * marginal.q1, u1 * eta_g];
2876 let h00 = u2 * eta_q * eta_q;
2878 let h01 = u2 * eta_q * eta_g + u1 * c1;
2879 let h11 = u2 * eta_g * eta_g + u1 * q * c2;
2880 let grad_q = u1 * eta_q;
2882 let hessian = [
2883 [
2884 h00 * marginal.q1 * marginal.q1 + grad_q * marginal.q2,
2885 h01 * marginal.q1,
2886 ],
2887 [h01 * marginal.q1, h11],
2888 ];
2889 (value, gradient, hessian)
2890 }
2891
2892 #[test]
2898 fn rigid_bernoulli_row_kernel_matches_hand_chain_witness() {
2899 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2900 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2901 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2902 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2903 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2904 let close = |a: f64, b: f64, label: &str| {
2905 let band = 1e-12 + 1e-9 * a.abs().max(b.abs());
2906 assert!(
2907 (a - b).abs() <= band,
2908 "{label}: jet {a:+.15e} vs hand {b:+.15e} (band {band:.3e})"
2909 );
2910 };
2911 for &probit_scale in &[1.0_f64, 0.8] {
2912 for r in 0..eta.len() {
2913 let marginal = bernoulli_marginal_link_map(
2914 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2915 eta[r],
2916 )
2917 .expect("link map");
2918 let (jv, jg, jh) = rigid_standard_normal_row_kernel(
2919 marginal,
2920 g[r],
2921 z[r],
2922 y[r],
2923 w[r],
2924 probit_scale,
2925 )
2926 .expect("jet kernel");
2927 let (hv, hg, hh) = hand_rigid_vgh(marginal, g[r], z[r], y[r], w[r], probit_scale);
2928 close(jv, hv, "value");
2929 for a in 0..2 {
2930 close(jg[a], hg[a], "grad");
2931 for b in 0..2 {
2932 close(jh[a][b], hh[a][b], "hess");
2933 }
2934 }
2935 }
2936 }
2937 }
2938
2939 #[test]
2948 fn release_measure_rigid_bernoulli_vgh_vs_hand_chain_932() {
2949 use std::time::Instant;
2950
2951 let cases = [
2955 (0.3_f64, 0.2_f64, 0.4_f64, 1.0_f64, 1.0_f64),
2956 (-0.7, -0.5, -1.1, 0.0, 0.8),
2957 ];
2958 let probit_scale = 0.8;
2959
2960 fn best_ns<F>(iterations: usize, base_g: f64, evaluate: F) -> f64
2966 where
2967 F: Fn(f64) -> (f64, [f64; 2], [[f64; 2]; 2]),
2968 {
2969 let mut best = f64::INFINITY;
2970 for _ in 0..5 {
2971 let mut checksum = 0.0_f64;
2972 let started = Instant::now();
2973 for _ in 0..iterations {
2974 let (value, gradient, hessian) = evaluate(base_g + checksum * 1e-18);
2975 checksum += value + gradient[0] + hessian[0][0];
2976 }
2977 assert!(
2978 checksum.is_finite(),
2979 "rigid Bernoulli release-measure checksum must stay finite"
2980 );
2981 best = best.min(started.elapsed().as_secs_f64());
2982 }
2983 best * 1e9 / iterations as f64
2984 }
2985
2986 let iterations = 2_000_000usize;
2987 for &(eta, g, z, y, w) in &cases {
2988 let marginal = bernoulli_marginal_link_map(
2989 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2990 eta,
2991 )
2992 .expect("link map");
2993
2994 let (jet_value, ..) =
2997 rigid_standard_normal_row_kernel(marginal, g, z, y, w, probit_scale)
2998 .expect("jet kernel");
2999 let (hand_value, ..) = hand_rigid_vgh(marginal, g, z, y, w, probit_scale);
3000 let band = 1e-12 + 1e-9 * jet_value.abs().max(hand_value.abs());
3001 assert!(
3002 (jet_value - hand_value).abs() <= band,
3003 "y={y:.0} value: jet {jet_value:+.15e} vs hand {hand_value:+.15e}"
3004 );
3005
3006 let production_ns = best_ns(iterations, g, |perturbed_g| {
3007 rigid_standard_normal_row_kernel(marginal, perturbed_g, z, y, w, probit_scale)
3008 .expect("jet kernel")
3009 });
3010 let hand_ns = best_ns(iterations, g, |perturbed_g| {
3011 hand_rigid_vgh(marginal, perturbed_g, z, y, w, probit_scale)
3012 });
3013 eprintln!(
3014 "RIGID-BERNOULLI-VGH-932 y={y:.0} production={production_ns:.2} ns/row \
3015 hand={hand_ns:.2} ns/row hand_over_production={:.6}",
3016 hand_ns / production_ns,
3017 );
3018 }
3019 }
3020}
3021
3022#[cfg(test)]
3023mod flex_primary_hessian_oracle_tests {
3024 use super::*;
3048 use super::family::*;
3055 use gam_linalg::matrix::DenseDesignMatrix;
3056 use ndarray::Array1;
3057 use ndarray::Array2;
3058 use std::sync::Arc;
3059 use std::sync::Mutex;
3060
3061 fn make_flex_oracle_family(
3067 n: usize,
3068 ) -> (BernoulliMarginalSlopeFamily, Vec<ParameterBlockState>) {
3069 let score_seed = Array1::linspace(-2.0, 2.0, n.max(6));
3070 let link_seed = Array1::linspace(-1.8, 1.8, n.max(6));
3071 let cfg = DeviationBlockConfig {
3072 num_internal_knots: 3,
3073 ..DeviationBlockConfig::default()
3074 };
3075 let score_prepared = build_score_warp_deviation_block_from_seed(&score_seed, &cfg)
3076 .expect("build score warp block");
3077 let link_prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
3078 &link_seed, &link_seed, &cfg,
3079 )
3080 .expect("build link deviation block");
3081
3082 let y: Array1<f64> =
3083 Array1::from_iter((0..n).map(|i| if (i * 17 + 3) % 7 >= 4 { 1.0 } else { 0.0 }));
3084 let weights: Array1<f64> =
3085 Array1::from_iter((0..n).map(|i| 0.75 + ((i * 11 + 5) % 5) as f64 * 0.05));
3086 let z: Array1<f64> =
3087 Array1::from_iter((0..n).map(|i| -1.7 + 3.4 * (i as f64 + 0.5) / n as f64));
3088 let marginal_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3089 if j == 0 {
3090 1.0
3091 } else {
3092 -0.4 + 0.8 * ((i * 19 + 7) % n) as f64 / n as f64
3093 }
3094 });
3095 let logslope_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3096 if j == 0 {
3097 1.0
3098 } else {
3099 0.3 - 0.6 * ((i * 23 + 11) % n) as f64 / n as f64
3100 }
3101 });
3102
3103 let family = BernoulliMarginalSlopeFamily {
3104 y: Arc::new(y),
3105 weights: Arc::new(weights),
3106 z: Arc::new(z.clone()),
3107 latent_measure: LatentMeasureKind::StandardNormal,
3108 gaussian_frailty_sd: Some(0.15),
3109 base_link: InverseLink::Standard(gam_problem::StandardLink::Probit),
3110 marginal_design: DesignMatrix::Dense(DenseDesignMatrix::from(marginal_x.clone())),
3111 logslope_design: DesignMatrix::Dense(DenseDesignMatrix::from(logslope_x.clone())),
3112 score_warp: Some(score_prepared.runtime.clone()),
3113 link_dev: Some(link_prepared.runtime.clone()),
3114 policy: gam_runtime::resource::ResourcePolicy::default_library(),
3115 cell_moment_lru: Arc::new(exact_kernel::CellMomentLruCache::new(1024)),
3116 cell_moment_cache_stats: Arc::new(exact_kernel::CellMomentCacheStats::default()),
3117 intercept_warm_starts: None,
3118 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
3119 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
3120 };
3121
3122 let beta_m = Array1::from_vec(vec![0.12, -0.04]);
3123 let beta_g = Array1::from_vec(vec![0.35, 0.03]);
3124 let beta_h = Array1::from_iter(
3125 (0..score_prepared.runtime.basis_dim()).map(|idx| 0.0015 * (idx as f64 + 1.0)),
3126 );
3127 let beta_w = Array1::from_iter(
3128 (0..link_prepared.runtime.basis_dim()).map(|idx| -0.001 * (idx as f64 + 1.0)),
3129 );
3130 let states = vec![
3131 ParameterBlockState {
3132 eta: marginal_x.dot(&beta_m),
3133 beta: beta_m,
3134 },
3135 ParameterBlockState {
3136 eta: logslope_x.dot(&beta_g),
3137 beta: beta_g,
3138 },
3139 ParameterBlockState {
3140 beta: beta_h,
3141 eta: Array1::zeros(z.len()),
3142 },
3143 ParameterBlockState {
3144 beta: beta_w,
3145 eta: Array1::zeros(z.len()),
3146 },
3147 ];
3148 (family, states)
3149 }
3150
3151 fn flex_gradient_at_perturbed(
3159 family: &BernoulliMarginalSlopeFamily,
3160 states: &[ParameterBlockState],
3161 primary: &super::super::hessian_paths::PrimarySlices,
3162 row: usize,
3163 u: usize,
3164 delta: f64,
3165 ) -> Array1<f64> {
3166 let mut states = states.to_vec();
3167 if u == primary.q {
3173 states[0].eta[row] += delta;
3174 } else if u == primary.logslope {
3175 states[1].eta[row] += delta;
3176 } else if let Some(h_range) = primary.h.as_ref()
3177 && h_range.contains(&u)
3178 {
3179 states[2].beta[u - h_range.start] += delta;
3180 } else if let Some(w_range) = primary.w.as_ref()
3181 && w_range.contains(&u)
3182 {
3183 states[3].beta[u - w_range.start] += delta;
3184 } else {
3185 panic!("primary coordinate {u} out of range for flex oracle");
3186 }
3187 let row_ctx = family
3188 .build_row_exact_context_with_stats_and_cell_cache(row, &states, None, false)
3189 .expect("perturbed row context");
3190 let (_neglog, grad, _hess) = family
3191 .compute_row_primary_gradient_hessian(row, &states, primary, &row_ctx)
3192 .expect("perturbed gradient");
3193 grad
3194 }
3195
3196 fn flex_nll_gradient_at_perturbed_z(
3202 family: &BernoulliMarginalSlopeFamily,
3203 states: &[ParameterBlockState],
3204 primary: &super::super::hessian_paths::PrimarySlices,
3205 row: usize,
3206 delta: f64,
3207 ) -> Array1<f64> {
3208 let mut perturbed = family.clone();
3209 let mut z = family.z.as_ref().clone();
3210 z[row] += delta;
3211 perturbed.z = Arc::new(z);
3212 let row_ctx = perturbed
3213 .build_row_exact_context_with_stats_and_cell_cache(row, states, None, false)
3214 .expect("z-perturbed row context");
3215 let mut scratch =
3216 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3217 perturbed
3218 .lower_bms_flex_row_order2(row, states, primary, &row_ctx, None, false, &mut scratch)
3219 .expect("z-perturbed flex gradient");
3220 scratch.grad
3221 }
3222
3223 #[test]
3229 fn flex_score_zeta_sensitivity_covers_all_active_deviation_blocks_2303() {
3230 let n = 12usize;
3231 let (family, states) = make_flex_oracle_family(n);
3232 let cache = family
3233 .build_exact_eval_cache(&states)
3234 .expect("flex exact eval cache");
3235 let primary = &cache.primary;
3236 let row = 5usize;
3237 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3238 let mut scratch =
3239 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3240 family
3241 .lower_bms_flex_row_order2(row, &states, primary, row_ctx, None, false, &mut scratch)
3242 .expect("analytic flex z-sensitivity");
3243 let analytic = scratch.score_zeta.clone();
3244
3245 let h = 1e-5_f64;
3246 let nll_plus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, h);
3247 let nll_minus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, -h);
3248 for u in 0..primary.total {
3249 let finite_difference = -(nll_plus[u] - nll_minus[u]) / (2.0 * h);
3251 let scale = 1.0 + analytic[u].abs().max(finite_difference.abs());
3252 let relative_error = (analytic[u] - finite_difference).abs() / scale;
3253 assert!(
3254 relative_error <= 2e-6,
3255 "flex score-zeta primary {u}: analytic={} FD={} relative_error={relative_error}",
3256 analytic[u],
3257 finite_difference
3258 );
3259 }
3260 let h_range = primary.h.as_ref().expect("score-warp primary range");
3261 let w_range = primary.w.as_ref().expect("link-deviation primary range");
3262 assert!(
3263 analytic
3264 .slice(s![h_range.start..h_range.end])
3265 .iter()
3266 .any(|value| value.abs() > 1e-10),
3267 "score-warp z-sensitivity must not be zero-filled"
3268 );
3269 assert!(
3270 analytic
3271 .slice(s![w_range.start..w_range.end])
3272 .iter()
3273 .any(|value| value.abs() > 1e-10),
3274 "link-deviation z-sensitivity must not be zero-filled"
3275 );
3276
3277 let coefficient = family
3278 .flex_score_zeta_sensitivity(
3279 &states,
3280 &BlockwiseFitOptions::default(),
3281 cache.slices.total,
3282 )
3283 .expect("full coefficient score-zeta sensitivity");
3284 assert_eq!(coefficient.dim(), (n, cache.slices.total));
3285 for range in [
3286 cache.slices.h.as_ref().expect("score-warp beta range"),
3287 cache.slices.w.as_ref().expect("link-deviation beta range"),
3288 ] {
3289 assert!(
3290 coefficient
3291 .slice(s![.., range.start..range.end])
3292 .iter()
3293 .any(|value| value.abs() > 1e-10),
3294 "active deviation coefficient range {range:?} must carry Murphy-Topel sensitivity"
3295 );
3296 }
3297
3298 let wrong_width = cache
3299 .slices
3300 .total
3301 .checked_sub(1)
3302 .expect("nonempty coefficient frame");
3303 let error = family
3304 .flex_score_zeta_sensitivity(&states, &BlockwiseFitOptions::default(), wrong_width)
3305 .expect_err("partial Murphy-Topel covariance frame must be rejected");
3306 assert!(
3307 error.contains("covariance/frame mismatch"),
3308 "unexpected partial-frame error: {error}"
3309 );
3310 }
3311
3312 #[test]
3315 fn flex_primary_hessian_matches_central_fd_of_gradient() {
3316 let n = 12usize;
3317 let (family, states) = make_flex_oracle_family(n);
3318 let cache = family
3319 .build_exact_eval_cache(&states)
3320 .expect("flex exact eval cache");
3321 let primary = &cache.primary;
3322 let r = primary.total;
3323 assert!(
3324 r >= 4,
3325 "flex fixture must carry q + logslope + deviation blocks"
3326 );
3327
3328 let h = 1e-4;
3332 let mut max_rel = 0.0_f64;
3333
3334 for &row in &[2usize, 5, 8] {
3337 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3338 let (_neglog, _grad, analytic_hess) = family
3339 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3340 .expect("analytic flex gradient + hessian");
3341
3342 for u in 0..r {
3343 let grad_plus = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3344 let grad_minus = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3345 for v in 0..r {
3346 let fd = (grad_plus[v] - grad_minus[v]) / (2.0 * h);
3347 let analytic = analytic_hess[[v, u]];
3348 let denom = 1.0 + analytic.abs().max(fd.abs());
3349 let rel = (analytic - fd).abs() / denom;
3350 max_rel = max_rel.max(rel);
3351 assert!(
3352 rel <= 1e-6,
3353 "flex hand Hessian H[{v}][{u}] = {analytic:.6e} disagrees with central \
3354 FD of the gradient {fd:.6e} at row {row} (rel {rel:.3e}); a product-rule \
3355 term is dropped or mis-signed"
3356 );
3357 }
3358 }
3359 }
3360 assert!(
3362 max_rel <= 1e-6,
3363 "flex Hessian FD oracle max rel {max_rel:.3e}"
3364 );
3365 }
3366
3367 #[test]
3376 fn arbiter_flex_hessian_h00_fd_step_scaling() {
3377 let n = 12usize;
3378 let (family, states) = make_flex_oracle_family(n);
3379 let cache = family
3380 .build_exact_eval_cache(&states)
3381 .expect("flex exact eval cache");
3382 let primary = &cache.primary;
3383 let row = 2usize;
3384 let u = primary.q; let v = primary.q;
3386
3387 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3388 let (_neglog, _grad, analytic_hess) = family
3389 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3390 .expect("analytic flex gradient + hessian");
3391 let analytic = analytic_hess[[v, u]];
3392
3393 let fd_at = |h: f64| -> f64 {
3394 let gp = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3395 let gm = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3396 (gp[v] - gm[v]) / (2.0 * h)
3397 };
3398
3399 let h = 1e-3_f64;
3406 let fd_h = fd_at(h);
3407 let fd_half = fd_at(h * 0.5);
3408 let fd_quarter = fd_at(h * 0.25);
3409 let gap_h = (analytic - fd_h).abs();
3410 let gap_half = (analytic - fd_half).abs();
3411 let gap_quarter = (analytic - fd_quarter).abs();
3412 let rich = (4.0 * fd_half - fd_h) / 3.0;
3413 let rich_gap = (analytic - rich).abs();
3414 let denom = analytic.abs().max(1.0);
3415
3416 let record = format!(
3418 "FLEX H[0][0] ARBITER row 2: analytic={analytic:+.12e} \
3419 fd(h)={fd_h:+.12e} fd(h/2)={fd_half:+.12e} fd(h/4)={fd_quarter:+.12e} \
3420 gap(h)={gap_h:.3e} gap(h/2)={gap_half:.3e} gap(h/4)={gap_quarter:.3e} \
3421 ratio_h_over_half={:.3} ratio_half_over_quarter={:.3} \
3422 richardson={rich:+.12e} richardson_gap={rich_gap:.3e} (rich_rel={:.3e})",
3423 gap_h / gap_half.max(f64::MIN_POSITIVE),
3424 gap_half / gap_quarter.max(f64::MIN_POSITIVE),
3425 rich_gap / denom,
3426 );
3427
3428 assert!(
3434 rich_gap / denom <= 1e-7,
3435 "{record}\nVERDICT: Richardson residual exceeds the FD-truncation floor — \
3436 the hand H[0][0] genuinely diverges (real dropped/mis-signed term), NOT FD noise"
3437 );
3438 }
3439}