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 mut square_root_factor = Array2::<f64>::zeros((dimension, dimension));
887 for (eigen_axis, &eigenvalue) in eigenvalues.iter().enumerate() {
888 if !(eigenvalue.is_finite() && eigenvalue >= 0.0) {
889 return Err(format!(
890 "marginal-slope full covariance must be positive semidefinite; eigenvalue {eigen_axis} is {eigenvalue}"
891 ));
892 }
893 let scale = eigenvalue.sqrt();
894 for axis in 0..dimension {
895 square_root_factor[[eigen_axis, axis]] = scale * eigenvectors[[axis, eigen_axis]];
896 }
897 }
898 let mut ones_quadratic_form = 0.0;
899 for factor_row in square_root_factor.rows() {
900 let projection = factor_row.sum();
901 ones_quadratic_form += projection * projection;
902 }
903 if !ones_quadratic_form.is_finite() {
904 return Err("marginal-slope full covariance geometry overflowed".to_string());
905 }
906 Ok(Self {
907 storage: MarginalSlopeCovarianceStorage::Full {
908 covariance,
909 square_root_factor,
910 },
911 ones_quadratic_form,
912 })
913 }
914
915 pub fn low_rank(factor: Array2<f64>) -> Result<Self, String> {
916 if factor.nrows() == 0 {
917 return Err("marginal-slope low-rank covariance factor has zero rows".to_string());
918 }
919 for ((row, column), &value) in factor.indexed_iter() {
920 if !value.is_finite() {
921 return Err(format!(
922 "marginal-slope low-rank covariance factor entry ({row},{column}) is non-finite"
923 ));
924 }
925 }
926 let mut ones_quadratic_form = 0.0;
927 for factor_column in factor.columns() {
928 let projection = factor_column.sum();
929 ones_quadratic_form += projection * projection;
930 }
931 if !ones_quadratic_form.is_finite() {
932 return Err("marginal-slope low-rank covariance geometry overflowed".to_string());
933 }
934 Ok(Self {
935 storage: MarginalSlopeCovarianceStorage::LowRank { factor },
936 ones_quadratic_form,
937 })
938 }
939
940 pub fn to_dense(&self) -> Array2<f64> {
941 match &self.storage {
942 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
943 Array2::from_diag(covariance)
944 }
945 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.clone(),
946 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.dot(&factor.t()),
947 }
948 }
949
950 pub fn shape(&self) -> MarginalSlopeCovarianceShape {
951 match &self.storage {
952 MarginalSlopeCovarianceStorage::Diagonal { .. } => {
953 MarginalSlopeCovarianceShape::Diagonal
954 }
955 MarginalSlopeCovarianceStorage::Full { .. } => MarginalSlopeCovarianceShape::Full,
956 MarginalSlopeCovarianceStorage::LowRank { .. } => MarginalSlopeCovarianceShape::LowRank,
957 }
958 }
959
960 pub fn dim(&self) -> usize {
961 match &self.storage {
962 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => covariance.len(),
963 MarginalSlopeCovarianceStorage::Full { covariance, .. } => covariance.nrows(),
964 MarginalSlopeCovarianceStorage::LowRank { factor } => factor.nrows(),
965 }
966 }
967
968 pub fn ones_quadratic_form(&self) -> f64 {
969 self.ones_quadratic_form
970 }
971
972 pub(crate) fn representation(&self) -> MarginalSlopeCovarianceRef<'_> {
973 match &self.storage {
974 MarginalSlopeCovarianceStorage::Diagonal { covariance, .. } => {
975 MarginalSlopeCovarianceRef::Diagonal(covariance)
976 }
977 MarginalSlopeCovarianceStorage::Full { covariance, .. } => {
978 MarginalSlopeCovarianceRef::Full(covariance)
979 }
980 MarginalSlopeCovarianceStorage::LowRank { factor } => {
981 MarginalSlopeCovarianceRef::LowRank(factor)
982 }
983 }
984 }
985
986 #[inline(always)]
987 pub(crate) fn quadratic_form_unchecked(&self, vector: &[f64]) -> f64 {
988 <Self as SymmetricQuadraticCoefficients>::quadratic_value(self, vector, |value| *value)
989 }
990
991 pub fn quadratic_form(&self, vector: &[f64]) -> Result<f64, String> {
992 if vector.len() != self.dim() {
993 return Err(format!(
994 "marginal-slope covariance dimension mismatch: vector={}, covariance={}",
995 vector.len(),
996 self.dim()
997 ));
998 }
999 if vector.iter().any(|value| !value.is_finite()) {
1000 return Err("marginal-slope covariance vector contains non-finite values".to_string());
1001 }
1002 let value = self.quadratic_form_unchecked(vector);
1003 if !value.is_finite() {
1004 return Err(format!(
1005 "marginal-slope covariance quadratic form is non-finite: {value}"
1006 ));
1007 }
1008 Ok(value)
1009 }
1010}
1011
1012enum VectorSupport {
1013 Zero,
1014 Singleton { axis: usize, value: f64 },
1015 Multiple,
1016}
1017
1018#[inline(always)]
1019fn vector_support(input: &[f64]) -> VectorSupport {
1020 let mut singleton = None;
1021 for (axis, &value) in input.iter().enumerate() {
1022 if value == 0.0 {
1023 continue;
1024 }
1025 if singleton.is_some() {
1026 return VectorSupport::Multiple;
1027 }
1028 singleton = Some((axis, value));
1029 }
1030 match singleton {
1031 None => VectorSupport::Zero,
1032 Some((axis, value)) => VectorSupport::Singleton { axis, value },
1033 }
1034}
1035
1036impl SymmetricQuadraticCoefficients for MarginalSlopeCovariance {
1037 fn dimension(&self) -> usize {
1038 self.dim()
1039 }
1040
1041 fn multiply(&self, input: &[f64], output: &mut [f64]) {
1042 assert_eq!(input.len(), self.dim());
1043 assert_eq!(output.len(), self.dim());
1044 match self.representation() {
1045 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1046 for axis in 0..input.len() {
1047 output[axis] = diagonal[axis] * input[axis];
1048 }
1049 }
1050 MarginalSlopeCovarianceRef::Full(matrix) => {
1051 match vector_support(input) {
1052 VectorSupport::Zero => {
1053 output.fill(0.0);
1054 return;
1055 }
1056 VectorSupport::Singleton { axis, value } => {
1057 for row in 0..input.len() {
1058 output[row] = matrix[[row, axis]] * value;
1059 }
1060 return;
1061 }
1062 VectorSupport::Multiple => {}
1063 }
1064 for row in 0..input.len() {
1065 let mut value = 0.0;
1066 for column in 0..input.len() {
1067 value += matrix[[row, column]] * input[column];
1068 }
1069 output[row] = value;
1070 }
1071 }
1072 MarginalSlopeCovarianceRef::LowRank(factor) => {
1073 output.fill(0.0);
1074 match vector_support(input) {
1075 VectorSupport::Zero => return,
1076 VectorSupport::Singleton { axis, value } => {
1077 for rank in 0..factor.ncols() {
1078 let projection = factor[[axis, rank]] * value;
1079 for row in 0..input.len() {
1080 output[row] += factor[[row, rank]] * projection;
1081 }
1082 }
1083 return;
1084 }
1085 VectorSupport::Multiple => {}
1086 }
1087 for rank in 0..factor.ncols() {
1088 let mut projection = 0.0;
1089 for row in 0..input.len() {
1090 projection += factor[[row, rank]] * input[row];
1091 }
1092 for row in 0..input.len() {
1093 output[row] += factor[[row, rank]] * projection;
1094 }
1095 }
1096 }
1097 }
1098 }
1099
1100 fn coefficient(&self, row: usize, column: usize) -> f64 {
1101 match self.representation() {
1102 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1103 if row == column {
1104 diagonal[row]
1105 } else {
1106 0.0
1107 }
1108 }
1109 MarginalSlopeCovarianceRef::Full(matrix) => matrix[[row, column]],
1110 MarginalSlopeCovarianceRef::LowRank(factor) => {
1111 let mut value = 0.0;
1112 for rank in 0..factor.ncols() {
1113 value += factor[[row, rank]] * factor[[column, rank]];
1114 }
1115 value
1116 }
1117 }
1118 }
1119
1120 fn visit_upper_triangle(
1121 &self,
1122 direction: &mut [f64],
1123 projected: &mut [f64],
1124 mut visit: impl FnMut(usize, usize, f64),
1125 ) {
1126 let dimension = self.dim();
1127 assert_eq!(direction.len(), dimension);
1128 assert_eq!(projected.len(), dimension);
1129 match self.representation() {
1130 MarginalSlopeCovarianceRef::Diagonal(diagonal) => {
1131 for column in 0..dimension {
1132 for row in 0..=column {
1133 visit(row, column, if row == column { diagonal[row] } else { 0.0 });
1134 }
1135 }
1136 }
1137 MarginalSlopeCovarianceRef::Full(matrix) => {
1138 for column in 0..dimension {
1139 for row in 0..=column {
1140 visit(row, column, matrix[[row, column]]);
1141 }
1142 }
1143 }
1144 MarginalSlopeCovarianceRef::LowRank(factor) => {
1145 for column in 0..dimension {
1146 for row in 0..=column {
1147 let mut value = 0.0;
1148 for rank in 0..factor.ncols() {
1149 value += factor[[row, rank]] * factor[[column, rank]];
1150 }
1151 visit(row, column, value);
1152 }
1153 }
1154 }
1155 }
1156 }
1157
1158 fn quadratic_value<T, F>(&self, input: &[T], value: F) -> f64
1159 where
1160 F: Fn(&T) -> f64,
1161 {
1162 assert_eq!(input.len(), self.dim());
1163 match &self.storage {
1164 MarginalSlopeCovarianceStorage::Diagonal { covariance } => input
1165 .iter()
1166 .zip(covariance)
1167 .map(|(input, &covariance)| {
1168 let input = value(input);
1169 covariance * input * input
1170 })
1171 .sum(),
1172 MarginalSlopeCovarianceStorage::Full {
1173 square_root_factor, ..
1174 } => {
1175 let mut total = 0.0;
1176 for factor_row in square_root_factor.rows() {
1177 let mut projection = 0.0;
1178 for axis in 0..input.len() {
1179 projection += factor_row[axis] * value(&input[axis]);
1180 }
1181 total += projection * projection;
1182 }
1183 total
1184 }
1185 MarginalSlopeCovarianceStorage::LowRank { factor } => {
1186 let mut total = 0.0;
1189 for rank in 0..factor.ncols() {
1190 let mut projection = 0.0;
1191 for row in 0..input.len() {
1192 projection += factor[[row, rank]] * value(&input[row]);
1193 }
1194 total += projection * projection;
1195 }
1196 total
1197 }
1198 }
1199 }
1200}
1201
1202pub fn marginal_slope_covariance_from_scores(
1229 scores: ArrayView2<'_, f64>,
1230 weights: &Array1<f64>,
1231) -> Result<MarginalSlopeCovariance, String> {
1232 let (n, k) = scores.dim();
1233 if k == 0 {
1234 return Err("marginal-slope score matrix must have at least one column".to_string());
1235 }
1236 if weights.len() != n {
1237 return Err(format!(
1238 "marginal-slope covariance weight length mismatch: weights={}, rows={n}",
1239 weights.len()
1240 ));
1241 }
1242 let total_weight = weights.iter().copied().sum::<f64>();
1243 if !(total_weight.is_finite() && total_weight > 0.0) {
1244 return Err("marginal-slope covariance needs positive finite total weight".to_string());
1245 }
1246 let mut mean = Array1::<f64>::zeros(k);
1247 for i in 0..n {
1248 let weight = weights[i];
1249 if !(weight.is_finite() && weight >= 0.0) {
1250 return Err(format!(
1251 "marginal-slope covariance weight {i} must be finite and non-negative, got {weight}"
1252 ));
1253 }
1254 for j in 0..k {
1255 let score = scores[[i, j]];
1256 if !score.is_finite() {
1257 return Err(format!(
1258 "marginal-slope covariance score ({i},{j}) is non-finite"
1259 ));
1260 }
1261 mean[j] += weight * score;
1262 }
1263 }
1264 mean.mapv_inplace(|value| value / total_weight);
1265
1266 let mut cov = Array2::<f64>::zeros((k, k));
1267 for i in 0..n {
1268 let weight = weights[i];
1269 for a in 0..k {
1270 let da = scores[[i, a]] - mean[a];
1271 for b in 0..=a {
1272 let value = weight * da * (scores[[i, b]] - mean[b]) / total_weight;
1273 cov[[a, b]] += value;
1274 if a != b {
1275 cov[[b, a]] += value;
1276 }
1277 }
1278 }
1279 }
1280
1281 let is_diagonal = (0..k).all(|row| ((row + 1)..k).all(|column| cov[[row, column]] == 0.0));
1285 if is_diagonal {
1286 MarginalSlopeCovariance::diagonal(cov.diag().to_owned())
1287 } else {
1288 MarginalSlopeCovariance::full(cov)
1289 }
1290}
1291
1292pub fn marginal_slope_preserving_scale(
1293 slopes: &[f64],
1294 covariance: &MarginalSlopeCovariance,
1295 probit_scale: f64,
1296) -> Result<f64, String> {
1297 if !probit_scale.is_finite() {
1298 return Err(format!(
1299 "marginal-slope probit scale must be finite, got {probit_scale}"
1300 ));
1301 }
1302 let variance = probit_scale * probit_scale * covariance.quadratic_form(slopes)?;
1303 if !variance.is_finite() {
1304 return Err("marginal-slope preserving variance is non-finite".to_string());
1305 }
1306 Ok((1.0 + variance).sqrt())
1307}
1308
1309pub fn marginal_slope_probit_eta(
1310 q: f64,
1311 z: &[f64],
1312 slopes: &[f64],
1313 covariance: &MarginalSlopeCovariance,
1314 probit_scale: f64,
1315) -> Result<f64, String> {
1316 if z.len() != slopes.len() {
1317 return Err(format!(
1318 "marginal-slope score/slope dimension mismatch: z={}, slopes={}",
1319 z.len(),
1320 slopes.len()
1321 ));
1322 }
1323 if slopes.len() != covariance.dim() {
1324 return Err(format!(
1325 "marginal-slope covariance dimension mismatch: slopes={}, covariance={}",
1326 slopes.len(),
1327 covariance.dim()
1328 ));
1329 }
1330 if !q.is_finite() || z.iter().any(|value| !value.is_finite()) {
1331 return Err("marginal-slope probit eta inputs must be finite".to_string());
1332 }
1333 let scale = marginal_slope_preserving_scale(slopes, covariance, probit_scale)?;
1334 let linear = z
1335 .iter()
1336 .zip(slopes.iter())
1337 .map(|(&score, &slope)| probit_scale * slope * score)
1338 .sum::<f64>();
1339 Ok(q * scale + linear)
1340}
1341
1342pub(super) fn empirical_rigid_calibration_eval(
1372 intercept: f64,
1373 log_target_mu: f64,
1374 slope: f64,
1375 probit_scale: f64,
1376 nodes: &[f64],
1377 weights: &[f64],
1378) -> Result<(f64, f64, f64), String> {
1379 if !intercept.is_finite() {
1380 return Err(format!(
1381 "empirical latent calibration: non-finite intercept {intercept}"
1382 ));
1383 }
1384 let observed_slope = rigid_observed_logslope(slope, probit_scale);
1385 const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_8; let mut log_max_phi = f64::NEG_INFINITY;
1389 let mut sum_phi = 0.0_f64;
1390 let mut log_max_cdf = f64::NEG_INFINITY;
1391 let mut sum_cdf = 0.0_f64;
1392
1393 let mut log_max_pos = f64::NEG_INFINITY;
1397 let mut sum_pos = 0.0_f64;
1398 let mut log_max_neg = f64::NEG_INFINITY;
1399 let mut sum_neg = 0.0_f64;
1400
1401 for (&node, &weight) in nodes.iter().zip(weights.iter()) {
1402 if !(weight.is_finite() && weight > 0.0) {
1403 continue;
1404 }
1405 let eta = intercept + observed_slope * node;
1406 if !eta.is_finite() {
1407 return Err(format!(
1408 "empirical latent calibration: non-finite η at intercept={intercept}, slope={slope}, node={node}"
1409 ));
1410 }
1411 let log_w = weight.ln();
1412 let log_phi = -0.5 * eta * eta - HALF_LOG_2PI;
1413 let log_term_phi = log_w + log_phi;
1414 let log_term_cdf = log_w + normal_logcdf(eta);
1415
1416 lse_accumulate(&mut log_max_phi, &mut sum_phi, log_term_phi);
1417 lse_accumulate(&mut log_max_cdf, &mut sum_cdf, log_term_cdf);
1418
1419 if eta != 0.0 {
1420 let log_term_eta_phi = log_term_phi + eta.abs().ln();
1421 if eta > 0.0 {
1422 lse_accumulate(&mut log_max_pos, &mut sum_pos, log_term_eta_phi);
1423 } else {
1424 lse_accumulate(&mut log_max_neg, &mut sum_neg, log_term_eta_phi);
1425 }
1426 }
1427 }
1428
1429 if !(sum_phi.is_finite() && sum_cdf.is_finite() && sum_phi > 0.0 && sum_cdf > 0.0) {
1430 return Err(format!(
1431 "empirical latent calibration: log-space accumulation failed (sum_phi={sum_phi}, sum_cdf={sum_cdf}, intercept={intercept})"
1432 ));
1433 }
1434
1435 let log_s_phi = log_max_phi + sum_phi.ln();
1436 let log_s_cdf = log_max_cdf + sum_cdf.ln();
1437
1438 let f = log_s_cdf - log_target_mu;
1440 let log_f_prime = log_s_phi - log_s_cdf;
1452 let f_prime = if log_f_prime > -740.0 {
1453 log_f_prime.exp()
1454 } else {
1455 f64::MIN_POSITIVE
1456 };
1457
1458 let exp_safe = |log_x: f64| -> f64 { if log_x > -740.0 { log_x.exp() } else { 0.0 } };
1466 let pos_over_cdf = if sum_pos > 0.0 {
1467 exp_safe(log_max_pos + sum_pos.ln() - log_s_cdf)
1468 } else {
1469 0.0
1470 };
1471 let neg_over_cdf = if sum_neg > 0.0 {
1472 exp_safe(log_max_neg + sum_neg.ln() - log_s_cdf)
1473 } else {
1474 0.0
1475 };
1476 let s_etaphi_over_s_cdf = pos_over_cdf - neg_over_cdf;
1477 let f_double_prime = -s_etaphi_over_s_cdf - f_prime * f_prime;
1478
1479 if !(f.is_finite() && f_prime.is_finite() && f_prime > 0.0 && f_double_prime.is_finite()) {
1480 return Err(format!(
1481 "empirical latent calibration: non-finite log-space state f={f}, f'={f_prime}, f''={f_double_prime} at intercept={intercept}"
1482 ));
1483 }
1484 Ok((f, f_prime, f_double_prime))
1485}
1486
1487pub(crate) fn empirical_intercept_from_marginal(
1488 target_mu: f64,
1489 target_q: f64,
1490 slope: f64,
1491 probit_scale: f64,
1492 nodes: &[f64],
1493 weights: &[f64],
1494 initial: Option<f64>,
1495) -> Result<f64, String> {
1496 if !(target_mu.is_finite() && target_mu > 0.0 && target_mu < 1.0) {
1497 return Err(format!(
1498 "empirical latent calibration requires target mu in (0,1), got {target_mu}"
1499 ));
1500 }
1501 let log_target_mu = target_mu.ln();
1502 let closed_form_seed = rigid_intercept_from_marginal(target_q, slope, probit_scale);
1503 let seed = initial.unwrap_or(closed_form_seed);
1504 let eval = |a: f64| {
1505 empirical_rigid_calibration_eval(a, log_target_mu, slope, probit_scale, nodes, weights)
1506 };
1507 let abs_tol = 1e-13_f64.max(4.0 * f64::EPSILON);
1514 let solve_from = |s: f64| {
1515 crate::monotone_root::solve_monotone_root(
1516 eval,
1517 s,
1518 "empirical latent intercept",
1519 abs_tol,
1520 64,
1521 48,
1522 )
1523 .map_err(|e| e.to_string())
1526 };
1527 let (root, _, f_best) = match solve_from(seed) {
1538 Ok(v) => v,
1539 Err(first_err) => {
1540 if seed == closed_form_seed {
1541 return Err(first_err);
1542 }
1543 solve_from(closed_form_seed).map_err(|retry_err| {
1544 format!("{first_err}; closed-form retry from a={closed_form_seed:.6}: {retry_err}")
1545 })?
1546 }
1547 };
1548 if f_best.abs() > abs_tol {
1549 return Err(format!(
1550 "empirical latent intercept solve failed: log-residual={f_best:.3e} at a={root:.6}, target mu={target_mu:.6}"
1551 ));
1552 }
1553 Ok(root)
1554}
1555
1556#[inline]
1557pub(super) fn rigid_standard_normal_neglog_only(
1558 q: f64,
1559 g: f64,
1560 z: f64,
1561 y: f64,
1562 w: f64,
1563 probit_scale: f64,
1564) -> Result<f64, String> {
1565 let s = 2.0 * y - 1.0;
1566 let eta = marginal_slope_standard_normal_scalar_eta(q, g, z, probit_scale);
1567 let m = s * eta;
1568 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
1569 if !logcdf.is_finite() {
1570 return Err(format!(
1571 "rigid probit neglog_only: non-finite log Φ at q={q}, g={g}, z={z}, y={y}"
1572 ));
1573 }
1574 Ok(-w * logcdf)
1575}
1576
1577#[inline]
1607pub(crate) fn rigid_standard_normal_row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1608 p: &[S; 2],
1609 marginal: BernoulliMarginalLinkMap,
1610 z: f64,
1611 y: f64,
1612 w: f64,
1613 probit_scale: f64,
1614) -> Result<S, String> {
1615 let signed = rigid_standard_normal_signed_margin(p, marginal, z, y, probit_scale);
1619 let m = signed.value();
1622 if !(m.is_finite() || m == f64::INFINITY) {
1623 return Err(format!(
1624 "non-finite signed margin in rigid probit row NLL: {m}"
1625 ));
1626 }
1627 Ok(signed.compose_unary(signed_probit_neglog_unary_stack(m, w)))
1629}
1630
1631#[inline]
1641pub(crate) fn rigid_standard_normal_signed_margin<S: gam_math::jet_scalar::JetScalar<2>>(
1642 p: &[S; 2],
1643 marginal: BernoulliMarginalLinkMap,
1644 z: f64,
1645 y: f64,
1646 probit_scale: f64,
1647) -> S {
1648 let q = p[0].compose_unary([
1650 marginal.q,
1651 marginal.q1,
1652 marginal.q2,
1653 marginal.q3,
1654 marginal.q4,
1655 ]);
1656 let slope = p[1];
1657 let observed_slope = slope.scale(probit_scale);
1659 let b2 = observed_slope.mul(&observed_slope);
1660 let c = b2.add(&S::constant(1.0)).sqrt();
1661 let eta = q.mul(&c).add(&observed_slope.scale(z));
1663 eta.scale(2.0 * y - 1.0)
1664}
1665
1666pub(crate) struct RigidStandardNormalRow {
1679 pub(crate) marginal: BernoulliMarginalLinkMap,
1680 pub(crate) g: f64,
1681 pub(crate) z: f64,
1682 pub(crate) y: f64,
1683 pub(crate) w: f64,
1684 pub(crate) probit_scale: f64,
1685}
1686
1687impl gam_math::jet_tower::RowProgram<2> for RigidStandardNormalRow {
1688 fn n_rows(&self) -> usize {
1689 1
1690 }
1691
1692 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
1693 if row != 0 {
1694 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1695 }
1696 Ok([self.marginal.eta_value(), self.g])
1697 }
1698
1699 fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
1700 &self,
1701 row: usize,
1702 p: &[S; 2],
1703 ) -> Result<S, String> {
1704 if row != 0 {
1705 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1706 }
1707 rigid_standard_normal_row_nll_generic(
1708 p,
1709 self.marginal,
1710 self.z,
1711 self.y,
1712 self.w,
1713 self.probit_scale,
1714 )
1715 }
1716}
1717
1718#[inline]
1719pub(crate) fn rigid_standard_normal_tower(
1720 marginal: BernoulliMarginalLinkMap,
1721 g: f64,
1722 z: f64,
1723 y: f64,
1724 w: f64,
1725 probit_scale: f64,
1726) -> Result<Tower4<2>, String> {
1727 let program = RigidStandardNormalRow {
1734 marginal,
1735 g,
1736 z,
1737 y,
1738 w,
1739 probit_scale,
1740 };
1741 gam_math::jet_tower::program_full_tower(&program, 0).map(|tower| *tower)
1742}
1743
1744#[inline]
1757fn rigid_standard_normal_signed_jet(
1758 marginal: BernoulliMarginalLinkMap,
1759 g: f64,
1760 z: f64,
1761 y: f64,
1762 probit_scale: f64,
1763) -> Tower4<2> {
1764 let p = [
1767 Tower4::<2>::variable(marginal.eta_value(), 0),
1768 Tower4::<2>::variable(g, 1),
1769 ];
1770 rigid_standard_normal_signed_margin(&p, marginal, z, y, probit_scale)
1771}
1772
1773#[inline]
1801pub(super) fn rigid_standard_normal_towers_batch<T>(
1802 marginals: &[BernoulliMarginalLinkMap],
1803 slopes: &[f64],
1804 zs: &[f64],
1805 ys: &[f64],
1806 weights: &[f64],
1807 probit_scale: f64,
1808 out: &mut [T],
1809 mut fill: impl FnMut(&Tower4<2>) -> Result<T, String>,
1810) -> Result<(), String> {
1811 let chunk = marginals.len();
1812 if slopes.len() != chunk
1813 || zs.len() != chunk
1814 || ys.len() != chunk
1815 || weights.len() != chunk
1816 || out.len() != chunk
1817 {
1818 return Err(format!(
1819 "rigid_standard_normal_towers_batch length mismatch: marginals={chunk}, \
1820 slopes={}, zs={}, ys={}, weights={}, out={}",
1821 slopes.len(),
1822 zs.len(),
1823 ys.len(),
1824 weights.len(),
1825 out.len()
1826 ));
1827 }
1828
1829 let mut signed: Vec<Tower4<2>> = Vec::with_capacity(chunk);
1831 let mut margins: Vec<f64> = Vec::with_capacity(chunk);
1832 for i in 0..chunk {
1833 let jet =
1834 rigid_standard_normal_signed_jet(marginals[i], slopes[i], zs[i], ys[i], probit_scale);
1835 margins.push(jet.v);
1836 signed.push(jet);
1837 }
1838
1839 let mut stacks: Vec<[f64; 5]> = Vec::with_capacity(chunk);
1843 for i in 0..chunk {
1844 let m = margins[i];
1845 if !(m.is_finite() || m == f64::INFINITY) {
1846 return Err(format!(
1847 "non-finite signed margin in rigid probit tower batch: {m}"
1848 ));
1849 }
1850 stacks.push(signed_probit_neglog_unary_stack(m, weights[i]));
1851 }
1852
1853 for i in 0..chunk {
1855 let tower = signed[i].compose_unary(stacks[i]);
1856 out[i] = fill(&tower)?;
1857 }
1858 Ok(())
1859}
1860
1861#[inline]
1862pub(super) fn rigid_standard_normal_row_kernel(
1863 marginal: BernoulliMarginalLinkMap,
1864 g: f64,
1865 z: f64,
1866 y: f64,
1867 w: f64,
1868 probit_scale: f64,
1869) -> Result<(f64, [f64; 2], [[f64; 2]; 2]), String> {
1870 let program = RigidStandardNormalRow {
1879 marginal,
1880 g,
1881 z,
1882 y,
1883 w,
1884 probit_scale,
1885 };
1886 gam_math::jet_tower::program_row_kernel(&program, 0)
1887}
1888
1889#[inline]
1927pub(super) fn rigid_standard_normal_mixed_z_sensitivity(
1928 marginal: BernoulliMarginalLinkMap,
1929 g: f64,
1930 z: f64,
1931 y: f64,
1932 w: f64,
1933 probit_scale: f64,
1934) -> Result<[f64; 2], String> {
1935 use gam_math::jet_tower::Tower2;
1946 let mut q = Tower2::<3>::constant(marginal.q);
1947 q.g[0] = marginal.q1;
1948 q.h[0][0] = marginal.q2;
1949 let slope = Tower2::<3>::variable(g, 1);
1950 let z_var = Tower2::<3>::variable(z, 2);
1951 let observed_logslope = slope * probit_scale;
1952 let c = (observed_logslope * observed_logslope + 1.0).sqrt();
1953 let eta = q * c + slope * (z_var * probit_scale);
1957 let signed = eta * (2.0 * y - 1.0);
1958 if !(signed.v.is_finite() || signed.v == f64::INFINITY) {
1960 return Err(format!(
1961 "rigid probit mixed-z sensitivity: non-finite signed margin {} at q={}, g={g}, z={z}, y={y}",
1962 signed.v, marginal.q
1963 ));
1964 }
1965 let stack = signed_probit_neglog_unary_stack(signed.v, w);
1966 if !stack[0].is_finite() {
1967 return Err(format!(
1968 "rigid probit mixed-z sensitivity: non-finite log Φ at q={}, g={g}, z={z}, y={y}",
1969 marginal.q
1970 ));
1971 }
1972 let tower = signed.compose_unary([stack[0], stack[1], stack[2]]);
1975 let s_q = -tower.h[0][2];
1982 let s_g = -tower.h[1][2];
1983 if !(s_q.is_finite() && s_g.is_finite()) {
1984 return Err(format!(
1985 "rigid probit mixed-z sensitivity: non-finite ∂²(log L)/∂(q,g)∂z = [{s_q}, {s_g}] at q={}, g={g}, z={z}",
1986 marginal.q
1987 ));
1988 }
1989 Ok([s_q, s_g])
1990}
1991
1992pub(super) fn rigid_standard_normal_score_zeta_sensitivity(
2016 base_link: &InverseLink,
2017 marginal_eta: &Array1<f64>,
2018 slope_eta: &Array1<f64>,
2019 z: &Array1<f64>,
2020 y: &Array1<f64>,
2021 weights: &Array1<f64>,
2022 probit_scale: f64,
2023 marginal_design: ArrayView2<'_, f64>,
2024 logslope_design: ArrayView2<'_, f64>,
2025 p_beta: usize,
2026) -> Result<Array2<f64>, String> {
2027 let n = marginal_eta.len();
2028 let p_m = marginal_design.ncols();
2029 let r = logslope_design.ncols();
2030 if slope_eta.len() != n
2031 || z.len() != n
2032 || y.len() != n
2033 || weights.len() != n
2034 || marginal_design.nrows() != n
2035 || logslope_design.nrows() != n
2036 {
2037 return Err(format!(
2038 "score_zeta_sensitivity row mismatch: marginal_eta={n}, slope_eta={}, z={}, y={}, \
2039 weights={}, marginal_design rows={}, logslope_design rows={}",
2040 slope_eta.len(),
2041 z.len(),
2042 y.len(),
2043 weights.len(),
2044 marginal_design.nrows(),
2045 logslope_design.nrows()
2046 ));
2047 }
2048 if p_m + r != p_beta {
2049 return Err(format!(
2050 "rigid score_zeta_sensitivity width mismatch: marginal({p_m}) + logslope({r}) != p_beta({p_beta})"
2051 ));
2052 }
2053 let mut s = Array2::<f64>::zeros((n, p_beta));
2054 for i in 0..n {
2055 let marginal = bernoulli_marginal_link_map(base_link, marginal_eta[i])?;
2056 let [s_q, s_g] = rigid_standard_normal_mixed_z_sensitivity(
2057 marginal,
2058 slope_eta[i],
2059 z[i],
2060 y[i],
2061 weights[i],
2062 probit_scale,
2063 )?;
2064 if s_q != 0.0 {
2067 let m_row = marginal_design.row(i);
2068 for (j, &mij) in m_row.iter().enumerate() {
2069 s[[i, j]] = s_q * mij;
2070 }
2071 }
2072 if s_g != 0.0 {
2073 let g_row = logslope_design.row(i);
2074 for (j, &gij) in g_row.iter().enumerate() {
2075 s[[i, p_m + j]] = s_g * gij;
2076 }
2077 }
2078 }
2079 Ok(s)
2080}
2081
2082#[inline]
2083pub(super) fn rigid_standard_normal_third_full(
2084 marginal: BernoulliMarginalLinkMap,
2085 g: f64,
2086 z: f64,
2087 y: f64,
2088 w: f64,
2089 probit_scale: f64,
2090) -> Result<[[[f64; 2]; 2]; 2], String> {
2091 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t3)
2092}
2093
2094#[inline]
2099pub(super) fn contract_third_full(t: &[[[f64; 2]; 2]; 2], d_eta: f64, d_g: f64) -> [[f64; 2]; 2] {
2100 [
2101 [
2102 t[0][0][0] * d_eta + t[0][0][1] * d_g,
2103 t[0][1][0] * d_eta + t[0][1][1] * d_g,
2104 ],
2105 [
2106 t[1][0][0] * d_eta + t[1][0][1] * d_g,
2107 t[1][1][0] * d_eta + t[1][1][1] * d_g,
2108 ],
2109 ]
2110}
2111
2112#[inline]
2113pub(super) fn rigid_standard_normal_fourth_full(
2114 marginal: BernoulliMarginalLinkMap,
2115 g: f64,
2116 z: f64,
2117 y: f64,
2118 w: f64,
2119 probit_scale: f64,
2120) -> Result<[[[[f64; 2]; 2]; 2]; 2], String> {
2121 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t4)
2135}
2136
2137#[inline]
2156pub(super) fn contract_fourth_full(
2157 t: &[[[[f64; 2]; 2]; 2]; 2],
2158 u_eta: f64,
2159 u_g: f64,
2160 v_eta: f64,
2161 v_g: f64,
2162) -> [[f64; 2]; 2] {
2163 let mut out = [[0.0; 2]; 2];
2164 for a in 0..2 {
2165 for b in 0..2 {
2166 let mut sum = 0.0;
2167 sum += t[a][b][0][0] * u_eta * v_eta;
2168 sum += t[a][b][0][1] * u_eta * v_g;
2169 sum += t[a][b][1][0] * u_g * v_eta;
2170 sum += t[a][b][1][1] * u_g * v_g;
2171 out[a][b] = sum;
2172 }
2173 }
2174 out
2175}
2176
2177pub(super) fn ensure_finite_third_full_cache_row(
2178 t: &[[[f64; 2]; 2]; 2],
2179 context: &str,
2180) -> Result<(), String> {
2181 if t.iter().flatten().flatten().all(|value| value.is_finite()) {
2182 Ok(())
2183 } else {
2184 Err(format!(
2185 "{context}: warmed third-derivative cache row contains a non-finite value"
2186 ))
2187 }
2188}
2189
2190pub(super) fn ensure_finite_fourth_full_cache_row(
2191 t: &[[[[f64; 2]; 2]; 2]; 2],
2192 context: &str,
2193) -> Result<(), String> {
2194 if t.iter()
2195 .flatten()
2196 .flatten()
2197 .flatten()
2198 .all(|value| value.is_finite())
2199 {
2200 Ok(())
2201 } else {
2202 Err(format!(
2203 "{context}: warmed fourth-derivative cache row contains a non-finite value"
2204 ))
2205 }
2206}
2207
2208pub(crate) fn unary_derivatives_sqrt(x: f64) -> [f64; 5] {
2209 let s = x.max(1e-300).sqrt();
2210 let x1 = x.max(1e-300);
2211 let x2 = x1 * x1;
2212 let x3 = x2 * x1;
2213 [
2214 s,
2215 0.5 / s,
2216 -0.25 / (x1 * s),
2217 3.0 / (8.0 * x2 * s),
2218 -15.0 / (16.0 * x3 * s),
2219 ]
2220}
2221pub(crate) fn unary_derivatives_neglog_phi(x: f64, weight: f64) -> [f64; 5] {
2222 signed_probit_neglog_unary_stack(x, weight)
2227}
2228
2229pub(crate) fn unary_derivatives_log(x: f64) -> [f64; 5] {
2247 let x2 = x * x;
2248 let x3 = x2 * x;
2249 let x4 = x3 * x;
2250 [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
2251}
2252
2253pub(crate) fn unary_derivatives_log_normal_pdf(x: f64) -> [f64; 5] {
2255 let c = 0.5 * (2.0 * std::f64::consts::PI).ln();
2256 [-0.5 * x * x - c, -x, -1.0, 0.0, 0.0]
2257}
2258
2259#[cfg(test)]
2260mod covariance_admission_tests {
2261 use super::*;
2262 use ndarray::array;
2263
2264 #[test]
2265 fn full_covariance_admission_rejects_one_ulp_asymmetry_932() {
2266 let upper = 0.25_f64;
2267 let lower = f64::from_bits(upper.to_bits() + 1);
2268 let error = MarginalSlopeCovariance::full(array![[1.0, upper], [lower, 1.0]])
2269 .expect_err("any asymmetric full operator must be rejected");
2270 assert!(error.contains("must be exactly symmetric"), "{error}");
2271 }
2272
2273 #[test]
2274 fn full_covariance_admission_rejects_indefinite_matrix_before_row_use_932() {
2275 let error = MarginalSlopeCovariance::full(array![[1.0, 2.0], [2.0, 1.0]])
2276 .expect_err("an indefinite full operator is not a covariance");
2277 assert!(error.contains("must be positive semidefinite"), "{error}");
2278 }
2279
2280 #[test]
2281 fn full_covariance_admission_accepts_exact_singular_psd_932() {
2282 MarginalSlopeCovariance::full(array![[1.0, 0.0], [0.0, 0.0]])
2283 .expect("an exact singular PSD covariance is admissible");
2284 }
2285
2286 #[test]
2287 fn full_covariance_admission_accepts_coupled_singular_psd_932() {
2288 let covariance = MarginalSlopeCovariance::full(array![[1.0, 1.0], [1.0, 1.0]])
2289 .expect("a coupled singular PSD covariance is admissible");
2290 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2291 assert_eq!(covariance.to_dense(), array![[1.0, 1.0], [1.0, 1.0]]);
2292 }
2293
2294 #[test]
2295 fn exact_nonzero_offdiagonal_classifier_retains_full_geometry_932() {
2296 let epsilon = 1.0e-14;
2297 let scores = array![[-1.0, -epsilon], [1.0, epsilon], [0.0, -1.0], [0.0, 1.0]];
2298 let covariance =
2299 marginal_slope_covariance_from_scores(scores.view(), &Array1::ones(4)).unwrap();
2300 let dense = covariance.to_dense();
2301 assert_eq!(covariance.shape(), MarginalSlopeCovarianceShape::Full);
2302 assert_ne!(dense[[0, 1]], 0.0);
2303 assert_eq!(dense[[0, 1]], dense[[1, 0]]);
2304 let direction = [0.75, -1.25];
2305 let expected = direction[0] * (dense[[0, 0]] * direction[0] + dense[[0, 1]] * direction[1])
2306 + direction[1] * (dense[[1, 0]] * direction[0] + dense[[1, 1]] * direction[1]);
2307 let actual = covariance.quadratic_form(&direction).unwrap();
2308 assert!((actual - expected).abs() <= 2.0e-15);
2309 }
2310
2311 #[test]
2312 fn diagonal_covariance_entries_are_the_exact_geometry_authority_932() {
2313 let covariance = MarginalSlopeCovariance::diagonal(array![3.75]).unwrap();
2314 assert_eq!(covariance.ones_quadratic_form(), 3.75);
2315 assert_eq!(covariance.quadratic_form(&[1.0]).unwrap(), 3.75);
2316 }
2317
2318 #[test]
2319 fn equal_dense_covariance_quadratic_forms_match_all_representations_932() {
2320 let diagonal = MarginalSlopeCovariance::diagonal(array![1.2, 0.7]).unwrap();
2321 let full = MarginalSlopeCovariance::full(array![[1.2, 0.0], [0.0, 0.7]]).unwrap();
2322 let low_rank =
2323 MarginalSlopeCovariance::low_rank(array![[1.2_f64.sqrt(), 0.0], [0.0, 0.7_f64.sqrt()]])
2324 .unwrap();
2325 let direction = [0.35, -0.8];
2326 let expected = diagonal.quadratic_form(&direction).unwrap();
2327 for covariance in [&full, &low_rank] {
2328 let actual = covariance.quadratic_form(&direction).unwrap();
2329 assert!((actual - expected).abs() <= 2.0e-15);
2330 assert!(
2331 (covariance.ones_quadratic_form() - diagonal.ones_quadratic_form()).abs()
2332 <= 2.0e-15
2333 );
2334 }
2335 }
2336}
2337
2338#[cfg(test)]
2339mod jet_tower_oracle_tests {
2340 use super::*;
2362
2363 #[test]
2364 fn signed_probit_stack_preserves_extreme_tail_derivatives_and_weight_sign() {
2365 let positive = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, 2.0);
2366 assert_eq!(
2367 positive,
2368 [f64::INFINITY, f64::NEG_INFINITY, 2.0, -0.0, -0.0]
2369 );
2370 let negative = signed_probit_neglog_unary_stack(f64::NEG_INFINITY, -2.0);
2371 assert_eq!(negative, [f64::NEG_INFINITY, f64::INFINITY, -2.0, 0.0, 0.0]);
2372
2373 let right = signed_probit_neglog_unary_stack(38.6, 1.0);
2374 assert_eq!(right[1], -0.0);
2375 assert!(right[2] > 0.0 && right[2].is_subnormal());
2376 assert!(right[3] < 0.0 && right[3].is_subnormal());
2377 assert!(right[4] > 0.0 && right[4].is_subnormal());
2378
2379 let left = signed_probit_neglog_unary_stack(-1.0e100, 1.0);
2380 assert_eq!(left[1], -1.0e100);
2381 assert_eq!(left[2], 1.0);
2382 assert!(left[3] < 0.0 && left[3].is_finite());
2383 assert_eq!(left[4], -0.0);
2384 }
2385
2386 fn rigid_standard_normal_third_and_fourth_full(
2394 marginal: BernoulliMarginalLinkMap,
2395 g: f64,
2396 z: f64,
2397 y: f64,
2398 w: f64,
2399 probit_scale: f64,
2400 ) -> Result<([[[f64; 2]; 2]; 2], [[[[f64; 2]; 2]; 2]; 2]), String> {
2401 let tower = rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?;
2402 Ok((tower.t3, tower.t4))
2403 }
2404 use gam_math::jet_tower::{
2405 KernelChannels, RowProgram, program_full_tower, verify_kernel_channels,
2406 };
2407
2408 struct BernoulliRigidStandardNormalNllProgram {
2411 primaries: Vec<[f64; 2]>,
2413 z: Vec<f64>,
2415 y: Vec<f64>,
2416 w: Vec<f64>,
2417 probit_scale: f64,
2418 }
2419
2420 impl RowProgram<2> for BernoulliRigidStandardNormalNllProgram {
2421 fn n_rows(&self) -> usize {
2422 self.primaries.len()
2423 }
2424
2425 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2426 self.primaries
2427 .get(row)
2428 .copied()
2429 .ok_or_else(|| format!("bernoulli rigid nll program: row {row} out of range"))
2430 }
2431
2432 fn eval<S: gam_math::jet_scalar::JetScalar<2>>(
2433 &self,
2434 row: usize,
2435 p: &[S; 2],
2436 ) -> Result<S, String> {
2437 let z = self.z[row];
2438 let y = self.y[row];
2439 let w = self.w[row];
2440 let s = self.probit_scale;
2441 let eta_marginal = p[0];
2445 let link = bernoulli_marginal_link_map(
2446 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2447 eta_marginal.value(),
2448 )?;
2449 let q = eta_marginal.compose_unary([link.q, link.q1, link.q2, link.q3, link.q4]);
2450 let g = p[1];
2451 let observed_slope = g.scale(s);
2453 let one_plus_slope_squared = observed_slope.mul(&observed_slope).add(&S::constant(1.0));
2454 let c = one_plus_slope_squared
2455 .compose_unary(unary_derivatives_sqrt(one_plus_slope_squared.value()));
2456 let eta = q.mul(&c).add(&observed_slope.scale(z));
2458 let signed = eta.scale(2.0 * y - 1.0);
2459 Ok(signed.compose_unary(unary_derivatives_neglog_phi(signed.value(), w)))
2461 }
2462 }
2463
2464 fn scalar_nll(eta_marginal: f64, g: f64, z: f64, y: f64, w: f64, s: f64) -> f64 {
2467 let link = bernoulli_marginal_link_map(
2468 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2469 eta_marginal,
2470 )
2471 .unwrap();
2472 let observed_slope = g * s;
2473 let c = (observed_slope * observed_slope + 1.0).sqrt();
2474 let eta = link.q * c + observed_slope * z;
2475 let signed = (2.0 * y - 1.0) * eta;
2476 let cdf = 0.5 * libm::erfc(-signed / std::f64::consts::SQRT_2);
2477 -w * cdf.max(1e-300).ln()
2478 }
2479
2480 #[test]
2481 fn rigid_bernoulli_row_kernel_agrees_with_jet_tower_program_all_channels() {
2482 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2486 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2487 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2488 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2489 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2490 let n = eta.len();
2491
2492 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2494
2495 for &probit_scale in &[1.0_f64, 0.8] {
2496 let program = BernoulliRigidStandardNormalNllProgram {
2497 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2498 z: z.to_vec(),
2499 y: y.to_vec(),
2500 w: w.to_vec(),
2501 probit_scale,
2502 };
2503
2504 for row in 0..n {
2505 let tower = program_full_tower(&program, row).expect("tower evaluation");
2506
2507 let marginal = bernoulli_marginal_link_map(
2509 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2510 eta[row],
2511 )
2512 .expect("link map");
2513 let (value, gradient, hessian) = rigid_standard_normal_row_kernel(
2514 marginal,
2515 g[row],
2516 z[row],
2517 y[row],
2518 w[row],
2519 probit_scale,
2520 )
2521 .expect("production row kernel");
2522
2523 let (third_full, fourth_full) = rigid_standard_normal_third_and_fourth_full(
2530 marginal,
2531 g[row],
2532 z[row],
2533 y[row],
2534 w[row],
2535 probit_scale,
2536 )
2537 .expect("production third+fourth");
2538 let third: Vec<([f64; 2], [[f64; 2]; 2])> = dirs
2539 .iter()
2540 .map(|d| (*d, contract_third_full(&third_full, d[0], d[1])))
2541 .collect();
2542
2543 let fourth: Vec<([f64; 2], [f64; 2], [[f64; 2]; 2])> = dirs
2544 .iter()
2545 .enumerate()
2546 .map(|(i, u)| {
2547 let v = dirs[(i + 1) % dirs.len()];
2548 (
2549 *u,
2550 v,
2551 contract_fourth_full(&fourth_full, u[0], u[1], v[0], v[1]),
2552 )
2553 })
2554 .collect();
2555
2556 let claims = KernelChannels {
2557 value,
2558 gradient,
2559 hessian,
2560 third,
2561 fourth,
2562 };
2563
2564 verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
2565 panic!(
2566 "probit_scale {probit_scale} row {row}: production rigid Bernoulli \
2567 RowKernel disagrees with #932 jet-tower truth: {e}"
2568 )
2569 });
2570
2571 let h = 1e-3;
2575 let f = |de: f64, dg: f64| {
2576 scalar_nll(
2577 eta[row] + de,
2578 g[row] + dg,
2579 z[row],
2580 y[row],
2581 w[row],
2582 probit_scale,
2583 )
2584 };
2585 let f0 = f(0.0, 0.0);
2586 assert!(
2587 (f0 - tower.v).abs() <= 1e-9 * f0.abs().max(1.0),
2588 "row {row}: independent scalar NLL {f0:+.12e} != tower value {:+.12e}",
2589 tower.v
2590 );
2591 let g_eta = (f(-2.0 * h, 0.0) - 8.0 * f(-h, 0.0) + 8.0 * f(h, 0.0)
2593 - f(2.0 * h, 0.0))
2594 / (12.0 * h);
2595 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))
2596 / (12.0 * h);
2597 for (label, fd, ad) in [("∂η", g_eta, tower.g[0]), ("∂g", g_g, tower.g[1])] {
2598 assert!(
2599 (fd - ad).abs() <= 1e-5 * ad.abs().max(1.0),
2600 "row {row} {label}: FD witness {fd:+.6e} != tower grad {ad:+.6e}"
2601 );
2602 }
2603 }
2604 }
2605 }
2606
2607 #[test]
2617 fn rigid_third_and_fourth_full_shares_one_tower_bit_identical() {
2618 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2619 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2620 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2621 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2622 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2623 for &probit_scale in &[1.0_f64, 0.8] {
2624 for r in 0..eta.len() {
2625 let marginal = bernoulli_marginal_link_map(
2626 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2627 eta[r],
2628 )
2629 .expect("link map");
2630 let t3_sep = rigid_standard_normal_third_full(
2631 marginal,
2632 g[r],
2633 z[r],
2634 y[r],
2635 w[r],
2636 probit_scale,
2637 )
2638 .expect("separate third");
2639 let t4_sep = rigid_standard_normal_fourth_full(
2640 marginal,
2641 g[r],
2642 z[r],
2643 y[r],
2644 w[r],
2645 probit_scale,
2646 )
2647 .expect("separate fourth");
2648 let (t3_comb, t4_comb) = rigid_standard_normal_third_and_fourth_full(
2649 marginal,
2650 g[r],
2651 z[r],
2652 y[r],
2653 w[r],
2654 probit_scale,
2655 )
2656 .expect("combined third+fourth");
2657 for a in 0..2 {
2659 for b in 0..2 {
2660 for c in 0..2 {
2661 assert_eq!(
2662 t3_comb[a][b][c], t3_sep[a][b][c],
2663 "t3[{a}][{b}][{c}] row {r} scale {probit_scale} not bit-identical"
2664 );
2665 for d in 0..2 {
2666 assert_eq!(
2667 t4_comb[a][b][c][d], t4_sep[a][b][c][d],
2668 "t4[{a}][{b}][{c}][{d}] row {r} scale {probit_scale} not bit-identical"
2669 );
2670 }
2671 }
2672 }
2673 }
2674 }
2675 }
2676 }
2677
2678 #[test]
2688 fn rigid_bernoulli_generic_program_matches_independent_program_all_channels() {
2689 use gam_math::jet_tower::{
2690 program_fourth_contracted, program_row_kernel, program_third_contracted,
2691 };
2692
2693 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2694 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2695 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2696 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2697 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2698 let n = eta.len();
2699 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2700
2701 let close = |a: f64, b: f64, label: &str| {
2702 let band = 1e-12 + 1e-12 * a.abs().max(b.abs());
2703 assert!(
2704 (a - b).abs() <= band,
2705 "{label}: generic {a:+.15e} vs Tower4-program {b:+.15e} (band {band:.3e})"
2706 );
2707 };
2708
2709 for &probit_scale in &[1.0_f64, 0.8] {
2710 let tower_program = BernoulliRigidStandardNormalNllProgram {
2712 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2713 z: z.to_vec(),
2714 y: y.to_vec(),
2715 w: w.to_vec(),
2716 probit_scale,
2717 };
2718
2719 for row in 0..n {
2720 let truth = program_full_tower(&tower_program, row).expect("program tower");
2721
2722 let marginal = bernoulli_marginal_link_map(
2723 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2724 eta[row],
2725 )
2726 .expect("link map");
2727 let program = RigidStandardNormalRow {
2728 marginal,
2729 g: g[row],
2730 z: z[row],
2731 y: y[row],
2732 w: w[row],
2733 probit_scale,
2734 };
2735
2736 let full = program_full_tower(&program, 0).expect("generic full tower");
2739 close(full.v, truth.v, "full value");
2740 for a in 0..2 {
2741 close(full.g[a], truth.g[a], "full grad");
2742 for b in 0..2 {
2743 close(full.h[a][b], truth.h[a][b], "full hess");
2744 for c in 0..2 {
2745 close(full.t3[a][b][c], truth.t3[a][b][c], "full t3");
2746 for d in 0..2 {
2747 close(full.t4[a][b][c][d], truth.t4[a][b][c][d], "full t4");
2748 }
2749 }
2750 }
2751 }
2752
2753 let (val, grad, hess) =
2755 program_row_kernel(&program, 0).expect("generic row kernel");
2756 close(val, truth.v, "order2 value");
2757 for a in 0..2 {
2758 close(grad[a], truth.g[a], "order2 grad");
2759 for b in 0..2 {
2760 close(hess[a][b], truth.h[a][b], "order2 hess");
2761 }
2762 }
2763
2764 for dir in &dirs {
2767 let third = program_third_contracted(&program, 0, dir)
2768 .expect("generic third contracted");
2769 let truth3 = truth.third_contracted(dir);
2770 for a in 0..2 {
2771 for b in 0..2 {
2772 close(third[a][b], truth3[a][b], "third contracted");
2773 }
2774 }
2775 }
2776
2777 for (i, u) in dirs.iter().enumerate() {
2780 let v = dirs[(i + 1) % dirs.len()];
2781 let fourth = program_fourth_contracted(&program, 0, u, &v)
2782 .expect("generic fourth contracted");
2783 let truth4 = truth.fourth_contracted(u, &v);
2784 for a in 0..2 {
2785 for b in 0..2 {
2786 close(fourth[a][b], truth4[a][b], "fourth contracted");
2787 }
2788 }
2789 }
2790 }
2791 }
2792 }
2793
2794 fn hand_rigid_vgh(
2803 marginal: BernoulliMarginalLinkMap,
2804 g: f64,
2805 z: f64,
2806 y: f64,
2807 w: f64,
2808 probit_scale: f64,
2809 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2810 let s = 2.0 * y - 1.0;
2811 let observed_logslope = probit_scale * g;
2812 let g2 = observed_logslope * observed_logslope;
2813 let c = (1.0 + g2).sqrt();
2814 let c1 = probit_scale * observed_logslope / c;
2815 let c_inv3 = 1.0 / (c * c * c);
2816 let c2 = probit_scale * probit_scale * c_inv3;
2817 let q = marginal.q;
2818 let eta = q * c + observed_logslope * z;
2820 let m = s * eta;
2821 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
2822 let (k1, k2, _k3, _k4) =
2826 signed_probit_neglog_derivatives_up_to_fourth(m, w).expect("hand kernel");
2827 let u1 = s * k1;
2828 let u2 = k2;
2829 let eta_q = c;
2830 let eta_g = q * c1 + probit_scale * z;
2831 let value = -w * logcdf;
2833 let gradient = [u1 * eta_q * marginal.q1, u1 * eta_g];
2835 let h00 = u2 * eta_q * eta_q;
2837 let h01 = u2 * eta_q * eta_g + u1 * c1;
2838 let h11 = u2 * eta_g * eta_g + u1 * q * c2;
2839 let grad_q = u1 * eta_q;
2841 let hessian = [
2842 [
2843 h00 * marginal.q1 * marginal.q1 + grad_q * marginal.q2,
2844 h01 * marginal.q1,
2845 ],
2846 [h01 * marginal.q1, h11],
2847 ];
2848 (value, gradient, hessian)
2849 }
2850
2851 #[test]
2857 fn rigid_bernoulli_row_kernel_matches_hand_chain_witness() {
2858 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2859 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2860 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2861 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2862 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2863 let close = |a: f64, b: f64, label: &str| {
2864 let band = 1e-12 + 1e-9 * a.abs().max(b.abs());
2865 assert!(
2866 (a - b).abs() <= band,
2867 "{label}: jet {a:+.15e} vs hand {b:+.15e} (band {band:.3e})"
2868 );
2869 };
2870 for &probit_scale in &[1.0_f64, 0.8] {
2871 for r in 0..eta.len() {
2872 let marginal = bernoulli_marginal_link_map(
2873 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2874 eta[r],
2875 )
2876 .expect("link map");
2877 let (jv, jg, jh) = rigid_standard_normal_row_kernel(
2878 marginal,
2879 g[r],
2880 z[r],
2881 y[r],
2882 w[r],
2883 probit_scale,
2884 )
2885 .expect("jet kernel");
2886 let (hv, hg, hh) = hand_rigid_vgh(marginal, g[r], z[r], y[r], w[r], probit_scale);
2887 close(jv, hv, "value");
2888 for a in 0..2 {
2889 close(jg[a], hg[a], "grad");
2890 for b in 0..2 {
2891 close(jh[a][b], hh[a][b], "hess");
2892 }
2893 }
2894 }
2895 }
2896 }
2897
2898 #[test]
2907 fn release_measure_rigid_bernoulli_vgh_vs_hand_chain_932() {
2908 use std::time::Instant;
2909
2910 let cases = [
2914 (0.3_f64, 0.2_f64, 0.4_f64, 1.0_f64, 1.0_f64),
2915 (-0.7, -0.5, -1.1, 0.0, 0.8),
2916 ];
2917 let probit_scale = 0.8;
2918
2919 fn best_ns<F>(iterations: usize, base_g: f64, evaluate: F) -> f64
2925 where
2926 F: Fn(f64) -> (f64, [f64; 2], [[f64; 2]; 2]),
2927 {
2928 let mut best = f64::INFINITY;
2929 for _ in 0..5 {
2930 let mut checksum = 0.0_f64;
2931 let started = Instant::now();
2932 for _ in 0..iterations {
2933 let (value, gradient, hessian) = evaluate(base_g + checksum * 1e-18);
2934 checksum += value + gradient[0] + hessian[0][0];
2935 }
2936 assert!(
2937 checksum.is_finite(),
2938 "rigid Bernoulli release-measure checksum must stay finite"
2939 );
2940 best = best.min(started.elapsed().as_secs_f64());
2941 }
2942 best * 1e9 / iterations as f64
2943 }
2944
2945 let iterations = 2_000_000usize;
2946 for &(eta, g, z, y, w) in &cases {
2947 let marginal = bernoulli_marginal_link_map(
2948 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2949 eta,
2950 )
2951 .expect("link map");
2952
2953 let (jet_value, ..) =
2956 rigid_standard_normal_row_kernel(marginal, g, z, y, w, probit_scale)
2957 .expect("jet kernel");
2958 let (hand_value, ..) = hand_rigid_vgh(marginal, g, z, y, w, probit_scale);
2959 let band = 1e-12 + 1e-9 * jet_value.abs().max(hand_value.abs());
2960 assert!(
2961 (jet_value - hand_value).abs() <= band,
2962 "y={y:.0} value: jet {jet_value:+.15e} vs hand {hand_value:+.15e}"
2963 );
2964
2965 let production_ns = best_ns(iterations, g, |perturbed_g| {
2966 rigid_standard_normal_row_kernel(marginal, perturbed_g, z, y, w, probit_scale)
2967 .expect("jet kernel")
2968 });
2969 let hand_ns = best_ns(iterations, g, |perturbed_g| {
2970 hand_rigid_vgh(marginal, perturbed_g, z, y, w, probit_scale)
2971 });
2972 eprintln!(
2973 "RIGID-BERNOULLI-VGH-932 y={y:.0} production={production_ns:.2} ns/row \
2974 hand={hand_ns:.2} ns/row hand_over_production={:.6}",
2975 hand_ns / production_ns,
2976 );
2977 }
2978 }
2979}
2980
2981#[cfg(test)]
2982mod flex_primary_hessian_oracle_tests {
2983 use super::*;
3007 use super::family::*;
3014 use gam_linalg::matrix::DenseDesignMatrix;
3015 use ndarray::Array1;
3016 use ndarray::Array2;
3017 use std::sync::Arc;
3018 use std::sync::Mutex;
3019
3020 fn make_flex_oracle_family(
3026 n: usize,
3027 ) -> (BernoulliMarginalSlopeFamily, Vec<ParameterBlockState>) {
3028 let score_seed = Array1::linspace(-2.0, 2.0, n.max(6));
3029 let link_seed = Array1::linspace(-1.8, 1.8, n.max(6));
3030 let cfg = DeviationBlockConfig {
3031 num_internal_knots: 3,
3032 ..DeviationBlockConfig::default()
3033 };
3034 let score_prepared = build_score_warp_deviation_block_from_seed(&score_seed, &cfg)
3035 .expect("build score warp block");
3036 let link_prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
3037 &link_seed, &link_seed, &cfg,
3038 )
3039 .expect("build link deviation block");
3040
3041 let y: Array1<f64> =
3042 Array1::from_iter((0..n).map(|i| if (i * 17 + 3) % 7 >= 4 { 1.0 } else { 0.0 }));
3043 let weights: Array1<f64> =
3044 Array1::from_iter((0..n).map(|i| 0.75 + ((i * 11 + 5) % 5) as f64 * 0.05));
3045 let z: Array1<f64> =
3046 Array1::from_iter((0..n).map(|i| -1.7 + 3.4 * (i as f64 + 0.5) / n as f64));
3047 let marginal_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3048 if j == 0 {
3049 1.0
3050 } else {
3051 -0.4 + 0.8 * ((i * 19 + 7) % n) as f64 / n as f64
3052 }
3053 });
3054 let logslope_x = Array2::from_shape_fn((n, 2), |(i, j)| {
3055 if j == 0 {
3056 1.0
3057 } else {
3058 0.3 - 0.6 * ((i * 23 + 11) % n) as f64 / n as f64
3059 }
3060 });
3061
3062 let family = BernoulliMarginalSlopeFamily {
3063 y: Arc::new(y),
3064 weights: Arc::new(weights),
3065 z: Arc::new(z.clone()),
3066 latent_measure: LatentMeasureKind::StandardNormal,
3067 gaussian_frailty_sd: Some(0.15),
3068 base_link: InverseLink::Standard(gam_problem::StandardLink::Probit),
3069 marginal_design: DesignMatrix::Dense(DenseDesignMatrix::from(marginal_x.clone())),
3070 logslope_design: DesignMatrix::Dense(DenseDesignMatrix::from(logslope_x.clone())),
3071 score_warp: Some(score_prepared.runtime.clone()),
3072 link_dev: Some(link_prepared.runtime.clone()),
3073 policy: gam_runtime::resource::ResourcePolicy::default_library(),
3074 cell_moment_lru: Arc::new(exact_kernel::CellMomentLruCache::new(1024)),
3075 cell_moment_cache_stats: Arc::new(exact_kernel::CellMomentCacheStats::default()),
3076 intercept_warm_starts: None,
3077 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
3078 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
3079 };
3080
3081 let beta_m = Array1::from_vec(vec![0.12, -0.04]);
3082 let beta_g = Array1::from_vec(vec![0.35, 0.03]);
3083 let beta_h = Array1::from_iter(
3084 (0..score_prepared.runtime.basis_dim()).map(|idx| 0.0015 * (idx as f64 + 1.0)),
3085 );
3086 let beta_w = Array1::from_iter(
3087 (0..link_prepared.runtime.basis_dim()).map(|idx| -0.001 * (idx as f64 + 1.0)),
3088 );
3089 let states = vec![
3090 ParameterBlockState {
3091 eta: marginal_x.dot(&beta_m),
3092 beta: beta_m,
3093 },
3094 ParameterBlockState {
3095 eta: logslope_x.dot(&beta_g),
3096 beta: beta_g,
3097 },
3098 ParameterBlockState {
3099 beta: beta_h,
3100 eta: Array1::zeros(z.len()),
3101 },
3102 ParameterBlockState {
3103 beta: beta_w,
3104 eta: Array1::zeros(z.len()),
3105 },
3106 ];
3107 (family, states)
3108 }
3109
3110 fn flex_gradient_at_perturbed(
3118 family: &BernoulliMarginalSlopeFamily,
3119 states: &[ParameterBlockState],
3120 primary: &super::super::hessian_paths::PrimarySlices,
3121 row: usize,
3122 u: usize,
3123 delta: f64,
3124 ) -> Array1<f64> {
3125 let mut states = states.to_vec();
3126 if u == primary.q {
3132 states[0].eta[row] += delta;
3133 } else if u == primary.logslope {
3134 states[1].eta[row] += delta;
3135 } else if let Some(h_range) = primary.h.as_ref()
3136 && h_range.contains(&u)
3137 {
3138 states[2].beta[u - h_range.start] += delta;
3139 } else if let Some(w_range) = primary.w.as_ref()
3140 && w_range.contains(&u)
3141 {
3142 states[3].beta[u - w_range.start] += delta;
3143 } else {
3144 panic!("primary coordinate {u} out of range for flex oracle");
3145 }
3146 let row_ctx = family
3147 .build_row_exact_context_with_stats_and_cell_cache(row, &states, None, false)
3148 .expect("perturbed row context");
3149 let (_neglog, grad, _hess) = family
3150 .compute_row_primary_gradient_hessian(row, &states, primary, &row_ctx)
3151 .expect("perturbed gradient");
3152 grad
3153 }
3154
3155 fn flex_nll_gradient_at_perturbed_z(
3161 family: &BernoulliMarginalSlopeFamily,
3162 states: &[ParameterBlockState],
3163 primary: &super::super::hessian_paths::PrimarySlices,
3164 row: usize,
3165 delta: f64,
3166 ) -> Array1<f64> {
3167 let mut perturbed = family.clone();
3168 let mut z = family.z.as_ref().clone();
3169 z[row] += delta;
3170 perturbed.z = Arc::new(z);
3171 let row_ctx = perturbed
3172 .build_row_exact_context_with_stats_and_cell_cache(row, states, None, false)
3173 .expect("z-perturbed row context");
3174 let mut scratch =
3175 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3176 perturbed
3177 .lower_bms_flex_row_order2(row, states, primary, &row_ctx, None, false, &mut scratch)
3178 .expect("z-perturbed flex gradient");
3179 scratch.grad
3180 }
3181
3182 #[test]
3188 fn flex_score_zeta_sensitivity_covers_all_active_deviation_blocks_2303() {
3189 let n = 12usize;
3190 let (family, states) = make_flex_oracle_family(n);
3191 let cache = family
3192 .build_exact_eval_cache(&states)
3193 .expect("flex exact eval cache");
3194 let primary = &cache.primary;
3195 let row = 5usize;
3196 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3197 let mut scratch =
3198 super::super::hessian_paths::BernoulliMarginalSlopeFlexRowScratch::new(primary.total);
3199 family
3200 .lower_bms_flex_row_order2(row, &states, primary, row_ctx, None, false, &mut scratch)
3201 .expect("analytic flex z-sensitivity");
3202 let analytic = scratch.score_zeta.clone();
3203
3204 let h = 1e-5_f64;
3205 let nll_plus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, h);
3206 let nll_minus = flex_nll_gradient_at_perturbed_z(&family, &states, primary, row, -h);
3207 for u in 0..primary.total {
3208 let finite_difference = -(nll_plus[u] - nll_minus[u]) / (2.0 * h);
3210 let scale = 1.0 + analytic[u].abs().max(finite_difference.abs());
3211 let relative_error = (analytic[u] - finite_difference).abs() / scale;
3212 assert!(
3213 relative_error <= 2e-6,
3214 "flex score-zeta primary {u}: analytic={} FD={} relative_error={relative_error}",
3215 analytic[u],
3216 finite_difference
3217 );
3218 }
3219 let h_range = primary.h.as_ref().expect("score-warp primary range");
3220 let w_range = primary.w.as_ref().expect("link-deviation primary range");
3221 assert!(
3222 analytic
3223 .slice(s![h_range.start..h_range.end])
3224 .iter()
3225 .any(|value| value.abs() > 1e-10),
3226 "score-warp z-sensitivity must not be zero-filled"
3227 );
3228 assert!(
3229 analytic
3230 .slice(s![w_range.start..w_range.end])
3231 .iter()
3232 .any(|value| value.abs() > 1e-10),
3233 "link-deviation z-sensitivity must not be zero-filled"
3234 );
3235
3236 let coefficient = family
3237 .flex_score_zeta_sensitivity(
3238 &states,
3239 &BlockwiseFitOptions::default(),
3240 cache.slices.total,
3241 )
3242 .expect("full coefficient score-zeta sensitivity");
3243 assert_eq!(coefficient.dim(), (n, cache.slices.total));
3244 for range in [
3245 cache.slices.h.as_ref().expect("score-warp beta range"),
3246 cache.slices.w.as_ref().expect("link-deviation beta range"),
3247 ] {
3248 assert!(
3249 coefficient
3250 .slice(s![.., range.start..range.end])
3251 .iter()
3252 .any(|value| value.abs() > 1e-10),
3253 "active deviation coefficient range {range:?} must carry Murphy-Topel sensitivity"
3254 );
3255 }
3256
3257 let wrong_width = cache
3258 .slices
3259 .total
3260 .checked_sub(1)
3261 .expect("nonempty coefficient frame");
3262 let error = family
3263 .flex_score_zeta_sensitivity(&states, &BlockwiseFitOptions::default(), wrong_width)
3264 .expect_err("partial Murphy-Topel covariance frame must be rejected");
3265 assert!(
3266 error.contains("covariance/frame mismatch"),
3267 "unexpected partial-frame error: {error}"
3268 );
3269 }
3270
3271 #[test]
3274 fn flex_primary_hessian_matches_central_fd_of_gradient() {
3275 let n = 12usize;
3276 let (family, states) = make_flex_oracle_family(n);
3277 let cache = family
3278 .build_exact_eval_cache(&states)
3279 .expect("flex exact eval cache");
3280 let primary = &cache.primary;
3281 let r = primary.total;
3282 assert!(
3283 r >= 4,
3284 "flex fixture must carry q + logslope + deviation blocks"
3285 );
3286
3287 let h = 1e-4;
3291 let mut max_rel = 0.0_f64;
3292
3293 for &row in &[2usize, 5, 8] {
3296 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3297 let (_neglog, _grad, analytic_hess) = family
3298 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3299 .expect("analytic flex gradient + hessian");
3300
3301 for u in 0..r {
3302 let grad_plus = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3303 let grad_minus = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3304 for v in 0..r {
3305 let fd = (grad_plus[v] - grad_minus[v]) / (2.0 * h);
3306 let analytic = analytic_hess[[v, u]];
3307 let denom = 1.0 + analytic.abs().max(fd.abs());
3308 let rel = (analytic - fd).abs() / denom;
3309 max_rel = max_rel.max(rel);
3310 assert!(
3311 rel <= 1e-6,
3312 "flex hand Hessian H[{v}][{u}] = {analytic:.6e} disagrees with central \
3313 FD of the gradient {fd:.6e} at row {row} (rel {rel:.3e}); a product-rule \
3314 term is dropped or mis-signed"
3315 );
3316 }
3317 }
3318 }
3319 assert!(
3321 max_rel <= 1e-6,
3322 "flex Hessian FD oracle max rel {max_rel:.3e}"
3323 );
3324 }
3325
3326 #[test]
3335 fn arbiter_flex_hessian_h00_fd_step_scaling() {
3336 let n = 12usize;
3337 let (family, states) = make_flex_oracle_family(n);
3338 let cache = family
3339 .build_exact_eval_cache(&states)
3340 .expect("flex exact eval cache");
3341 let primary = &cache.primary;
3342 let row = 2usize;
3343 let u = primary.q; let v = primary.q;
3345
3346 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
3347 let (_neglog, _grad, analytic_hess) = family
3348 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
3349 .expect("analytic flex gradient + hessian");
3350 let analytic = analytic_hess[[v, u]];
3351
3352 let fd_at = |h: f64| -> f64 {
3353 let gp = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
3354 let gm = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
3355 (gp[v] - gm[v]) / (2.0 * h)
3356 };
3357
3358 let h = 1e-3_f64;
3365 let fd_h = fd_at(h);
3366 let fd_half = fd_at(h * 0.5);
3367 let fd_quarter = fd_at(h * 0.25);
3368 let gap_h = (analytic - fd_h).abs();
3369 let gap_half = (analytic - fd_half).abs();
3370 let gap_quarter = (analytic - fd_quarter).abs();
3371 let rich = (4.0 * fd_half - fd_h) / 3.0;
3372 let rich_gap = (analytic - rich).abs();
3373 let denom = analytic.abs().max(1.0);
3374
3375 let record = format!(
3377 "FLEX H[0][0] ARBITER row 2: analytic={analytic:+.12e} \
3378 fd(h)={fd_h:+.12e} fd(h/2)={fd_half:+.12e} fd(h/4)={fd_quarter:+.12e} \
3379 gap(h)={gap_h:.3e} gap(h/2)={gap_half:.3e} gap(h/4)={gap_quarter:.3e} \
3380 ratio_h_over_half={:.3} ratio_half_over_quarter={:.3} \
3381 richardson={rich:+.12e} richardson_gap={rich_gap:.3e} (rich_rel={:.3e})",
3382 gap_h / gap_half.max(f64::MIN_POSITIVE),
3383 gap_half / gap_quarter.max(f64::MIN_POSITIVE),
3384 rich_gap / denom,
3385 );
3386
3387 assert!(
3393 rich_gap / denom <= 1e-7,
3394 "{record}\nVERDICT: Richardson residual exceeds the FD-truncation floor — \
3395 the hand H[0][0] genuinely diverges (real dropped/mis-signed term), NOT FD noise"
3396 );
3397 }
3398}