1use super::family::clamp_bernoulli_link_probability;
2use super::*;
3use gam_linalg::matrix::{LinearOperator, SignedWeightsView};
4use gam_math::jet_tower::Tower4;
5use opt::{BacktrackConfig, RidgeSchedule, backtracking_line_search, escalate_ridge};
6
7pub(crate) fn standardize_latent_z_with_policy(
8 z: &Array1<f64>,
9 weights: &Array1<f64>,
10 context: &str,
11 policy: &LatentZPolicy,
12) -> Result<(Array1<f64>, LatentZNormalization), String> {
13 if z.len() != weights.len() {
14 return Err(format!(
15 "{context} latent-score normalization length mismatch: z={}, weights={}",
16 z.len(),
17 weights.len()
18 ));
19 }
20 let weight_sum = weights.iter().copied().sum::<f64>();
21 let weight_sq_sum = weights.iter().map(|&w| w * w).sum::<f64>();
22 if !(weight_sum.is_finite()
23 && weight_sum > 0.0
24 && weight_sq_sum.is_finite()
25 && weight_sq_sum > 0.0)
26 {
27 return Err(format!("{context} requires positive finite total weight"));
28 }
29 let effective_n = weight_sum * weight_sum / weight_sq_sum;
30 if !(effective_n.is_finite() && effective_n > 1.0) {
31 return Err(format!(
32 "{context} requires at least two effective observations for latent-score normalization"
33 ));
34 }
35 let mean = z
36 .iter()
37 .zip(weights.iter())
38 .map(|(&zi, &wi)| wi * zi)
39 .sum::<f64>()
40 / weight_sum;
41 let var = z
42 .iter()
43 .zip(weights.iter())
44 .map(|(&zi, &wi)| wi * (zi - mean) * (zi - mean))
45 .sum::<f64>()
46 / weight_sum;
47 let sd = var.sqrt();
48 if !(sd.is_finite() && sd > BMS_VARIANCE_FLOOR) {
49 return Err(format!(
50 "{context} requires z with positive finite weighted standard deviation"
51 ));
52 }
53 let target_norm = match policy.normalization {
54 LatentZNormalizationMode::None => LatentZNormalization { mean: 0.0, sd: 1.0 },
55 LatentZNormalizationMode::FitWeighted => LatentZNormalization { mean, sd },
56 LatentZNormalizationMode::Frozen {
57 mean: frozen_mean,
58 sd: frozen_sd,
59 } => LatentZNormalization {
60 mean: frozen_mean,
61 sd: frozen_sd,
62 },
63 };
64 let mean_tol = policy.mean_tol_multiplier / effective_n.sqrt();
65 let sd_tol = policy.sd_tol_multiplier / (2.0 * (effective_n - 1.0).max(1.0)).sqrt();
66 let check_msg = || {
67 format!(
68 "{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}"
69 )
70 };
71 if mean.abs() > mean_tol || (sd - 1.0).abs() > sd_tol {
72 match policy.check_mode {
73 LatentZCheckMode::Strict => return Err(check_msg()),
74 LatentZCheckMode::WarnOnly => log::warn!("{}", check_msg()),
75 LatentZCheckMode::Off => {}
76 }
77 }
78
79 let normalization = target_norm;
80 let z_std = normalization.apply(z, context)?;
81 let std_mean = z_std
87 .iter()
88 .zip(weights.iter())
89 .map(|(&zi, &wi)| wi * zi)
90 .sum::<f64>()
91 / weight_sum;
92 let std_var = (z_std
93 .iter()
94 .zip(weights.iter())
95 .map(|(&zi, &wi)| wi * (zi - std_mean) * (zi - std_mean))
96 .sum::<f64>()
97 / weight_sum)
98 .max(f64::MIN_POSITIVE);
99 let skew = z_std
100 .iter()
101 .zip(weights.iter())
102 .map(|(&zi, &wi)| wi * (zi - std_mean).powi(3))
103 .sum::<f64>()
104 / weight_sum
105 / std_var.powf(1.5);
106 let kurt = z_std
107 .iter()
108 .zip(weights.iter())
109 .map(|(&zi, &wi)| wi * (zi - std_mean).powi(4))
110 .sum::<f64>()
111 / weight_sum
112 / (std_var * std_var)
113 - 3.0;
114 if skew.abs() > policy.max_abs_skew || kurt.abs() > policy.max_abs_excess_kurtosis {
115 let msg = format!(
116 "{context} requires z to be approximately Gaussian after identification normalization; got skewness={skew:.3}, excess_kurtosis={kurt:.3}"
117 );
118 match policy.check_mode {
119 LatentZCheckMode::Strict => return Err(msg),
120 LatentZCheckMode::WarnOnly => log::warn!("{}", msg),
121 LatentZCheckMode::Off => {}
122 }
123 }
124 if skew.abs() > 0.75 || kurt.abs() > 2.0 {
125 log::warn!(
126 "{context}: z has skewness={skew:.3} and excess kurtosis={kurt:.3}; latent-measure auto-selection will use empirical calibration unless stricter diagnostics pass"
127 );
128 }
129 Ok((z_std, normalization))
130}
131
132pub fn padded_deviation_seed(seed: &Array1<f64>, min_iqr: f64, pad_fraction: f64) -> Array1<f64> {
133 let mut sorted = seed.to_vec();
134 sorted.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
135
136 if sorted.len() < 4 {
137 return seed.clone();
138 }
139
140 let n = sorted.len();
141 let q1 = sorted[n / 4];
142 let q3 = sorted[3 * n / 4];
143 let iqr = (q3 - q1).max(min_iqr);
144 let pad = pad_fraction * iqr;
145
146 let mut out = seed.to_vec();
147 out.push(sorted[0] - pad);
148 out.push(sorted[n - 1] + pad);
149 Array1::from_vec(out)
150}
151
152const POOLED_PILOT_MAX_NEWTON_ITERS: usize = 50;
163pub(crate) const POOLED_PILOT_RIDGE_INIT: f64 = 1e-8;
165pub(crate) const POOLED_PILOT_DET_FLOOR: f64 = 1e-18;
168pub(crate) const POOLED_PILOT_RIDGE_GROWTH: f64 = 10.0;
170pub(crate) const POOLED_PILOT_RIDGE_MAX: f64 = 1e6;
173const POOLED_PILOT_MAX_BACKTRACKS: usize = 25;
175pub(crate) const POOLED_PILOT_BACKTRACK_SHRINK: f64 = 0.5;
177pub(crate) const POOLED_PILOT_STALL_TOL: f64 = 1e-10;
180pub(crate) const POOLED_PILOT_MIN_ABS_SLOPE: f64 = 1e-6;
183
184pub(super) fn pooled_probit_baseline(
185 y: &Array1<f64>,
186 z: &Array1<f64>,
187 weights: &Array1<f64>,
188) -> Result<(f64, f64), String> {
189 if y.len() != z.len() || y.len() != weights.len() {
190 return Err(format!(
191 "pooled bernoulli-marginal-slope pilot length mismatch: y={}, z={}, weights={}",
192 y.len(),
193 z.len(),
194 weights.len()
195 ));
196 }
197 let weight_sum = weights.iter().copied().sum::<f64>();
198 if !weight_sum.is_finite() || weight_sum <= 0.0 {
199 return Err(
200 "pooled bernoulli-marginal-slope pilot requires positive finite total weight"
201 .to_string(),
202 );
203 }
204 let prevalence = y
205 .iter()
206 .zip(weights.iter())
207 .map(|(&yi, &wi)| yi * wi)
208 .sum::<f64>()
209 / weight_sum;
210 let prevalence = prevalence.clamp(1e-6, 1.0 - 1e-6);
211 let z_mean = z
212 .iter()
213 .zip(weights.iter())
214 .map(|(&zi, &wi)| zi * wi)
215 .sum::<f64>()
216 / weight_sum;
217 let z_var = z
218 .iter()
219 .zip(weights.iter())
220 .map(|(&zi, &wi)| wi * (zi - z_mean) * (zi - z_mean))
221 .sum::<f64>()
222 / weight_sum;
223 let yz_cov = y
224 .iter()
225 .zip(z.iter())
226 .zip(weights.iter())
227 .map(|((&yi, &zi), &wi)| wi * (yi - prevalence) * (zi - z_mean))
228 .sum::<f64>()
229 / weight_sum;
230 let mut beta0 = standard_normal_quantile(prevalence).map_err(|e| {
231 format!("failed to initialize pooled bernoulli-marginal-slope pilot intercept: {e}")
232 })?;
233 let mut beta1 = if z_var > BMS_VARIANCE_FLOOR {
234 yz_cov / z_var
235 } else {
236 0.0
237 };
238
239 let objective_grad_hess =
240 |intercept: f64, slope: f64| -> Result<(f64, f64, f64, f64, f64, f64), String> {
241 let mut obj = 0.0;
242 let mut g0 = 0.0;
243 let mut g1 = 0.0;
244 let mut h00 = 0.0;
245 let mut h01 = 0.0;
246 let mut h11 = 0.0;
247 for ((&yi, &zi), &wi) in y.iter().zip(z.iter()).zip(weights.iter()) {
248 if wi == 0.0 {
249 continue;
250 }
251 let eta = intercept + slope * zi;
252 let s = 2.0 * yi - 1.0;
253 let margin = s * eta;
254 let (logcdf, lambda) = signed_probit_logcdf_and_mills_ratio(margin);
255 let g_eta = -wi * s * lambda;
256 let h_eta = wi * lambda * (margin + lambda);
257 obj -= wi * logcdf;
258 g0 += g_eta;
259 g1 += g_eta * zi;
260 h00 += h_eta;
261 h01 += h_eta * zi;
262 h11 += h_eta * zi * zi;
263 }
264 Ok((obj, g0, g1, h00, h01, h11))
265 };
266
267 let mut obj_prev = f64::INFINITY;
268 for _ in 0..POOLED_PILOT_MAX_NEWTON_ITERS {
269 let (obj, g0, g1, h00, h01, h11) = objective_grad_hess(beta0, beta1)?;
270 if !obj.is_finite() || !g0.is_finite() || !g1.is_finite() {
271 return Err(
272 "pooled bernoulli-marginal-slope pilot produced non-finite objective or gradient"
273 .to_string(),
274 );
275 }
276 let grad_max = g0.abs().max(g1.abs());
277 if grad_max < BMS_DERIV_TOL {
278 break;
279 }
280 let ridge_trials = (POOLED_PILOT_RIDGE_MAX / POOLED_PILOT_RIDGE_INIT)
284 .log10()
285 .ceil() as usize
286 + 1;
287 let (step0, step1) = escalate_ridge(
288 RidgeSchedule {
289 initial: POOLED_PILOT_RIDGE_INIT,
290 growth: POOLED_PILOT_RIDGE_GROWTH,
291 max_escalations: ridge_trials,
292 },
293 |ridge| {
294 let h00_r = h00 + ridge;
295 let h11_r = h11 + ridge;
296 let det = h00_r * h11_r - h01 * h01;
297 if !(det.is_finite() && det.abs() > POOLED_PILOT_DET_FLOOR) {
298 return None;
299 }
300 let s0 = (h11_r * g0 - h01 * g1) / det;
301 let s1 = (-h01 * g0 + h00_r * g1) / det;
302 (s0.is_finite() && s1.is_finite()).then_some((s0, s1))
303 },
304 )
305 .map(|success| success.value)
306 .map_err(|_| "pooled bernoulli-marginal-slope pilot Hessian solve failed".to_string())?;
307 let accepted = backtracking_line_search::<_, String>(
308 BacktrackConfig {
309 contraction: POOLED_PILOT_BACKTRACK_SHRINK,
310 max_steps: POOLED_PILOT_MAX_BACKTRACKS,
311 ..BacktrackConfig::default()
312 },
313 |step_scale| {
314 let cand0 = beta0 - step_scale * step0;
315 let cand1 = beta1 - step_scale * step1;
316 let (cand_obj, _, _, _, _, _) = objective_grad_hess(cand0, cand1)?;
317 Ok(Some((cand_obj, (cand0, cand1))))
318 },
319 |_scale, cand_obj| cand_obj.is_finite() && cand_obj <= obj,
320 )?;
321 match accepted {
322 Some(step) => {
323 (beta0, beta1) = step.payload;
324 obj_prev = step.value;
325 }
326 None => {
327 if (obj_prev - obj).abs() < POOLED_PILOT_STALL_TOL {
328 break;
329 }
330 return Err("pooled bernoulli-marginal-slope pilot line search failed".to_string());
331 }
332 }
333 }
334 let a = beta0;
335 let b = if beta1.abs() < POOLED_PILOT_MIN_ABS_SLOPE {
337 if beta1.is_sign_negative() {
338 -POOLED_PILOT_MIN_ABS_SLOPE
339 } else {
340 POOLED_PILOT_MIN_ABS_SLOPE
341 }
342 } else {
343 beta1
344 };
345 Ok((a / (1.0 + b * b).sqrt(), b))
346}
347
348pub(super) fn pilot_irls_hessian_row_metric_at_eta(
388 eta_pilot: &Array1<f64>,
389 sample_weights: &Array1<f64>,
390) -> Array1<f64> {
391 let n = eta_pilot.len();
392 let mut w = Array1::<f64>::zeros(n);
393 for i in 0..n {
394 let eta = eta_pilot[i];
395 let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
396 let phi = normal_pdf(eta).max(1e-300);
397 let var = (mu * (1.0 - mu)).max(1e-300);
398 w[i] = sample_weights[i] * (phi * phi) / var;
399 }
400 w
401}
402
403pub(super) fn rigid_pooled_probit_pilot_eta(
410 base_link: &InverseLink,
411 z: &Array1<f64>,
412 marginal_offset: &Array1<f64>,
413 logslope_offset: &Array1<f64>,
414 baseline_marginal: f64,
415 baseline_logslope: f64,
416 probit_scale: f64,
417) -> Result<Array1<f64>, String> {
418 let n = z.len();
419 let mut out = Array1::<f64>::zeros(n);
420 for i in 0..n {
421 let a_pre = baseline_marginal + marginal_offset[i];
422 let b_pre = baseline_logslope + logslope_offset[i];
423 let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
424 .map_err(|e| format!("rigid_pooled_probit_pilot_eta marginal link map: {e}"))?
425 .q;
426 out[i] = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
427 }
428 Ok(out)
429}
430
431pub(crate) const PILOT_RIDGE_DIAG_FRACTION: f64 = 1e-6;
437pub(crate) const PILOT_RIDGE_DIAG_FLOOR: f64 = 1e-12;
440
441pub(super) fn pilot_eta_for_link_dev_orthogonalisation(
442 base_link: &InverseLink,
443 y: &Array1<f64>,
444 z: &Array1<f64>,
445 weights: &Array1<f64>,
446 marginal_design: &DesignMatrix,
447 marginal_offset: &Array1<f64>,
448 logslope_offset: &Array1<f64>,
449 baseline_marginal: f64,
450 baseline_logslope: f64,
451 probit_scale: f64,
452) -> Result<Array1<f64>, String> {
453 use gam_linalg::faer_ndarray::FaerCholesky;
454
455 let n = y.len();
456 if marginal_design.nrows() != n {
457 return Err(format!(
458 "pilot_eta_for_link_dev_orthogonalisation: marginal design has {} rows, expected {}",
459 marginal_design.nrows(),
460 n,
461 ));
462 }
463 let mut working_eta = Array1::<f64>::zeros(n);
464 let mut w_irls = Array1::<f64>::zeros(n);
465 let mut residual = Array1::<f64>::zeros(n);
466 for i in 0..n {
467 let a_pre = baseline_marginal + marginal_offset[i];
468 let b_pre = baseline_logslope + logslope_offset[i];
469 let q_marg = bernoulli_marginal_link_map(base_link, a_pre)
470 .map_err(|e| {
471 format!("pilot_eta_for_link_dev_orthogonalisation marginal link map: {e}")
472 })?
473 .q;
474 let eta = rigid_observed_eta(q_marg, b_pre, z[i], probit_scale);
475 working_eta[i] = eta;
476 let mu = clamp_bernoulli_link_probability(normal_cdf(eta));
477 let phi = normal_pdf(eta).max(1e-300);
478 let var = (mu * (1.0 - mu)).max(1e-300);
479 w_irls[i] = weights[i] * (phi * phi) / var;
480 residual[i] = (y[i] - mu) / phi;
481 }
482 let p_marg = marginal_design.ncols();
483 if p_marg == 0 {
484 return Ok(working_eta);
485 }
486 let xtwr = marginal_design.compute_xtwy(&w_irls, &residual)?;
487 let mut xtwx = marginal_design.xt_diag_x_signed_op(SignedWeightsView::from_array(&w_irls))?;
488 let trace_diag: f64 = (0..p_marg).map(|i| xtwx[[i, i]]).sum();
489 let ridge =
490 (trace_diag / p_marg as f64).max(PILOT_RIDGE_DIAG_FLOOR) * PILOT_RIDGE_DIAG_FRACTION;
491 for i in 0..p_marg {
492 xtwx[[i, i]] += ridge;
493 }
494 let factor = xtwx
495 .cholesky(faer::Side::Lower)
496 .map_err(|e| format!("pilot_eta_for_link_dev_orthogonalisation Cholesky failed: {e}"))?;
497 let delta_beta_marg = factor.solvevec(&xtwr);
498 let marg_contrib = marginal_design.dot(&delta_beta_marg);
499 Ok(&working_eta + &marg_contrib)
500}
501
502pub(super) fn joint_setup(
503 data: ArrayView2<'_, f64>,
504 marginalspec: &TermCollectionSpec,
505 logslopespec: &TermCollectionSpec,
506 marginal_penalties: usize,
507 logslope_penalties: usize,
508 absorber_rho0: Option<f64>,
509 extra_rho0: &[f64],
510 kappa_options: &SpatialLengthScaleOptimizationOptions,
511) -> ExactJointHyperSetup {
512 let marginal_terms = spatial_length_scale_term_indices(marginalspec);
513 let logslope_terms = spatial_length_scale_term_indices(logslopespec);
514 let rho_dim = marginal_penalties + logslope_penalties + extra_rho0.len();
515 let mut rho0vec = Array1::<f64>::zeros(rho_dim);
516 if let Some(seed) = absorber_rho0 {
520 assert!(
521 marginal_penalties > 0,
522 "an absorber rho0 seed requires at least one marginal penalty to land in"
523 );
524 rho0vec[marginal_penalties - 1] = seed;
525 }
526 for (idx, &value) in extra_rho0.iter().enumerate() {
527 rho0vec[marginal_penalties + logslope_penalties + idx] = value;
528 }
529 let rho_lower = Array1::<f64>::from_elem(rho_dim, -12.0);
530 let rho_upper = Array1::<f64>::from_elem(rho_dim, 12.0);
531 let marginal_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
532 marginalspec,
533 &marginal_terms,
534 kappa_options,
535 )
536 .reseed_from_data(data, marginalspec, &marginal_terms, kappa_options);
537 let logslope_kappa = SpatialLogKappaCoords::from_length_scales_aniso(
538 logslopespec,
539 &logslope_terms,
540 kappa_options,
541 )
542 .reseed_from_data(data, logslopespec, &logslope_terms, kappa_options);
543 let mut values = marginal_kappa.as_array().to_vec();
544 values.extend(logslope_kappa.as_array().iter());
545 let marginal_dims = marginal_kappa.dims_per_term().to_vec();
546 let logslope_dims = logslope_kappa.dims_per_term().to_vec();
547 let mut dims = marginal_dims.clone();
548 dims.extend(logslope_dims.iter().copied());
549 let log_kappa0 = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(values), dims.clone());
550 let marginal_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
552 data,
553 marginalspec,
554 &marginal_terms,
555 &marginal_dims,
556 kappa_options,
557 );
558 let logslope_lower = SpatialLogKappaCoords::lower_bounds_aniso_from_data(
559 data,
560 logslopespec,
561 &logslope_terms,
562 &logslope_dims,
563 kappa_options,
564 );
565 let mut lower_vals = marginal_lower.as_array().to_vec();
566 lower_vals.extend(logslope_lower.as_array().iter());
567 let log_kappa_lower =
568 SpatialLogKappaCoords::new_with_dims(Array1::from_vec(lower_vals), dims.clone());
569 let marginal_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
570 data,
571 marginalspec,
572 &marginal_terms,
573 &marginal_dims,
574 kappa_options,
575 );
576 let logslope_upper = SpatialLogKappaCoords::upper_bounds_aniso_from_data(
577 data,
578 logslopespec,
579 &logslope_terms,
580 &logslope_dims,
581 kappa_options,
582 );
583 let mut upper_vals = marginal_upper.as_array().to_vec();
584 upper_vals.extend(logslope_upper.as_array().iter());
585 let log_kappa_upper = SpatialLogKappaCoords::new_with_dims(Array1::from_vec(upper_vals), dims);
586 let log_kappa0 = log_kappa0.clamp_to_bounds(&log_kappa_lower, &log_kappa_upper);
589 ExactJointHyperSetup::new(
590 rho0vec,
591 rho_lower,
592 rho_upper,
593 log_kappa0,
594 log_kappa_lower,
595 log_kappa_upper,
596 )
597}
598
599#[inline]
600pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth_numeric(
601 signed_margin: f64,
602 weight: f64,
603) -> (f64, f64, f64, f64) {
604 if weight == 0.0 || signed_margin == f64::INFINITY {
605 return (0.0, 0.0, 0.0, 0.0);
606 }
607 if signed_margin == f64::NEG_INFINITY {
608 return (f64::NEG_INFINITY, weight, 0.0, 0.0);
609 }
610 if signed_margin.is_nan() {
611 return (f64::NAN, f64::NAN, f64::NAN, f64::NAN);
612 }
613 let (_, lambda) = signed_probit_logcdf_and_mills_ratio(signed_margin);
614 let k1 = -lambda;
615 let k2 = lambda * (signed_margin + lambda);
616 let k3 = lambda
617 * (1.0
618 - signed_margin * signed_margin
619 - 3.0 * signed_margin * lambda
620 - 2.0 * lambda * lambda);
621 let k4 = lambda
622 * ((signed_margin.powi(3) - 3.0 * signed_margin)
623 + (7.0 * signed_margin * signed_margin - 4.0) * lambda
624 + 12.0 * signed_margin * lambda * lambda
625 + 6.0 * lambda.powi(3));
626 (weight * k1, weight * k2, weight * k3, weight * k4)
627}
628
629pub(crate) fn signed_probit_neglog_derivatives_up_to_fourth(
637 signed_margin: f64,
638 weight: f64,
639) -> Result<(f64, f64, f64, f64), String> {
640 if weight == 0.0 || signed_margin == f64::INFINITY {
641 return Ok((0.0, 0.0, 0.0, 0.0));
642 }
643 if !signed_margin.is_finite() {
644 return Err(format!(
645 "non-finite signed margin in exact probit derivative helper: {signed_margin}"
646 ));
647 }
648 Ok(signed_probit_neglog_derivatives_up_to_fourth_numeric(
649 signed_margin,
650 weight,
651 ))
652}
653
654#[inline]
680pub(crate) fn signed_probit_neglog_unary_stack(signed_margin: f64, weight: f64) -> [f64; 5] {
681 if weight == 0.0 || signed_margin == f64::INFINITY {
682 return [0.0; 5];
683 }
684 if signed_margin == f64::NEG_INFINITY {
685 return [f64::INFINITY, f64::NEG_INFINITY, weight, 0.0, 0.0];
688 }
689 if signed_margin.is_nan() {
690 return [f64::NAN; 5];
691 }
692 let (logcdf, lambda) = signed_probit_logcdf_and_mills_ratio(signed_margin);
695 let m = signed_margin;
696 let k1 = -lambda;
697 let k2 = lambda * (m + lambda);
698 let k3 = lambda * (1.0 - m * m - 3.0 * m * lambda - 2.0 * lambda * lambda);
699 let k4 = lambda
700 * ((m * m * m - 3.0 * m)
701 + (7.0 * m * m - 4.0) * lambda
702 + 12.0 * m * lambda * lambda
703 + 6.0 * lambda * lambda * lambda);
704 [
705 -weight * logcdf,
706 weight * k1,
707 weight * k2,
708 weight * k3,
709 weight * k4,
710 ]
711}
712
713#[inline]
714pub(super) fn rigid_observed_logslope(logslope: f64, probit_scale: f64) -> f64 {
715 probit_scale * logslope
716}
717
718#[inline]
719pub(super) fn rigid_observed_scale(logslope: f64, probit_scale: f64) -> f64 {
720 let observed_logslope = rigid_observed_logslope(logslope, probit_scale);
721 (1.0 + observed_logslope * observed_logslope).sqrt()
722}
723
724#[inline]
725pub(super) fn rigid_intercept_from_marginal(
726 marginal_eta: f64,
727 logslope: f64,
728 probit_scale: f64,
729) -> f64 {
730 marginal_eta * rigid_observed_scale(logslope, probit_scale)
731}
732
733#[inline]
734pub(super) fn rigid_prescale_intercept_from_marginal(
735 marginal_eta: f64,
736 logslope: f64,
737 probit_scale: f64,
738) -> f64 {
739 rigid_intercept_from_marginal(marginal_eta, logslope, probit_scale) / probit_scale
740}
741
742#[inline]
743pub(super) fn rigid_prescale_intercept_derivative_abs(
744 marginal_eta: f64,
745 logslope: f64,
746 probit_scale: f64,
747) -> f64 {
748 let c = rigid_observed_scale(logslope, probit_scale);
749 probit_scale * normal_pdf(marginal_eta) / c
750}
751
752#[inline]
753pub(super) fn rigid_observed_eta(
754 marginal_eta: f64,
755 logslope: f64,
756 z: f64,
757 probit_scale: f64,
758) -> f64 {
759 marginal_slope_standard_normal_scalar_eta(marginal_eta, logslope, z, probit_scale)
760}
761
762#[inline]
763pub(super) fn marginal_slope_standard_normal_scalar_eta(
764 q: f64,
765 slope: f64,
766 z: f64,
767 probit_scale: f64,
768) -> f64 {
769 let observed_slope = rigid_observed_logslope(slope, probit_scale);
770 q * (1.0 + observed_slope * observed_slope).sqrt() + observed_slope * z
771}
772
773pub(super) fn unary_derivatives_normal_cdf(x: f64) -> [f64; 5] {
774 let pdf = normal_pdf(x);
775 [
776 normal_cdf(x),
777 pdf,
778 -x * pdf,
779 (x * x - 1.0) * pdf,
780 (-x.powi(3) + 3.0 * x) * pdf,
781 ]
782}
783
784pub(super) fn unary_derivatives_normal_pdf(x: f64) -> [f64; 5] {
785 let pdf = normal_pdf(x);
786 [
787 pdf,
788 -x * pdf,
789 (x * x - 1.0) * pdf,
790 (-x.powi(3) + 3.0 * x) * pdf,
791 (x.powi(4) - 6.0 * x * x + 3.0) * pdf,
792 ]
793}
794
795#[inline]
802pub(super) fn lse_accumulate(log_max: &mut f64, sum: &mut f64, log_term: f64) {
803 if !log_term.is_finite() {
804 return;
805 }
806 if log_term > *log_max {
807 if log_max.is_finite() {
808 *sum = *sum * (*log_max - log_term).exp() + 1.0;
809 } else {
810 *sum = 1.0;
811 }
812 *log_max = log_term;
813 } else {
814 *sum += (log_term - *log_max).exp();
815 }
816}
817
818#[derive(Clone, Copy, Debug, PartialEq, Eq)]
819pub enum MarginalSlopeCovarianceShape {
820 Diagonal,
821 Full,
822 LowRank,
823}
824
825#[derive(Clone, Debug, PartialEq)]
826pub enum MarginalSlopeCovariance {
827 Diagonal(Array1<f64>),
828 Full(Array2<f64>),
829 LowRank(Array2<f64>),
831}
832
833pub(crate) const COVARIANCE_QUADRATIC_FORM_PSD_TOL: f64 = -1e-10;
838
839impl MarginalSlopeCovariance {
840 pub fn shape(&self) -> MarginalSlopeCovarianceShape {
841 match self {
842 Self::Diagonal(_) => MarginalSlopeCovarianceShape::Diagonal,
843 Self::Full(_) => MarginalSlopeCovarianceShape::Full,
844 Self::LowRank(_) => MarginalSlopeCovarianceShape::LowRank,
845 }
846 }
847
848 pub fn dim(&self) -> usize {
849 match self {
850 Self::Diagonal(diag) => diag.len(),
851 Self::Full(cov) => cov.nrows(),
852 Self::LowRank(factor) => factor.nrows(),
853 }
854 }
855
856 pub fn validate(&self, context: &str) -> Result<(), String> {
857 match self {
858 Self::Diagonal(diag) => {
859 if diag.is_empty() {
860 return Err(format!("{context} diagonal covariance is empty"));
861 }
862 for (idx, &value) in diag.iter().enumerate() {
863 if !(value.is_finite() && value >= 0.0) {
864 return Err(format!(
865 "{context} diagonal covariance entry {idx} must be finite and non-negative, got {value}"
866 ));
867 }
868 }
869 }
870 Self::Full(cov) => {
871 if cov.nrows() == 0 || cov.nrows() != cov.ncols() {
872 return Err(format!(
873 "{context} full covariance must be non-empty and square, got {}x{}",
874 cov.nrows(),
875 cov.ncols()
876 ));
877 }
878 for i in 0..cov.nrows() {
879 for j in 0..cov.ncols() {
880 let value = cov[[i, j]];
881 if !value.is_finite() {
882 return Err(format!(
883 "{context} full covariance entry ({i},{j}) is non-finite"
884 ));
885 }
886 if (value - cov[[j, i]]).abs()
887 > 1e-10 * (1.0 + value.abs().max(cov[[j, i]].abs()))
888 {
889 return Err(format!(
890 "{context} full covariance must be symmetric at ({i},{j})"
891 ));
892 }
893 }
894 }
895 }
896 Self::LowRank(factor) => {
897 if factor.nrows() == 0 {
898 return Err(format!(
899 "{context} low-rank covariance factor has zero rows"
900 ));
901 }
902 for ((i, j), &value) in factor.indexed_iter() {
903 if !value.is_finite() {
904 return Err(format!(
905 "{context} low-rank covariance factor entry ({i},{j}) is non-finite"
906 ));
907 }
908 }
909 }
910 }
911 Ok(())
912 }
913
914 pub fn quadratic_form(&self, vector: &[f64]) -> Result<f64, String> {
915 self.validate("marginal-slope covariance")?;
916 if vector.len() != self.dim() {
917 return Err(format!(
918 "marginal-slope covariance dimension mismatch: vector={}, covariance={}",
919 vector.len(),
920 self.dim()
921 ));
922 }
923 if vector.iter().any(|value| !value.is_finite()) {
924 return Err("marginal-slope covariance vector contains non-finite values".to_string());
925 }
926 let value = match self {
927 Self::Diagonal(diag) => vector
928 .iter()
929 .zip(diag.iter())
930 .map(|(&v, &sigma)| v * v * sigma)
931 .sum::<f64>(),
932 Self::Full(cov) => {
933 let mut total = 0.0;
934 for i in 0..cov.nrows() {
935 let mut row_dot = 0.0;
936 for j in 0..cov.ncols() {
937 row_dot += cov[[i, j]] * vector[j];
938 }
939 total += vector[i] * row_dot;
940 }
941 total
942 }
943 Self::LowRank(factor) => {
944 let mut total = 0.0;
950 for r in 0..factor.ncols() {
951 let mut projection = 0.0;
952 for k in 0..factor.nrows() {
953 projection += factor[[k, r]] * vector[k];
954 }
955 total += projection * projection;
956 }
957 total
958 }
959 };
960 if value.is_finite() && value >= COVARIANCE_QUADRATIC_FORM_PSD_TOL {
961 Ok(value.max(0.0))
962 } else {
963 Err(format!(
964 "marginal-slope covariance quadratic form must be non-negative, got {value}"
965 ))
966 }
967 }
968}
969
970pub fn marginal_slope_covariance_from_scores(
997 scores: ArrayView2<'_, f64>,
998 weights: &Array1<f64>,
999) -> Result<MarginalSlopeCovariance, String> {
1000 let (n, k) = scores.dim();
1001 if k == 0 {
1002 return Err("marginal-slope score matrix must have at least one column".to_string());
1003 }
1004 if weights.len() != n {
1005 return Err(format!(
1006 "marginal-slope covariance weight length mismatch: weights={}, rows={n}",
1007 weights.len()
1008 ));
1009 }
1010 let total_weight = weights.iter().copied().sum::<f64>();
1011 if !(total_weight.is_finite() && total_weight > 0.0) {
1012 return Err("marginal-slope covariance needs positive finite total weight".to_string());
1013 }
1014 let mut mean = Array1::<f64>::zeros(k);
1015 for i in 0..n {
1016 let weight = weights[i];
1017 if !(weight.is_finite() && weight >= 0.0) {
1018 return Err(format!(
1019 "marginal-slope covariance weight {i} must be finite and non-negative, got {weight}"
1020 ));
1021 }
1022 for j in 0..k {
1023 let score = scores[[i, j]];
1024 if !score.is_finite() {
1025 return Err(format!(
1026 "marginal-slope covariance score ({i},{j}) is non-finite"
1027 ));
1028 }
1029 mean[j] += weight * score;
1030 }
1031 }
1032 mean.mapv_inplace(|value| value / total_weight);
1033
1034 let mut cov = Array2::<f64>::zeros((k, k));
1035 for i in 0..n {
1036 let weight = weights[i];
1037 for a in 0..k {
1038 let da = scores[[i, a]] - mean[a];
1039 for b in 0..=a {
1040 let value = weight * da * (scores[[i, b]] - mean[b]) / total_weight;
1041 cov[[a, b]] += value;
1042 if a != b {
1043 cov[[b, a]] += value;
1044 }
1045 }
1046 }
1047 }
1048
1049 if k == 1 {
1074 return Ok(MarginalSlopeCovariance::Diagonal(cov.diag().to_owned()));
1075 }
1076
1077 let diag: Vec<f64> = (0..k).map(|i| cov[[i, i]]).collect();
1078 let diag_max = diag.iter().fold(0.0_f64, |acc, &v| acc.max(v.abs()));
1079 let numerical_floor = 1e-10 * (1.0 + diag_max);
1080
1081 let mut is_strict_diagonal = true;
1082 'strict: for a in 0..k {
1083 for b in (a + 1)..k {
1084 if cov[[a, b]].abs() > numerical_floor {
1085 is_strict_diagonal = false;
1086 break 'strict;
1087 }
1088 }
1089 }
1090 if is_strict_diagonal {
1091 return Ok(MarginalSlopeCovariance::Diagonal(cov.diag().to_owned()));
1092 }
1093
1094 use gam_linalg::faer_ndarray::FaerEigh;
1095 let (evals, evecs) = cov
1096 .eigh(faer::Side::Lower)
1097 .map_err(|err| format!("marginal-slope covariance eigendecomposition failed: {err}"))?;
1098 let max_eval = evals
1099 .iter()
1100 .fold(0.0_f64, |acc, &value| acc.max(value.abs()));
1101 let rank_tol = 1e-10 * max_eval.max(1.0);
1102 let positive: Vec<(usize, f64)> = evals
1103 .iter()
1104 .enumerate()
1105 .filter_map(|(idx, &value)| (value > rank_tol).then_some((idx, value)))
1106 .collect();
1107
1108 if positive.len() < k {
1109 let mut factor = Array2::<f64>::zeros((k, positive.len()));
1112 for (col, (idx, value)) in positive.iter().enumerate() {
1113 let scale = value.sqrt();
1114 for row in 0..k {
1115 factor[[row, col]] = evecs[[row, *idx]] * scale;
1116 }
1117 }
1118 return Ok(MarginalSlopeCovariance::LowRank(factor));
1119 }
1120
1121 let sum_w_sq = weights.iter().map(|&w| w * w).sum::<f64>();
1123 let n_eff = if sum_w_sq > 0.0 {
1124 (total_weight * total_weight) / sum_w_sq
1125 } else {
1126 1.0
1127 };
1128 const OFFDIAG_Z_THRESHOLD: f64 = 4.0;
1129 let mut is_stat_diagonal = true;
1130 'stat: for a in 0..k {
1131 for b in (a + 1)..k {
1132 let stat_se = (diag[a].max(0.0) * diag[b].max(0.0) / n_eff)
1133 .max(0.0)
1134 .sqrt();
1135 let threshold = numerical_floor.max(OFFDIAG_Z_THRESHOLD * stat_se);
1136 if cov[[a, b]].abs() > threshold {
1137 is_stat_diagonal = false;
1138 break 'stat;
1139 }
1140 }
1141 }
1142 if is_stat_diagonal {
1143 Ok(MarginalSlopeCovariance::Diagonal(cov.diag().to_owned()))
1144 } else {
1145 Ok(MarginalSlopeCovariance::Full(cov))
1146 }
1147}
1148
1149pub fn marginal_slope_preserving_scale(
1150 slopes: &[f64],
1151 covariance: &MarginalSlopeCovariance,
1152 probit_scale: f64,
1153) -> Result<f64, String> {
1154 if !probit_scale.is_finite() {
1155 return Err(format!(
1156 "marginal-slope probit scale must be finite, got {probit_scale}"
1157 ));
1158 }
1159 let observed_slopes = slopes
1160 .iter()
1161 .map(|&slope| probit_scale * slope)
1162 .collect::<Vec<_>>();
1163 let variance = covariance.quadratic_form(&observed_slopes)?;
1164 Ok((1.0 + variance).sqrt())
1165}
1166
1167pub fn marginal_slope_probit_eta(
1168 q: f64,
1169 z: &[f64],
1170 slopes: &[f64],
1171 covariance: &MarginalSlopeCovariance,
1172 probit_scale: f64,
1173) -> Result<f64, String> {
1174 if z.len() != slopes.len() {
1175 return Err(format!(
1176 "marginal-slope score/slope dimension mismatch: z={}, slopes={}",
1177 z.len(),
1178 slopes.len()
1179 ));
1180 }
1181 if slopes.len() != covariance.dim() {
1182 return Err(format!(
1183 "marginal-slope covariance dimension mismatch: slopes={}, covariance={}",
1184 slopes.len(),
1185 covariance.dim()
1186 ));
1187 }
1188 if !q.is_finite() || z.iter().any(|value| !value.is_finite()) {
1189 return Err("marginal-slope probit eta inputs must be finite".to_string());
1190 }
1191 let scale = marginal_slope_preserving_scale(slopes, covariance, probit_scale)?;
1192 let linear = z
1193 .iter()
1194 .zip(slopes.iter())
1195 .map(|(&score, &slope)| probit_scale * slope * score)
1196 .sum::<f64>();
1197 Ok(q * scale + linear)
1198}
1199
1200pub(super) fn empirical_rigid_calibration_eval(
1230 intercept: f64,
1231 log_target_mu: f64,
1232 slope: f64,
1233 probit_scale: f64,
1234 nodes: &[f64],
1235 weights: &[f64],
1236) -> Result<(f64, f64, f64), String> {
1237 if !intercept.is_finite() {
1238 return Err(format!(
1239 "empirical latent calibration: non-finite intercept {intercept}"
1240 ));
1241 }
1242 let observed_slope = rigid_observed_logslope(slope, probit_scale);
1243 const HALF_LOG_2PI: f64 = 0.918_938_533_204_672_8; let mut log_max_phi = f64::NEG_INFINITY;
1247 let mut sum_phi = 0.0_f64;
1248 let mut log_max_cdf = f64::NEG_INFINITY;
1249 let mut sum_cdf = 0.0_f64;
1250
1251 let mut log_max_pos = f64::NEG_INFINITY;
1255 let mut sum_pos = 0.0_f64;
1256 let mut log_max_neg = f64::NEG_INFINITY;
1257 let mut sum_neg = 0.0_f64;
1258
1259 for (&node, &weight) in nodes.iter().zip(weights.iter()) {
1260 if !(weight.is_finite() && weight > 0.0) {
1261 continue;
1262 }
1263 let eta = intercept + observed_slope * node;
1264 if !eta.is_finite() {
1265 return Err(format!(
1266 "empirical latent calibration: non-finite η at intercept={intercept}, slope={slope}, node={node}"
1267 ));
1268 }
1269 let log_w = weight.ln();
1270 let log_phi = -0.5 * eta * eta - HALF_LOG_2PI;
1271 let log_term_phi = log_w + log_phi;
1272 let log_term_cdf = log_w + normal_logcdf(eta);
1273
1274 lse_accumulate(&mut log_max_phi, &mut sum_phi, log_term_phi);
1275 lse_accumulate(&mut log_max_cdf, &mut sum_cdf, log_term_cdf);
1276
1277 if eta != 0.0 {
1278 let log_term_eta_phi = log_term_phi + eta.abs().ln();
1279 if eta > 0.0 {
1280 lse_accumulate(&mut log_max_pos, &mut sum_pos, log_term_eta_phi);
1281 } else {
1282 lse_accumulate(&mut log_max_neg, &mut sum_neg, log_term_eta_phi);
1283 }
1284 }
1285 }
1286
1287 if !(sum_phi.is_finite() && sum_cdf.is_finite() && sum_phi > 0.0 && sum_cdf > 0.0) {
1288 return Err(format!(
1289 "empirical latent calibration: log-space accumulation failed (sum_phi={sum_phi}, sum_cdf={sum_cdf}, intercept={intercept})"
1290 ));
1291 }
1292
1293 let log_s_phi = log_max_phi + sum_phi.ln();
1294 let log_s_cdf = log_max_cdf + sum_cdf.ln();
1295
1296 let f = log_s_cdf - log_target_mu;
1298 let log_f_prime = log_s_phi - log_s_cdf;
1310 let f_prime = if log_f_prime > -740.0 {
1311 log_f_prime.exp()
1312 } else {
1313 f64::MIN_POSITIVE
1314 };
1315
1316 let exp_safe = |log_x: f64| -> f64 { if log_x > -740.0 { log_x.exp() } else { 0.0 } };
1324 let pos_over_cdf = if sum_pos > 0.0 {
1325 exp_safe(log_max_pos + sum_pos.ln() - log_s_cdf)
1326 } else {
1327 0.0
1328 };
1329 let neg_over_cdf = if sum_neg > 0.0 {
1330 exp_safe(log_max_neg + sum_neg.ln() - log_s_cdf)
1331 } else {
1332 0.0
1333 };
1334 let s_etaphi_over_s_cdf = pos_over_cdf - neg_over_cdf;
1335 let f_double_prime = -s_etaphi_over_s_cdf - f_prime * f_prime;
1336
1337 if !(f.is_finite() && f_prime.is_finite() && f_prime > 0.0 && f_double_prime.is_finite()) {
1338 return Err(format!(
1339 "empirical latent calibration: non-finite log-space state f={f}, f'={f_prime}, f''={f_double_prime} at intercept={intercept}"
1340 ));
1341 }
1342 Ok((f, f_prime, f_double_prime))
1343}
1344
1345pub(crate) fn empirical_intercept_from_marginal(
1346 target_mu: f64,
1347 target_q: f64,
1348 slope: f64,
1349 probit_scale: f64,
1350 nodes: &[f64],
1351 weights: &[f64],
1352 initial: Option<f64>,
1353) -> Result<f64, String> {
1354 if !(target_mu.is_finite() && target_mu > 0.0 && target_mu < 1.0) {
1355 return Err(format!(
1356 "empirical latent calibration requires target mu in (0,1), got {target_mu}"
1357 ));
1358 }
1359 let log_target_mu = target_mu.ln();
1360 let closed_form_seed = rigid_intercept_from_marginal(target_q, slope, probit_scale);
1361 let seed = initial.unwrap_or(closed_form_seed);
1362 let eval = |a: f64| {
1363 empirical_rigid_calibration_eval(a, log_target_mu, slope, probit_scale, nodes, weights)
1364 };
1365 let abs_tol = 1e-13_f64.max(4.0 * f64::EPSILON);
1372 let solve_from = |s: f64| {
1373 crate::monotone_root::solve_monotone_root(
1374 eval,
1375 s,
1376 "empirical latent intercept",
1377 abs_tol,
1378 64,
1379 48,
1380 )
1381 .map_err(|e| e.to_string())
1384 };
1385 let (root, _, f_best) = match solve_from(seed) {
1396 Ok(v) => v,
1397 Err(first_err) => {
1398 if seed == closed_form_seed {
1399 return Err(first_err);
1400 }
1401 solve_from(closed_form_seed).map_err(|retry_err| {
1402 format!("{first_err}; closed-form retry from a={closed_form_seed:.6}: {retry_err}")
1403 })?
1404 }
1405 };
1406 if f_best.abs() > abs_tol {
1407 return Err(format!(
1408 "empirical latent intercept solve failed: log-residual={f_best:.3e} at a={root:.6}, target mu={target_mu:.6}"
1409 ));
1410 }
1411 Ok(root)
1412}
1413
1414#[inline]
1415pub(super) fn rigid_standard_normal_neglog_only(
1416 q: f64,
1417 g: f64,
1418 z: f64,
1419 y: f64,
1420 w: f64,
1421 probit_scale: f64,
1422) -> Result<f64, String> {
1423 let s = 2.0 * y - 1.0;
1424 let eta = marginal_slope_standard_normal_scalar_eta(q, g, z, probit_scale);
1425 let m = s * eta;
1426 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
1427 if !logcdf.is_finite() {
1428 return Err(format!(
1429 "rigid probit neglog_only: non-finite log Φ at q={q}, g={g}, z={z}, y={y}"
1430 ));
1431 }
1432 Ok(-w * logcdf)
1433}
1434
1435#[inline]
1465pub(crate) fn rigid_standard_normal_row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1466 p: &[S; 2],
1467 marginal: BernoulliMarginalLinkMap,
1468 z: f64,
1469 y: f64,
1470 w: f64,
1471 probit_scale: f64,
1472) -> Result<S, String> {
1473 let signed = rigid_standard_normal_signed_margin(p, marginal, z, y, probit_scale);
1477 let m = signed.value();
1480 if !(m.is_finite() || m == f64::INFINITY) {
1481 return Err(format!(
1482 "non-finite signed margin in rigid probit row NLL: {m}"
1483 ));
1484 }
1485 Ok(signed.compose_unary(signed_probit_neglog_unary_stack(m, w)))
1487}
1488
1489#[inline]
1499pub(crate) fn rigid_standard_normal_signed_margin<S: gam_math::jet_scalar::JetScalar<2>>(
1500 p: &[S; 2],
1501 marginal: BernoulliMarginalLinkMap,
1502 z: f64,
1503 y: f64,
1504 probit_scale: f64,
1505) -> S {
1506 let q = p[0].compose_unary([
1508 marginal.q,
1509 marginal.q1,
1510 marginal.q2,
1511 marginal.q3,
1512 marginal.q4,
1513 ]);
1514 let slope = p[1];
1515 let observed_slope = slope.scale(probit_scale);
1517 let b2 = observed_slope.mul(&observed_slope);
1518 let c = b2.add(&S::constant(1.0)).sqrt();
1519 let eta = q.mul(&c).add(&observed_slope.scale(z));
1521 eta.scale(2.0 * y - 1.0)
1522}
1523
1524pub(crate) struct RigidStandardNormalRow {
1537 pub(crate) marginal: BernoulliMarginalLinkMap,
1538 pub(crate) g: f64,
1539 pub(crate) z: f64,
1540 pub(crate) y: f64,
1541 pub(crate) w: f64,
1542 pub(crate) probit_scale: f64,
1543}
1544
1545impl gam_math::jet_tower::RowNllProgramGeneric<2> for RigidStandardNormalRow {
1546 fn n_rows(&self) -> usize {
1547 1
1548 }
1549
1550 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
1551 if row != 0 {
1552 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1553 }
1554 Ok([self.marginal.eta_value(), self.g])
1555 }
1556
1557 fn row_nll_generic<S: gam_math::jet_scalar::JetScalar<2>>(
1558 &self,
1559 row: usize,
1560 p: &[S; 2],
1561 ) -> Result<S, String> {
1562 if row != 0 {
1563 return Err(format!("RigidStandardNormalRow: row {row} out of range"));
1564 }
1565 rigid_standard_normal_row_nll_generic(
1566 p,
1567 self.marginal,
1568 self.z,
1569 self.y,
1570 self.w,
1571 self.probit_scale,
1572 )
1573 }
1574}
1575
1576#[inline]
1577pub(crate) fn rigid_standard_normal_tower(
1578 marginal: BernoulliMarginalLinkMap,
1579 g: f64,
1580 z: f64,
1581 y: f64,
1582 w: f64,
1583 probit_scale: f64,
1584) -> Result<Tower4<2>, String> {
1585 let program = RigidStandardNormalRow {
1592 marginal,
1593 g,
1594 z,
1595 y,
1596 w,
1597 probit_scale,
1598 };
1599 gam_math::jet_tower::generic_full_tower(&program, 0)
1600}
1601
1602#[inline]
1615fn rigid_standard_normal_signed_jet(
1616 marginal: BernoulliMarginalLinkMap,
1617 g: f64,
1618 z: f64,
1619 y: f64,
1620 probit_scale: f64,
1621) -> Tower4<2> {
1622 let p = [
1625 Tower4::<2>::variable(marginal.eta_value(), 0),
1626 Tower4::<2>::variable(g, 1),
1627 ];
1628 rigid_standard_normal_signed_margin(&p, marginal, z, y, probit_scale)
1629}
1630
1631#[inline]
1659pub(super) fn rigid_standard_normal_towers_batch<T>(
1660 marginals: &[BernoulliMarginalLinkMap],
1661 slopes: &[f64],
1662 zs: &[f64],
1663 ys: &[f64],
1664 weights: &[f64],
1665 probit_scale: f64,
1666 out: &mut [T],
1667 mut fill: impl FnMut(&Tower4<2>) -> Result<T, String>,
1668) -> Result<(), String> {
1669 let chunk = marginals.len();
1670 if slopes.len() != chunk
1671 || zs.len() != chunk
1672 || ys.len() != chunk
1673 || weights.len() != chunk
1674 || out.len() != chunk
1675 {
1676 return Err(format!(
1677 "rigid_standard_normal_towers_batch length mismatch: marginals={chunk}, \
1678 slopes={}, zs={}, ys={}, weights={}, out={}",
1679 slopes.len(),
1680 zs.len(),
1681 ys.len(),
1682 weights.len(),
1683 out.len()
1684 ));
1685 }
1686
1687 let mut signed: Vec<Tower4<2>> = Vec::with_capacity(chunk);
1689 let mut margins: Vec<f64> = Vec::with_capacity(chunk);
1690 for i in 0..chunk {
1691 let jet =
1692 rigid_standard_normal_signed_jet(marginals[i], slopes[i], zs[i], ys[i], probit_scale);
1693 margins.push(jet.v);
1694 signed.push(jet);
1695 }
1696
1697 let mut stacks: Vec<[f64; 5]> = Vec::with_capacity(chunk);
1701 for i in 0..chunk {
1702 let m = margins[i];
1703 if !(m.is_finite() || m == f64::INFINITY) {
1704 return Err(format!(
1705 "non-finite signed margin in rigid probit tower batch: {m}"
1706 ));
1707 }
1708 stacks.push(signed_probit_neglog_unary_stack(m, weights[i]));
1709 }
1710
1711 for i in 0..chunk {
1713 let tower = signed[i].compose_unary(stacks[i]);
1714 out[i] = fill(&tower)?;
1715 }
1716 Ok(())
1717}
1718
1719#[inline]
1720pub(super) fn rigid_standard_normal_row_kernel(
1721 marginal: BernoulliMarginalLinkMap,
1722 g: f64,
1723 z: f64,
1724 y: f64,
1725 w: f64,
1726 probit_scale: f64,
1727) -> Result<(f64, [f64; 2], [[f64; 2]; 2]), String> {
1728 let program = RigidStandardNormalRow {
1737 marginal,
1738 g,
1739 z,
1740 y,
1741 w,
1742 probit_scale,
1743 };
1744 gam_math::jet_tower::generic_row_kernel(&program, 0)
1745}
1746
1747#[inline]
1785pub(super) fn rigid_standard_normal_mixed_z_sensitivity(
1786 marginal: BernoulliMarginalLinkMap,
1787 g: f64,
1788 z: f64,
1789 y: f64,
1790 w: f64,
1791 probit_scale: f64,
1792) -> Result<[f64; 2], String> {
1793 use gam_math::jet_tower::Tower2;
1804 let mut q = Tower2::<3>::constant(marginal.q);
1805 q.g[0] = marginal.q1;
1806 q.h[0][0] = marginal.q2;
1807 let slope = Tower2::<3>::variable(g, 1);
1808 let z_var = Tower2::<3>::variable(z, 2);
1809 let observed_logslope = slope * probit_scale;
1810 let c = (observed_logslope * observed_logslope + 1.0).sqrt();
1811 let eta = q * c + slope * (z_var * probit_scale);
1815 let signed = eta * (2.0 * y - 1.0);
1816 if !(signed.v.is_finite() || signed.v == f64::INFINITY) {
1818 return Err(format!(
1819 "rigid probit mixed-z sensitivity: non-finite signed margin {} at q={}, g={g}, z={z}, y={y}",
1820 signed.v, marginal.q
1821 ));
1822 }
1823 let stack = signed_probit_neglog_unary_stack(signed.v, w);
1824 if !stack[0].is_finite() {
1825 return Err(format!(
1826 "rigid probit mixed-z sensitivity: non-finite log Φ at q={}, g={g}, z={z}, y={y}",
1827 marginal.q
1828 ));
1829 }
1830 let tower = signed.compose_unary([stack[0], stack[1], stack[2]]);
1833 let s_q = -tower.h[0][2];
1840 let s_g = -tower.h[1][2];
1841 if !(s_q.is_finite() && s_g.is_finite()) {
1842 return Err(format!(
1843 "rigid probit mixed-z sensitivity: non-finite ∂²(log L)/∂(q,g)∂z = [{s_q}, {s_g}] at q={}, g={g}, z={z}",
1844 marginal.q
1845 ));
1846 }
1847 Ok([s_q, s_g])
1848}
1849
1850pub(super) fn rigid_standard_normal_score_zeta_sensitivity(
1877 base_link: &InverseLink,
1878 marginal_eta: &Array1<f64>,
1879 slope_eta: &Array1<f64>,
1880 z: &Array1<f64>,
1881 y: &Array1<f64>,
1882 weights: &Array1<f64>,
1883 probit_scale: f64,
1884 marginal_design: ArrayView2<'_, f64>,
1885 logslope_design: ArrayView2<'_, f64>,
1886 p_beta: usize,
1887) -> Result<Array2<f64>, String> {
1888 let n = marginal_eta.len();
1889 let p_m = marginal_design.ncols();
1890 let r = logslope_design.ncols();
1891 if slope_eta.len() != n
1892 || z.len() != n
1893 || y.len() != n
1894 || weights.len() != n
1895 || marginal_design.nrows() != n
1896 || logslope_design.nrows() != n
1897 {
1898 return Err(format!(
1899 "score_zeta_sensitivity row mismatch: marginal_eta={n}, slope_eta={}, z={}, y={}, \
1900 weights={}, marginal_design rows={}, logslope_design rows={}",
1901 slope_eta.len(),
1902 z.len(),
1903 y.len(),
1904 weights.len(),
1905 marginal_design.nrows(),
1906 logslope_design.nrows()
1907 ));
1908 }
1909 if p_m + r > p_beta {
1910 return Err(format!(
1911 "score_zeta_sensitivity width overflow: marginal({p_m}) + logslope({r}) > p_beta({p_beta})"
1912 ));
1913 }
1914 let mut s = Array2::<f64>::zeros((n, p_beta));
1915 for i in 0..n {
1916 let marginal = bernoulli_marginal_link_map(base_link, marginal_eta[i])?;
1917 let [s_q, s_g] = rigid_standard_normal_mixed_z_sensitivity(
1918 marginal,
1919 slope_eta[i],
1920 z[i],
1921 y[i],
1922 weights[i],
1923 probit_scale,
1924 )?;
1925 if s_q != 0.0 {
1928 let m_row = marginal_design.row(i);
1929 for (j, &mij) in m_row.iter().enumerate() {
1930 s[[i, j]] = s_q * mij;
1931 }
1932 }
1933 if s_g != 0.0 {
1934 let g_row = logslope_design.row(i);
1935 for (j, &gij) in g_row.iter().enumerate() {
1936 s[[i, p_m + j]] = s_g * gij;
1937 }
1938 }
1939 }
1940 Ok(s)
1941}
1942
1943#[inline]
1944pub(super) fn rigid_standard_normal_third_full(
1945 marginal: BernoulliMarginalLinkMap,
1946 g: f64,
1947 z: f64,
1948 y: f64,
1949 w: f64,
1950 probit_scale: f64,
1951) -> Result<[[[f64; 2]; 2]; 2], String> {
1952 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t3)
1953}
1954
1955#[inline]
1960pub(super) fn contract_third_full(t: &[[[f64; 2]; 2]; 2], d_eta: f64, d_g: f64) -> [[f64; 2]; 2] {
1961 [
1962 [
1963 t[0][0][0] * d_eta + t[0][0][1] * d_g,
1964 t[0][1][0] * d_eta + t[0][1][1] * d_g,
1965 ],
1966 [
1967 t[1][0][0] * d_eta + t[1][0][1] * d_g,
1968 t[1][1][0] * d_eta + t[1][1][1] * d_g,
1969 ],
1970 ]
1971}
1972
1973#[inline]
1974pub(super) fn rigid_standard_normal_fourth_full(
1975 marginal: BernoulliMarginalLinkMap,
1976 g: f64,
1977 z: f64,
1978 y: f64,
1979 w: f64,
1980 probit_scale: f64,
1981) -> Result<[[[[f64; 2]; 2]; 2]; 2], String> {
1982 Ok(rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?.t4)
1996}
1997
1998#[inline]
2017pub(super) fn contract_fourth_full(
2018 t: &[[[[f64; 2]; 2]; 2]; 2],
2019 u_eta: f64,
2020 u_g: f64,
2021 v_eta: f64,
2022 v_g: f64,
2023) -> [[f64; 2]; 2] {
2024 let mut out = [[0.0; 2]; 2];
2025 for a in 0..2 {
2026 for b in 0..2 {
2027 let mut sum = 0.0;
2028 sum += t[a][b][0][0] * u_eta * v_eta;
2029 sum += t[a][b][0][1] * u_eta * v_g;
2030 sum += t[a][b][1][0] * u_g * v_eta;
2031 sum += t[a][b][1][1] * u_g * v_g;
2032 out[a][b] = sum;
2033 }
2034 }
2035 out
2036}
2037
2038pub(super) fn ensure_finite_third_full_cache_row(
2039 t: &[[[f64; 2]; 2]; 2],
2040 context: &str,
2041) -> Result<(), String> {
2042 if t.iter().flatten().flatten().all(|value| value.is_finite()) {
2043 Ok(())
2044 } else {
2045 Err(format!(
2046 "{context}: warmed third-derivative cache row contains a non-finite value"
2047 ))
2048 }
2049}
2050
2051pub(super) fn ensure_finite_fourth_full_cache_row(
2052 t: &[[[[f64; 2]; 2]; 2]; 2],
2053 context: &str,
2054) -> Result<(), String> {
2055 if t.iter()
2056 .flatten()
2057 .flatten()
2058 .flatten()
2059 .all(|value| value.is_finite())
2060 {
2061 Ok(())
2062 } else {
2063 Err(format!(
2064 "{context}: warmed fourth-derivative cache row contains a non-finite value"
2065 ))
2066 }
2067}
2068
2069pub(crate) fn unary_derivatives_sqrt(x: f64) -> [f64; 5] {
2070 let s = x.max(1e-300).sqrt();
2071 let x1 = x.max(1e-300);
2072 let x2 = x1 * x1;
2073 let x3 = x2 * x1;
2074 [
2075 s,
2076 0.5 / s,
2077 -0.25 / (x1 * s),
2078 3.0 / (8.0 * x2 * s),
2079 -15.0 / (16.0 * x3 * s),
2080 ]
2081}
2082pub(crate) fn unary_derivatives_neglog_phi(x: f64, weight: f64) -> [f64; 5] {
2083 signed_probit_neglog_unary_stack(x, weight)
2088}
2089
2090pub(crate) fn unary_derivatives_log(x: f64) -> [f64; 5] {
2108 let x2 = x * x;
2109 let x3 = x2 * x;
2110 let x4 = x3 * x;
2111 [x.ln(), 1.0 / x, -1.0 / x2, 2.0 / x3, -6.0 / x4]
2112}
2113
2114pub(crate) fn unary_derivatives_log_normal_pdf(x: f64) -> [f64; 5] {
2116 let c = 0.5 * (2.0 * std::f64::consts::PI).ln();
2117 [-0.5 * x * x - c, -x, -1.0, 0.0, 0.0]
2118}
2119
2120#[cfg(test)]
2121mod jet_tower_oracle_tests {
2122 use super::*;
2144
2145 fn rigid_standard_normal_third_and_fourth_full(
2153 marginal: BernoulliMarginalLinkMap,
2154 g: f64,
2155 z: f64,
2156 y: f64,
2157 w: f64,
2158 probit_scale: f64,
2159 ) -> Result<([[[f64; 2]; 2]; 2], [[[[f64; 2]; 2]; 2]; 2]), String> {
2160 let tower = rigid_standard_normal_tower(marginal, g, z, y, w, probit_scale)?;
2161 Ok((tower.t3, tower.t4))
2162 }
2163 use gam_math::jet_tower::{
2164 KernelChannels, RowNllProgram, evaluate_program, verify_kernel_channels,
2165 };
2166
2167 struct BernoulliRigidStandardNormalNllProgram {
2170 primaries: Vec<[f64; 2]>,
2172 z: Vec<f64>,
2174 y: Vec<f64>,
2175 w: Vec<f64>,
2176 probit_scale: f64,
2177 }
2178
2179 impl RowNllProgram<2> for BernoulliRigidStandardNormalNllProgram {
2180 fn n_rows(&self) -> usize {
2181 self.primaries.len()
2182 }
2183
2184 fn primaries(&self, row: usize) -> Result<[f64; 2], String> {
2185 self.primaries
2186 .get(row)
2187 .copied()
2188 .ok_or_else(|| format!("bernoulli rigid nll program: row {row} out of range"))
2189 }
2190
2191 fn row_nll(&self, row: usize, p: &[Tower4<2>; 2]) -> Result<Tower4<2>, String> {
2192 let z = self.z[row];
2193 let y = self.y[row];
2194 let w = self.w[row];
2195 let s = self.probit_scale;
2196 let eta_marginal = p[0];
2200 let link = bernoulli_marginal_link_map(
2201 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2202 eta_marginal.v,
2203 )?;
2204 let q = eta_marginal.compose_unary([link.q, link.q1, link.q2, link.q3, link.q4]);
2205 let g = p[1];
2206 let observed_slope = g * s;
2208 let c = (observed_slope * observed_slope + 1.0).compose_unary(unary_derivatives_sqrt(
2209 observed_slope.v * observed_slope.v + 1.0,
2210 ));
2211 let eta = q * c + observed_slope * z;
2213 let signed = eta * (2.0 * y - 1.0);
2214 Ok(signed.compose_unary(unary_derivatives_neglog_phi(signed.v, w)))
2216 }
2217 }
2218
2219 fn scalar_nll(eta_marginal: f64, g: f64, z: f64, y: f64, w: f64, s: f64) -> f64 {
2222 let link = bernoulli_marginal_link_map(
2223 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2224 eta_marginal,
2225 )
2226 .unwrap();
2227 let observed_slope = g * s;
2228 let c = (observed_slope * observed_slope + 1.0).sqrt();
2229 let eta = link.q * c + observed_slope * z;
2230 let signed = (2.0 * y - 1.0) * eta;
2231 let cdf = 0.5 * libm::erfc(-signed / std::f64::consts::SQRT_2);
2232 -w * cdf.max(1e-300).ln()
2233 }
2234
2235 #[test]
2236 fn rigid_bernoulli_row_kernel_agrees_with_jet_tower_program_all_channels() {
2237 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2241 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2242 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2243 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2244 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2245 let n = eta.len();
2246
2247 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2249
2250 for &probit_scale in &[1.0_f64, 0.8] {
2251 let program = BernoulliRigidStandardNormalNllProgram {
2252 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2253 z: z.to_vec(),
2254 y: y.to_vec(),
2255 w: w.to_vec(),
2256 probit_scale,
2257 };
2258
2259 for row in 0..n {
2260 let tower = evaluate_program(&program, row).expect("tower evaluation");
2261
2262 let marginal = bernoulli_marginal_link_map(
2264 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2265 eta[row],
2266 )
2267 .expect("link map");
2268 let (value, gradient, hessian) = rigid_standard_normal_row_kernel(
2269 marginal,
2270 g[row],
2271 z[row],
2272 y[row],
2273 w[row],
2274 probit_scale,
2275 )
2276 .expect("production row kernel");
2277
2278 let (third_full, fourth_full) = rigid_standard_normal_third_and_fourth_full(
2285 marginal,
2286 g[row],
2287 z[row],
2288 y[row],
2289 w[row],
2290 probit_scale,
2291 )
2292 .expect("production third+fourth");
2293 let third: Vec<([f64; 2], [[f64; 2]; 2])> = dirs
2294 .iter()
2295 .map(|d| (*d, contract_third_full(&third_full, d[0], d[1])))
2296 .collect();
2297
2298 let fourth: Vec<([f64; 2], [f64; 2], [[f64; 2]; 2])> = dirs
2299 .iter()
2300 .enumerate()
2301 .map(|(i, u)| {
2302 let v = dirs[(i + 1) % dirs.len()];
2303 (
2304 *u,
2305 v,
2306 contract_fourth_full(&fourth_full, u[0], u[1], v[0], v[1]),
2307 )
2308 })
2309 .collect();
2310
2311 let claims = KernelChannels {
2312 value,
2313 gradient,
2314 hessian,
2315 third,
2316 fourth,
2317 };
2318
2319 verify_kernel_channels(&tower, &claims, 1e-9).unwrap_or_else(|e| {
2320 panic!(
2321 "probit_scale {probit_scale} row {row}: production rigid Bernoulli \
2322 RowKernel disagrees with #932 jet-tower truth: {e}"
2323 )
2324 });
2325
2326 let h = 1e-3;
2330 let f = |de: f64, dg: f64| {
2331 scalar_nll(
2332 eta[row] + de,
2333 g[row] + dg,
2334 z[row],
2335 y[row],
2336 w[row],
2337 probit_scale,
2338 )
2339 };
2340 let f0 = f(0.0, 0.0);
2341 assert!(
2342 (f0 - tower.v).abs() <= 1e-9 * f0.abs().max(1.0),
2343 "row {row}: independent scalar NLL {f0:+.12e} != tower value {:+.12e}",
2344 tower.v
2345 );
2346 let g_eta = (f(-2.0 * h, 0.0) - 8.0 * f(-h, 0.0) + 8.0 * f(h, 0.0)
2348 - f(2.0 * h, 0.0))
2349 / (12.0 * h);
2350 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))
2351 / (12.0 * h);
2352 for (label, fd, ad) in [("∂η", g_eta, tower.g[0]), ("∂g", g_g, tower.g[1])] {
2353 assert!(
2354 (fd - ad).abs() <= 1e-5 * ad.abs().max(1.0),
2355 "row {row} {label}: FD witness {fd:+.6e} != tower grad {ad:+.6e}"
2356 );
2357 }
2358 }
2359 }
2360 }
2361
2362 #[test]
2372 fn rigid_third_and_fourth_full_shares_one_tower_bit_identical() {
2373 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2374 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2375 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2376 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2377 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2378 for &probit_scale in &[1.0_f64, 0.8] {
2379 for r in 0..eta.len() {
2380 let marginal = bernoulli_marginal_link_map(
2381 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2382 eta[r],
2383 )
2384 .expect("link map");
2385 let t3_sep = rigid_standard_normal_third_full(
2386 marginal,
2387 g[r],
2388 z[r],
2389 y[r],
2390 w[r],
2391 probit_scale,
2392 )
2393 .expect("separate third");
2394 let t4_sep = rigid_standard_normal_fourth_full(
2395 marginal,
2396 g[r],
2397 z[r],
2398 y[r],
2399 w[r],
2400 probit_scale,
2401 )
2402 .expect("separate fourth");
2403 let (t3_comb, t4_comb) = rigid_standard_normal_third_and_fourth_full(
2404 marginal,
2405 g[r],
2406 z[r],
2407 y[r],
2408 w[r],
2409 probit_scale,
2410 )
2411 .expect("combined third+fourth");
2412 for a in 0..2 {
2414 for b in 0..2 {
2415 for c in 0..2 {
2416 assert_eq!(
2417 t3_comb[a][b][c], t3_sep[a][b][c],
2418 "t3[{a}][{b}][{c}] row {r} scale {probit_scale} not bit-identical"
2419 );
2420 for d in 0..2 {
2421 assert_eq!(
2422 t4_comb[a][b][c][d], t4_sep[a][b][c][d],
2423 "t4[{a}][{b}][{c}][{d}] row {r} scale {probit_scale} not bit-identical"
2424 );
2425 }
2426 }
2427 }
2428 }
2429 }
2430 }
2431 }
2432
2433 #[test]
2444 fn rigid_bernoulli_generic_program_matches_tower4_program_all_channels() {
2445 use gam_math::jet_tower::{
2446 generic_fourth_contracted, generic_full_tower, generic_row_kernel,
2447 generic_third_contracted,
2448 };
2449
2450 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2451 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2452 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2453 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2454 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2455 let n = eta.len();
2456 let dirs: [[f64; 2]; 3] = [[0.7, -1.3], [-0.4, 0.6], [1.2, 0.2]];
2457
2458 let close = |a: f64, b: f64, label: &str| {
2459 let band = 1e-12 + 1e-12 * a.abs().max(b.abs());
2460 assert!(
2461 (a - b).abs() <= band,
2462 "{label}: generic {a:+.15e} vs Tower4-program {b:+.15e} (band {band:.3e})"
2463 );
2464 };
2465
2466 for &probit_scale in &[1.0_f64, 0.8] {
2467 let tower_program = BernoulliRigidStandardNormalNllProgram {
2469 primaries: (0..n).map(|r| [eta[r], g[r]]).collect(),
2470 z: z.to_vec(),
2471 y: y.to_vec(),
2472 w: w.to_vec(),
2473 probit_scale,
2474 };
2475
2476 for row in 0..n {
2477 let truth = evaluate_program(&tower_program, row).expect("Tower4 program tower");
2478
2479 let marginal = bernoulli_marginal_link_map(
2480 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2481 eta[row],
2482 )
2483 .expect("link map");
2484 let program = RigidStandardNormalRow {
2485 marginal,
2486 g: g[row],
2487 z: z[row],
2488 y: y[row],
2489 w: w[row],
2490 probit_scale,
2491 };
2492
2493 let full = generic_full_tower(&program, 0).expect("generic full tower");
2496 close(full.v, truth.v, "full value");
2497 for a in 0..2 {
2498 close(full.g[a], truth.g[a], "full grad");
2499 for b in 0..2 {
2500 close(full.h[a][b], truth.h[a][b], "full hess");
2501 for c in 0..2 {
2502 close(full.t3[a][b][c], truth.t3[a][b][c], "full t3");
2503 for d in 0..2 {
2504 close(full.t4[a][b][c][d], truth.t4[a][b][c][d], "full t4");
2505 }
2506 }
2507 }
2508 }
2509
2510 let (val, grad, hess) =
2512 generic_row_kernel(&program, 0).expect("generic row kernel");
2513 close(val, truth.v, "order2 value");
2514 for a in 0..2 {
2515 close(grad[a], truth.g[a], "order2 grad");
2516 for b in 0..2 {
2517 close(hess[a][b], truth.h[a][b], "order2 hess");
2518 }
2519 }
2520
2521 for dir in &dirs {
2524 let third = generic_third_contracted(&program, 0, dir)
2525 .expect("generic third contracted");
2526 let truth3 = truth.third_contracted(dir);
2527 for a in 0..2 {
2528 for b in 0..2 {
2529 close(third[a][b], truth3[a][b], "third contracted");
2530 }
2531 }
2532 }
2533
2534 for (i, u) in dirs.iter().enumerate() {
2537 let v = dirs[(i + 1) % dirs.len()];
2538 let fourth = generic_fourth_contracted(&program, 0, u, &v)
2539 .expect("generic fourth contracted");
2540 let truth4 = truth.fourth_contracted(u, &v);
2541 for a in 0..2 {
2542 for b in 0..2 {
2543 close(fourth[a][b], truth4[a][b], "fourth contracted");
2544 }
2545 }
2546 }
2547 }
2548 }
2549 }
2550
2551 fn hand_rigid_vgh(
2560 marginal: BernoulliMarginalLinkMap,
2561 g: f64,
2562 z: f64,
2563 y: f64,
2564 w: f64,
2565 probit_scale: f64,
2566 ) -> (f64, [f64; 2], [[f64; 2]; 2]) {
2567 let s = 2.0 * y - 1.0;
2568 let observed_logslope = probit_scale * g;
2569 let g2 = observed_logslope * observed_logslope;
2570 let c = (1.0 + g2).sqrt();
2571 let c1 = probit_scale * observed_logslope / c;
2572 let c_inv3 = 1.0 / (c * c * c);
2573 let c2 = probit_scale * probit_scale * c_inv3;
2574 let q = marginal.q;
2575 let eta = q * c + observed_logslope * z;
2577 let m = s * eta;
2578 let (logcdf, _) = signed_probit_logcdf_and_mills_ratio(m);
2579 let (k1, k2, _k3, _k4) =
2583 signed_probit_neglog_derivatives_up_to_fourth(m, w).expect("hand kernel");
2584 let u1 = s * k1;
2585 let u2 = k2;
2586 let eta_q = c;
2587 let eta_g = q * c1 + probit_scale * z;
2588 let value = -w * logcdf;
2590 let gradient = [u1 * eta_q * marginal.q1, u1 * eta_g];
2592 let h00 = u2 * eta_q * eta_q;
2594 let h01 = u2 * eta_q * eta_g + u1 * c1;
2595 let h11 = u2 * eta_g * eta_g + u1 * q * c2;
2596 let grad_q = u1 * eta_q;
2598 let hessian = [
2599 [
2600 h00 * marginal.q1 * marginal.q1 + grad_q * marginal.q2,
2601 h01 * marginal.q1,
2602 ],
2603 [h01 * marginal.q1, h11],
2604 ];
2605 (value, gradient, hessian)
2606 }
2607
2608 #[test]
2614 fn rigid_bernoulli_row_kernel_matches_hand_chain_witness() {
2615 let eta = [0.3_f64, -0.7, 0.05, 0.9, -1.2, 2.1, -2.4];
2616 let g = [0.2_f64, -0.5, 0.35, -0.15, 0.6, 0.45, -0.55];
2617 let z = [0.4_f64, -1.1, 0.0, 0.7, -0.3, 1.6, -1.4];
2618 let y = [1.0_f64, 0.0, 0.0, 1.0, 1.0, 0.0, 1.0];
2619 let w = [1.0_f64, 0.8, 1.3, 0.9, 1.1, 0.7, 1.4];
2620 let close = |a: f64, b: f64, label: &str| {
2621 let band = 1e-12 + 1e-9 * a.abs().max(b.abs());
2622 assert!(
2623 (a - b).abs() <= band,
2624 "{label}: jet {a:+.15e} vs hand {b:+.15e} (band {band:.3e})"
2625 );
2626 };
2627 for &probit_scale in &[1.0_f64, 0.8] {
2628 for r in 0..eta.len() {
2629 let marginal = bernoulli_marginal_link_map(
2630 &InverseLink::Standard(gam_problem::StandardLink::Probit),
2631 eta[r],
2632 )
2633 .expect("link map");
2634 let (jv, jg, jh) = rigid_standard_normal_row_kernel(
2635 marginal,
2636 g[r],
2637 z[r],
2638 y[r],
2639 w[r],
2640 probit_scale,
2641 )
2642 .expect("jet kernel");
2643 let (hv, hg, hh) = hand_rigid_vgh(marginal, g[r], z[r], y[r], w[r], probit_scale);
2644 close(jv, hv, "value");
2645 for a in 0..2 {
2646 close(jg[a], hg[a], "grad");
2647 for b in 0..2 {
2648 close(jh[a][b], hh[a][b], "hess");
2649 }
2650 }
2651 }
2652 }
2653 }
2654
2655 }
2661
2662#[cfg(test)]
2663mod flex_primary_hessian_oracle_tests {
2664 use super::*;
2688 use super::family::*;
2695 use gam_linalg::matrix::DenseDesignMatrix;
2696 use ndarray::Array1;
2697 use ndarray::Array2;
2698 use std::sync::Arc;
2699 use std::sync::Mutex;
2700
2701 fn make_flex_oracle_family(
2707 n: usize,
2708 ) -> (BernoulliMarginalSlopeFamily, Vec<ParameterBlockState>) {
2709 let score_seed = Array1::linspace(-2.0, 2.0, n.max(6));
2710 let link_seed = Array1::linspace(-1.8, 1.8, n.max(6));
2711 let cfg = DeviationBlockConfig {
2712 num_internal_knots: 3,
2713 ..DeviationBlockConfig::default()
2714 };
2715 let score_prepared = build_score_warp_deviation_block_from_seed(&score_seed, &cfg)
2716 .expect("build score warp block");
2717 let link_prepared = build_link_deviation_block_from_knots_design_seed_and_weights(
2718 &link_seed, &link_seed, &cfg,
2719 )
2720 .expect("build link deviation block");
2721
2722 let y: Array1<f64> =
2723 Array1::from_iter((0..n).map(|i| if (i * 17 + 3) % 7 >= 4 { 1.0 } else { 0.0 }));
2724 let weights: Array1<f64> =
2725 Array1::from_iter((0..n).map(|i| 0.75 + ((i * 11 + 5) % 5) as f64 * 0.05));
2726 let z: Array1<f64> =
2727 Array1::from_iter((0..n).map(|i| -1.7 + 3.4 * (i as f64 + 0.5) / n as f64));
2728 let marginal_x = Array2::from_shape_fn((n, 2), |(i, j)| {
2729 if j == 0 {
2730 1.0
2731 } else {
2732 -0.4 + 0.8 * ((i * 19 + 7) % n) as f64 / n as f64
2733 }
2734 });
2735 let logslope_x = Array2::from_shape_fn((n, 2), |(i, j)| {
2736 if j == 0 {
2737 1.0
2738 } else {
2739 0.3 - 0.6 * ((i * 23 + 11) % n) as f64 / n as f64
2740 }
2741 });
2742
2743 let family = BernoulliMarginalSlopeFamily {
2744 y: Arc::new(y),
2745 weights: Arc::new(weights),
2746 z: Arc::new(z.clone()),
2747 latent_measure: LatentMeasureKind::StandardNormal,
2748 gaussian_frailty_sd: Some(0.15),
2749 base_link: InverseLink::Standard(gam_problem::StandardLink::Probit),
2750 marginal_design: DesignMatrix::Dense(DenseDesignMatrix::from(marginal_x.clone())),
2751 logslope_design: DesignMatrix::Dense(DenseDesignMatrix::from(logslope_x.clone())),
2752 score_warp: Some(score_prepared.runtime.clone()),
2753 link_dev: Some(link_prepared.runtime.clone()),
2754 policy: gam_runtime::resource::ResourcePolicy::default_library(),
2755 cell_moment_lru: Arc::new(exact_kernel::CellMomentLruCache::new(1024)),
2756 cell_moment_cache_stats: Arc::new(exact_kernel::CellMomentCacheStats::default()),
2757 intercept_warm_starts: None,
2758 auto_subsample_phase_counter: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
2759 auto_subsample_last_rho: Arc::new(Mutex::new(None)),
2760 };
2761
2762 let beta_m = Array1::from_vec(vec![0.12, -0.04]);
2763 let beta_g = Array1::from_vec(vec![0.35, 0.03]);
2764 let beta_h = Array1::from_iter(
2765 (0..score_prepared.runtime.basis_dim()).map(|idx| 0.0015 * (idx as f64 + 1.0)),
2766 );
2767 let beta_w = Array1::from_iter(
2768 (0..link_prepared.runtime.basis_dim()).map(|idx| -0.001 * (idx as f64 + 1.0)),
2769 );
2770 let states = vec![
2771 ParameterBlockState {
2772 eta: marginal_x.dot(&beta_m),
2773 beta: beta_m,
2774 },
2775 ParameterBlockState {
2776 eta: logslope_x.dot(&beta_g),
2777 beta: beta_g,
2778 },
2779 ParameterBlockState {
2780 beta: beta_h,
2781 eta: Array1::zeros(z.len()),
2782 },
2783 ParameterBlockState {
2784 beta: beta_w,
2785 eta: Array1::zeros(z.len()),
2786 },
2787 ];
2788 (family, states)
2789 }
2790
2791 fn flex_gradient_at_perturbed(
2799 family: &BernoulliMarginalSlopeFamily,
2800 states: &[ParameterBlockState],
2801 primary: &super::super::hessian_paths::PrimarySlices,
2802 row: usize,
2803 u: usize,
2804 delta: f64,
2805 ) -> Array1<f64> {
2806 let mut states = states.to_vec();
2807 if u == primary.q {
2813 states[0].eta[row] += delta;
2814 } else if u == primary.logslope {
2815 states[1].eta[row] += delta;
2816 } else if let Some(h_range) = primary.h.as_ref()
2817 && h_range.contains(&u)
2818 {
2819 states[2].beta[u - h_range.start] += delta;
2820 } else if let Some(w_range) = primary.w.as_ref()
2821 && w_range.contains(&u)
2822 {
2823 states[3].beta[u - w_range.start] += delta;
2824 } else {
2825 panic!("primary coordinate {u} out of range for flex oracle");
2826 }
2827 let row_ctx = family
2828 .build_row_exact_context_with_stats_and_cell_cache(row, &states, None, false)
2829 .expect("perturbed row context");
2830 let (_neglog, grad, _hess) = family
2831 .compute_row_primary_gradient_hessian(row, &states, primary, &row_ctx)
2832 .expect("perturbed gradient");
2833 grad
2834 }
2835
2836 #[test]
2839 fn flex_primary_hessian_matches_central_fd_of_gradient() {
2840 let n = 12usize;
2841 let (family, states) = make_flex_oracle_family(n);
2842 let cache = family
2843 .build_exact_eval_cache(&states)
2844 .expect("flex exact eval cache");
2845 let primary = &cache.primary;
2846 let r = primary.total;
2847 assert!(
2848 r >= 4,
2849 "flex fixture must carry q + logslope + deviation blocks"
2850 );
2851
2852 let h = 1e-4;
2856 let mut max_rel = 0.0_f64;
2857
2858 for &row in &[2usize, 5, 8] {
2861 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
2862 let (_neglog, _grad, analytic_hess) = family
2863 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
2864 .expect("analytic flex gradient + hessian");
2865
2866 for u in 0..r {
2867 let grad_plus = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
2868 let grad_minus = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
2869 for v in 0..r {
2870 let fd = (grad_plus[v] - grad_minus[v]) / (2.0 * h);
2871 let analytic = analytic_hess[[v, u]];
2872 let denom = 1.0 + analytic.abs().max(fd.abs());
2873 let rel = (analytic - fd).abs() / denom;
2874 max_rel = max_rel.max(rel);
2875 assert!(
2876 rel <= 1e-6,
2877 "flex hand Hessian H[{v}][{u}] = {analytic:.6e} disagrees with central \
2878 FD of the gradient {fd:.6e} at row {row} (rel {rel:.3e}); a product-rule \
2879 term is dropped or mis-signed"
2880 );
2881 }
2882 }
2883 }
2884 assert!(
2886 max_rel <= 1e-6,
2887 "flex Hessian FD oracle max rel {max_rel:.3e}"
2888 );
2889 }
2890
2891 #[test]
2900 fn arbiter_flex_hessian_h00_fd_step_scaling() {
2901 let n = 12usize;
2902 let (family, states) = make_flex_oracle_family(n);
2903 let cache = family
2904 .build_exact_eval_cache(&states)
2905 .expect("flex exact eval cache");
2906 let primary = &cache.primary;
2907 let row = 2usize;
2908 let u = primary.q; let v = primary.q;
2910
2911 let row_ctx = BernoulliMarginalSlopeFamily::row_ctx(&cache, row);
2912 let (_neglog, _grad, analytic_hess) = family
2913 .compute_row_primary_gradient_hessian(row, &states, primary, row_ctx)
2914 .expect("analytic flex gradient + hessian");
2915 let analytic = analytic_hess[[v, u]];
2916
2917 let fd_at = |h: f64| -> f64 {
2918 let gp = flex_gradient_at_perturbed(&family, &states, primary, row, u, h);
2919 let gm = flex_gradient_at_perturbed(&family, &states, primary, row, u, -h);
2920 (gp[v] - gm[v]) / (2.0 * h)
2921 };
2922
2923 let h = 1e-3_f64;
2930 let fd_h = fd_at(h);
2931 let fd_half = fd_at(h * 0.5);
2932 let fd_quarter = fd_at(h * 0.25);
2933 let gap_h = (analytic - fd_h).abs();
2934 let gap_half = (analytic - fd_half).abs();
2935 let gap_quarter = (analytic - fd_quarter).abs();
2936 let rich = (4.0 * fd_half - fd_h) / 3.0;
2937 let rich_gap = (analytic - rich).abs();
2938 let denom = analytic.abs().max(1.0);
2939
2940 let record = format!(
2942 "FLEX H[0][0] ARBITER row 2: analytic={analytic:+.12e} \
2943 fd(h)={fd_h:+.12e} fd(h/2)={fd_half:+.12e} fd(h/4)={fd_quarter:+.12e} \
2944 gap(h)={gap_h:.3e} gap(h/2)={gap_half:.3e} gap(h/4)={gap_quarter:.3e} \
2945 ratio_h_over_half={:.3} ratio_half_over_quarter={:.3} \
2946 richardson={rich:+.12e} richardson_gap={rich_gap:.3e} (rich_rel={:.3e})",
2947 gap_h / gap_half.max(f64::MIN_POSITIVE),
2948 gap_half / gap_quarter.max(f64::MIN_POSITIVE),
2949 rich_gap / denom,
2950 );
2951
2952 assert!(
2958 rich_gap / denom <= 1e-7,
2959 "{record}\nVERDICT: Richardson residual exceeds the FD-truncation floor — \
2960 the hand H[0][0] genuinely diverges (real dropped/mis-signed term), NOT FD noise"
2961 );
2962 }
2963}