1use crate::estimate::EstimationError;
2use crate::rho_optimizer::{FallbackPolicy, OuterProblem};
3use faer::Side;
4use gam_linalg::faer_ndarray::{
5 FaerCholesky, FaerEigh, default_rrqr_rank_alpha, fast_ab, fast_atb, fast_xt_diag_x,
6 fast_xt_diag_y, rrqr_with_permutation,
7};
8use gam_problem::{
9 DeclaredHessianForm, Derivative, HessianValue, OuterEval, StationarityStandard,
10};
11use gam_terms::construction::CanonicalPenalty;
12use gam_terms::smooth::BlockwisePenalty;
13use ndarray::{Array1, Array2, Array3, ArrayView1, ArrayView2, ArrayView3, Axis, s};
14use opt::{RidgeSchedule, escalate_ridge};
15use rayon::prelude::*;
16use std::sync::Once;
17
18static ILL_CONDITIONED_BACKWARD_WARNED: Once = Once::new();
26
27fn warn_ill_conditioned_backward_once(p: usize, d: usize, condition_number: f64) {
28 ILL_CONDITIONED_BACKWARD_WARNED.call_once(|| {
29 log::warn!(
30 "gaussian_reml_fit_backward: K = XᵀWX + λS is near-singular \
31 (p={p}, d={d}, cond≈{condition_number:.2e}); returning zero gradients \
32 for this fit (λ has saturated, atom is effectively unused). \
33 Further occurrences are silent."
34 );
35 });
36}
37
38fn zero_backward_result(n: usize, p: usize, d: usize) -> GaussianRemlBackwardResult {
39 GaussianRemlBackwardResult {
40 grad_x: Array2::<f64>::zeros((n, p)),
41 grad_y: Array2::<f64>::zeros((n, d)),
42 grad_penalty: Array2::<f64>::zeros((p, p)),
43 grad_weights: Array1::<f64>::zeros(n),
44 }
45}
46
47pub const RHO_LOWER: f64 = -30.0;
53pub const RHO_UPPER: f64 = 30.0;
54const EIGEN_REL_TOL: f64 = 1.0e-10;
55const BLOCK_ORTHOGONAL_SCORE_TOL: f64 = 1.0e-7;
61const BLOCK_ORTHOGONAL_MAX_OUTER_PASSES: usize = 200;
65const BLOCK_ORTHOGONAL_BLOCK_UPDATES_PER_PASS: usize = 32;
69
70#[derive(Clone)]
87pub struct GaussianRemlBlocksDomain {
88 p_total: usize,
89 canonical_penalties: Vec<CanonicalPenalty>,
90 nullspace_dims: Vec<usize>,
91}
92
93impl GaussianRemlBlocksDomain {
94 pub fn from_blockwise_penalties(
95 p_total: usize,
96 penalties: &[BlockwisePenalty],
97 ) -> Result<Self, EstimationError> {
98 if p_total == 0 || penalties.is_empty() {
99 return Err(EstimationError::InvalidInput(
100 "block Gaussian REML domain requires at least one coefficient and one penalty block"
101 .to_string(),
102 ));
103 }
104
105 let mut canonical_penalties = Vec::with_capacity(penalties.len());
106 let mut nullspace_dims = Vec::with_capacity(penalties.len());
107 let mut expected_start = 0_usize;
108 for (block, penalty) in penalties.iter().enumerate() {
109 if penalty.col_range.start != expected_start
110 || penalty.col_range.end <= penalty.col_range.start
111 {
112 return Err(EstimationError::InvalidInput(format!(
113 "block Gaussian REML penalties must form a non-empty contiguous partition: \
114 block {block} has range {:?}, expected start {expected_start}",
115 penalty.col_range
116 )));
117 }
118 expected_start = penalty.col_range.end;
119
120 let spec = gam_terms::PenaltySpec::from_blockwise_ref(penalty);
121 let canonical = gam_terms::construction::canonicalize_penalty_spec(
122 &spec,
123 p_total,
124 block,
125 "block Gaussian REML domain",
126 )?
127 .ok_or_else(|| {
128 EstimationError::InvalidInput(format!(
129 "block Gaussian REML penalty {block} has no positive-curvature direction"
130 ))
131 })?;
132 let block_dim = canonical.block_dim();
133 let rank = canonical.rank();
134 if rank + canonical.nullity != block_dim {
135 return Err(EstimationError::InvalidInput(format!(
136 "block Gaussian REML penalty {block} is not positive semidefinite under the \
137 canonical spectral classification: rank={rank}, nullity={}, dimension={block_dim}",
138 canonical.nullity
139 )));
140 }
141 if canonical.positive_eigenvalues.len() != rank {
142 return Err(EstimationError::InvalidInput(format!(
143 "block Gaussian REML penalty {block} canonical root/eigenspectrum mismatch: \
144 root rank={rank}, positive eigenvalues={}",
145 canonical.positive_eigenvalues.len()
146 )));
147 }
148 nullspace_dims.push(canonical.nullity);
149 canonical_penalties.push(canonical);
150 }
151 if expected_start != p_total {
152 return Err(EstimationError::InvalidInput(format!(
153 "block Gaussian REML penalty partition ends at {expected_start}, \
154 but the joint design has {p_total} columns"
155 )));
156 }
157
158 Ok(Self {
159 p_total,
160 canonical_penalties,
161 nullspace_dims,
162 })
163 }
164
165 #[inline]
166 pub fn nullspace_dims(&self) -> &[usize] {
167 &self.nullspace_dims
168 }
169
170 fn local_penalties(&self) -> Vec<Array2<f64>> {
171 self.canonical_penalties
172 .iter()
173 .map(CanonicalPenalty::local_penalty)
174 .collect()
175 }
176
177 fn normal_matrix(
178 &self,
179 xtwx: &Array2<f64>,
180 lambdas: ArrayView1<'_, f64>,
181 ) -> Result<Array2<f64>, EstimationError> {
182 if xtwx.dim() != (self.p_total, self.p_total) {
183 return Err(EstimationError::InvalidInput(format!(
184 "block Gaussian REML Gram shape mismatch: expected {}x{}, got {}x{}",
185 self.p_total,
186 self.p_total,
187 xtwx.nrows(),
188 xtwx.ncols()
189 )));
190 }
191 if lambdas.len() != self.canonical_penalties.len() {
192 return Err(EstimationError::InvalidInput(format!(
193 "block Gaussian REML lambda count mismatch: expected {}, got {}",
194 self.canonical_penalties.len(),
195 lambdas.len()
196 )));
197 }
198 let mut normal = xtwx.clone();
199 for (block, penalty) in self.canonical_penalties.iter().enumerate() {
200 let lambda = lambdas[block];
201 if !lambda.is_finite() || lambda <= 0.0 {
202 return Err(EstimationError::InvalidInput(format!(
203 "block Gaussian REML lambda[{block}] must be finite and positive; got {lambda}"
204 )));
205 }
206 penalty.accumulate_weighted(&mut normal, lambda);
207 }
208 gam_linalg::matrix::symmetrize_in_place(&mut normal);
209 Ok(normal)
210 }
211
212 fn penalty_pseudoinverses(&self) -> Result<Vec<Array2<f64>>, EstimationError> {
219 let mut out = Vec::with_capacity(self.canonical_penalties.len());
220 for (block, penalty) in self.canonical_penalties.iter().enumerate() {
221 let k = penalty.block_dim();
222 let mut pinv = Array2::<f64>::zeros((k, k));
223 for (row, &eigenvalue) in penalty.positive_eigenvalues.iter().enumerate() {
224 if !eigenvalue.is_finite() || eigenvalue <= 0.0 {
225 return Err(EstimationError::InvalidInput(format!(
226 "block Gaussian REML penalty {block} has invalid canonical positive \
227 eigenvalue {row}: {eigenvalue}"
228 )));
229 }
230 let scale = 1.0 / (eigenvalue * eigenvalue);
231 for i in 0..k {
232 for j in 0..k {
233 pinv[[i, j]] += scale * penalty.root[[row, i]] * penalty.root[[row, j]];
234 }
235 }
236 }
237 if pinv.iter().any(|value| !value.is_finite()) {
238 return Err(EstimationError::InvalidInput(format!(
239 "block Gaussian REML penalty {block} canonical pseudoinverse is not representable"
240 )));
241 }
242 out.push(pinv);
243 }
244 Ok(out)
245 }
246
247 pub fn certify_joint_coefficient_map(
251 &self,
252 design: ArrayView2<'_, f64>,
253 weights: ArrayView1<'_, f64>,
254 lambdas: ArrayView1<'_, f64>,
255 ) -> Result<Array2<f64>, EstimationError> {
256 if design.ncols() != self.p_total || weights.len() != design.nrows() {
257 return Err(EstimationError::InvalidInput(format!(
258 "block Gaussian REML domain shape mismatch: design={}x{}, weights={}, coefficients={}",
259 design.nrows(),
260 design.ncols(),
261 weights.len(),
262 self.p_total
263 )));
264 }
265 if lambdas.len() != self.canonical_penalties.len() {
266 return Err(EstimationError::InvalidInput(format!(
267 "block Gaussian REML lambda count mismatch: expected {}, got {}",
268 self.canonical_penalties.len(),
269 lambdas.len()
270 )));
271 }
272 if let Some(((row, col), value)) =
273 design.indexed_iter().find(|(_, value)| !value.is_finite())
274 {
275 return Err(EstimationError::InvalidInput(format!(
276 "block Gaussian REML design[{row},{col}] must be finite; got {value}"
277 )));
278 }
279 if let Some((row, value)) = weights
280 .iter()
281 .enumerate()
282 .find(|(_, value)| !value.is_finite() || **value < 0.0)
283 {
284 return Err(EstimationError::InvalidInput(format!(
285 "block Gaussian REML weights[{row}] must be finite and non-negative; got {value}"
286 )));
287 }
288 if let Some((block, value)) = lambdas
289 .iter()
290 .enumerate()
291 .find(|(_, value)| !value.is_finite() || **value <= 0.0)
292 {
293 return Err(EstimationError::InvalidInput(format!(
294 "block Gaussian REML lambda[{block}] must be finite and positive; got {value}"
295 )));
296 }
297
298 let augmented_rows = design.nrows()
299 + self
300 .canonical_penalties
301 .iter()
302 .map(CanonicalPenalty::rank)
303 .sum::<usize>();
304 let mut augmented = Array2::<f64>::zeros((augmented_rows, self.p_total));
305 for row in 0..design.nrows() {
306 let scale = weights[row].sqrt();
307 for col in 0..self.p_total {
308 augmented[[row, col]] = scale * design[[row, col]];
309 }
310 }
311 let mut augmented_row = design.nrows();
312 for (block, penalty) in self.canonical_penalties.iter().enumerate() {
313 let scale = lambdas[block].sqrt();
314 for root_row in 0..penalty.rank() {
315 for local_col in 0..penalty.block_dim() {
316 augmented[[
317 augmented_row + root_row,
318 penalty.col_range.start + local_col,
319 ]] = scale * penalty.root[[root_row, local_col]];
320 }
321 }
322 augmented_row += penalty.rank();
323 }
324
325 let rank = rrqr_with_permutation(&augmented, default_rrqr_rank_alpha())
326 .map_err(|error| {
327 EstimationError::InvalidInput(format!(
328 "block Gaussian REML augmented-rank certificate failed: {error}"
329 ))
330 })?
331 .rank;
332 if rank != self.p_total {
333 return Err(EstimationError::InvalidInput(format!(
334 "block Gaussian REML joint coefficient map is not identified: \
335 augmented operator [sqrt(W)X; sqrt(lambda_k)R_k] has numerical \
336 rank {rank} < {}; constrain shared design/penalty-null directions \
337 before fitting",
338 self.p_total
339 )));
340 }
341
342 let xtwx = fast_xt_diag_x(&design, &weights);
343 let normal = self.normal_matrix(&xtwx, lambdas)?;
344 gam_linalg::utils::certified_spd_factorize(
345 &normal,
346 "block Gaussian REML penalized normal matrix",
347 )
348 .map_err(|error| {
349 EstimationError::InvalidInput(format!(
350 "block Gaussian REML requires an exact SPD penalized normal matrix: {error}"
351 ))
352 })?;
353 Ok(normal)
354 }
355}
356
357#[derive(Clone, Debug)]
358pub struct GaussianRemlBlocksResult {
359 pub coefficients: Array2<f64>,
360 pub fitted: Array2<f64>,
361 pub lambdas: Array1<f64>,
362 pub log_lambdas: Array1<f64>,
363 pub reml_score: f64,
364 pub edf: Array1<f64>,
365}
366
367struct GaussianRemlBlocksProfile {
368 domain: GaussianRemlBlocksDomain,
369 design: Array2<f64>,
370 weights: Array1<f64>,
371 y: Array1<f64>,
372 xtwx: Array2<f64>,
373 xtwy: Array1<f64>,
374 nu: f64,
375 observation_measure: TermDerivs,
376}
377
378struct GaussianRemlBlocksProfileEval {
379 cost: f64,
380 gradient: Array1<f64>,
381 hessian: Array2<f64>,
382 lambdas: Array1<f64>,
383 coefficients: Array1<f64>,
384 fitted: Array1<f64>,
385 edf: Array1<f64>,
386}
387
388impl GaussianRemlBlocksProfile {
389 fn evaluate(
390 &self,
391 rhos: ArrayView1<'_, f64>,
392 ) -> Result<GaussianRemlBlocksProfileEval, EstimationError> {
393 let f_blocks = self.domain.canonical_penalties.len();
394 if rhos.len() != f_blocks {
395 return Err(EstimationError::InvalidInput(format!(
396 "block Gaussian REML rho count mismatch: expected {f_blocks}, got {}",
397 rhos.len()
398 )));
399 }
400 let lambdas = Array1::from_vec(
401 gam_problem::checked_exp_log_strengths(rhos.iter().copied())
402 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?,
403 );
404 let normal = self.domain.normal_matrix(&self.xtwx, lambdas.view())?;
405 let inverse = gam_linalg::utils::certified_spd_inverse(
406 &normal,
407 "block Gaussian REML penalized normal matrix",
408 )
409 .map(gam_linalg::utils::CertifiedSpdInverse::into_inverse)
410 .map_err(|error| {
411 EstimationError::InvalidInput(format!(
412 "block Gaussian REML requires an exact SPD penalized normal matrix: {error}"
413 ))
414 })?;
415 let lower = normal
419 .cholesky(Side::Lower)
420 .map_err(|error| {
421 EstimationError::InvalidInput(format!(
422 "block Gaussian REML penalized normal log-determinant failed: {error}"
423 ))
424 })?
425 .lower_triangular();
426 let logdet_normal = 2.0 * lower.diag().iter().map(|value| value.ln()).sum::<f64>();
427 if !logdet_normal.is_finite() {
428 return Err(EstimationError::InvalidInput(
429 "block Gaussian REML penalized normal log-determinant is not finite".to_string(),
430 ));
431 }
432
433 let coefficients = inverse.dot(&self.xtwy);
434 let fitted = self.design.dot(&coefficients);
435 let residual = &self.y - &fitted;
436
437 let mut q = residual
442 .iter()
443 .zip(self.weights.iter())
444 .map(|(&value, &weight)| weight * value * value)
445 .sum::<f64>();
446 let mut logdet_penalty = 0.0_f64;
447 let mut p_betas = Vec::with_capacity(f_blocks);
448 let mut rp_matrices = Vec::with_capacity(f_blocks);
449 let mut b_values = Array1::<f64>::zeros(f_blocks);
450 let mut t_values = Array1::<f64>::zeros(f_blocks);
451 let mut edf = Array1::<f64>::zeros(f_blocks);
452 for (block, penalty) in self.domain.canonical_penalties.iter().enumerate() {
453 let start = penalty.col_range.start;
454 let end = penalty.col_range.end;
455 let beta_block = coefficients.slice(s![start..end]);
456 let local_p_beta = penalty.local.dot(&beta_block);
457 let lambda = lambdas[block];
458 let mut p_beta = Array1::<f64>::zeros(self.domain.p_total);
459 for local in 0..penalty.block_dim() {
460 p_beta[start + local] = lambda * local_p_beta[local];
461 }
462 let b_value = coefficients.dot(&p_beta);
463 q += b_value;
464 b_values[block] = b_value;
465
466 let weighted_penalty = penalty.local.mapv(|value| lambda * value);
467 let rp_block = inverse
468 .slice(s![.., start..end])
469 .dot(&weighted_penalty);
470 let mut rp = Array2::<f64>::zeros((self.domain.p_total, self.domain.p_total));
471 rp.slice_mut(s![.., start..end]).assign(&rp_block);
472 let trace = (0..penalty.block_dim())
473 .map(|local| rp_block[[start + local, local]])
474 .sum::<f64>();
475 t_values[block] = trace;
476 edf[block] = penalty.block_dim() as f64 - trace;
477 logdet_penalty += penalty
478 .positive_eigenvalues
479 .iter()
480 .map(|eigenvalue| eigenvalue.ln())
481 .sum::<f64>()
482 + penalty.rank() as f64 * rhos[block];
483 p_betas.push(p_beta);
484 rp_matrices.push(rp);
485 }
486 if !q.is_finite() || q <= 0.0 {
487 return Err(EstimationError::InvalidInput(format!(
488 "block Gaussian REML profiled residual quadratic form must be finite and positive; got {q}"
489 )));
490 }
491 if !logdet_penalty.is_finite() {
492 return Err(EstimationError::InvalidInput(
493 "block Gaussian REML penalty pseudo-log-determinant is not finite".to_string(),
494 ));
495 }
496
497 let tau = self.nu / q;
498 let tau_q = -self.nu / (q * q);
499 let cost = 0.5
500 * (self.nu
501 * (1.0 + (2.0 * std::f64::consts::PI * q / self.nu).ln())
502 + logdet_normal
503 - logdet_penalty)
504 + self.observation_measure.value;
505 let mut gradient = Array1::<f64>::zeros(f_blocks);
506 for block in 0..f_blocks {
507 gradient[block] = 0.5
508 * (t_values[block]
509 - self.domain.canonical_penalties[block].rank() as f64
510 + tau * b_values[block]);
511 }
512
513 let mut hessian = Array2::<f64>::zeros((f_blocks, f_blocks));
514 for k in 0..f_blocks {
515 for j in 0..f_blocks {
516 let trace_pair = gam_linalg::utils::trace_of_product(
517 rp_matrices[k].view(),
518 rp_matrices[j].view(),
519 );
520 let beta_pk_r_pj_beta = p_betas[k].dot(&inverse.dot(&p_betas[j]));
521 hessian[[k, j]] = 0.5
522 * ((if k == j { t_values[k] } else { 0.0 }) - trace_pair
523 + tau_q * b_values[k] * b_values[j]
524 + tau
525 * ((if k == j { b_values[k] } else { 0.0 })
526 - 2.0 * beta_pk_r_pj_beta));
527 }
528 }
529 gam_linalg::matrix::symmetrize_in_place(&mut hessian);
530 if !cost.is_finite()
531 || coefficients.iter().any(|value| !value.is_finite())
532 || fitted.iter().any(|value| !value.is_finite())
533 || edf.iter().any(|value| !value.is_finite())
534 || gradient.iter().any(|value| !value.is_finite())
535 || hessian.iter().any(|value| !value.is_finite())
536 {
537 return Err(EstimationError::InvalidInput(
538 "block Gaussian REML profile evaluation produced a non-finite value".to_string(),
539 ));
540 }
541
542 Ok(GaussianRemlBlocksProfileEval {
543 cost,
544 gradient,
545 hessian,
546 lambdas,
547 coefficients,
548 fitted,
549 edf,
550 })
551 }
552}
553
554fn gaussian_reml_blocks_profile_cost(
555 state: &mut GaussianRemlBlocksProfile,
556 rhos: &Array1<f64>,
557) -> Result<f64, EstimationError> {
558 Ok(state.evaluate(rhos.view())?.cost)
559}
560
561fn gaussian_reml_blocks_profile_outer_eval(
562 state: &mut GaussianRemlBlocksProfile,
563 rhos: &Array1<f64>,
564) -> Result<OuterEval, EstimationError> {
565 let evaluated = state.evaluate(rhos.view())?;
566 Ok(OuterEval {
567 cost: evaluated.cost,
568 gradient: evaluated.gradient,
569 hessian: HessianValue::Dense(evaluated.hessian),
570 inner_beta_hint: Some(evaluated.coefficients),
571 })
572}
573
574pub fn gaussian_reml_fit_blocks_exact(
587 designs: &[Array2<f64>],
588 penalties: &[Array2<f64>],
589 y: ArrayView1<'_, f64>,
590 weights: Option<ArrayView1<'_, f64>>,
591 init_rhos: Option<&[f64]>,
592) -> Result<GaussianRemlBlocksResult, EstimationError> {
593 let f_blocks = designs.len();
594 if f_blocks == 0 || penalties.len() != f_blocks {
595 return Err(EstimationError::InvalidInput(format!(
596 "exact block Gaussian REML requires equal non-zero design and penalty block counts; \
597 got designs={}, penalties={}",
598 f_blocks,
599 penalties.len()
600 )));
601 }
602 if let Some(rhos) = init_rhos {
603 if rhos.len() != f_blocks {
604 return Err(EstimationError::InvalidInput(format!(
605 "exact block Gaussian REML init_rhos length mismatch: expected {f_blocks}, got {}",
606 rhos.len()
607 )));
608 }
609 if let Some((block, value)) = rhos
610 .iter()
611 .enumerate()
612 .find(|(_, value)| !value.is_finite())
613 {
614 return Err(EstimationError::InvalidInput(format!(
615 "exact block Gaussian REML init_rhos[{block}] must be finite; got {value}"
616 )));
617 }
618 }
619
620 let n = y.len();
621 if n == 0 {
622 return Err(EstimationError::InvalidInput(
623 "exact block Gaussian REML requires at least one observation".to_string(),
624 ));
625 }
626 if let Some((row, value)) = y.iter().enumerate().find(|(_, value)| !value.is_finite()) {
627 return Err(EstimationError::InvalidInput(format!(
628 "exact block Gaussian REML y[{row}] must be finite; got {value}"
629 )));
630 }
631
632 let mut offsets = Vec::with_capacity(f_blocks + 1);
633 offsets.push(0_usize);
634 let mut p_total = 0_usize;
635 for (block, (design, penalty)) in designs.iter().zip(penalties.iter()).enumerate() {
636 if design.nrows() != n {
637 return Err(EstimationError::InvalidInput(format!(
638 "exact block Gaussian REML designs[{block}] has {} rows, expected {n}",
639 design.nrows()
640 )));
641 }
642 if design.ncols() == 0 || penalty.dim() != (design.ncols(), design.ncols()) {
643 return Err(EstimationError::InvalidInput(format!(
644 "exact block Gaussian REML block {block} requires a non-empty square penalty \
645 matching its {} design columns; got {}x{}",
646 design.ncols(),
647 penalty.nrows(),
648 penalty.ncols()
649 )));
650 }
651 if let Some(((row, col), value)) =
652 design.indexed_iter().find(|(_, value)| !value.is_finite())
653 {
654 return Err(EstimationError::InvalidInput(format!(
655 "exact block Gaussian REML designs[{block}][{row},{col}] must be finite; got {value}"
656 )));
657 }
658 if let Some(((row, col), value)) =
659 penalty.indexed_iter().find(|(_, value)| !value.is_finite())
660 {
661 return Err(EstimationError::InvalidInput(format!(
662 "exact block Gaussian REML penalties[{block}][{row},{col}] must be finite; got {value}"
663 )));
664 }
665 p_total += design.ncols();
666 offsets.push(p_total);
667 }
668
669 let weight = gaussian_reml_weights(n, weights)?;
670 let mut design = Array2::<f64>::zeros((n, p_total));
671 let mut blockwise_penalties = Vec::with_capacity(f_blocks);
672 let mut canonical_keys = Vec::with_capacity(f_blocks);
673 for block in 0..f_blocks {
674 design
675 .slice_mut(s![.., offsets[block]..offsets[block + 1]])
676 .assign(&designs[block]);
677 blockwise_penalties.push(BlockwisePenalty::new(
678 offsets[block]..offsets[block + 1],
679 penalties[block].clone(),
680 ));
681 canonical_keys.push(fnv1a_mix(
682 matrix_fingerprint(designs[block].view()),
683 matrix_fingerprint(penalties[block].view()),
684 ));
685 }
686 let domain =
687 GaussianRemlBlocksDomain::from_blockwise_penalties(p_total, &blockwise_penalties)?;
688 let unit_lambdas = Array1::<f64>::ones(f_blocks);
689 domain.certify_joint_coefficient_map(design.view(), weight.view(), unit_lambdas.view())?;
690
691 let n_effective = effective_observation_count(weight.view());
692 let nullity = domain.nullspace_dims.iter().sum::<usize>();
693 if n_effective <= nullity {
694 return Err(EstimationError::InvalidInput(format!(
695 "exact block Gaussian REML requires more positive-weight rows than total penalty \
696 nullity; got n_effective={n_effective}, nullity={nullity}"
697 )));
698 }
699
700 if f_blocks == 1 {
701 let xtwx = fast_xt_diag_x(&design.view(), &weight.view());
705 gam_linalg::utils::certified_spd_factorize(
706 &xtwx,
707 "one-block Gaussian REML unpenalized normal matrix",
708 )
709 .map_err(|error| {
710 EstimationError::InvalidInput(format!(
711 "one-block Gaussian REML requires an exact SPD unpenalized normal matrix: {error}"
712 ))
713 })?;
714 let scalar = gaussian_reml_closed_form(
715 design.view(),
716 y,
717 penalties[0].view(),
718 Some(weight.view()),
719 init_rhos.map(|rhos| rhos[0]),
720 )?;
721 let lambdas = Array1::from_elem(1, scalar.lambda);
722 domain.certify_joint_coefficient_map(design.view(), weight.view(), lambdas.view())?;
723 return Ok(GaussianRemlBlocksResult {
724 coefficients: scalar.coefficients.insert_axis(Axis(1)),
725 fitted: scalar.fitted.insert_axis(Axis(1)),
726 lambdas,
727 log_lambdas: Array1::from_elem(1, scalar.rho),
728 reml_score: scalar.reml_score,
729 edf: Array1::from_elem(1, scalar.edf),
730 });
731 }
732
733 let xtwx = fast_xt_diag_x(&design.view(), &weight.view());
734 let y_owned = y.to_owned();
735 let y_matrix = y_owned.view().insert_axis(Axis(1));
736 let xtwy = fast_xt_diag_y(&design.view(), &weight.view(), &y_matrix)
737 .column(0)
738 .to_owned();
739 let profile = GaussianRemlBlocksProfile {
740 domain,
741 design,
742 observation_measure: gaussian_reml_observation_measure(weight.view(), 1),
743 weights: weight,
744 y: y_owned,
745 xtwx,
746 xtwy,
747 nu: (n_effective - nullity) as f64,
748 };
749
750 let mut seed_config = gam_problem::SeedConfig::default();
751 seed_config.bounds = (RHO_LOWER, RHO_UPPER);
752 seed_config.risk_profile = gam_problem::SeedRiskProfile::Gaussian;
753 let mut problem = OuterProblem::new(f_blocks)
754 .with_gradient(Derivative::Analytic)
755 .with_hessian(DeclaredHessianForm::Dense)
756 .with_prefer_gradient_only(false)
757 .with_disable_fixed_point(true)
758 .with_tolerance(1.0e-10)
759 .with_required_projected_gradient_norm(Some(1.0e-8))
760 .with_max_iter(200)
761 .with_bounds(
762 Array1::from_elem(f_blocks, RHO_LOWER),
763 Array1::from_elem(f_blocks, RHO_UPPER),
764 )
765 .with_rho_bound(RHO_UPPER)
766 .with_seed_config(seed_config)
767 .with_rho_canonical_keys(Some(canonical_keys))
768 .with_fallback_policy(FallbackPolicy::Disabled)
769 .with_problem_size(n, p_total);
770 if let Some(rhos) = init_rhos {
771 problem = problem
772 .with_initial_rho(Array1::from_iter(
773 rhos
774 .iter()
775 .map(|rho| rho.clamp(RHO_LOWER, RHO_UPPER)),
776 ))
777 .with_screen_initial_rho(true);
778 }
779 let mut objective = problem.build_objective(
780 profile,
781 gaussian_reml_blocks_profile_cost,
782 gaussian_reml_blocks_profile_outer_eval,
783 None::<fn(&mut GaussianRemlBlocksProfile)>,
784 None::<
785 fn(
786 &mut GaussianRemlBlocksProfile,
787 &Array1<f64>,
788 ) -> Result<gam_problem::EfsEval, EstimationError>,
789 >,
790 );
791 let optimum = problem.run(&mut objective, "exact block Gaussian REML")?;
792 let final_eval = objective.state.evaluate(optimum.rho.view())?;
793 objective.state.domain.certify_joint_coefficient_map(
794 objective.state.design.view(),
795 objective.state.weights.view(),
796 final_eval.lambdas.view(),
797 )?;
798
799 Ok(GaussianRemlBlocksResult {
800 coefficients: final_eval.coefficients.insert_axis(Axis(1)),
801 fitted: final_eval.fitted.insert_axis(Axis(1)),
802 lambdas: final_eval.lambdas,
803 log_lambdas: optimum.rho,
804 reml_score: final_eval.cost,
805 edf: final_eval.edf,
806 })
807}
808
809#[derive(Clone, Copy)]
810struct BlockOrthogonalControls {
811 score_tol: f64,
812 max_outer_passes: usize,
813 block_updates_per_pass: usize,
814}
815
816impl Default for BlockOrthogonalControls {
817 fn default() -> Self {
818 Self {
819 score_tol: BLOCK_ORTHOGONAL_SCORE_TOL,
820 max_outer_passes: BLOCK_ORTHOGONAL_MAX_OUTER_PASSES,
821 block_updates_per_pass: BLOCK_ORTHOGONAL_BLOCK_UPDATES_PER_PASS,
822 }
823 }
824}
825
826fn canonicalize_penalty(penalty: ArrayView2<'_, f64>) -> Array2<f64> {
837 let p = penalty.nrows();
838 let mut out = penalty.to_owned();
839 for i in 0..p {
840 for j in (i + 1)..p {
841 let avg = 0.5 * (out[[i, j]] + out[[j, i]]);
842 out[[i, j]] = avg;
843 out[[j, i]] = avg;
844 }
845 }
846 out
847}
848
849#[derive(Clone, Debug)]
850pub struct GaussianRemlEigenCache {
851 pub penalty_eigenvalues: Array1<f64>,
852 pub eigenvectors: Array2<f64>,
853 pub coefficient_basis: Array2<f64>,
854 pub xtwx_fingerprint: u64,
855 pub penalty_fingerprint: u64,
856 pub logdet_xtwx: f64,
857 pub logdet_penalty_positive: f64,
858 pub penalty_rank: usize,
859 pub nullity: usize,
860}
861
862#[derive(Clone, Debug, Default)]
863pub struct GaussianRemlWarmStart {
864 pub lambda: Option<f64>,
865 pub eigen_cache: Option<GaussianRemlEigenCache>,
866}
867
868#[derive(Clone, Debug)]
869pub struct GaussianRemlResult {
870 pub lambda: f64,
871 pub rho: f64,
872 pub coefficients: Array1<f64>,
873 pub fitted: Array1<f64>,
874 pub reml_score: f64,
875 pub reml_grad_lambda: f64,
876 pub reml_hess_lambda: f64,
877 pub reml_grad_rho: f64,
878 pub reml_hess_rho: f64,
879 pub edf: f64,
880 pub sigma2: f64,
881 pub cache: GaussianRemlEigenCache,
882}
883
884#[derive(Clone, Debug)]
885pub struct GaussianRemlMultiResult {
886 pub lambda: f64,
887 pub rho: f64,
888 pub coefficients: Array2<f64>,
889 pub fitted: Array2<f64>,
890 pub reml_score: f64,
891 pub reml_score_roundoff: Option<f64>,
901 pub reml_grad_lambda: f64,
902 pub reml_hess_lambda: f64,
903 pub reml_grad_rho: f64,
904 pub reml_hess_rho: f64,
905 pub edf: f64,
906 pub sigma2: Array1<f64>,
907 pub cache: GaussianRemlEigenCache,
908}
909
910#[derive(Clone, Debug)]
911pub struct GaussianRemlFreeBScore {
912 pub reml_score: f64,
913 pub grad_coefficients: Array2<f64>,
914 pub grad_penalty: Array2<f64>,
915 pub grad_log_lambda: f64,
916 pub fitted: Array2<f64>,
917 pub sigma2: Array1<f64>,
918 pub edf: f64,
919}
920
921#[derive(Clone, Debug)]
922pub struct GaussianRemlBackwardResult {
923 pub grad_x: Array2<f64>,
924 pub grad_y: Array2<f64>,
925 pub grad_penalty: Array2<f64>,
926 pub grad_weights: Array1<f64>,
930}
931
932#[derive(Clone, Debug)]
933pub struct GaussianRemlMultiBackwardProblem<'a> {
934 pub x: ArrayView2<'a, f64>,
935 pub y: ArrayView2<'a, f64>,
936 pub weights: Option<ArrayView1<'a, f64>>,
937 pub fit: &'a GaussianRemlMultiResult,
938 pub grad_lambda: f64,
939 pub grad_coefficients: Option<ArrayView2<'a, f64>>,
940 pub grad_fitted: Option<ArrayView2<'a, f64>>,
941 pub grad_reml_score: f64,
942 pub grad_edf: f64,
943}
944
945#[derive(Clone, Debug)]
946pub struct GaussianRemlNoAllocWorkspace {
947 pub xtwy: Array2<f64>,
948 pub ywy: Array1<f64>,
949 pub projected_rhs: Array2<f64>,
950 pub projected_rhs_squared: Array2<f64>,
951 pub scaled_projected_rhs: Array2<f64>,
952}
953
954impl GaussianRemlNoAllocWorkspace {
955 pub fn new(n_coefficients: usize, n_outputs: usize) -> Self {
956 Self {
957 xtwy: Array2::zeros((n_coefficients, n_outputs)),
958 ywy: Array1::zeros(n_outputs),
959 projected_rhs: Array2::zeros((n_coefficients, n_outputs)),
960 projected_rhs_squared: Array2::zeros((n_coefficients, n_outputs)),
961 scaled_projected_rhs: Array2::zeros((n_coefficients, n_outputs)),
962 }
963 }
964
965}
966
967#[derive(Clone, Copy, Debug)]
968pub struct GaussianRemlNoAllocFit {
969 pub lambda: f64,
970 pub rho: f64,
971 pub reml_score: f64,
972 pub reml_grad_lambda: f64,
973 pub reml_hess_lambda: f64,
974 pub reml_grad_rho: f64,
975 pub reml_hess_rho: f64,
976 pub edf: f64,
977}
978
979#[derive(Clone, Debug)]
980pub struct GaussianRemlMultiBatchProblem<'a> {
981 pub x: ArrayView2<'a, f64>,
982 pub y: ArrayView2<'a, f64>,
983 pub weights: Option<ArrayView1<'a, f64>>,
984 pub init_rho: Option<f64>,
985}
986
987#[derive(Clone, Debug)]
988pub struct GaussianRemlBlockOrthogonalResult {
989 pub coefficients: Vec<Array2<f64>>,
990 pub fitted: Array2<f64>,
991 pub lambdas: Array1<f64>,
992 pub log_lambdas: Array1<f64>,
993 pub reml_score: f64,
994 pub edf: Array1<f64>,
995}
996
997#[derive(Clone)]
998struct GaussianRemlPrepared {
999 cache: GaussianRemlEigenCache,
1000 ywy: Array1<f64>,
1001 projected_rhs_squared: Array2<f64>,
1002 projected_rhs: Array2<f64>,
1003 n_effective: usize,
1007 n_outputs: usize,
1008 observation_measure: TermDerivs,
1011}
1012
1013#[derive(Clone, Copy)]
1014struct ObjectiveEval {
1015 cost: f64,
1016 grad: f64,
1017 hess: f64,
1018 edf: f64,
1019 cost_roundoff: f64,
1025}
1026
1027const UNIT_ROUNDOFF: f64 = 0.5 * f64::EPSILON;
1030
1031fn roundoff_growth(operation_count: usize) -> f64 {
1035 let accumulated = operation_count as f64 * UNIT_ROUNDOFF;
1036 if accumulated < 1.0 {
1037 accumulated / (1.0 - accumulated)
1038 } else {
1039 f64::INFINITY
1040 }
1041}
1042
1043#[derive(Clone, Copy)]
1054struct TermDerivs {
1055 value: f64,
1056 grad: f64,
1057 hess: f64,
1058 roundoff: f64,
1063}
1064
1065fn gaussian_reml_observation_measure(weights: ArrayView1<'_, f64>, n_outputs: usize) -> TermDerivs {
1071 let mut logdet = 0.0;
1072 let mut magnitude = 0.0;
1073 let mut active = 0_usize;
1074 for &weight in weights {
1075 if weight > 0.0 {
1076 let term = weight.ln();
1077 logdet += term;
1078 magnitude += term.abs();
1079 active += 1;
1080 }
1081 }
1082 let scale = -0.5 * n_outputs as f64;
1083 TermDerivs {
1084 value: scale * logdet,
1085 grad: 0.0,
1086 hess: 0.0,
1087 roundoff: scale.abs()
1088 * roundoff_growth(active.saturating_mul(2).saturating_add(2))
1089 * magnitude,
1090 }
1091}
1092
1093fn finish_gaussian_reml_weight_vjp(
1098 weights: ArrayView1<'_, f64>,
1099 n_outputs: usize,
1100 upstream_score: f64,
1101 gradient: &mut Array1<f64>,
1102) {
1103 let scale = -0.5 * n_outputs as f64 * upstream_score;
1104 for (&weight, value) in weights.iter().zip(gradient.iter_mut()) {
1105 if weight > 0.0 {
1106 *value += scale / weight;
1107 } else {
1108 *value = 0.0;
1109 }
1110 }
1111}
1112
1113#[derive(Clone, Copy)]
1121struct ModalKernels {
1122 log_one_plus_t: f64,
1123 u: f64,
1125 v: f64,
1127 w: f64,
1129 k: f64,
1131}
1132
1133fn modal_kernels(rho: f64, delta: f64) -> ModalKernels {
1134 if delta == 0.0 {
1135 return ModalKernels {
1136 log_one_plus_t: 0.0,
1137 u: 0.0,
1138 v: 1.0,
1139 w: 0.0,
1140 k: 0.0,
1141 };
1142 }
1143 let log_t = rho + delta.ln();
1144 let (log_one_plus_t, u, v) = if log_t >= 0.0 {
1145 let reciprocal_t = (-log_t).exp();
1146 let v = reciprocal_t / (1.0 + reciprocal_t);
1147 (log_t + reciprocal_t.ln_1p(), 1.0 - v, v)
1148 } else {
1149 let t = log_t.exp();
1150 let u = t / (1.0 + t);
1151 (t.ln_1p(), u, 1.0 - u)
1152 };
1153 let w = u * v;
1154 ModalKernels {
1155 log_one_plus_t,
1156 u,
1157 v,
1158 w,
1159 k: w * (v - u),
1160 }
1161}
1162
1163impl std::ops::AddAssign<TermDerivs> for ObjectiveEval {
1164 fn add_assign(&mut self, rhs: TermDerivs) {
1167 self.cost += rhs.value;
1168 self.grad += rhs.grad;
1169 self.hess += rhs.hess;
1170 self.cost_roundoff += rhs.roundoff + UNIT_ROUNDOFF * self.cost.abs();
1173 }
1174}
1175
1176fn gaussian_reml_logdet_term(
1182 cache: &GaussianRemlEigenCache,
1183 rho: f64,
1184 n_outputs: f64,
1185) -> (TermDerivs, f64) {
1186 let mut logdet_h = cache.logdet_xtwx;
1187 let mut trace_h = 0.0;
1188 let mut trace_h_deriv = 0.0;
1189 let mut edf = 0.0;
1190 let mut logdet_magnitude = cache.logdet_xtwx.abs();
1196 for delta in PenaltyRangeSpectrum::of(cache).iter() {
1200 let mode = modal_kernels(rho, delta);
1201 logdet_h += mode.log_one_plus_t;
1202 logdet_magnitude += mode.log_one_plus_t.abs();
1203 if delta > 0.0 {
1204 trace_h += mode.u;
1205 trace_h_deriv += mode.w;
1206 }
1207 edf += mode.v;
1208 }
1209 let logdet_s = cache.logdet_penalty_positive + (cache.penalty_rank as f64) * rho;
1210 logdet_magnitude += cache.logdet_penalty_positive.abs() + logdet_s.abs();
1211 let value = 0.5 * n_outputs * (logdet_h - logdet_s);
1212 let operation_count = cache
1215 .penalty_eigenvalues
1216 .len()
1217 .saturating_mul(2)
1218 .saturating_add(5);
1219 let term = TermDerivs {
1220 value,
1221 grad: 0.5 * n_outputs * (trace_h - cache.penalty_rank as f64),
1222 hess: 0.5 * n_outputs * trace_h_deriv,
1223 roundoff: 0.5 * n_outputs * roundoff_growth(operation_count) * logdet_magnitude,
1224 };
1225 (term, edf)
1226}
1227
1228#[inline]
1256fn dispersion_residual_parts(
1257 cache: &GaussianRemlEigenCache,
1258 ywy: ArrayView1<'_, f64>,
1259 projected_rhs_squared: ArrayView2<'_, f64>,
1260 output: usize,
1261 rho: f64,
1262) -> DispersionResidualParts {
1263 let mut total_c2 = 0.0;
1264 let mut penalized_residual = 0.0;
1265 let mut dp_grad = 0.0;
1266 let mut dp_hess = 0.0;
1267 let spectrum = PenaltyRangeSpectrum::of(cache);
1268 for eig in 0..spectrum.len() {
1269 let c2 = projected_rhs_squared[[eig, output]];
1270 let mode = modal_kernels(rho, spectrum.get(eig));
1271 total_c2 += c2;
1272 penalized_residual += c2 * mode.u;
1273 dp_grad += c2 * mode.w;
1274 dp_hess += c2 * mode.k;
1275 }
1276 let unpenalized_residual = (ywy[output] - total_c2).max(0.0);
1280 DispersionResidualParts {
1281 unpenalized_residual,
1282 penalized_residual,
1283 dp_grad,
1284 dp_hess,
1285 total_c2,
1286 }
1287}
1288
1289#[derive(Clone, Copy)]
1297struct DispersionResidualParts {
1298 unpenalized_residual: f64,
1299 penalized_residual: f64,
1300 dp_grad: f64,
1301 dp_hess: f64,
1302 total_c2: f64,
1303}
1304
1305fn gaussian_reml_dispersion_term(
1312 cache: &GaussianRemlEigenCache,
1313 ywy: ArrayView1<'_, f64>,
1314 projected_rhs_squared: ArrayView2<'_, f64>,
1315 output: usize,
1316 nu: f64,
1317 rho: f64,
1318) -> TermDerivs {
1319 let parts = dispersion_residual_parts(cache, ywy, projected_rhs_squared, output, rho);
1320 let dp = parts.unpenalized_residual + parts.penalized_residual;
1321 let value = 0.5 * nu * (1.0 + (2.0 * std::f64::consts::PI * dp / nu).ln());
1322 let operation_count = cache
1327 .penalty_eigenvalues
1328 .len()
1329 .saturating_mul(3)
1330 .saturating_add(3);
1331 let dp_magnitude = ywy[output].abs() + parts.total_c2.abs() + parts.penalized_residual.abs();
1332 let dp_roundoff = roundoff_growth(operation_count) * dp_magnitude;
1333 TermDerivs {
1338 value,
1339 grad: 0.5 * nu * parts.dp_grad / dp,
1340 hess: 0.5 * nu * (parts.dp_hess / dp - (parts.dp_grad * parts.dp_grad) / (dp * dp)),
1341 roundoff: 0.5 * nu * (dp_roundoff / dp) + roundoff_growth(4) * value.abs(),
1342 }
1343}
1344
1345pub fn gaussian_reml_closed_form(
1346 x: ArrayView2<'_, f64>,
1347 y: ArrayView1<'_, f64>,
1348 penalty: ArrayView2<'_, f64>,
1349 weights: Option<ArrayView1<'_, f64>>,
1350 init_rho: Option<f64>,
1351) -> Result<GaussianRemlResult, EstimationError> {
1352 gaussian_reml_closed_form_with_nullspace_dim(x, y, penalty, None, weights, init_rho)
1353}
1354
1355pub fn gaussian_reml_closed_form_with_nullspace_dim(
1356 x: ArrayView2<'_, f64>,
1357 y: ArrayView1<'_, f64>,
1358 penalty: ArrayView2<'_, f64>,
1359 nullspace_dim: Option<usize>,
1360 weights: Option<ArrayView1<'_, f64>>,
1361 init_rho: Option<f64>,
1362) -> Result<GaussianRemlResult, EstimationError> {
1363 let y2 = y.insert_axis(Axis(1));
1364 let result = gaussian_reml_multi_closed_form_with_nullspace_dim(
1365 x,
1366 y2,
1367 penalty,
1368 nullspace_dim,
1369 weights,
1370 init_rho,
1371 )?;
1372 scalar_result_from_multi(result)
1373}
1374
1375fn scalar_result_from_multi(
1376 result: GaussianRemlMultiResult,
1377) -> Result<GaussianRemlResult, EstimationError> {
1378 Ok(GaussianRemlResult {
1379 lambda: result.lambda,
1380 rho: result.rho,
1381 coefficients: result.coefficients.column(0).to_owned(),
1382 fitted: result.fitted.column(0).to_owned(),
1383 reml_score: result.reml_score,
1384 reml_grad_lambda: result.reml_grad_lambda,
1385 reml_hess_lambda: result.reml_hess_lambda,
1386 reml_grad_rho: result.reml_grad_rho,
1387 reml_hess_rho: result.reml_hess_rho,
1388 edf: result.edf,
1389 sigma2: result.sigma2[0],
1390 cache: result.cache,
1391 })
1392}
1393
1394#[derive(Clone, Debug)]
1400pub struct GaussianRemlPointEval {
1401 pub rho: f64,
1402 pub lambda: f64,
1403 pub reml_score: f64,
1404 pub edf: f64,
1405 pub sigma2: f64,
1406 pub coefficients: Array1<f64>,
1407}
1408
1409#[derive(Clone, Debug)]
1424pub struct GaussianRemlStationarySet {
1425 pub roots: Vec<f64>,
1426 pub root_brackets: Vec<[f64; 2]>,
1427 pub root_gradients: Vec<f64>,
1428 pub selected_rho: f64,
1429 pub endpoint_costs: [f64; 2],
1430 pub rho_window: [f64; 2],
1431 pub root_location_resolution: f64,
1432}
1433
1434pub fn gaussian_reml_stationary_set(
1440 x: ArrayView2<'_, f64>,
1441 y: ArrayView1<'_, f64>,
1442 penalty: ArrayView2<'_, f64>,
1443 nullspace_dim: Option<usize>,
1444 weights: Option<ArrayView1<'_, f64>>,
1445 init_rho: Option<f64>,
1446) -> Result<GaussianRemlStationarySet, EstimationError> {
1447 if init_rho.is_some_and(|rho| !rho.is_finite()) {
1448 crate::bail_invalid_estim!("Gaussian REML stationary search requires a finite rho hint");
1449 }
1450 let y2 = y.insert_axis(Axis(1));
1451 let prepared = prepare_gaussian_reml(x, y2.view(), penalty, nullspace_dim, weights, None)?;
1452 let endpoint_costs = [
1453 prepared.evaluate(RHO_LOWER).cost,
1454 prepared.evaluate(RHO_UPPER).cost,
1455 ];
1456 validate_reml_profile_residuals(
1457 &prepared.cache,
1458 prepared.ywy.view(),
1459 prepared.projected_rhs_squared.view(),
1460 RHO_LOWER,
1461 )?;
1462 if prepared.cache.penalty_rank == 0 {
1463 return Ok(GaussianRemlStationarySet {
1464 roots: Vec::new(),
1465 root_brackets: Vec::new(),
1466 root_gradients: Vec::new(),
1467 selected_rho: init_rho.unwrap_or(0.0).clamp(RHO_LOWER, RHO_UPPER),
1468 endpoint_costs,
1469 rho_window: [RHO_LOWER, RHO_UPPER],
1470 root_location_resolution: RHO_BRACKET_RESOLUTION,
1471 });
1472 }
1473 let eval = |rho: f64| prepared.evaluate(rho);
1474 let enclose = |a: f64, b: f64| {
1475 reml_deriv_enclosure(
1476 &prepared.cache,
1477 prepared.ywy.view(),
1478 prepared.projected_rhs_squared.view(),
1479 prepared.n_effective,
1480 prepared.n_outputs,
1481 a,
1482 b,
1483 )
1484 };
1485 let mut roots = Vec::new();
1486 let mut root_brackets = Vec::new();
1487 let mut root_gradients = Vec::new();
1488 let selection = {
1489 let mut observer = |root: StationaryRoot, e: &ObjectiveEval| {
1490 roots.push(root.rho);
1491 root_brackets.push(root.bracket);
1492 root_gradients.push(e.grad);
1493 };
1494 enumerate_and_select_rho(&eval, &enclose, init_rho, Some(&mut observer))?
1495 };
1496 Ok(GaussianRemlStationarySet {
1497 roots,
1498 root_brackets,
1499 root_gradients,
1500 selected_rho: selection.rho,
1501 endpoint_costs,
1502 rho_window: [RHO_LOWER, RHO_UPPER],
1503 root_location_resolution: RHO_BRACKET_RESOLUTION,
1504 })
1505}
1506
1507pub fn gaussian_reml_multi_closed_form(
1508 x: ArrayView2<'_, f64>,
1509 y: ArrayView2<'_, f64>,
1510 penalty: ArrayView2<'_, f64>,
1511 weights: Option<ArrayView1<'_, f64>>,
1512 init_rho: Option<f64>,
1513) -> Result<GaussianRemlMultiResult, EstimationError> {
1514 gaussian_reml_multi_closed_form_with_nullspace_dim(x, y, penalty, None, weights, init_rho)
1515}
1516
1517pub fn gaussian_reml_multi_shared_dispersion_closed_form(
1535 x: ArrayView2<'_, f64>,
1536 y: ArrayView2<'_, f64>,
1537 penalty: ArrayView2<'_, f64>,
1538 weights: Option<ArrayView1<'_, f64>>,
1539 init_rho: Option<f64>,
1540) -> Result<GaussianRemlMultiResult, EstimationError> {
1541 if y.ncols() == 0 {
1542 crate::bail_invalid_estim!(
1543 "shared-dispersion Gaussian REML requires at least one response column"
1544 );
1545 }
1546 let prepared = prepare_gaussian_reml(x, y, penalty, None, weights, None)?;
1547 let init_rho = init_rho
1548 .map(f64::exp)
1549 .map(validate_initial_lambda)
1550 .transpose()?
1551 .map(f64::ln);
1552 let d = prepared.n_outputs;
1553 let mut pooled_ywy = Array1::<f64>::zeros(1);
1554 pooled_ywy[0] = prepared.ywy.iter().copied().sum();
1555 let mut pooled_projected_rhs_squared =
1556 Array2::<f64>::zeros((prepared.cache.penalty_eigenvalues.len(), 1));
1557 for eig in 0..prepared.cache.penalty_eigenvalues.len() {
1558 pooled_projected_rhs_squared[[eig, 0]] = prepared
1559 .projected_rhs_squared
1560 .row(eig)
1561 .iter()
1562 .copied()
1563 .sum();
1564 }
1565 let per_output_nu = prepared.n_effective as f64 - prepared.cache.nullity as f64;
1566 let shared_nu = (d as f64) * per_output_nu;
1567 validate_reml_profile_residuals(
1568 &prepared.cache,
1569 pooled_ywy.view(),
1570 pooled_projected_rhs_squared.view(),
1571 RHO_LOWER,
1572 )?;
1573 let eval = |rho: f64| {
1574 let mut value = evaluate_reml_profile(
1575 &prepared.cache,
1576 pooled_ywy.view(),
1577 pooled_projected_rhs_squared.view(),
1578 d,
1579 shared_nu,
1580 rho,
1581 );
1582 value += prepared.observation_measure;
1583 value
1584 };
1585 let rho = if prepared.cache.penalty_rank == 0 {
1586 init_rho.unwrap_or(0.0).clamp(RHO_LOWER, RHO_UPPER)
1587 } else {
1588 let enclose = |a: f64, b: f64| {
1589 reml_deriv_enclosure_profile(
1590 &prepared.cache,
1591 pooled_ywy.view(),
1592 pooled_projected_rhs_squared.view(),
1593 d,
1594 shared_nu,
1595 a,
1596 b,
1597 )
1598 };
1599 enumerate_and_select_rho(eval, enclose, init_rho, None)?.rho
1600 };
1601 let objective = eval(rho);
1602 let lambda = gam_problem::checked_exp_log_strength(rho)
1603 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
1604 let coefficients = prepared.coefficients(lambda);
1605 let fitted = dense_ab(x, coefficients.view());
1606 let mut fitted_quadratic = 0.0_f64;
1607 let spectrum = PenaltyRangeSpectrum::of(&prepared.cache);
1611 for eig in 0..spectrum.len() {
1612 let denom = 1.0 + lambda * spectrum.get(eig);
1613 fitted_quadratic += pooled_projected_rhs_squared[[eig, 0]] / denom;
1614 }
1615 let shared_sigma2 = (pooled_ywy[0] - fitted_quadratic) / shared_nu;
1616 let (reml_grad_lambda, reml_hess_lambda) =
1617 rho_derivatives_to_lambda(lambda, objective.grad, objective.hess);
1618 Ok(GaussianRemlMultiResult {
1619 lambda,
1620 rho,
1621 coefficients,
1622 fitted,
1623 reml_score: objective.cost,
1624 reml_score_roundoff: Some(objective.cost_roundoff),
1625 reml_grad_lambda,
1626 reml_hess_lambda,
1627 reml_grad_rho: objective.grad,
1628 reml_hess_rho: objective.hess,
1629 edf: objective.edf,
1630 sigma2: Array1::from_elem(d, shared_sigma2),
1631 cache: prepared.cache,
1632 })
1633}
1634
1635pub fn gaussian_reml_multi_closed_form_with_nullspace_dim(
1636 x: ArrayView2<'_, f64>,
1637 y: ArrayView2<'_, f64>,
1638 penalty: ArrayView2<'_, f64>,
1639 nullspace_dim: Option<usize>,
1640 weights: Option<ArrayView1<'_, f64>>,
1641 init_rho: Option<f64>,
1642) -> Result<GaussianRemlMultiResult, EstimationError> {
1643 let init_lambda = init_rho.map(f64::exp);
1644 gaussian_reml_multi_closed_form_from_parts(
1645 x,
1646 y,
1647 penalty,
1648 nullspace_dim,
1649 weights,
1650 init_lambda,
1651 None,
1652 )
1653}
1654
1655pub fn gaussian_reml_multi_closed_form_with_cache(
1656 x: ArrayView2<'_, f64>,
1657 y: ArrayView2<'_, f64>,
1658 penalty: ArrayView2<'_, f64>,
1659 weights: Option<ArrayView1<'_, f64>>,
1660 init_lambda: Option<f64>,
1661 eigen_cache: Option<&GaussianRemlEigenCache>,
1662) -> Result<GaussianRemlMultiResult, EstimationError> {
1663 gaussian_reml_multi_closed_form_from_parts(
1664 x,
1665 y,
1666 penalty,
1667 None,
1668 weights,
1669 init_lambda,
1670 eigen_cache,
1671 )
1672}
1673
1674struct BlockOrthogonalEval {
1675 beta: Array2<f64>,
1676 logdet: f64,
1677 trace: f64,
1678 trace_pair: f64,
1679 fitted_energy: Array1<f64>,
1680 penalty_energy: Array1<f64>,
1681 curvature_energy: Array1<f64>,
1682 edf: f64,
1683}
1684
1685fn block_penalty_rank_logdet(
1686 penalty: ArrayView2<'_, f64>,
1687) -> Result<(usize, f64), EstimationError> {
1688 let eigs = penalty
1689 .to_owned()
1690 .eigh(Side::Lower)
1691 .map_err(|_| EstimationError::ModelIsIllConditioned {
1692 condition_number: f64::INFINITY,
1693 })?
1694 .0;
1695 let max_abs = eigs.iter().fold(0.0_f64, |m, &v| m.max(v.abs()));
1696 let tol = (EIGEN_REL_TOL * max_abs).max(1.0e-14);
1697 let mut rank = 0_usize;
1698 let mut logdet = 0.0;
1699 for eig in eigs.iter().copied() {
1700 if eig > tol {
1701 rank += 1;
1702 logdet += eig.ln();
1703 }
1704 }
1705 Ok((rank, logdet))
1706}
1707
1708fn block_orthogonal_eval(
1709 gram: &Array2<f64>,
1710 rhs: &Array2<f64>,
1711 penalty: &Array2<f64>,
1712 rho: f64,
1713) -> Result<BlockOrthogonalEval, EstimationError> {
1714 let lambda = gam_problem::checked_exp_log_strength(rho)
1715 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
1716 validate_initial_lambda(lambda)?;
1717 let scaled_penalty = penalty * lambda;
1718 let hessian = canonicalize_penalty((gram + &scaled_penalty).view());
1719 let chol = gaussian_reml_cholesky_lower(hessian)?;
1720 let beta = solve_spd_from_lower_factor(&chol, rhs)?;
1721 let solved_penalty = solve_spd_from_lower_factor(&chol, &scaled_penalty)?;
1722 let logdet = 2.0 * chol.diag().iter().map(|value| value.ln()).sum::<f64>();
1723 let trace = (0..solved_penalty.nrows())
1724 .map(|i| solved_penalty[[i, i]])
1725 .sum::<f64>();
1726 let trace_pair =
1727 gam_linalg::utils::trace_of_product(solved_penalty.view(), solved_penalty.view());
1728 let fitted_energy = (rhs * &beta).sum_axis(Axis(0));
1729 let p_beta = scaled_penalty.dot(&beta);
1730 let penalty_energy = (&beta * &p_beta).sum_axis(Axis(0));
1731 let solved_p_beta = solve_spd_from_lower_factor(&chol, &p_beta)?;
1732 let curvature_energy = (&p_beta * &solved_p_beta).sum_axis(Axis(0));
1733 Ok(BlockOrthogonalEval {
1734 beta,
1735 logdet,
1736 trace,
1737 trace_pair,
1738 fitted_energy,
1739 penalty_energy,
1740 curvature_energy,
1741 edf: penalty.nrows() as f64 - trace,
1742 })
1743}
1744
1745struct BlockOrthogonalScaleDerivs {
1756 value: f64,
1757 value_roundoff: f64,
1768 grad: f64,
1769 hess: f64,
1770}
1771
1772fn block_orthogonal_scale_objective(
1773 eval: &BlockOrthogonalEval,
1774 rho: f64,
1775 scale_precision: ArrayView1<'_, f64>,
1776 rank: usize,
1777) -> BlockOrthogonalScaleDerivs {
1778 let d = scale_precision.len() as f64;
1779 let fit_term = scale_precision
1780 .iter()
1781 .zip(eval.fitted_energy.iter())
1782 .map(|(scale, energy)| scale * energy)
1783 .sum::<f64>();
1784 let logdet_term = 0.5 * d * eval.logdet;
1786 let rank_term = 0.5 * d * (rank as f64) * rho;
1787 let value = logdet_term - 0.5 * fit_term - rank_term;
1788 let value_roundoff =
1792 f64::EPSILON * (logdet_term.abs() + 0.5 * fit_term.abs() + rank_term.abs());
1793 let grad = 0.5 * d * (eval.trace - rank as f64)
1797 + 0.5
1798 * scale_precision
1799 .iter()
1800 .zip(eval.penalty_energy.iter())
1801 .map(|(scale, energy)| scale * energy)
1802 .sum::<f64>();
1803 let hess = 0.5 * d * (eval.trace - eval.trace_pair)
1806 + 0.5
1807 * scale_precision
1808 .iter()
1809 .zip(eval.penalty_energy.iter().zip(eval.curvature_energy.iter()))
1810 .map(|(scale, (energy, curvature))| scale * (energy - 2.0 * curvature))
1811 .sum::<f64>();
1812 BlockOrthogonalScaleDerivs {
1813 value,
1814 value_roundoff,
1815 grad,
1816 hess,
1817 }
1818}
1819
1820fn solve_block_orthogonal_rho(
1827 gram: &Array2<f64>,
1828 rhs: &Array2<f64>,
1829 penalty: &Array2<f64>,
1830 rho0: f64,
1831 scale_precision: ArrayView1<'_, f64>,
1832 rank: usize,
1833 max_iter: usize,
1834) -> Result<(f64, BlockOrthogonalEval), EstimationError> {
1835 let mut rho = rho0;
1836 let mut current = block_orthogonal_eval(gram, rhs, penalty, rho)?;
1837 for _ in 0..max_iter {
1838 let derivs = block_orthogonal_scale_objective(¤t, rho, scale_precision, rank);
1841 let grad = derivs.grad;
1842 let hess = derivs.hess;
1843 if !(grad.is_finite() && hess.is_finite()) {
1844 return Err(EstimationError::ModelIsIllConditioned {
1845 condition_number: f64::INFINITY,
1846 });
1847 }
1848 if grad == 0.0 {
1849 break;
1850 }
1851 let direction = if hess > 0.0 { -grad / hess } else { -grad };
1857 if !direction.is_finite() || grad * direction >= 0.0 {
1858 return Err(EstimationError::ModelIsIllConditioned {
1859 condition_number: f64::INFINITY,
1860 });
1861 }
1862 let current_value = derivs.value;
1863 let model_decrease = -grad * direction - 0.5 * hess * direction * direction;
1873 let value_decides = model_decrease.is_finite() && model_decrease > derivs.value_roundoff;
1874 let accepted = if value_decides {
1875 let mut step_scale = 1.0_f64;
1876 loop {
1877 let candidate_rho = rho + step_scale * direction;
1878 if candidate_rho == rho {
1879 break None;
1880 }
1881 if let Ok(candidate_eval) = block_orthogonal_eval(gram, rhs, penalty, candidate_rho)
1882 {
1883 let candidate_value = block_orthogonal_scale_objective(
1884 &candidate_eval,
1885 candidate_rho,
1886 scale_precision,
1887 rank,
1888 )
1889 .value;
1890 if candidate_value.is_finite() && candidate_value < current_value {
1891 break Some((candidate_rho, candidate_eval));
1892 }
1893 }
1894 step_scale *= 0.5;
1898 }
1899 } else {
1900 None
1901 };
1902 let accepted = accepted.or_else(|| {
1912 if hess <= 0.0 {
1913 return None;
1914 }
1915 let mut step_scale = 1.0_f64;
1916 loop {
1917 let candidate_rho = rho + step_scale * direction;
1918 if candidate_rho == rho {
1919 break None;
1920 }
1921 if let Ok(candidate_eval) = block_orthogonal_eval(gram, rhs, penalty, candidate_rho)
1922 {
1923 let candidate = block_orthogonal_scale_objective(
1924 &candidate_eval,
1925 candidate_rho,
1926 scale_precision,
1927 rank,
1928 );
1929 if candidate.grad.is_finite() && candidate.grad.abs() < grad.abs() {
1930 break Some((candidate_rho, candidate_eval));
1931 }
1932 }
1933 step_scale *= 0.5;
1934 }
1935 });
1936 let Some((next_rho, next_eval)) = accepted else {
1937 break;
1938 };
1939 rho = next_rho;
1940 current = next_eval;
1941 }
1942 Ok((rho, current))
1943}
1944
1945fn block_orthogonal_conditional_scale(
1946 evals: &[BlockOrthogonalEval],
1947 ywy: ArrayView1<'_, f64>,
1948 nu: f64,
1949) -> Result<Array1<f64>, EstimationError> {
1950 let mut explained = Array1::<f64>::zeros(ywy.len());
1951 for eval in evals {
1952 explained += &eval.fitted_energy;
1953 }
1954 let q = &ywy - &explained;
1955 if q.iter().any(|value| !value.is_finite() || *value <= 0.0) {
1956 return Err(EstimationError::ModelIsIllConditioned {
1957 condition_number: f64::INFINITY,
1958 });
1959 }
1960 let scale = q.mapv(|value| nu / value);
1961 if scale
1962 .iter()
1963 .any(|value| !value.is_finite() || *value <= 0.0)
1964 {
1965 return Err(EstimationError::ModelIsIllConditioned {
1966 condition_number: f64::INFINITY,
1967 });
1968 }
1969 Ok(scale)
1970}
1971
1972fn validate_weighted_block_orthogonality(
1978 designs: &[Array2<f64>],
1979 weight: ArrayView1<'_, f64>,
1980) -> Result<(), EstimationError> {
1981 let unit_roundoff = 0.5 * f64::EPSILON;
1982 let operation_count = weight.len().saturating_mul(4);
1983 let accumulated = operation_count as f64 * unit_roundoff;
1984 if accumulated >= 1.0 {
1985 crate::bail_invalid_estim!(
1986 "block-orthogonality verification has no finite floating-point error bound for {} rows",
1987 weight.len()
1988 );
1989 }
1990 let gamma = accumulated / (1.0 - accumulated);
1991 for left_block in 0..designs.len() {
1992 for right_block in (left_block + 1)..designs.len() {
1993 let left = &designs[left_block];
1994 let right = &designs[right_block];
1995 for left_col in 0..left.ncols() {
1996 for right_col in 0..right.ncols() {
1997 let mut cross_product = 0.0_f64;
1998 let mut magnitude_sum = 0.0_f64;
1999 for row in 0..weight.len() {
2000 let term = weight[row] * left[[row, left_col]] * right[[row, right_col]];
2001 cross_product += term;
2002 magnitude_sum += term.abs();
2003 }
2004 let roundoff = gamma * magnitude_sum;
2005 if !cross_product.is_finite()
2006 || !roundoff.is_finite()
2007 || cross_product.abs() > roundoff
2008 {
2009 crate::bail_invalid_estim!(
2010 "block-orthogonal Gaussian REML requires X[{left_block}]' W X[{right_block}] = 0, but columns ({left_col}, {right_col}) have weighted cross-product {cross_product:.6e} beyond the arithmetic bound {roundoff:.3e}"
2011 );
2012 }
2013 }
2014 }
2015 }
2016 }
2017 Ok(())
2018}
2019
2020#[derive(Clone, Copy, Debug)]
2021struct BlockOrthogonalProfileCurvature {
2022 min_eigenvalue: f64,
2023 roundoff: f64,
2024}
2025
2026fn block_orthogonal_profile_hessian(
2036 evals: &[BlockOrthogonalEval],
2037 rhos: ArrayView1<'_, f64>,
2038 scale_precision: ArrayView1<'_, f64>,
2039 ranks: &[usize],
2040 nu: f64,
2041) -> Result<Array2<f64>, EstimationError> {
2042 let blocks = evals.len();
2043 let mut hessian = Array2::<f64>::zeros((blocks, blocks));
2044 for block in 0..blocks {
2045 hessian[[block, block]] = block_orthogonal_scale_objective(
2046 &evals[block],
2047 rhos[block],
2048 scale_precision.view(),
2049 ranks[block],
2050 )
2051 .hess;
2052 }
2053 for left in 0..blocks {
2054 for right in 0..=left {
2055 let correction = evals[left]
2056 .penalty_energy
2057 .iter()
2058 .zip(evals[right].penalty_energy.iter())
2059 .zip(scale_precision.iter())
2060 .map(|((&left_energy, &right_energy), &scale)| {
2061 0.5 * scale * scale * left_energy * right_energy / nu
2062 })
2063 .sum::<f64>();
2064 hessian[[left, right]] -= correction;
2065 if left != right {
2066 hessian[[right, left]] -= correction;
2067 }
2068 }
2069 }
2070 if hessian.iter().any(|value| !value.is_finite()) {
2071 return Err(EstimationError::ModelIsIllConditioned {
2072 condition_number: f64::INFINITY,
2073 });
2074 }
2075 Ok(hessian)
2076}
2077
2078struct BlockOrthogonalProfileSpectrum {
2086 curvature: BlockOrthogonalProfileCurvature,
2087 eigenvalues: Array1<f64>,
2088 eigenvectors: Array2<f64>,
2089}
2090
2091fn block_orthogonal_profile_spectrum(
2092 hessian: &Array2<f64>,
2093) -> Result<BlockOrthogonalProfileSpectrum, EstimationError> {
2094 let blocks = hessian.nrows();
2095 let (eigenvalues, eigenvectors) =
2096 hessian
2097 .clone()
2098 .eigh(Side::Lower)
2099 .map_err(|_| EstimationError::ModelIsIllConditioned {
2100 condition_number: f64::INFINITY,
2101 })?;
2102 let min_eigenvalue = eigenvalues.iter().copied().fold(f64::INFINITY, f64::min);
2103 let spectral_scale = eigenvalues
2104 .iter()
2105 .copied()
2106 .map(f64::abs)
2107 .fold(0.0_f64, f64::max);
2108 let roundoff = f64::EPSILON * blocks.max(1) as f64 * spectral_scale.max(f64::MIN_POSITIVE);
2109 Ok(BlockOrthogonalProfileSpectrum {
2110 curvature: BlockOrthogonalProfileCurvature {
2111 min_eigenvalue,
2112 roundoff,
2113 },
2114 eigenvalues,
2115 eigenvectors,
2116 })
2117}
2118
2119impl BlockOrthogonalProfileSpectrum {
2120 fn newton_direction(&self, gradient: ArrayView1<'_, f64>) -> Option<Array1<f64>> {
2124 if self.curvature.min_eigenvalue.is_nan() || self.curvature.min_eigenvalue <= 0.0 {
2125 return None;
2126 }
2127 let projected = self.eigenvectors.t().dot(&gradient);
2128 let scaled = Array1::from_iter(
2129 projected
2130 .iter()
2131 .zip(self.eigenvalues.iter())
2132 .map(|(component, eigenvalue)| -component / eigenvalue),
2133 );
2134 let direction = self.eigenvectors.dot(&scaled);
2135 direction
2136 .iter()
2137 .all(|value| value.is_finite())
2138 .then_some(direction)
2139 }
2140}
2141
2142struct BlockOrthogonalProfileValue {
2153 value: f64,
2154 roundoff: f64,
2155}
2156
2157fn block_orthogonal_profile_value(
2158 evals: &[BlockOrthogonalEval],
2159 rhos: ArrayView1<'_, f64>,
2160 ranks: &[usize],
2161 ywy: ArrayView1<'_, f64>,
2162 nu: f64,
2163 d: usize,
2164) -> Option<BlockOrthogonalProfileValue> {
2165 let mut explained = Array1::<f64>::zeros(ywy.len());
2166 for eval in evals {
2167 explained += &eval.fitted_energy;
2168 }
2169 let mut q = ywy.to_owned();
2170 q -= &explained;
2171 if q.iter().any(|value| !value.is_finite() || *value <= 0.0) {
2172 return None;
2173 }
2174 let determinant_term = 0.5
2175 * d as f64
2176 * evals
2177 .iter()
2178 .enumerate()
2179 .map(|(block, eval)| eval.logdet - ranks[block] as f64 * rhos[block])
2180 .sum::<f64>();
2181 let deviance_term = 0.5 * nu * q.iter().map(|value| value.ln()).sum::<f64>();
2182 let value = determinant_term + deviance_term;
2183 if !value.is_finite() {
2184 return None;
2185 }
2186 Some(BlockOrthogonalProfileValue {
2187 value,
2188 roundoff: f64::EPSILON * (determinant_term.abs() + deviance_term.abs()),
2189 })
2190}
2191
2192struct BlockOrthogonalStateMeasurement {
2197 score_residual: f64,
2198 gradient: Array1<f64>,
2199 spectrum: BlockOrthogonalProfileSpectrum,
2200}
2201
2202fn measure_block_orthogonal_state(
2203 evals: &[BlockOrthogonalEval],
2204 rhos: ArrayView1<'_, f64>,
2205 scale_precision: ArrayView1<'_, f64>,
2206 ranks: &[usize],
2207 nu: f64,
2208 d: usize,
2209) -> Result<BlockOrthogonalStateMeasurement, EstimationError> {
2210 let mut gradient = Array1::<f64>::zeros(evals.len());
2211 let mut score_residual = 0.0_f64;
2212 for (block, eval) in evals.iter().enumerate() {
2213 let derivs =
2214 block_orthogonal_scale_objective(eval, rhos[block], scale_precision, ranks[block]);
2215 let residual = derivs.grad.abs() / ((d as f64) * (ranks[block].max(1) as f64));
2216 if !residual.is_finite() {
2217 return Err(EstimationError::ModelIsIllConditioned {
2218 condition_number: f64::INFINITY,
2219 });
2220 }
2221 gradient[block] = derivs.grad;
2222 score_residual = score_residual.max(residual);
2223 }
2224 let hessian = block_orthogonal_profile_hessian(evals, rhos, scale_precision, ranks, nu)?;
2225 Ok(BlockOrthogonalStateMeasurement {
2226 score_residual,
2227 gradient,
2228 spectrum: block_orthogonal_profile_spectrum(&hessian)?,
2229 })
2230}
2231
2232pub fn gaussian_reml_blocks_orthogonal_shared_scale(
2233 designs: &[Array2<f64>],
2234 penalties: &[Array2<f64>],
2235 y: ArrayView2<'_, f64>,
2236 weights: Option<ArrayView1<'_, f64>>,
2237 init_rhos: Option<&[f64]>,
2238) -> Result<GaussianRemlBlockOrthogonalResult, EstimationError> {
2239 gaussian_reml_blocks_orthogonal_shared_scale_with_controls(
2240 designs,
2241 penalties,
2242 y,
2243 weights,
2244 init_rhos,
2245 BlockOrthogonalControls::default(),
2246 )
2247}
2248
2249fn gaussian_reml_blocks_orthogonal_shared_scale_with_controls(
2250 designs: &[Array2<f64>],
2251 penalties: &[Array2<f64>],
2252 y: ArrayView2<'_, f64>,
2253 weights: Option<ArrayView1<'_, f64>>,
2254 init_rhos: Option<&[f64]>,
2255 controls: BlockOrthogonalControls,
2256) -> Result<GaussianRemlBlockOrthogonalResult, EstimationError> {
2257 if designs.is_empty() {
2258 crate::bail_invalid_estim!("block-orthogonal Gaussian REML requires at least one block");
2259 }
2260 if designs.len() != penalties.len() {
2261 crate::bail_invalid_estim!(
2262 "block-orthogonal Gaussian REML block mismatch: {} designs, {} penalties",
2263 designs.len(),
2264 penalties.len()
2265 );
2266 }
2267 let n = y.nrows();
2268 let d = y.ncols();
2269 if d == 0 {
2270 crate::bail_invalid_estim!("block-orthogonal Gaussian REML requires at least one output");
2271 }
2272 if y.iter().any(|value| !value.is_finite()) {
2273 crate::bail_invalid_estim!("block-orthogonal Gaussian REML response must be finite");
2274 }
2275 let weight = gaussian_reml_weights(n, weights)?;
2276 if let Some(rhos) = init_rhos {
2277 if rhos.len() != designs.len() {
2278 crate::bail_invalid_estim!(
2279 "block-orthogonal Gaussian REML init_rhos length mismatch: expected {}, got {}",
2280 designs.len(),
2281 rhos.len()
2282 );
2283 }
2284 if rhos.iter().any(|value| !value.is_finite()) {
2285 crate::bail_invalid_estim!("block-orthogonal Gaussian REML init_rhos must be finite");
2286 }
2287 }
2288
2289 let mut ywy = Array1::<f64>::zeros(d);
2290 for row in 0..n {
2291 for output in 0..d {
2292 ywy[output] += weight[row] * y[[row, output]] * y[[row, output]];
2293 }
2294 }
2295 let mut grams = Vec::with_capacity(designs.len());
2296 let mut rhs_blocks = Vec::with_capacity(designs.len());
2297 let mut penalties_owned = Vec::with_capacity(penalties.len());
2298 let mut ranks = Vec::with_capacity(penalties.len());
2299 let mut penalty_logdets = Vec::with_capacity(penalties.len());
2300 let mut nullity_total = 0_usize;
2301 for (block, (design, penalty)) in designs.iter().zip(penalties.iter()).enumerate() {
2302 let penalty_owned = canonicalize_penalty(penalty.view());
2303 validate_gaussian_reml_design(design.view(), penalty_owned.view(), Some(weight.view()))?;
2304 if design.nrows() != n {
2305 crate::bail_invalid_estim!(
2306 "block-orthogonal Gaussian REML designs[{block}] has {} rows, expected {n}",
2307 design.nrows()
2308 );
2309 }
2310 let gram = dense_xt_diag_x(design.view(), weight.view());
2311 let rhs = dense_xt_diag_y(design.view(), weight.view(), y);
2312 let (rank, logdet) = block_penalty_rank_logdet(penalty_owned.view())?;
2313 nullity_total += penalty_owned.nrows().saturating_sub(rank);
2314 grams.push(canonicalize_penalty(gram.view()));
2315 rhs_blocks.push(rhs);
2316 penalties_owned.push(penalty_owned);
2317 ranks.push(rank);
2318 penalty_logdets.push(logdet);
2319 }
2320 validate_weighted_block_orthogonality(designs, weight.view())?;
2321 let n_effective = effective_observation_count(weight.view());
2322 if n_effective <= nullity_total {
2323 crate::bail_invalid_estim!(
2324 "block-orthogonal Gaussian REML requires more positive-weight rows than the total penalty nullity; got n_effective={n_effective}, nullity={nullity_total}"
2325 );
2326 }
2327 let nu = (n_effective - nullity_total) as f64;
2328 let mut rhos = match init_rhos {
2329 Some(values) => Array1::from_vec(values.to_vec()),
2330 None => Array1::zeros(designs.len()),
2331 };
2332 let mut evals = (0..designs.len())
2337 .map(|block| {
2338 block_orthogonal_eval(
2339 &grams[block],
2340 &rhs_blocks[block],
2341 &penalties_owned[block],
2342 rhos[block],
2343 )
2344 })
2345 .collect::<Result<Vec<_>, _>>()?;
2346 let mut scale_precision = block_orthogonal_conditional_scale(&evals, ywy.view(), nu)?;
2347 let mut converged = false;
2380 let mut cycle_detected = false;
2381 let mut outer_passes = 0usize;
2382 let mut last_score_residual = f64::INFINITY;
2383 let mut last_min_profile_curvature = f64::NEG_INFINITY;
2384 let mut last_profile_curvature_roundoff = 0.0_f64;
2385 let mut last_scale_step = f64::INFINITY;
2386 let mut recent_states: [Option<(Array1<f64>, Array1<f64>)>; 2] = [None, None];
2387 while outer_passes < controls.max_outer_passes {
2388 outer_passes += 1;
2389 let scale_at_pass_start = scale_precision.clone();
2390 evals.clear();
2391 for block in 0..designs.len() {
2392 let (rho, eval) = solve_block_orthogonal_rho(
2393 &grams[block],
2394 &rhs_blocks[block],
2395 &penalties_owned[block],
2396 rhos[block],
2397 scale_precision.view(),
2398 ranks[block],
2399 controls.block_updates_per_pass,
2400 )?;
2401 rhos[block] = rho;
2402 evals.push(eval);
2403 }
2404 scale_precision = block_orthogonal_conditional_scale(&evals, ywy.view(), nu)?;
2405 let mut measured = measure_block_orthogonal_state(
2406 &evals,
2407 rhos.view(),
2408 scale_precision.view(),
2409 &ranks,
2410 nu,
2411 d,
2412 )?;
2413 let alternation_certified = measured.score_residual <= controls.score_tol
2416 && measured.spectrum.curvature.min_eigenvalue >= -measured.spectrum.curvature.roundoff;
2417 let newton_step = if alternation_certified {
2418 None
2419 } else {
2420 measured
2421 .spectrum
2422 .newton_direction(measured.gradient.view())
2423 .zip(block_orthogonal_profile_value(
2424 &evals,
2425 rhos.view(),
2426 &ranks,
2427 ywy.view(),
2428 nu,
2429 d,
2430 ))
2431 };
2432 if let Some((direction, current_profile)) = newton_step {
2433 let model_decrease = -0.5 * measured.gradient.dot(&direction);
2439 let value_decides =
2440 model_decrease.is_finite() && model_decrease > current_profile.roundoff;
2441 let mut step_scale = 1.0_f64;
2442 let accepted = loop {
2443 let candidate_rhos = &rhos + &direction.mapv(|value| step_scale * value);
2444 if candidate_rhos == rhos {
2445 break None;
2446 }
2447 let candidate = (0..designs.len())
2448 .map(|block| {
2449 block_orthogonal_eval(
2450 &grams[block],
2451 &rhs_blocks[block],
2452 &penalties_owned[block],
2453 candidate_rhos[block],
2454 )
2455 })
2456 .collect::<Result<Vec<_>, _>>()
2457 .ok()
2458 .and_then(|candidate_evals| {
2459 let candidate_scale =
2460 block_orthogonal_conditional_scale(&candidate_evals, ywy.view(), nu)
2461 .ok()?;
2462 let candidate_measured = measure_block_orthogonal_state(
2463 &candidate_evals,
2464 candidate_rhos.view(),
2465 candidate_scale.view(),
2466 &ranks,
2467 nu,
2468 d,
2469 )
2470 .ok()?;
2471 let improves = if value_decides {
2472 block_orthogonal_profile_value(
2473 &candidate_evals,
2474 candidate_rhos.view(),
2475 &ranks,
2476 ywy.view(),
2477 nu,
2478 d,
2479 )
2480 .is_some_and(|profile| profile.value < current_profile.value)
2481 } else {
2482 candidate_measured.score_residual < measured.score_residual
2483 };
2484 improves.then_some((candidate_evals, candidate_scale, candidate_measured))
2485 });
2486 if let Some((candidate_evals, candidate_scale, candidate_measured)) = candidate {
2487 break Some((
2488 candidate_rhos,
2489 candidate_evals,
2490 candidate_scale,
2491 candidate_measured,
2492 ));
2493 }
2494 step_scale *= 0.5;
2497 };
2498 if let Some((next_rhos, next_evals, next_scale, next_measured)) = accepted {
2499 rhos = next_rhos;
2500 evals = next_evals;
2501 scale_precision = next_scale;
2502 measured = next_measured;
2503 }
2504 }
2505 last_scale_step = scale_precision
2506 .iter()
2507 .zip(scale_at_pass_start.iter())
2508 .map(|(next, old)| (next.ln() - old.ln()).abs())
2509 .fold(0.0_f64, f64::max);
2510 last_score_residual = measured.score_residual;
2511 last_min_profile_curvature = measured.spectrum.curvature.min_eigenvalue;
2512 last_profile_curvature_roundoff = measured.spectrum.curvature.roundoff;
2513 if last_score_residual <= controls.score_tol
2514 && last_min_profile_curvature >= -last_profile_curvature_roundoff
2515 {
2516 converged = true;
2517 break;
2518 }
2519 let state = (rhos.clone(), scale_precision.clone());
2525 if recent_states
2526 .iter()
2527 .flatten()
2528 .any(|prev| prev.0 == state.0 && prev.1 == state.1)
2529 {
2530 cycle_detected = true;
2531 break;
2532 }
2533 recent_states[1] = recent_states[0].take();
2534 recent_states[0] = Some(state);
2535 }
2536 if !converged {
2537 return Err(EstimationError::BlockOrthogonalRemlDidNotConverge {
2538 iterations: outer_passes,
2539 max_score_residual: last_score_residual,
2540 score_tol: controls.score_tol,
2541 min_profile_curvature: last_min_profile_curvature,
2542 profile_curvature_roundoff: last_profile_curvature_roundoff,
2543 last_scale_step,
2544 cycle_detected,
2545 rho_checkpoint: rhos.to_vec(),
2546 });
2547 }
2548
2549 let coefficients = evals
2550 .iter()
2551 .map(|eval| eval.beta.clone())
2552 .collect::<Vec<_>>();
2553 let mut fitted = Array2::<f64>::zeros((n, d));
2554 for (design, coef) in designs.iter().zip(coefficients.iter()) {
2555 fitted += &fast_ab(&design.view(), &coef.view());
2556 }
2557 let mut explained = Array1::<f64>::zeros(d);
2558 for eval in evals.iter() {
2559 explained += &eval.fitted_energy;
2560 }
2561 let q = &ywy - &explained;
2562 if q.iter().any(|value| !value.is_finite() || *value <= 0.0) {
2563 return Err(EstimationError::ModelIsIllConditioned {
2564 condition_number: f64::INFINITY,
2565 });
2566 }
2567 let lambdas = Array1::from_vec(gam_problem::checked_exp_log_strengths(
2568 rhos.iter().copied(),
2569 )?);
2570 let edf = Array1::from_iter(evals.iter().map(|eval| eval.edf));
2571 let logdet_term = evals
2572 .iter()
2573 .enumerate()
2574 .map(|(block, eval)| {
2575 eval.logdet - penalty_logdets[block] - (ranks[block] as f64) * rhos[block]
2576 })
2577 .sum::<f64>();
2578 let scale_term = q
2579 .iter()
2580 .map(|value| nu * (1.0 + (2.0 * std::f64::consts::PI * value / nu).ln()))
2581 .sum::<f64>();
2582 Ok(GaussianRemlBlockOrthogonalResult {
2583 coefficients,
2584 fitted,
2585 lambdas,
2586 log_lambdas: rhos,
2587 reml_score: 0.5 * (d as f64) * logdet_term + 0.5 * scale_term
2588 + gaussian_reml_observation_measure(weight.view(), d).value,
2589 edf,
2590 })
2591}
2592
2593pub fn gaussian_reml_multi_shared_dispersion_penalty_gradient_from_fit(
2604 x: ArrayView2<'_, f64>,
2605 y: ArrayView2<'_, f64>,
2606 penalty: ArrayView2<'_, f64>,
2607 weights: Option<ArrayView1<'_, f64>>,
2608 fit: &GaussianRemlMultiResult,
2609) -> Result<Array2<f64>, EstimationError> {
2610 validate_gaussian_reml_forward_fit(x, y, penalty, weights, fit)?;
2611 let n = x.nrows();
2612 let p = x.ncols();
2613 let d = y.ncols();
2614 if d == 0 {
2615 crate::bail_invalid_estim!(
2616 "shared-dispersion REML penalty gradient requires at least one response column"
2617 );
2618 }
2619 let weight = gaussian_reml_weights(n, weights)?;
2620 let n_effective = effective_observation_count(weight.view());
2621 let per_output_nu = n_effective.checked_sub(fit.cache.nullity).ok_or_else(|| {
2622 EstimationError::InvalidInput(
2623 "shared-dispersion REML penalty gradient has non-positive residual degrees of freedom"
2624 .to_string(),
2625 )
2626 })?;
2627 if per_output_nu == 0 {
2628 crate::bail_invalid_estim!(
2629 "shared-dispersion REML penalty gradient requires positive residual degrees of freedom"
2630 );
2631 }
2632 let shared_nu = (d as f64) * (per_output_nu as f64);
2633 let shared_sigma2 = fit.sigma2[0];
2644 if fit
2645 .sigma2
2646 .iter()
2647 .any(|sigma2| sigma2.to_bits() != shared_sigma2.to_bits())
2648 {
2649 crate::bail_invalid_estim!(
2650 "shared-dispersion REML penalty gradient requires one shared forward dispersion"
2651 );
2652 }
2653 let pooled_deviance = shared_sigma2 * shared_nu;
2654 let mut pooled_response_energy = 0.0_f64;
2679 for output in 0..d {
2680 for row in 0..n {
2681 let value = y[[row, output]];
2682 pooled_response_energy += weight[row] * value * value;
2683 }
2684 }
2685 let unit_roundoff = 0.5 * f64::EPSILON;
2686 let operation_count = n
2687 .saturating_mul(d)
2688 .saturating_mul(3)
2689 .saturating_add(p.saturating_mul(2))
2690 .saturating_add(1);
2691 let accumulated = operation_count as f64 * unit_roundoff;
2692 if accumulated >= 1.0 {
2693 crate::bail_invalid_estim!(
2694 "shared-dispersion REML penalty gradient has no finite floating-point error bound for {n} rows, {d} responses and {p} coefficients"
2695 );
2696 }
2697 let deviance_roundoff = (accumulated / (1.0 - accumulated)) * pooled_response_energy;
2698 if !(pooled_deviance.is_finite()
2699 && deviance_roundoff.is_finite()
2700 && pooled_deviance > deviance_roundoff)
2701 {
2702 crate::bail_invalid_estim!(
2703 "shared-dispersion REML penalty gradient requires a forward deviance resolved above the roundoff of its own formation; the chart is interpolating to arithmetic precision: pooled deviance {pooled_deviance:.6e} does not exceed the forward bound {deviance_roundoff:.6e} on the cancellation that produced it from pooled response energy {pooled_response_energy:.6e}"
2704 );
2705 }
2706
2707 let inverse_hessian = gaussian_reml_inverse_hessian_from_cache(&fit.cache, fit.lambda)?;
2708 let penalty_pseudoinverse = gaussian_reml_penalty_pseudoinverse_from_cache(&fit.cache)?;
2709 let mut gradient = Array2::<f64>::zeros((p, p));
2710 for row in 0..p {
2711 for col in 0..p {
2712 gradient[[row, col]] = 0.5
2713 * (d as f64)
2714 * (fit.lambda * inverse_hessian[[col, row]] - penalty_pseudoinverse[[col, row]]);
2715 }
2716 }
2717 let deviance_scale = 0.5 * shared_nu * fit.lambda / pooled_deviance;
2718 for output in 0..d {
2719 add_rank_one_penalty_vjp(
2720 deviance_scale,
2721 fit.coefficients.column(output),
2722 &mut gradient,
2723 );
2724 }
2725 for row in 0..p {
2726 for col in (row + 1)..p {
2727 let mean = 0.5 * (gradient[[row, col]] + gradient[[col, row]]);
2728 gradient[[row, col]] = mean;
2729 gradient[[col, row]] = mean;
2730 }
2731 }
2732 if gradient.iter().any(|value| !value.is_finite()) {
2733 crate::bail_invalid_estim!(
2734 "shared-dispersion REML penalty gradient produced a non-finite value"
2735 );
2736 }
2737 Ok(gradient)
2738}
2739
2740fn gaussian_reml_multi_closed_form_from_parts(
2741 x: ArrayView2<'_, f64>,
2742 y: ArrayView2<'_, f64>,
2743 penalty: ArrayView2<'_, f64>,
2744 nullspace_dim: Option<usize>,
2745 weights: Option<ArrayView1<'_, f64>>,
2746 init_lambda: Option<f64>,
2747 eigen_cache: Option<&GaussianRemlEigenCache>,
2748) -> Result<GaussianRemlMultiResult, EstimationError> {
2749 let prepared = prepare_gaussian_reml(x, y, penalty, nullspace_dim, weights, eigen_cache)?;
2750 let init_rho = init_lambda
2751 .map(validate_initial_lambda)
2752 .transpose()?
2753 .map(f64::ln);
2754 let rho = optimize_rho(&prepared, init_rho)?;
2755 let eval = prepared.evaluate(rho);
2756 let lambda = gam_problem::checked_exp_log_strength(rho)
2757 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
2758 let coefficients = prepared.coefficients(lambda);
2759 let fitted = dense_ab(x, coefficients.view());
2760 let sigma2 = prepared.sigma2(rho);
2761 let (reml_grad_lambda, reml_hess_lambda) =
2762 rho_derivatives_to_lambda(lambda, eval.grad, eval.hess);
2763 Ok(GaussianRemlMultiResult {
2764 lambda,
2765 rho,
2766 coefficients,
2767 fitted,
2768 reml_score: eval.cost,
2769 reml_score_roundoff: Some(eval.cost_roundoff),
2770 reml_grad_lambda,
2771 reml_hess_lambda,
2772 reml_grad_rho: eval.grad,
2773 reml_hess_rho: eval.hess,
2774 edf: eval.edf,
2775 sigma2,
2776 cache: prepared.cache,
2777 })
2778}
2779
2780pub fn gaussian_reml_free_b_score(
2781 x: ArrayView2<'_, f64>,
2782 y: ArrayView2<'_, f64>,
2783 coefficients: ArrayView2<'_, f64>,
2784 log_lambda: f64,
2785 penalty: ArrayView2<'_, f64>,
2786 weights: Option<ArrayView1<'_, f64>>,
2787) -> Result<GaussianRemlFreeBScore, EstimationError> {
2788 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
2789 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?;
2790 let penalty_owned = canonicalize_penalty(penalty);
2791 let penalty = penalty_owned.view();
2792 let n = x.nrows();
2793 let p = x.ncols();
2794 let d = y.ncols();
2795 validate_gaussian_reml_design(x, penalty, weights)?;
2796 if y.nrows() != n {
2797 crate::bail_invalid_estim!(
2798 "Gaussian REML row mismatch: X has {n} rows but Y has {}",
2799 y.nrows()
2800 );
2801 }
2802 if coefficients.dim() != (p, d) {
2803 crate::bail_invalid_estim!(
2804 "Gaussian REML coefficient shape mismatch: expected {p}x{d}, got {}x{}",
2805 coefficients.nrows(),
2806 coefficients.ncols()
2807 );
2808 }
2809 if y.iter().chain(coefficients.iter()).any(|v| !v.is_finite()) {
2810 crate::bail_invalid_estim!("Gaussian REML inputs must be finite");
2811 }
2812
2813 let weight = gaussian_reml_weights(n, weights)?;
2814 let n_effective = effective_observation_count(weight.view());
2815 let cache =
2816 build_gaussian_reml_eigen_cache_with_nullspace_dim(x, penalty, None, Some(weight.view()))?;
2817 if n_effective <= cache.nullity {
2818 crate::bail_invalid_estim!(
2819 "Gaussian REML requires more positive-weight rows than the nullspace dimension; got n_effective={n_effective}, nullity={}",
2820 cache.nullity
2821 );
2822 }
2823 let nu = n_effective as f64 - cache.nullity as f64;
2824 let fitted = dense_ab(x, coefficients);
2825 let residual = y.to_owned() - &fitted;
2826 let xtw_residual = dense_xt_diag_y(x, weight.view(), residual.view());
2827 let s_beta = dense_ab(penalty, coefficients);
2828
2829 let mut logdet_h = cache.logdet_xtwx;
2830 let mut trace_h = 0.0;
2831 let mut edf = 0.0;
2832 for delta in PenaltyRangeSpectrum::of(&cache).iter() {
2835 let t = lambda * delta;
2836 logdet_h += (1.0 + t).ln();
2837 if delta > 0.0 {
2838 trace_h += t / (1.0 + t);
2839 }
2840 edf += 1.0 / (1.0 + t);
2841 }
2842 let logdet_s = cache.logdet_penalty_positive + (cache.penalty_rank as f64) * log_lambda;
2843 let mut reml_score = 0.5 * (d as f64) * (logdet_h - logdet_s)
2844 + gaussian_reml_observation_measure(weight.view(), d).value;
2845 let mut grad_log_lambda = 0.5 * (d as f64) * (trace_h - cache.penalty_rank as f64);
2846 let mut grad_coefficients = Array2::<f64>::zeros((p, d));
2847 let inverse_hessian = {
2848 let xtwx = dense_xt_diag_x(x, weight.view());
2849 let mut hessian = xtwx;
2850 hessian += &(penalty.to_owned() * lambda);
2851 hessian
2852 .cholesky(Side::Lower)
2853 .map_err(EstimationError::LinearSystemSolveFailed)?
2854 .solve_mat(&Array2::<f64>::eye(p))
2855 };
2856 let penalty_pinv = gaussian_reml_penalty_pseudoinverse_from_cache(&cache)?;
2857 let mut grad_penalty = Array2::<f64>::zeros((p, p));
2858 for row in 0..p {
2859 for col in 0..p {
2860 grad_penalty[[row, col]] += 0.5
2861 * (d as f64)
2862 * (lambda * inverse_hessian[[col, row]] - penalty_pinv[[col, row]]);
2863 }
2864 }
2865 let mut sigma2 = Array1::<f64>::zeros(d);
2866
2867 for output in 0..d {
2868 let mut weighted_rss = 0.0;
2869 for row in 0..n {
2870 let r = residual[[row, output]];
2871 weighted_rss += weight[row] * r * r;
2872 }
2873 let beta_col = coefficients.column(output);
2874 let s_beta_col = s_beta.column(output);
2875 let penalty_quadratic = beta_col.dot(&s_beta_col);
2876 let dp = weighted_rss + lambda * penalty_quadratic;
2877 if !(dp > 0.0) {
2884 crate::bail_invalid_estim!(
2885 "Gaussian REML output {output} has a non-positive penalized deviance {dp}: the \
2886 profiled scale is not identifiable (interpolating fit), so the REML criterion \
2887 is undefined there"
2888 );
2889 }
2890 sigma2[output] = dp / nu;
2891 reml_score += 0.5 * nu * (1.0 + (2.0 * std::f64::consts::PI * dp / nu).ln());
2892 grad_log_lambda += 0.5 * nu * lambda * penalty_quadratic / dp;
2893 let scale = nu / dp;
2894 for coeff in 0..p {
2895 grad_coefficients[[coeff, output]] =
2896 scale * (-xtw_residual[[coeff, output]] + lambda * s_beta[[coeff, output]]);
2897 }
2898 add_rank_one_penalty_vjp(0.5 * scale * lambda, beta_col, &mut grad_penalty);
2899 }
2900 for i in 0..p {
2901 for j in (i + 1)..p {
2902 let avg = 0.5 * (grad_penalty[[i, j]] + grad_penalty[[j, i]]);
2903 grad_penalty[[i, j]] = avg;
2904 grad_penalty[[j, i]] = avg;
2905 }
2906 }
2907
2908 Ok(GaussianRemlFreeBScore {
2909 reml_score,
2910 grad_coefficients,
2911 grad_penalty,
2912 grad_log_lambda,
2913 fitted,
2914 sigma2,
2915 edf,
2916 })
2917}
2918
2919pub fn gaussian_reml_multi_closed_form_backward(
2920 x: ArrayView2<'_, f64>,
2921 y: ArrayView2<'_, f64>,
2922 penalty: ArrayView2<'_, f64>,
2923 weights: Option<ArrayView1<'_, f64>>,
2924 init_lambda: Option<f64>,
2925 upstream_lambda: f64,
2926 upstream_coefficients: Option<ArrayView2<'_, f64>>,
2927 upstream_fitted: Option<ArrayView2<'_, f64>>,
2928 upstream_reml_score: f64,
2929 upstream_edf: f64,
2930) -> Result<GaussianRemlBackwardResult, EstimationError> {
2931 let fit =
2932 gaussian_reml_multi_closed_form_with_cache(x, y, penalty, weights, init_lambda, None)?;
2933 gaussian_reml_multi_closed_form_backward_from_fit(
2934 x,
2935 y,
2936 penalty,
2937 weights,
2938 &fit,
2939 upstream_lambda,
2940 upstream_coefficients,
2941 upstream_fitted,
2942 upstream_reml_score,
2943 upstream_edf,
2944 )
2945}
2946
2947pub fn gaussian_reml_multi_closed_form_backward_from_fit(
2948 x: ArrayView2<'_, f64>,
2949 y: ArrayView2<'_, f64>,
2950 penalty: ArrayView2<'_, f64>,
2951 weights: Option<ArrayView1<'_, f64>>,
2952 fit: &GaussianRemlMultiResult,
2953 upstream_lambda: f64,
2954 upstream_coefficients: Option<ArrayView2<'_, f64>>,
2955 upstream_fitted: Option<ArrayView2<'_, f64>>,
2956 upstream_reml_score: f64,
2957 upstream_edf: f64,
2958) -> Result<GaussianRemlBackwardResult, EstimationError> {
2959 validate_gaussian_reml_backward_upstreams(
2960 x,
2961 y,
2962 penalty,
2963 upstream_lambda,
2964 upstream_coefficients,
2965 upstream_fitted,
2966 upstream_reml_score,
2967 upstream_edf,
2968 )?;
2969 validate_gaussian_reml_forward_fit(x, y, penalty, weights, fit)?;
2970 let lambda = fit.lambda;
2971 let n = x.nrows();
2972 let p = x.ncols();
2973 let d = y.ncols();
2974 let rho_hat = lambda.ln();
2985 let rho_at_bound =
2986 (rho_hat - RHO_UPPER).abs() <= 1.0e-9 || (rho_hat - RHO_LOWER).abs() <= 1.0e-9;
2987 let implicit_rho_usable =
2988 fit.reml_hess_rho.is_finite() && fit.reml_hess_rho.abs() > 1.0e-14 && !rho_at_bound;
2989 let weight = gaussian_reml_weights(n, weights)?;
2990 let inverse_hessian = match gaussian_reml_inverse_hessian_from_cache(&fit.cache, lambda) {
2991 Ok(inv) => inv,
2992 Err(EstimationError::ModelIsIllConditioned { condition_number }) => {
2993 warn_ill_conditioned_backward_once(p, d, condition_number);
2994 return Ok(zero_backward_result(n, p, d));
2995 }
2996 Err(err) => return Err(err),
2997 };
2998 gaussian_reml_multi_closed_form_backward_from_fit_with_inverse_hessian_impl(
2999 x,
3000 y,
3001 penalty,
3002 weight,
3003 fit,
3004 inverse_hessian,
3005 upstream_lambda,
3006 upstream_coefficients,
3007 upstream_fitted,
3008 upstream_reml_score,
3009 upstream_edf,
3010 implicit_rho_usable,
3011 n,
3012 p,
3013 d,
3014 )
3015}
3016
3017fn gaussian_reml_multi_closed_form_backward_from_fit_with_inverse_hessian_impl(
3018 x: ArrayView2<'_, f64>,
3019 y: ArrayView2<'_, f64>,
3020 penalty: ArrayView2<'_, f64>,
3021 weight: Array1<f64>,
3022 fit: &GaussianRemlMultiResult,
3023 inverse_hessian: Array2<f64>,
3024 upstream_lambda: f64,
3025 upstream_coefficients: Option<ArrayView2<'_, f64>>,
3026 upstream_fitted: Option<ArrayView2<'_, f64>>,
3027 upstream_reml_score: f64,
3028 upstream_edf: f64,
3029 implicit_rho_usable: bool,
3030 n: usize,
3031 p: usize,
3032 d: usize,
3033) -> Result<GaussianRemlBackwardResult, EstimationError> {
3034 let penalty_owned = canonicalize_penalty(penalty);
3038 let penalty = penalty_owned.view();
3039 let lambda = fit.lambda;
3040 let beta = &fit.coefficients;
3041 let residual = y.to_owned() - &fit.fitted;
3042 let nu = effective_observation_count(weight.view()) as f64 - fit.cache.nullity as f64;
3046
3047 let mut grad_x = Array2::<f64>::zeros((n, p));
3048 let mut grad_y = Array2::<f64>::zeros((n, d));
3049 let mut grad_penalty = Array2::<f64>::zeros((p, p));
3050 let mut grad_weights = Array1::<f64>::zeros(n);
3051
3052 let mut upstream_beta = Array2::<f64>::zeros((p, d));
3053 if let Some(upstream_coefficients) = upstream_coefficients {
3054 upstream_beta += &upstream_coefficients;
3055 }
3056 if let Some(upstream_fitted) = upstream_fitted {
3057 upstream_beta += &dense_atb(x, upstream_fitted);
3058 grad_x += &dense_ab(upstream_fitted, beta.t());
3059 }
3060
3061 let mut lambda_adjoint = upstream_lambda;
3062 if upstream_beta.iter().any(|value| *value != 0.0) {
3063 add_ridge_profile_vjp_with_lambda_grad(
3068 1.0,
3069 x,
3070 y,
3071 penalty,
3072 &weight,
3073 lambda,
3074 &inverse_hessian,
3075 beta,
3076 upstream_beta.view(),
3077 &mut grad_x,
3078 &mut grad_y,
3079 &mut grad_penalty,
3080 &mut grad_weights,
3081 &mut lambda_adjoint,
3082 );
3083 }
3084
3085 if upstream_reml_score != 0.0 {
3086 add_reml_score_vjp(
3087 upstream_reml_score,
3088 x,
3089 &weight,
3090 &inverse_hessian,
3091 beta,
3092 &residual,
3093 &fit.sigma2,
3094 nu,
3095 lambda,
3096 &fit.cache,
3097 &mut grad_x,
3098 &mut grad_y,
3099 &mut grad_penalty,
3100 &mut grad_weights,
3101 )?;
3102 lambda_adjoint += upstream_reml_score * fit.reml_grad_lambda;
3103 }
3104
3105 if upstream_edf != 0.0 {
3106 lambda_adjoint += add_edf_vjp(
3107 upstream_edf,
3108 x,
3109 penalty,
3110 &weight,
3111 lambda,
3112 &inverse_hessian,
3113 &mut grad_x,
3114 &mut grad_penalty,
3115 &mut grad_weights,
3116 );
3117 }
3118
3119 if lambda_adjoint != 0.0 && implicit_rho_usable {
3120 let root_scale = -lambda_adjoint * lambda / fit.reml_hess_rho;
3121 add_reml_rho_gradient_vjp(
3122 root_scale,
3123 x,
3124 y,
3125 penalty,
3126 &weight,
3127 lambda,
3128 &inverse_hessian,
3129 beta,
3130 &residual,
3131 &fit.sigma2,
3132 nu,
3133 &mut grad_x,
3134 &mut grad_y,
3135 &mut grad_penalty,
3136 &mut grad_weights,
3137 );
3138 }
3139
3140 let p = grad_penalty.nrows();
3149 for i in 0..p {
3150 for j in (i + 1)..p {
3151 let avg = 0.5 * (grad_penalty[[i, j]] + grad_penalty[[j, i]]);
3152 grad_penalty[[i, j]] = avg;
3153 grad_penalty[[j, i]] = avg;
3154 }
3155 }
3156 finish_gaussian_reml_weight_vjp(weight.view(), d, upstream_reml_score, &mut grad_weights);
3157 Ok(GaussianRemlBackwardResult {
3158 grad_x,
3159 grad_y,
3160 grad_penalty,
3161 grad_weights,
3162 })
3163}
3164
3165pub fn gaussian_reml_multi_closed_form_backward_batch<'a>(
3166 problems: &[GaussianRemlMultiBackwardProblem<'a>],
3167 penalty: ArrayView2<'a, f64>,
3168) -> Vec<Result<GaussianRemlBackwardResult, EstimationError>> {
3169 let inverse_hessians = batched_inverse_hessians_from_caches(problems);
3170 let results: Vec<Result<GaussianRemlBackwardResult, EstimationError>> = problems
3171 .par_iter()
3172 .zip(inverse_hessians.into_par_iter())
3173 .map(|(problem, inverse_hessian_result)| {
3174 validate_gaussian_reml_backward_upstreams(
3175 problem.x.view(),
3176 problem.y.view(),
3177 penalty,
3178 problem.grad_lambda,
3179 problem.grad_coefficients.as_ref().map(|g| g.view()),
3180 problem.grad_fitted.as_ref().map(|g| g.view()),
3181 problem.grad_reml_score,
3182 problem.grad_edf,
3183 )?;
3184 validate_gaussian_reml_forward_fit(
3185 problem.x.view(),
3186 problem.y.view(),
3187 penalty,
3188 problem.weights.as_ref().map(|w| w.view()),
3189 problem.fit,
3190 )?;
3191 let n = problem.x.nrows();
3192 let p = problem.x.ncols();
3193 let d = problem.y.ncols();
3194 if !(problem.fit.reml_hess_rho.is_finite() && problem.fit.reml_hess_rho.abs() > 1.0e-14)
3195 {
3196 warn_ill_conditioned_backward_once(p, d, f64::INFINITY);
3198 return Ok(zero_backward_result(n, p, d));
3199 }
3200 let weight = gaussian_reml_weights(n, problem.weights.as_ref().map(|w| w.view()))?;
3201 let inverse_hessian = match inverse_hessian_result {
3202 Ok(inv) => inv,
3203 Err(EstimationError::ModelIsIllConditioned { condition_number }) => {
3204 warn_ill_conditioned_backward_once(p, d, condition_number);
3205 return Ok(zero_backward_result(n, p, d));
3206 }
3207 Err(err) => return Err(err),
3208 };
3209 let rho_hat = problem.fit.lambda.ln();
3214 let rho_at_bound =
3215 (rho_hat - RHO_UPPER).abs() <= 1.0e-9 || (rho_hat - RHO_LOWER).abs() <= 1.0e-9;
3216 let implicit_rho_usable = problem.fit.reml_hess_rho.is_finite()
3217 && problem.fit.reml_hess_rho.abs() > 1.0e-14
3218 && !rho_at_bound;
3219 gaussian_reml_multi_closed_form_backward_from_fit_with_inverse_hessian_impl(
3220 problem.x.view(),
3221 problem.y.view(),
3222 penalty,
3223 weight,
3224 problem.fit,
3225 inverse_hessian,
3226 problem.grad_lambda,
3227 problem.grad_coefficients.as_ref().map(|g| g.view()),
3228 problem.grad_fitted.as_ref().map(|g| g.view()),
3229 problem.grad_reml_score,
3230 problem.grad_edf,
3231 implicit_rho_usable,
3232 n,
3233 p,
3234 d,
3235 )
3236 })
3237 .collect();
3238 results
3239}
3240
3241fn rho_derivatives_to_lambda(lambda: f64, grad_rho: f64, hess_rho: f64) -> (f64, f64) {
3242 (grad_rho / lambda, (hess_rho - grad_rho) / (lambda * lambda))
3243}
3244
3245fn validate_gaussian_reml_backward_upstreams(
3246 x: ArrayView2<'_, f64>,
3247 y: ArrayView2<'_, f64>,
3248 penalty: ArrayView2<'_, f64>,
3249 upstream_lambda: f64,
3250 upstream_coefficients: Option<ArrayView2<'_, f64>>,
3251 upstream_fitted: Option<ArrayView2<'_, f64>>,
3252 upstream_reml_score: f64,
3253 upstream_edf: f64,
3254) -> Result<(), EstimationError> {
3255 if !(upstream_lambda.is_finite() && upstream_reml_score.is_finite() && upstream_edf.is_finite())
3256 {
3257 crate::bail_invalid_estim!("Gaussian REML backward upstream scalars must be finite");
3258 }
3259 if let Some(upstream_coefficients) = upstream_coefficients {
3260 if upstream_coefficients.dim() != (x.ncols(), y.ncols()) {
3261 crate::bail_invalid_estim!(
3262 "Gaussian REML backward coefficient upstream shape mismatch: expected {}x{}, got {}x{}",
3263 x.ncols(),
3264 y.ncols(),
3265 upstream_coefficients.nrows(),
3266 upstream_coefficients.ncols()
3267 );
3268 }
3269 if upstream_coefficients.iter().any(|value| !value.is_finite()) {
3270 crate::bail_invalid_estim!(
3271 "Gaussian REML backward coefficient upstream must be finite"
3272 );
3273 }
3274 }
3275 if let Some(upstream_fitted) = upstream_fitted {
3276 if upstream_fitted.dim() != y.dim() {
3277 crate::bail_invalid_estim!(
3278 "Gaussian REML backward fitted upstream shape mismatch: expected {}x{}, got {}x{}",
3279 y.nrows(),
3280 y.ncols(),
3281 upstream_fitted.nrows(),
3282 upstream_fitted.ncols()
3283 );
3284 }
3285 if upstream_fitted.iter().any(|value| !value.is_finite()) {
3286 crate::bail_invalid_estim!("Gaussian REML backward fitted upstream must be finite");
3287 }
3288 }
3289 validate_gaussian_reml_design(x, penalty, None)?;
3290 Ok(())
3291}
3292
3293fn validate_gaussian_reml_forward_fit(
3294 x: ArrayView2<'_, f64>,
3295 y: ArrayView2<'_, f64>,
3296 penalty: ArrayView2<'_, f64>,
3297 weights: Option<ArrayView1<'_, f64>>,
3298 fit: &GaussianRemlMultiResult,
3299) -> Result<(), EstimationError> {
3300 let penalty_owned = canonicalize_penalty(penalty);
3304 let penalty = penalty_owned.view();
3305 let n = x.nrows();
3306 let p = x.ncols();
3307 let d = y.ncols();
3308 validate_gaussian_reml_design(x, penalty, weights)?;
3309 validate_gaussian_reml_eigen_cache(&fit.cache, p)?;
3310 if y.nrows() != n
3311 || fit.coefficients.dim() != (p, d)
3312 || fit.fitted.dim() != (n, d)
3313 || fit.sigma2.len() != d
3314 {
3315 crate::bail_invalid_estim!(
3316 "Gaussian REML backward forward-state shape mismatch: expected coefficients=({p},{d}), fitted=({n},{d}), sigma2={d}"
3317 );
3318 }
3319 if !(fit.lambda.is_finite()
3320 && fit.lambda > 0.0
3321 && fit.rho.is_finite()
3322 && fit.reml_score.is_finite()
3323 && fit.reml_hess_rho.is_finite()
3324 && fit.edf.is_finite())
3325 || fit.coefficients.iter().any(|value| !value.is_finite())
3326 || fit.fitted.iter().any(|value| !value.is_finite())
3327 || fit.sigma2.iter().any(|value| !(value.is_finite() && *value > 0.0))
3328 {
3329 crate::bail_invalid_estim!(
3330 "Gaussian REML backward forward state must be finite with positive profiled scales"
3331 );
3332 }
3333 let penalty_fingerprint = matrix_fingerprint(penalty);
3334 if fit.cache.penalty_fingerprint != penalty_fingerprint {
3335 crate::bail_invalid_estim!("Gaussian REML backward forward-state penalty mismatch");
3336 }
3337 let weight = gaussian_reml_weights(n, weights)?;
3338 let xtwx = dense_xt_diag_x(x, weight.view());
3339 if fit.cache.xtwx_fingerprint != matrix_fingerprint(xtwx.view()) {
3340 crate::bail_invalid_estim!("Gaussian REML backward forward-state X'WX mismatch");
3341 }
3342 Ok(())
3343}
3344
3345fn gaussian_reml_inverse_hessian_from_cache(
3346 cache: &GaussianRemlEigenCache,
3347 lambda: f64,
3348) -> Result<Array2<f64>, EstimationError> {
3349 if !(lambda.is_finite() && lambda > 0.0) {
3350 crate::bail_invalid_estim!(
3351 "Gaussian REML lambda must be finite and positive; got {lambda}"
3352 );
3353 }
3354 let p = cache.penalty_eigenvalues.len();
3355 let spectrum = PenaltyRangeSpectrum::of(cache);
3356 let mut scaled_basis = cache.coefficient_basis.clone();
3357 for eig in 0..p {
3358 let scale = 1.0 / (1.0 + lambda * spectrum.get(eig));
3362 for row in 0..p {
3363 scaled_basis[[row, eig]] *= scale;
3364 }
3365 }
3366 let inverse = dense_ab(scaled_basis.view(), cache.coefficient_basis.t());
3367 if inverse.iter().any(|value| !value.is_finite()) {
3368 return Err(EstimationError::ModelIsIllConditioned {
3369 condition_number: f64::INFINITY,
3370 });
3371 }
3372 Ok(inverse)
3373}
3374
3375fn batched_inverse_hessians_from_caches(
3376 problems: &[GaussianRemlMultiBackwardProblem<'_>],
3377) -> Vec<Result<Array2<f64>, EstimationError>> {
3378 if problems.is_empty() {
3379 return Vec::new();
3380 }
3381 let p = problems[0].fit.cache.coefficient_basis.nrows();
3382 let uniform = p > 0
3383 && problems.iter().all(|problem| {
3384 let cache = &problem.fit.cache;
3385 cache.coefficient_basis.dim() == (p, p) && cache.penalty_eigenvalues.len() == p
3386 });
3387 if uniform && problems.len() > 1 {
3388 let mut scaled_basis = Array3::<f64>::zeros((problems.len(), p, p));
3389 let mut basis = Array3::<f64>::zeros((problems.len(), p, p));
3390 let mut valid = true;
3391 for (idx, problem) in problems.iter().enumerate() {
3392 let lambda = problem.fit.lambda;
3393 if !(lambda.is_finite() && lambda > 0.0) {
3394 valid = false;
3395 break;
3396 }
3397 let cache = &problem.fit.cache;
3398 let spectrum = PenaltyRangeSpectrum::of(cache);
3399 basis
3400 .slice_mut(s![idx, .., ..])
3401 .assign(&cache.coefficient_basis);
3402 for eig in 0..p {
3403 let scale = 1.0 / (1.0 + lambda * spectrum.get(eig));
3404 for row in 0..p {
3405 scaled_basis[[idx, row, eig]] = cache.coefficient_basis[[row, eig]] * scale;
3406 }
3407 }
3408 }
3409 if valid
3410 && let Some(inverses) =
3411 gam_gpu::try_fast_abt_strided_batched(scaled_basis.view(), basis.view())
3412 {
3413 return inverses
3414 .axis_iter(Axis(0))
3415 .map(|inverse| Ok(inverse.to_owned()))
3416 .collect();
3417 }
3418 }
3419 problems
3420 .iter()
3421 .map(|problem| {
3422 gaussian_reml_inverse_hessian_from_cache(&problem.fit.cache, problem.fit.lambda)
3423 })
3424 .collect()
3425}
3426
3427fn ridge_profile_vjp_data_partials(
3434 scale: f64,
3435 x: ArrayView2<'_, f64>,
3436 y: ArrayView2<'_, f64>,
3437 penalty: ArrayView2<'_, f64>,
3438 weights: &Array1<f64>,
3439 lambda: f64,
3440 inverse_hessian: &Array2<f64>,
3441 beta: &Array2<f64>,
3442 upstream_beta: ArrayView2<'_, f64>,
3443 grad_x: &mut Array2<f64>,
3444 grad_y: &mut Array2<f64>,
3445 grad_penalty: &mut Array2<f64>,
3446 grad_weights: &mut Array1<f64>,
3447) -> Array2<f64> {
3448 let m = dense_ab(inverse_hessian.view(), upstream_beta);
3449 let c = dense_ab(m.view(), beta.t());
3450 let c_sym = &c + &c.t();
3451 let ymt = dense_ab(y, m.t());
3452 let xcs = dense_ab(x, c_sym.view());
3453 for i in 0..x.nrows() {
3454 let wi = weights[i] * scale;
3455 for k in 0..x.ncols() {
3456 grad_x[[i, k]] += wi * (ymt[[i, k]] - xcs[[i, k]]);
3457 }
3458 }
3459
3460 let xm = dense_ab(x, m.view());
3461 for i in 0..x.nrows() {
3462 let wi = weights[i] * scale;
3463 for j in 0..y.ncols() {
3464 grad_y[[i, j]] += wi * xm[[i, j]];
3465 }
3466 }
3467
3468 let xc = dense_ab(x, c.view());
3469 for i in 0..x.nrows() {
3470 let mut from_b = 0.0;
3471 for j in 0..y.ncols() {
3472 from_b += y[[i, j]] * xm[[i, j]];
3473 }
3474 let mut from_a = 0.0;
3475 for k in 0..x.ncols() {
3476 from_a += x[[i, k]] * xc[[i, k]];
3477 }
3478 grad_weights[i] += scale * (from_b - from_a);
3479 }
3480
3481 for row in 0..penalty.nrows() {
3482 for col in 0..penalty.ncols() {
3483 let mut value = 0.0;
3484 for output in 0..beta.ncols() {
3485 value += m[[row, output]] * beta[[col, output]];
3486 }
3487 grad_penalty[[row, col]] -= scale * lambda * value;
3488 }
3489 }
3490 m
3491}
3492
3493fn add_ridge_profile_vjp_with_lambda_grad(
3498 scale: f64,
3499 x: ArrayView2<'_, f64>,
3500 y: ArrayView2<'_, f64>,
3501 penalty: ArrayView2<'_, f64>,
3502 weights: &Array1<f64>,
3503 lambda: f64,
3504 inverse_hessian: &Array2<f64>,
3505 beta: &Array2<f64>,
3506 upstream_beta: ArrayView2<'_, f64>,
3507 grad_x: &mut Array2<f64>,
3508 grad_y: &mut Array2<f64>,
3509 grad_penalty: &mut Array2<f64>,
3510 grad_weights: &mut Array1<f64>,
3511 lambda_adjoint_out: &mut f64,
3512) {
3513 let m = ridge_profile_vjp_data_partials(
3514 scale,
3515 x,
3516 y,
3517 penalty,
3518 weights,
3519 lambda,
3520 inverse_hessian,
3521 beta,
3522 upstream_beta,
3523 grad_x,
3524 grad_y,
3525 grad_penalty,
3526 grad_weights,
3527 );
3528 let penalty_beta = dense_ab(penalty, beta.view());
3529 let dot = m
3530 .iter()
3531 .zip(penalty_beta.iter())
3532 .map(|(left, right)| left * right)
3533 .sum::<f64>();
3534 *lambda_adjoint_out += -scale * dot;
3535}
3536
3537fn add_ridge_profile_vjp_fixed_lambda(
3541 scale: f64,
3542 x: ArrayView2<'_, f64>,
3543 y: ArrayView2<'_, f64>,
3544 penalty: ArrayView2<'_, f64>,
3545 weights: &Array1<f64>,
3546 lambda: f64,
3547 inverse_hessian: &Array2<f64>,
3548 beta: &Array2<f64>,
3549 upstream_beta: ArrayView2<'_, f64>,
3550 grad_x: &mut Array2<f64>,
3551 grad_y: &mut Array2<f64>,
3552 grad_penalty: &mut Array2<f64>,
3553 grad_weights: &mut Array1<f64>,
3554) {
3555 ridge_profile_vjp_data_partials(
3556 scale,
3557 x,
3558 y,
3559 penalty,
3560 weights,
3561 lambda,
3562 inverse_hessian,
3563 beta,
3564 upstream_beta,
3565 grad_x,
3566 grad_y,
3567 grad_penalty,
3568 grad_weights,
3569 );
3570}
3571
3572fn add_reml_score_vjp(
3573 scale: f64,
3574 x: ArrayView2<'_, f64>,
3575 weights: &Array1<f64>,
3576 inverse_hessian: &Array2<f64>,
3577 beta: &Array2<f64>,
3578 residual: &Array2<f64>,
3579 sigma2: &Array1<f64>,
3580 nu: f64,
3581 lambda: f64,
3582 cache: &GaussianRemlEigenCache,
3583 grad_x: &mut Array2<f64>,
3584 grad_y: &mut Array2<f64>,
3585 grad_penalty: &mut Array2<f64>,
3586 grad_weights: &mut Array1<f64>,
3587) -> Result<(), EstimationError> {
3588 let d = beta.ncols() as f64;
3589 let xp = dense_ab(x, inverse_hessian.view());
3590 let penalty_pinv = gaussian_reml_penalty_pseudoinverse_from_cache(cache)?;
3591 for row in 0..grad_penalty.nrows() {
3592 for col in 0..grad_penalty.ncols() {
3593 grad_penalty[[row, col]] +=
3594 scale * 0.5 * d * (lambda * inverse_hessian[[col, row]] - penalty_pinv[[col, row]]);
3595 }
3596 }
3597 for i in 0..x.nrows() {
3598 let wi = weights[i] * scale * d;
3599 for k in 0..x.ncols() {
3600 grad_x[[i, k]] += wi * xp[[i, k]];
3601 }
3602 let mut leverage = 0.0;
3603 for k in 0..x.ncols() {
3604 leverage += x[[i, k]] * xp[[i, k]];
3605 }
3606 grad_weights[i] += scale * 0.5 * d * leverage;
3607 }
3608
3609 for j in 0..beta.ncols() {
3610 let dp = sigma2[j] * nu;
3611 let coef = scale * 0.5 * nu / dp;
3612 add_deviance_profile_vjp(
3613 coef,
3614 j,
3615 x,
3616 weights,
3617 beta,
3618 residual,
3619 grad_x,
3620 grad_y,
3621 grad_weights,
3622 );
3623 add_rank_one_penalty_vjp(coef * lambda, beta.column(j), grad_penalty);
3624 }
3625 Ok(())
3626}
3627
3628fn add_edf_vjp(
3639 scale: f64,
3640 x: ArrayView2<'_, f64>,
3641 penalty: ArrayView2<'_, f64>,
3642 weights: &Array1<f64>,
3643 lambda: f64,
3644 inverse_hessian: &Array2<f64>,
3645 grad_x: &mut Array2<f64>,
3646 grad_penalty: &mut Array2<f64>,
3647 grad_weights: &mut Array1<f64>,
3648) -> f64 {
3649 let m_inv_s = dense_ab(inverse_hessian.view(), penalty);
3651 let mut g_a = dense_ab(m_inv_s.view(), inverse_hessian.view());
3652 g_a.mapv_inplace(|v| v * lambda);
3653
3654 let xg = dense_ab(x, g_a.view());
3658 let leading_scale = 2.0 * scale;
3662 for i in 0..xg.nrows() {
3663 let row_scale = leading_scale * weights[i];
3664 for k in 0..xg.ncols() {
3665 grad_x[[i, k]] += row_scale * xg[[i, k]];
3666 }
3667 }
3668 for i in 0..x.nrows() {
3669 let mut quad = 0.0;
3670 for k in 0..x.ncols() {
3671 quad += x[[i, k]] * xg[[i, k]];
3672 }
3673 grad_weights[i] += scale * quad;
3674 }
3675
3676 for row in 0..grad_penalty.nrows() {
3679 for col in 0..grad_penalty.ncols() {
3680 grad_penalty[[row, col]] +=
3681 scale * (-lambda * inverse_hessian[[row, col]] + lambda * g_a[[row, col]]);
3682 }
3683 }
3684
3685 let p_dim = m_inv_s.nrows();
3687 let mut tr_m_inv_s = 0.0;
3688 for i in 0..p_dim {
3689 tr_m_inv_s += m_inv_s[[i, i]];
3690 }
3691 let mut tr_squared = 0.0;
3692 for i in 0..p_dim {
3693 for j in 0..p_dim {
3694 tr_squared += m_inv_s[[i, j]] * m_inv_s[[j, i]];
3695 }
3696 }
3697 scale * (-tr_m_inv_s + lambda * tr_squared)
3698}
3699
3700fn add_reml_rho_gradient_vjp(
3701 scale: f64,
3702 x: ArrayView2<'_, f64>,
3703 y: ArrayView2<'_, f64>,
3704 penalty: ArrayView2<'_, f64>,
3705 weights: &Array1<f64>,
3706 lambda: f64,
3707 inverse_hessian: &Array2<f64>,
3708 beta: &Array2<f64>,
3709 residual: &Array2<f64>,
3710 sigma2: &Array1<f64>,
3711 nu: f64,
3712 grad_x: &mut Array2<f64>,
3713 grad_y: &mut Array2<f64>,
3714 grad_penalty: &mut Array2<f64>,
3715 grad_weights: &mut Array1<f64>,
3716) {
3717 let d = beta.ncols() as f64;
3718 let inverse_s = dense_ab(inverse_hessian.view(), penalty);
3719 let trace_kernel = dense_ab(inverse_s.view(), inverse_hessian.view());
3720 for row in 0..grad_penalty.nrows() {
3721 for col in 0..grad_penalty.ncols() {
3722 grad_penalty[[row, col]] += scale
3723 * 0.5
3724 * d
3725 * lambda
3726 * (inverse_hessian[[col, row]] - lambda * trace_kernel[[col, row]]);
3727 }
3728 }
3729 let xt = dense_ab(x, trace_kernel.view());
3730 for i in 0..x.nrows() {
3731 let wi = -scale * d * lambda * weights[i];
3732 for k in 0..x.ncols() {
3733 grad_x[[i, k]] += wi * xt[[i, k]];
3734 }
3735 let mut quad = 0.0;
3736 for k in 0..x.ncols() {
3737 quad += x[[i, k]] * xt[[i, k]];
3738 }
3739 grad_weights[i] -= scale * 0.5 * d * lambda * quad;
3740 }
3741
3742 let s_beta = dense_ab(penalty, beta.view());
3743 let mut upstream_beta = Array2::<f64>::zeros(beta.dim());
3744 for j in 0..beta.ncols() {
3745 let dp = sigma2[j] * nu;
3746 let q = lambda * beta.column(j).dot(&s_beta.column(j));
3747 let q_coef = scale * nu / dp;
3748 for row in 0..beta.nrows() {
3749 upstream_beta[[row, j]] = q_coef * lambda * s_beta[[row, j]];
3750 }
3751 let dp_coef = -scale * 0.5 * nu * q / (dp * dp);
3752 add_rank_one_penalty_vjp(
3753 (0.5 * q_coef + dp_coef) * lambda,
3754 beta.column(j),
3755 grad_penalty,
3756 );
3757 add_deviance_profile_vjp(
3758 dp_coef,
3759 j,
3760 x,
3761 weights,
3762 beta,
3763 residual,
3764 grad_x,
3765 grad_y,
3766 grad_weights,
3767 );
3768 }
3769 add_ridge_profile_vjp_fixed_lambda(
3772 1.0,
3773 x,
3774 y,
3775 penalty,
3776 weights,
3777 lambda,
3778 inverse_hessian,
3779 beta,
3780 upstream_beta.view(),
3781 grad_x,
3782 grad_y,
3783 grad_penalty,
3784 grad_weights,
3785 );
3786}
3787
3788fn add_rank_one_penalty_vjp(
3789 scale: f64,
3790 beta_col: ArrayView1<'_, f64>,
3791 grad_penalty: &mut Array2<f64>,
3792) {
3793 for row in 0..beta_col.len() {
3794 for col in 0..beta_col.len() {
3795 grad_penalty[[row, col]] += scale * beta_col[row] * beta_col[col];
3796 }
3797 }
3798}
3799
3800fn penalty_range_tolerance(eigenvalues: ArrayView1<'_, f64>) -> f64 {
3817 let max_abs = eigenvalues
3818 .iter()
3819 .fold(0.0_f64, |acc, &value| acc.max(value.abs()));
3820 max_abs * EIGEN_REL_TOL
3821}
3822
3823#[derive(Clone, Copy)]
3852struct PenaltyRangeSpectrum<'a> {
3853 eigenvalues: &'a Array1<f64>,
3854 tolerance: f64,
3855}
3856
3857impl<'a> PenaltyRangeSpectrum<'a> {
3858 fn of(cache: &'a GaussianRemlEigenCache) -> Self {
3859 Self {
3860 eigenvalues: &cache.penalty_eigenvalues,
3861 tolerance: penalty_range_tolerance(cache.penalty_eigenvalues.view()),
3862 }
3863 }
3864
3865 fn len(&self) -> usize {
3866 self.eigenvalues.len()
3867 }
3868
3869 #[inline]
3872 fn get(&self, index: usize) -> f64 {
3873 let delta = self.eigenvalues[index];
3874 if delta > self.tolerance { delta } else { 0.0 }
3875 }
3876
3877 fn iter(&self) -> impl Iterator<Item = f64> + '_ {
3878 (0..self.len()).map(move |index| self.get(index))
3879 }
3880
3881 fn rank(&self) -> usize {
3886 self.eigenvalues
3887 .iter()
3888 .filter(|&&delta| delta > self.tolerance)
3889 .count()
3890 }
3891}
3892
3893fn gaussian_reml_penalty_pseudoinverse_from_cache(
3894 cache: &GaussianRemlEigenCache,
3895) -> Result<Array2<f64>, EstimationError> {
3896 let p = cache.penalty_eigenvalues.len();
3897 let spectrum = PenaltyRangeSpectrum::of(cache);
3926 let tolerance = spectrum.tolerance;
3927 let selected: Vec<usize> = (0..p).filter(|eig| spectrum.get(*eig) > 0.0).collect();
3928 if selected.len() != cache.penalty_rank {
3929 crate::bail_invalid_estim!(
3930 "Gaussian REML penalty pseudoinverse: the cache reports penalty_rank={} but {} of its \
3931 {p} eigenvalues exceed the range tolerance {tolerance:e}; the pseudoinverse divides by \
3932 each selected eigenvalue, so it cannot reconcile a rank it did not derive",
3933 cache.penalty_rank,
3934 selected.len()
3935 );
3936 }
3937 let mut scaled_basis = Array2::<f64>::zeros((p, p));
3938 for eig in selected {
3939 let delta = spectrum.get(eig);
3940 for row in 0..p {
3941 scaled_basis[[row, eig]] = cache.coefficient_basis[[row, eig]] / delta;
3942 }
3943 }
3944 Ok(dense_ab(scaled_basis.view(), cache.coefficient_basis.t()))
3945}
3946
3947fn add_deviance_profile_vjp(
3948 scale: f64,
3949 output: usize,
3950 x: ArrayView2<'_, f64>,
3951 weights: &Array1<f64>,
3952 beta: &Array2<f64>,
3953 residual: &Array2<f64>,
3954 grad_x: &mut Array2<f64>,
3955 grad_y: &mut Array2<f64>,
3956 grad_weights: &mut Array1<f64>,
3957) {
3958 for i in 0..x.nrows() {
3959 let r = residual[[i, output]];
3960 let wr_scale = scale * weights[i] * r;
3961 grad_y[[i, output]] += 2.0 * wr_scale;
3962 for k in 0..x.ncols() {
3963 grad_x[[i, k]] -= 2.0 * wr_scale * beta[[k, output]];
3964 }
3965 grad_weights[i] += scale * r * r;
3966 }
3967}
3968
3969fn validate_initial_lambda(lambda: f64) -> Result<f64, EstimationError> {
3970 if lambda.is_finite() && lambda > 0.0 {
3971 Ok(lambda)
3972 } else {
3973 Err(EstimationError::InvalidInput(format!(
3974 "Gaussian REML initial lambda must be finite and positive; got {lambda}"
3975 )))
3976 }
3977}
3978
3979fn dense_ab(a: ArrayView2<'_, f64>, b: ArrayView2<'_, f64>) -> Array2<f64> {
3980 fast_ab(&a, &b)
3981}
3982
3983fn dense_atb(a: ArrayView2<'_, f64>, b: ArrayView2<'_, f64>) -> Array2<f64> {
3984 fast_atb(&a, &b)
3985}
3986
3987fn dense_xt_diag_x(x: ArrayView2<'_, f64>, w: ArrayView1<'_, f64>) -> Array2<f64> {
3988 fast_xt_diag_x(&x, &w)
3989}
3990
3991fn dense_xt_diag_y(
3992 x: ArrayView2<'_, f64>,
3993 w: ArrayView1<'_, f64>,
3994 y: ArrayView2<'_, f64>,
3995) -> Array2<f64> {
3996 fast_xt_diag_y(&x, &w, &y)
3997}
3998
3999fn matrix_fingerprint(matrix: ArrayView2<'_, f64>) -> u64 {
4000 let mut hash = 0xcbf29ce484222325_u64;
4001 hash = fnv1a_mix(hash, matrix.nrows() as u64);
4002 hash = fnv1a_mix(hash, matrix.ncols() as u64);
4003 for &value in matrix {
4004 hash = fnv1a_mix(hash, value.to_bits());
4005 }
4006 hash
4007}
4008
4009fn fnv1a_mix(hash: u64, value: u64) -> u64 {
4010 (hash ^ value).wrapping_mul(0x100000001b3)
4011}
4012
4013pub fn build_gaussian_reml_eigen_cache_batched(
4018 xtwx_matrices: Vec<Array2<f64>>,
4019 penalty: ArrayView2<'_, f64>,
4020 nullspace_dim: Option<usize>,
4021) -> Vec<Result<GaussianRemlEigenCache, EstimationError>> {
4022 let penalty_owned = canonicalize_penalty(penalty);
4023 let penalty = penalty_owned.view();
4024 let k = xtwx_matrices.len();
4025 if k == 0 {
4026 return Vec::new();
4027 }
4028 let fingerprints: Vec<u64> = xtwx_matrices
4029 .iter()
4030 .map(|m| matrix_fingerprint(m.view()))
4031 .collect();
4032
4033 let p = xtwx_matrices[0].nrows();
4034 let uniform_square = p > 0 && xtwx_matrices.iter().all(|matrix| matrix.dim() == (p, p));
4035 if uniform_square && k > 1 {
4036 let mut lower_matrices = xtwx_matrices.clone();
4037 if gam_gpu::try_cholesky_batched_lower_inplace(&mut lower_matrices).is_some() {
4038 let transforms = batched_whitened_penalty_transforms(&lower_matrices, penalty);
4046 return lower_matrices
4047 .into_iter()
4048 .enumerate()
4049 .map(|(b, lower)| {
4050 let precomputed_transform = transforms.as_ref().map(|t| t[b].clone());
4051 gaussian_reml_eigen_cache_from_lower_with_transform(
4052 lower,
4053 penalty,
4054 nullspace_dim,
4055 fingerprints[b],
4056 precomputed_transform,
4057 )
4058 })
4059 .collect();
4060 }
4061 }
4062
4063 let mut results = Vec::with_capacity(k);
4064 for (b, xtwx) in xtwx_matrices.into_iter().enumerate() {
4065 let lower = match gaussian_reml_cholesky_lower(xtwx) {
4066 Ok(l) => l,
4067 Err(err) => {
4068 results.push(Err(err));
4069 continue;
4070 }
4071 };
4072 results.push(gaussian_reml_eigen_cache_from_lower_with_transform(
4073 lower,
4074 penalty,
4075 nullspace_dim,
4076 fingerprints[b],
4077 None,
4078 ));
4079 }
4080 results
4081}
4082
4083fn batched_whitened_penalty_transforms(
4084 lowers: &[Array2<f64>],
4085 penalty: ArrayView2<'_, f64>,
4086) -> Option<Vec<Array2<f64>>> {
4087 let first = lowers.first()?;
4088 let p = first.nrows();
4089 if p == 0 || first.ncols() != p || lowers.iter().any(|lower| lower.dim() != (p, p)) {
4090 return None;
4091 }
4092 let mut linv_stack = Array3::<f64>::zeros((lowers.len(), p, p));
4093 for (idx, lower) in lowers.iter().enumerate() {
4094 let l_inv = invert_lower_triangular(lower).ok()?;
4095 linv_stack.slice_mut(s![idx, .., ..]).assign(&l_inv);
4096 }
4097 let penalty_in_metric = gam_gpu::try_fast_ab_broadcast_b_batched(linv_stack.view(), penalty)?;
4098 let transformed =
4099 gam_gpu::try_fast_abt_strided_batched(penalty_in_metric.view(), linv_stack.view())?;
4100 Some(
4101 transformed
4102 .axis_iter(Axis(0))
4103 .map(|matrix| matrix.to_owned())
4104 .collect(),
4105 )
4106}
4107
4108pub fn build_gaussian_reml_eigen_cache_with_nullspace_dim(
4109 x: ArrayView2<'_, f64>,
4110 penalty: ArrayView2<'_, f64>,
4111 nullspace_dim: Option<usize>,
4112 weights: Option<ArrayView1<'_, f64>>,
4113) -> Result<GaussianRemlEigenCache, EstimationError> {
4114 let penalty_owned = canonicalize_penalty(penalty);
4115 let penalty = penalty_owned.view();
4116 let n = x.nrows();
4117 validate_gaussian_reml_design(x, penalty, weights)?;
4118 let weight = gaussian_reml_weights(n, weights)?;
4119
4120 let xtwx = dense_xt_diag_x(x, weight.view());
4121 gaussian_reml_eigen_cache_from_xtwx(xtwx, penalty, nullspace_dim)
4122}
4123
4124fn validate_gaussian_reml_design(
4125 x: ArrayView2<'_, f64>,
4126 penalty: ArrayView2<'_, f64>,
4127 weights: Option<ArrayView1<'_, f64>>,
4128) -> Result<(), EstimationError> {
4129 let n = x.nrows();
4130 let p = x.ncols();
4131 if penalty.nrows() != p || penalty.ncols() != p {
4132 crate::bail_invalid_estim!(
4133 "Gaussian REML penalty shape mismatch: expected {p}x{p}, got {}x{}",
4134 penalty.nrows(),
4135 penalty.ncols()
4136 );
4137 }
4138 if x.iter().chain(penalty.iter()).any(|v| !v.is_finite()) {
4139 crate::bail_invalid_estim!("Gaussian REML inputs must be finite");
4140 }
4141 if let Some(w) = weights {
4142 if w.len() != n {
4143 crate::bail_invalid_estim!(
4144 "Gaussian REML weights length mismatch: expected {n}, got {}",
4145 w.len()
4146 );
4147 }
4148 if w.iter().any(|value| !value.is_finite() || *value < 0.0) {
4149 crate::bail_invalid_estim!("Gaussian REML weights must be finite and non-negative");
4150 }
4151 }
4152 Ok(())
4153}
4154
4155fn effective_observation_count(weight: ArrayView1<'_, f64>) -> usize {
4169 weight.iter().filter(|&&w| w > 0.0).count()
4170}
4171
4172fn gaussian_reml_weights(
4173 n: usize,
4174 weights: Option<ArrayView1<'_, f64>>,
4175) -> Result<Array1<f64>, EstimationError> {
4176 match weights {
4177 Some(w) => {
4178 if w.len() != n {
4179 crate::bail_invalid_estim!(
4180 "Gaussian REML weights length mismatch: expected {n}, got {}",
4181 w.len()
4182 );
4183 }
4184 if w.iter().any(|value| !value.is_finite() || *value < 0.0) {
4185 crate::bail_invalid_estim!("Gaussian REML weights must be finite and non-negative");
4186 }
4187 Ok(w.to_owned())
4188 }
4189 None => Ok(Array1::ones(n)),
4190 }
4191}
4192
4193fn gaussian_reml_eigen_cache_from_xtwx(
4194 xtwx: Array2<f64>,
4195 penalty: ArrayView2<'_, f64>,
4196 nullspace_dim: Option<usize>,
4197) -> Result<GaussianRemlEigenCache, EstimationError> {
4198 let xtwx_fingerprint = matrix_fingerprint(xtwx.view());
4199 let lower = gaussian_reml_cholesky_lower(xtwx)?;
4200 gaussian_reml_eigen_cache_from_lower(lower, penalty, nullspace_dim, xtwx_fingerprint)
4201}
4202
4203fn gaussian_reml_eigen_cache_from_lower(
4208 lower: Array2<f64>,
4209 penalty: ArrayView2<'_, f64>,
4210 nullspace_dim: Option<usize>,
4211 xtwx_fingerprint: u64,
4212) -> Result<GaussianRemlEigenCache, EstimationError> {
4213 gaussian_reml_eigen_cache_from_lower_with_transform(
4214 lower,
4215 penalty,
4216 nullspace_dim,
4217 xtwx_fingerprint,
4218 None,
4219 )
4220}
4221
4222fn gaussian_reml_eigen_cache_from_lower_with_transform(
4225 lower: Array2<f64>,
4226 penalty: ArrayView2<'_, f64>,
4227 nullspace_dim: Option<usize>,
4228 xtwx_fingerprint: u64,
4229 precomputed_transform: Option<Array2<f64>>,
4230) -> Result<GaussianRemlEigenCache, EstimationError> {
4231 let p = lower.nrows();
4232 if lower.ncols() != p {
4233 crate::bail_invalid_estim!("Gaussian REML Cholesky factor must be square");
4234 }
4235 let penalty_fingerprint = matrix_fingerprint(penalty);
4236 let logdet_xtwx = 2.0 * lower.diag().iter().map(|v| v.ln()).sum::<f64>();
4237 let transformed_penalty = match precomputed_transform {
4238 Some(transformed) => transformed,
4239 None => {
4240 let l_inv = invert_lower_triangular(&lower)?;
4241 let penalty_in_metric = dense_ab(l_inv.view(), penalty);
4242 dense_ab(penalty_in_metric.view(), l_inv.t())
4243 }
4244 };
4245 let (mut penalty_eigenvalues, eigenvectors) =
4246 transformed_penalty.eigh(Side::Lower).map_err(|_| {
4247 EstimationError::ModelIsIllConditioned {
4248 condition_number: f64::INFINITY,
4249 }
4250 })?;
4251 let eig_tol = penalty_range_tolerance(penalty_eigenvalues.view());
4261 for value in &mut penalty_eigenvalues {
4262 if *value < 0.0 && value.abs() <= eig_tol {
4263 *value = 0.0;
4264 }
4265 if *value < 0.0 {
4266 crate::bail_invalid_estim!(
4267 "Gaussian REML penalty is not positive semidefinite; eigenvalue={value:.3e}"
4268 );
4269 }
4270 }
4271 let penalty_rank = penalty_eigenvalues
4272 .iter()
4273 .filter(|&&value| value > eig_tol)
4274 .count();
4275 let nullity = p - penalty_rank;
4276 if let Some(expected_nullity) = nullspace_dim
4277 && expected_nullity != nullity
4278 {
4279 crate::bail_invalid_estim!(
4280 "Gaussian REML penalty nullspace mismatch: expected {expected_nullity}, inferred {nullity}"
4281 );
4282 }
4283 let logdet_penalty_positive = gaussian_penalty_positive_logdet(penalty, penalty_rank)?;
4284 let coefficient_basis = solve_upper_triangular_matrix(&lower.t().to_owned(), &eigenvectors)?;
4285
4286 Ok(GaussianRemlEigenCache {
4287 penalty_eigenvalues,
4288 eigenvectors,
4289 coefficient_basis,
4290 xtwx_fingerprint,
4291 penalty_fingerprint,
4292 logdet_xtwx,
4293 logdet_penalty_positive,
4294 penalty_rank,
4295 nullity,
4296 })
4297}
4298
4299fn gaussian_reml_cholesky_lower(xtwx: Array2<f64>) -> Result<Array2<f64>, EstimationError> {
4300 let mut gpu_candidate = xtwx.clone();
4310 if gam_gpu::try_cholesky_lower_inplace(&mut gpu_candidate).is_some() {
4311 return Ok(gpu_candidate);
4312 }
4313 if let Ok(chol) = xtwx.cholesky(Side::Lower) {
4314 return Ok(chol.lower_triangular());
4315 }
4316 let p = xtwx.nrows();
4317 let trace: f64 = (0..p).map(|i| xtwx[[i, i]]).sum();
4318 if !trace.is_finite() || trace <= 0.0 {
4319 return Err(EstimationError::ModelIsIllConditioned {
4320 condition_number: f64::INFINITY,
4321 });
4322 }
4323 let schedule = RidgeSchedule::geometric(1e-12 * trace / (p as f64), 6);
4324 escalate_ridge(
4325 schedule,
4326 |jitter| {
4327 let mut jittered = xtwx.clone();
4328 for i in 0..p {
4329 jittered[[i, i]] += jitter;
4330 }
4331 let mut gpu_candidate = jittered.clone();
4332 if gam_gpu::try_cholesky_lower_inplace(&mut gpu_candidate).is_some() {
4333 return Some(gpu_candidate);
4334 }
4335 jittered
4336 .cholesky(Side::Lower)
4337 .ok()
4338 .map(|chol| chol.lower_triangular())
4339 },
4340 )
4341 .map(|success| success.value)
4342 .map_err(|exhausted| {
4343 let last_attempted = exhausted.next_ridge / schedule.growth;
4349 EstimationError::ModelIsIllConditioned {
4350 condition_number: if last_attempted > 0.0 && last_attempted.is_finite() {
4351 trace / last_attempted
4352 } else {
4353 f64::INFINITY
4354 },
4355 }
4356 })
4357}
4358
4359fn gaussian_penalty_positive_logdet(
4360 penalty: ArrayView2<'_, f64>,
4361 penalty_rank: usize,
4362) -> Result<f64, EstimationError> {
4363 if penalty_rank == 0 {
4364 return Ok(0.0);
4365 }
4366 let (pen_eigs, _) = penalty.to_owned().eigh(Side::Lower).map_err(|_| {
4367 EstimationError::ModelIsIllConditioned {
4368 condition_number: f64::INFINITY,
4369 }
4370 })?;
4371 let pen_tol = penalty_range_tolerance(pen_eigs.view());
4378 let mut positive_eigs: Vec<f64> = pen_eigs
4379 .iter()
4380 .copied()
4381 .filter(|&value| value > pen_tol)
4382 .collect();
4383 if positive_eigs.len() != penalty_rank {
4384 positive_eigs = pen_eigs
4385 .iter()
4386 .copied()
4387 .filter(|&value| value > 0.0)
4388 .collect();
4389 positive_eigs.sort_by(|a, b| b.total_cmp(a));
4390 if positive_eigs.len() < penalty_rank {
4391 return Err(EstimationError::ModelIsIllConditioned {
4392 condition_number: f64::INFINITY,
4393 });
4394 }
4395 positive_eigs.truncate(penalty_rank);
4396 }
4397 Ok(positive_eigs.iter().map(|value| value.ln()).sum())
4398}
4399
4400fn validate_gaussian_reml_eigen_cache(
4401 cache: &GaussianRemlEigenCache,
4402 p: usize,
4403) -> Result<(), EstimationError> {
4404 if cache.penalty_eigenvalues.len() != p
4405 || cache.eigenvectors.dim() != (p, p)
4406 || cache.coefficient_basis.dim() != (p, p)
4407 {
4408 crate::bail_invalid_estim!(
4409 "Gaussian REML eigen cache dimension mismatch: expected {p} coefficients"
4410 );
4411 }
4412 if cache.penalty_rank > p || cache.nullity > p || cache.penalty_rank + cache.nullity != p {
4413 crate::bail_invalid_estim!(
4414 "Gaussian REML eigen cache rank/nullity mismatch: rank={}, nullity={}, p={p}",
4415 cache.penalty_rank,
4416 cache.nullity
4417 );
4418 }
4419 if !(cache.logdet_xtwx.is_finite() && cache.logdet_penalty_positive.is_finite()) {
4420 crate::bail_invalid_estim!("Gaussian REML eigen cache log-determinants must be finite");
4421 }
4422 if cache
4423 .penalty_eigenvalues
4424 .iter()
4425 .any(|value| !value.is_finite() || *value < 0.0)
4426 || cache.eigenvectors.iter().any(|value| !value.is_finite())
4427 || cache
4428 .coefficient_basis
4429 .iter()
4430 .any(|value| !value.is_finite())
4431 {
4432 crate::bail_invalid_estim!(
4433 "Gaussian REML eigen cache entries must be finite with non-negative eigenvalues"
4434 .to_string(),
4435 );
4436 }
4437 let spectrum = PenaltyRangeSpectrum::of(cache);
4448 let classified_rank = spectrum.rank();
4449 if classified_rank != cache.penalty_rank {
4450 crate::bail_invalid_estim!(
4451 "Gaussian REML eigen cache reports penalty_rank={} but {classified_rank} of its {p} \
4452 eigenvalues clear the range tolerance {:e}; the log-determinant sums run over the \
4453 directions that clear it while log|S|₊ and the gradient offset are denominated in \
4454 penalty_rank, so the two must be the same count",
4455 cache.penalty_rank,
4456 spectrum.tolerance
4457 );
4458 }
4459 Ok::<(), _>(())
4460}
4461
4462fn prepare_gaussian_reml(
4463 x: ArrayView2<'_, f64>,
4464 y: ArrayView2<'_, f64>,
4465 penalty: ArrayView2<'_, f64>,
4466 nullspace_dim: Option<usize>,
4467 weights: Option<ArrayView1<'_, f64>>,
4468 eigen_cache: Option<&GaussianRemlEigenCache>,
4469) -> Result<GaussianRemlPrepared, EstimationError> {
4470 let penalty_owned = canonicalize_penalty(penalty);
4473 let penalty = penalty_owned.view();
4474 let n = x.nrows();
4475 let p = x.ncols();
4476 let d = y.ncols();
4477 validate_gaussian_reml_design(x, penalty, weights)?;
4478 if y.nrows() != n {
4479 crate::bail_invalid_estim!(
4480 "Gaussian REML row mismatch: X has {n} rows but Y has {}",
4481 y.nrows()
4482 );
4483 }
4484 if y.iter().any(|v| !v.is_finite()) {
4485 crate::bail_invalid_estim!("Gaussian REML inputs must be finite");
4486 }
4487 let weight = gaussian_reml_weights(n, weights)?;
4488 let n_effective = effective_observation_count(weight.view());
4489
4490 let xtwy = dense_xt_diag_y(x, weight.view(), y);
4491 let ywy = Array1::from_iter((0..d).map(|j| {
4492 let mut value = 0.0;
4493 for row in 0..n {
4494 value += weight[row] * y[[row, j]] * y[[row, j]];
4495 }
4496 value
4497 }));
4498 let xtwx = dense_xt_diag_x(x, weight.view());
4499
4500 if let Some(cache) = eigen_cache {
4501 validate_gaussian_reml_eigen_cache(cache, p)?;
4502 let xtwx_fingerprint = matrix_fingerprint(xtwx.view());
4503 if cache.xtwx_fingerprint != xtwx_fingerprint {
4504 crate::bail_invalid_estim!("Gaussian REML eigen cache X'WX mismatch");
4505 }
4506 let penalty_fingerprint = matrix_fingerprint(penalty);
4507 if cache.penalty_fingerprint != penalty_fingerprint {
4508 crate::bail_invalid_estim!("Gaussian REML eigen cache penalty mismatch");
4509 }
4510 if let Some(expected_nullity) = nullspace_dim
4511 && expected_nullity != cache.nullity
4512 {
4513 crate::bail_invalid_estim!(
4514 "Gaussian REML eigen cache nullspace mismatch: expected {expected_nullity}, got {}",
4515 cache.nullity
4516 );
4517 }
4518 if n_effective <= cache.nullity {
4519 crate::bail_invalid_estim!(
4520 "Gaussian REML requires more positive-weight rows than the nullspace dimension; got n_effective={n_effective}, nullity={}",
4521 cache.nullity
4522 );
4523 }
4524 let projected_rhs = dense_atb(cache.coefficient_basis.view(), xtwy.view());
4525 let projected_rhs_squared = projected_rhs.mapv(|value| value * value);
4526 return Ok(GaussianRemlPrepared {
4527 cache: cache.clone(),
4528 ywy,
4529 projected_rhs_squared,
4530 projected_rhs,
4531 n_effective,
4532 n_outputs: d,
4533 observation_measure: gaussian_reml_observation_measure(weight.view(), d),
4534 });
4535 }
4536
4537 let cache = gaussian_reml_eigen_cache_from_xtwx(xtwx, penalty, nullspace_dim)?;
4538 if n_effective <= cache.nullity {
4539 crate::bail_invalid_estim!(
4540 "Gaussian REML requires more positive-weight rows than the nullspace dimension; got n_effective={n_effective}, nullity={}",
4541 cache.nullity
4542 );
4543 }
4544 let projected_rhs = dense_atb(cache.coefficient_basis.view(), xtwy.view());
4545 let projected_rhs_squared = projected_rhs.mapv(|value| value * value);
4546
4547 Ok(GaussianRemlPrepared {
4548 cache,
4549 ywy,
4550 projected_rhs_squared,
4551 projected_rhs,
4552 n_effective,
4553 n_outputs: d,
4554 observation_measure: gaussian_reml_observation_measure(weight.view(), d),
4555 })
4556}
4557
4558impl GaussianRemlPrepared {
4559 fn nu(&self) -> f64 {
4560 self.n_effective as f64 - self.cache.nullity as f64
4561 }
4562
4563 fn evaluate(&self, rho: f64) -> ObjectiveEval {
4564 let mut value = evaluate_reml_parts(
4565 &self.cache,
4566 self.ywy.view(),
4567 self.projected_rhs_squared.view(),
4568 self.n_effective,
4569 self.n_outputs,
4570 rho,
4571 );
4572 value += self.observation_measure;
4573 value
4574 }
4575
4576 fn coefficients(&self, lambda: f64) -> Array2<f64> {
4577 let mut scaled = self.projected_rhs.clone();
4578 let spectrum = PenaltyRangeSpectrum::of(&self.cache);
4579 for i in 0..spectrum.len() {
4580 let scale = 1.0 / (1.0 + lambda * spectrum.get(i));
4581 for value in scaled.row_mut(i) {
4582 *value *= scale;
4583 }
4584 }
4585 dense_ab(self.cache.coefficient_basis.view(), scaled.view())
4586 }
4587
4588 fn sigma2(&self, rho: f64) -> Array1<f64> {
4594 let nu = self.nu();
4595 Array1::from_iter((0..self.n_outputs).map(|j| {
4596 let DispersionResidualParts {
4597 unpenalized_residual,
4598 penalized_residual,
4599 ..
4600 } = dispersion_residual_parts(
4601 &self.cache,
4602 self.ywy.view(),
4603 self.projected_rhs_squared.view(),
4604 j,
4605 rho,
4606 );
4607 (unpenalized_residual + penalized_residual) / nu
4608 }))
4609 }
4610}
4611
4612fn profile_residual_resolution(cache: &GaussianRemlEigenCache, ywy_output: f64) -> f64 {
4632 let operations =
4633 64usize.saturating_add(32usize.saturating_mul(cache.penalty_eigenvalues.len()));
4634 let n_eps = (operations as f64) * f64::EPSILON;
4635 if !(ywy_output.is_finite() && ywy_output >= 0.0) || n_eps >= 1.0 {
4636 return f64::INFINITY;
4637 }
4638 (n_eps / (1.0 - n_eps)) * 2.0 * ywy_output
4639}
4640
4641fn validate_reml_profile_residuals(
4675 cache: &GaussianRemlEigenCache,
4676 ywy: ArrayView1<'_, f64>,
4677 projected_rhs_squared: ArrayView2<'_, f64>,
4678 rho: f64,
4679) -> Result<(), EstimationError> {
4680 for output in 0..ywy.len() {
4681 let DispersionResidualParts {
4686 unpenalized_residual,
4687 penalized_residual,
4688 ..
4689 } = dispersion_residual_parts(cache, ywy, projected_rhs_squared, output, rho);
4690 let residual = unpenalized_residual + penalized_residual;
4691 let resolution = profile_residual_resolution(cache, ywy[output]);
4692 if !(residual.is_finite() && residual > resolution) {
4693 return Err(EstimationError::InvalidInput(format!(
4694 "Gaussian REML profiled residual {output} is not resolvably positive at rho={rho}: {residual} against its own arithmetic resolution {resolution} (gamma_m * 2 * ywy, ywy={}); the design interpolates its response, so the profiled dispersion has no finite value",
4695 ywy[output]
4696 )));
4697 }
4698 }
4699 Ok(())
4700}
4701
4702const RHO_BRACKET_RESOLUTION: f64 = 1.0e-12;
4790
4791const fn dfs_max_depth(range: f64, resolution: f64) -> usize {
4796 let mut width = range;
4797 let mut depth = 0usize;
4798 while width > resolution {
4799 width *= 0.5;
4800 depth += 1;
4801 }
4802 depth
4803}
4804
4805const MAX_DEPTH: usize = dfs_max_depth(RHO_UPPER - RHO_LOWER, RHO_BRACKET_RESOLUTION);
4807
4808#[derive(Clone, Copy)]
4810struct Interval {
4811 lo: f64,
4812 hi: f64,
4813}
4814
4815impl Interval {
4816 fn entire() -> Self {
4817 Self {
4818 lo: f64::NEG_INFINITY,
4819 hi: f64::INFINITY,
4820 }
4821 }
4822}
4823
4824fn round_down(x: f64) -> f64 {
4827 if x.is_nan() || x == f64::NEG_INFINITY {
4828 return x;
4829 }
4830 if x == 0.0 {
4831 return -f64::from_bits(1);
4832 }
4833 let bits = x.to_bits();
4834 let next = if x > 0.0 { bits - 1 } else { bits + 1 };
4835 f64::from_bits(next)
4836}
4837
4838fn round_up(x: f64) -> f64 {
4841 if x.is_nan() || x == f64::INFINITY {
4842 return x;
4843 }
4844 if x == 0.0 {
4845 return f64::from_bits(1);
4846 }
4847 let bits = x.to_bits();
4848 let next = if x > 0.0 { bits + 1 } else { bits - 1 };
4849 f64::from_bits(next)
4850}
4851
4852fn add_down(lhs: f64, rhs: f64) -> f64 {
4853 round_down(lhs + rhs)
4854}
4855
4856fn add_up(lhs: f64, rhs: f64) -> f64 {
4857 round_up(lhs + rhs)
4858}
4859
4860fn nonnegative_product_interval(lhs: f64, rhs: Interval) -> Option<Interval> {
4864 if !(lhs.is_finite()
4865 && lhs >= 0.0
4866 && rhs.lo.is_finite()
4867 && rhs.hi.is_finite()
4868 && rhs.lo >= 0.0
4869 && rhs.hi >= rhs.lo)
4870 {
4871 return None;
4872 }
4873 Some(Interval {
4874 lo: round_down(lhs * rhs.lo).max(0.0),
4875 hi: round_up(lhs * rhs.hi),
4876 })
4877}
4878
4879fn nonnegative_square_interval(bounds: Interval) -> Option<Interval> {
4881 if !(bounds.lo.is_finite()
4882 && bounds.hi.is_finite()
4883 && bounds.lo >= 0.0
4884 && bounds.hi >= bounds.lo)
4885 {
4886 return None;
4887 }
4888 Some(Interval {
4889 lo: round_down(bounds.lo * bounds.lo).max(0.0),
4890 hi: round_up(bounds.hi * bounds.hi),
4891 })
4892}
4893
4894fn conservative_interval(lo: f64, hi: f64, magnitude: f64, operations: usize) -> Interval {
4900 if !(lo.is_finite() && hi.is_finite() && magnitude.is_finite() && lo <= hi) {
4901 return Interval::entire();
4902 }
4903 let n_eps = (operations as f64) * f64::EPSILON;
4904 if n_eps >= 1.0 {
4905 return Interval::entire();
4906 }
4907 let pad =
4908 (n_eps / (1.0 - n_eps)) * magnitude.max(lo.abs()).max(hi.abs()).max(f64::MIN_POSITIVE);
4909 Interval {
4910 lo: round_down(lo - pad),
4911 hi: round_up(hi + pad),
4912 }
4913}
4914
4915#[derive(Clone, Copy)]
4925struct KernelRange {
4926 u_lo: f64,
4927 u_hi: f64,
4928 w_lo: f64,
4929 w_hi: f64,
4930 k_lo: f64,
4931 k_hi: f64,
4932}
4933
4934fn kernel_ranges(log_t_lo: f64, log_t_hi: f64) -> KernelRange {
4935 let kernels = |log_t: f64| modal_kernels(log_t, 1.0);
4936 let left = kernels(log_t_lo);
4937 let right = kernels(log_t_hi);
4938
4939 let u_lo = left.u;
4942 let u_hi = right.u;
4943
4944 let w_a = left.w;
4946 let w_b = right.w;
4947 let w_lo = w_a.min(w_b);
4948 let w_hi = if log_t_lo <= 0.0 && 0.0 <= log_t_hi {
4949 0.25
4950 } else {
4951 w_a.max(w_b)
4952 };
4953
4954 let sqrt3 = 3.0_f64.sqrt();
4957 let cp_lo = (2.0 - sqrt3).ln();
4958 let cp_hi = (2.0 + sqrt3).ln();
4959 let mut k_lo = left.k.min(right.k);
4960 let mut k_hi = left.k.max(right.k);
4961 if log_t_lo < cp_lo && cp_lo < log_t_hi {
4962 let kc = kernels(cp_lo).k;
4963 k_lo = k_lo.min(kc);
4964 k_hi = k_hi.max(kc);
4965 }
4966 if log_t_lo < cp_hi && cp_hi < log_t_hi {
4967 let kc = kernels(cp_hi).k;
4968 k_lo = k_lo.min(kc);
4969 k_hi = k_hi.max(kc);
4970 }
4971
4972 KernelRange {
4973 u_lo: round_down(u_lo).max(0.0),
4974 u_hi: round_up(u_hi),
4975 w_lo: round_down(w_lo).max(0.0),
4976 w_hi: round_up(w_hi),
4977 k_lo: round_down(k_lo),
4978 k_hi: round_up(k_hi),
4979 }
4980}
4981
4982fn reml_deriv_enclosure(
4990 cache: &GaussianRemlEigenCache,
4991 ywy: ArrayView1<'_, f64>,
4992 projected_rhs_squared: ArrayView2<'_, f64>,
4993 n_effective: usize,
4994 n_outputs: usize,
4995 a: f64,
4996 b: f64,
4997) -> (Interval, Interval) {
4998 reml_deriv_enclosure_profile(
4999 cache,
5000 ywy,
5001 projected_rhs_squared,
5002 n_outputs,
5003 n_effective as f64 - cache.nullity as f64,
5004 a,
5005 b,
5006 )
5007}
5008
5009fn reml_deriv_enclosure_profile(
5016 cache: &GaussianRemlEigenCache,
5017 ywy: ArrayView1<'_, f64>,
5018 projected_rhs_squared: ArrayView2<'_, f64>,
5019 logdet_output_count: usize,
5020 dispersion_dof: f64,
5021 a: f64,
5022 b: f64,
5023) -> (Interval, Interval) {
5024 let d = logdet_output_count as f64;
5025 let spectrum = PenaltyRangeSpectrum::of(cache);
5026 let rank = cache.penalty_rank as f64;
5027 let half_d = 0.5 * d;
5028 let half_nu = 0.5 * dispersion_dof;
5029 let mut sum_u_lo = 0.0;
5032 let mut sum_u_hi = 0.0;
5033 let mut sum_w_lo = 0.0;
5034 let mut sum_w_hi = 0.0;
5035 for delta in spectrum.iter() {
5039 if delta > 0.0 {
5040 let log_delta = delta.ln();
5041 let kr = kernel_ranges(a + log_delta, b + log_delta);
5042 sum_u_lo = add_down(sum_u_lo, kr.u_lo);
5043 sum_u_hi = add_up(sum_u_hi, kr.u_hi);
5044 sum_w_lo = add_down(sum_w_lo, kr.w_lo);
5045 sum_w_hi = add_up(sum_w_hi, kr.w_hi);
5046 }
5047 }
5048 let g1_lo = round_down(half_d * round_down(sum_u_lo - rank));
5049 let g1_hi = round_up(half_d * round_up(sum_u_hi - rank));
5050
5051 let mut g2_lo = 0.0;
5054 let mut g2_hi = 0.0;
5055 let mut vpp_disp_lo = 0.0;
5056 let mut vpp_disp_hi = 0.0;
5057 for j in 0..ywy.len() {
5058 let mut num_lo = 0.0; let mut num_hi = 0.0;
5060 let mut su_lo = 0.0; let mut su_hi = 0.0;
5062 let mut c2_point = 0.0;
5067 let mut dph_lo = 0.0; let mut dph_hi = 0.0;
5069 for eig in 0..spectrum.len() {
5070 let delta = spectrum.get(eig);
5071 let c2 = projected_rhs_squared[[eig, j]];
5072 let log_delta = if delta == 0.0 {
5073 f64::NEG_INFINITY
5074 } else {
5075 delta.ln()
5076 };
5077 let kr = kernel_ranges(a + log_delta, b + log_delta);
5078 let Some(w_product) = nonnegative_product_interval(
5079 c2,
5080 Interval {
5081 lo: kr.w_lo,
5082 hi: kr.w_hi,
5083 },
5084 ) else {
5085 return (Interval::entire(), Interval::entire());
5086 };
5087 let Some(u_product) = nonnegative_product_interval(
5088 c2,
5089 Interval {
5090 lo: kr.u_lo,
5091 hi: kr.u_hi,
5092 },
5093 ) else {
5094 return (Interval::entire(), Interval::entire());
5095 };
5096 num_lo = add_down(num_lo, w_product.lo);
5097 num_hi = add_up(num_hi, w_product.hi);
5098 su_lo = add_down(su_lo, u_product.lo);
5099 su_hi = add_up(su_hi, u_product.hi);
5100 c2_point += c2;
5101 dph_lo = add_down(dph_lo, round_down(c2 * kr.k_lo));
5102 dph_hi = add_up(dph_hi, round_up(c2 * kr.k_hi));
5103 }
5104 let r0 = (ywy[j] - c2_point).max(0.0);
5146 let dp_lo = add_down(r0, su_lo);
5147 let dp_hi = add_up(r0, su_hi);
5148 if !(dp_lo.is_finite() && dp_hi.is_finite() && dp_lo > 0.0 && dp_hi >= dp_lo) {
5149 return (Interval::entire(), Interval::entire());
5150 }
5151
5152 let ratio_lo = round_down(num_lo / dp_hi).max(0.0);
5154 let ratio_hi = round_up(num_hi / dp_lo);
5155 g2_lo = add_down(g2_lo, ratio_lo);
5156 g2_hi = add_up(g2_hi, ratio_hi);
5157
5158 let quotients = [
5161 dph_lo / dp_lo,
5162 dph_lo / dp_hi,
5163 dph_hi / dp_lo,
5164 dph_hi / dp_hi,
5165 ];
5166 let adp_lo = round_down(quotients.iter().copied().fold(f64::INFINITY, f64::min));
5167 let adp_hi = round_up(quotients.iter().copied().fold(f64::NEG_INFINITY, f64::max));
5168
5169 let bl = round_down(num_lo / dp_hi).max(0.0);
5171 let bh = round_up(num_hi / dp_lo);
5172 let Some(squared_ratio) = nonnegative_square_interval(Interval { lo: bl, hi: bh }) else {
5173 return (Interval::entire(), Interval::entire());
5174 };
5175
5176 vpp_disp_lo = add_down(vpp_disp_lo, round_down(adp_lo - squared_ratio.hi));
5178 vpp_disp_hi = add_up(vpp_disp_hi, round_up(adp_hi - squared_ratio.lo));
5179 }
5180
5181 let vp_lo = add_down(g1_lo, round_down(half_nu * g2_lo));
5182 let vp_hi = add_up(g1_hi, round_up(half_nu * g2_hi));
5183 let vpp_lo = add_down(
5184 round_down(half_d * sum_w_lo),
5185 round_down(half_nu * vpp_disp_lo),
5186 );
5187 let vpp_hi = add_up(round_up(half_d * sum_w_hi), round_up(half_nu * vpp_disp_hi));
5188
5189 let operations = 64usize.saturating_add(
5190 32usize.saturating_mul(
5191 cache
5192 .penalty_eigenvalues
5193 .len()
5194 .saturating_mul(ywy.len().max(1)),
5195 ),
5196 );
5197 let vp_magnitude = g1_lo.abs() + g1_hi.abs() + half_nu.abs() * (g2_lo.abs() + g2_hi.abs());
5198 let vpp_magnitude = half_d.abs() * (sum_w_lo.abs() + sum_w_hi.abs())
5199 + half_nu.abs() * (vpp_disp_lo.abs() + vpp_disp_hi.abs());
5200 (
5201 conservative_interval(vp_lo, vp_hi, vp_magnitude, operations),
5202 conservative_interval(vpp_lo, vpp_hi, vpp_magnitude, operations),
5203 )
5204}
5205
5206#[derive(Clone, Copy, Debug)]
5207struct StationaryRoot {
5208 rho: f64,
5209 bracket: [f64; 2],
5210}
5211
5212#[derive(Clone, Copy, Debug)]
5213struct ProfileSelection {
5214 rho: f64,
5215}
5216
5217#[derive(Clone, Copy)]
5218struct ProfileSearchControls {
5219 lower: f64,
5220 upper: f64,
5221 resolution: f64,
5222 max_depth: usize,
5223}
5224
5225impl ProfileSearchControls {
5226 const PRODUCTION: Self = Self {
5227 lower: RHO_LOWER,
5228 upper: RHO_UPPER,
5229 resolution: RHO_BRACKET_RESOLUTION,
5230 max_depth: MAX_DEPTH,
5231 };
5232}
5233
5234fn profile_search_refusal(
5235 eval: &impl Fn(f64) -> ObjectiveEval,
5236 checkpoint: f64,
5237 reason: String,
5238) -> EstimationError {
5239 let e = eval(checkpoint);
5240 EstimationError::RemlDidNotConverge {
5241 context: "closed-form Gaussian profiled REML stationary search".to_string(),
5242 reason,
5243 iterations: 0,
5244 final_value: e.cost,
5245 projected_grad_norm: e.grad.is_finite().then_some(e.grad.abs()),
5246 stationarity_standard: StationarityStandard::NoComparison,
5259 rho_checkpoint: vec![checkpoint],
5260 }
5261}
5262
5263fn refine_stationary_rho_core(
5269 eval: &impl Fn(f64) -> ObjectiveEval,
5270 mut lo: f64,
5271 mut hi: f64,
5272 resolution: f64,
5273 mut hint: Option<f64>,
5274) -> Result<StationaryRoot, EstimationError> {
5275 let mut left = eval(lo);
5276 let mut right = eval(hi);
5277 if left.grad == 0.0 {
5278 return Ok(StationaryRoot {
5279 rho: lo,
5280 bracket: [lo, lo],
5281 });
5282 }
5283 if right.grad == 0.0 {
5284 return Ok(StationaryRoot {
5285 rho: hi,
5286 bracket: [hi, hi],
5287 });
5288 }
5289 if left.grad.is_sign_positive() == right.grad.is_sign_positive() {
5290 return Err(profile_search_refusal(
5291 eval,
5292 0.5 * (lo + hi),
5293 format!("stationary refinement received a non-bracketing cell [{lo}, {hi}]"),
5294 ));
5295 }
5296
5297 loop {
5298 let width = hi - lo;
5299 let scale = 1.0 + lo.abs().max(hi.abs());
5300 if width <= resolution * scale {
5301 let midpoint = lo + 0.5 * width;
5302 let middle = if midpoint > lo && midpoint < hi {
5303 Some((midpoint, eval(midpoint)))
5304 } else {
5305 None
5306 };
5307 let mut representative = (lo, left);
5308 if right.grad.abs() < representative.1.grad.abs() {
5309 representative = (hi, right);
5310 }
5311 if let Some(candidate) = middle
5312 && candidate.1.grad.abs() < representative.1.grad.abs()
5313 {
5314 representative = candidate;
5315 }
5316 return Ok(StationaryRoot {
5317 rho: representative.0,
5318 bracket: [lo, hi],
5319 });
5320 }
5321
5322 let midpoint = lo + 0.5 * width;
5323 if !(midpoint > lo && midpoint < hi) {
5324 return Err(profile_search_refusal(
5325 eval,
5326 midpoint,
5327 format!(
5328 "stationary root on [{lo}, {hi}] reached floating-point spacing before rho resolution {resolution}"
5329 ),
5330 ));
5331 }
5332 let guard = 0.25 * width;
5333 let base = if left.grad.abs() <= right.grad.abs() {
5334 (lo, left)
5335 } else {
5336 (hi, right)
5337 };
5338 let newton = if base.1.hess != 0.0 {
5339 base.0 - base.1.grad / base.1.hess
5340 } else {
5341 f64::NAN
5342 };
5343 let candidate = hint
5344 .take()
5345 .filter(|&rho| rho >= lo + guard && rho <= hi - guard)
5346 .or_else(|| {
5347 (newton.is_finite() && newton >= lo + guard && newton <= hi - guard)
5348 .then_some(newton)
5349 })
5350 .unwrap_or(midpoint);
5351 if !(candidate > lo && candidate < hi) {
5352 return Err(profile_search_refusal(
5353 eval,
5354 midpoint,
5355 format!(
5356 "stationary refinement could not represent an interior point on [{lo}, {hi}]"
5357 ),
5358 ));
5359 }
5360 let current = eval(candidate);
5361 if current.grad == 0.0 {
5362 return Ok(StationaryRoot {
5363 rho: candidate,
5364 bracket: [candidate, candidate],
5365 });
5366 }
5367 if current.grad.is_sign_positive() == left.grad.is_sign_positive() {
5368 lo = candidate;
5369 left = current;
5370 } else {
5371 hi = candidate;
5372 right = current;
5373 }
5374 }
5375}
5376
5377fn intersect_intervals(left: Interval, right: Interval) -> Interval {
5382 let lo = if right.lo.is_nan() { left.lo } else { left.lo.max(right.lo) };
5383 let hi = if right.hi.is_nan() { left.hi } else { left.hi.min(right.hi) };
5384 if lo > hi { left } else { Interval { lo, hi } }
5385}
5386
5387fn mean_value_derivative_enclosure(
5399 at_a: Interval,
5400 at_b: Interval,
5401 curvature: Interval,
5402 h: f64,
5403) -> Interval {
5404 if !(h.is_finite()
5405 && h >= 0.0
5406 && curvature.lo.is_finite()
5407 && curvature.hi.is_finite()
5408 && at_a.lo.is_finite()
5409 && at_a.hi.is_finite()
5410 && at_b.lo.is_finite()
5411 && at_b.hi.is_finite())
5412 {
5413 return Interval::entire();
5414 }
5415 let down = round_down(curvature.lo * h).min(0.0);
5416 let up = round_up(curvature.hi * h).max(0.0);
5417 let from_a = Interval {
5418 lo: round_down(at_a.lo + down),
5419 hi: round_up(at_a.hi + up),
5420 };
5421 let from_b = Interval {
5422 lo: round_down(at_b.lo - up),
5423 hi: round_up(at_b.hi - down),
5424 };
5425 intersect_intervals(from_a, from_b)
5426}
5427
5428fn widen_to_include(interval: Interval, first: f64, second: f64) -> Interval {
5437 let mut out = interval;
5438 for value in [first, second] {
5439 if value.is_finite() {
5440 out.lo = out.lo.min(value);
5441 out.hi = out.hi.max(value);
5442 }
5443 }
5444 out
5445}
5446
5447fn interval_contains(interval: Interval, value: f64) -> bool {
5448 value.is_finite() && interval.lo <= value && value <= interval.hi
5449}
5450
5451fn enumerate_and_select_rho_with_controls(
5456 eval: impl Fn(f64) -> ObjectiveEval,
5457 enclose: impl Fn(f64, f64) -> (Interval, Interval),
5458 init_rho: Option<f64>,
5459 controls: ProfileSearchControls,
5460 mut visit: Option<&mut dyn FnMut(StationaryRoot, &ObjectiveEval)>,
5461) -> Result<ProfileSelection, EstimationError> {
5462 const CAP: usize = MAX_DEPTH + 4;
5463 let lower_eval = eval(controls.lower);
5464 let upper_eval = eval(controls.upper);
5465 let lower_point = enclose(controls.lower, controls.lower).0;
5479 let upper_point = enclose(controls.upper, controls.upper).0;
5480 let mut stack = [(
5481 controls.lower,
5482 lower_eval,
5483 lower_point,
5484 controls.upper,
5485 upper_eval,
5486 upper_point,
5487 0usize,
5488 ); CAP];
5489 let mut top = 1usize;
5490
5491 let (mut best_rho, mut best_eval) = if upper_eval.cost < lower_eval.cost {
5492 (controls.upper, upper_eval)
5493 } else {
5494 (controls.lower, lower_eval)
5495 };
5496 let mut last_root: Option<StationaryRoot> = None;
5497 let mut cells_visited = 0usize;
5502 let mut evaluations = 2usize;
5503 let mut deepest = 0usize;
5504 let mut unbounded_enclosures = 0usize;
5505
5506 while top > 0 {
5507 top -= 1;
5508 let (a, ea, pa, b, eb, pb, depth) = stack[top];
5509 cells_visited += 1;
5510 deepest = deepest.max(depth);
5511 let (direct_dv, dvv) = enclose(a, b);
5512 let dv = widen_to_include(
5535 intersect_intervals(
5536 direct_dv,
5537 mean_value_derivative_enclosure(pa, pb, dvv, b - a),
5538 ),
5539 ea.grad,
5540 eb.grad,
5541 );
5542 if !(dv.lo.is_finite() && dv.hi.is_finite()) {
5543 unbounded_enclosures += 1;
5544 }
5545 if !(interval_contains(dv, ea.grad)
5546 && interval_contains(dv, eb.grad)
5547 && interval_contains(dvv, ea.hess)
5548 && interval_contains(dvv, eb.hess))
5549 {
5550 return Err(profile_search_refusal(
5551 &eval,
5552 0.5 * (a + b),
5553 format!(
5554 "analytic derivative enclosure [{}, {}] / curvature enclosure [{}, {}] missed an endpoint jet on [{a}, {b}]",
5555 dv.lo, dv.hi, dvv.lo, dvv.hi
5556 ),
5557 ));
5558 }
5559 if dv.lo > 0.0 || dv.hi < 0.0 {
5560 continue;
5561 }
5562
5563 let monotone = dvv.lo > 0.0 || dvv.hi < 0.0;
5564 let at_floor = depth >= controls.max_depth
5565 || (b - a) <= controls.resolution * (1.0 + a.abs().max(b.abs()));
5566 if !monotone && at_floor {
5567 return Err(profile_search_refusal(
5568 &eval,
5569 0.5 * (a + b),
5570 format!(
5571 "stationary structure remained non-monotone on [{a}, {b}] at rho resolution {} \
5572 after {cells_visited} branch-and-bound cells ({evaluations} objective \
5573 evaluations, deepest bisection {deepest}, {unbounded_enclosures} cells whose \
5574 derivative enclosure was unbounded)",
5575 controls.resolution
5576 ),
5577 ));
5578 }
5579
5580 if monotone {
5581 let crosses = (ea.grad <= 0.0 && eb.grad >= 0.0) || (ea.grad >= 0.0 && eb.grad <= 0.0);
5582 if crosses {
5583 let hint = init_rho.filter(|rho| rho.is_finite() && *rho >= a && *rho <= b);
5584 let root = refine_stationary_rho_core(&eval, a, b, controls.resolution, hint)?;
5585 let duplicate = last_root.is_some_and(|previous| {
5586 root.rho.to_bits() == previous.rho.to_bits()
5587 || (root.bracket[0] <= previous.bracket[1]
5588 && previous.bracket[0] <= root.bracket[1])
5589 });
5590 if !duplicate {
5591 let e = eval(root.rho);
5592 if e.cost < best_eval.cost {
5593 best_rho = root.rho;
5594 best_eval = e;
5595 }
5596 if let Some(observer) = visit.as_deref_mut() {
5597 observer(root, &e);
5598 }
5599 last_root = Some(root);
5600 }
5601 }
5602 continue;
5603 }
5604
5605 let mid = a + 0.5 * (b - a);
5606 if !(mid > a && mid < b) || top + 2 > CAP {
5607 return Err(profile_search_refusal(
5608 &eval,
5609 mid,
5610 format!("stationary subdivision could not continue on [{a}, {b}]"),
5611 ));
5612 }
5613 let emid = eval(mid);
5614 let pmid = enclose(mid, mid).0;
5615 evaluations += 1;
5616 stack[top] = (mid, emid, pmid, b, eb, pb, depth + 1);
5617 top += 1;
5618 stack[top] = (a, ea, pa, mid, emid, pmid, depth + 1);
5619 top += 1;
5620 }
5621 log::info!(
5622 "[REML-BNB] certified 1-D rho search over [{}, {}]: {cells_visited} cells, \
5623 {evaluations} objective evaluations, deepest bisection {deepest}/{}, \
5624 {unbounded_enclosures} unbounded enclosures",
5625 controls.lower,
5626 controls.upper,
5627 controls.max_depth,
5628 );
5629
5630 if !(best_eval.cost.is_finite() && best_eval.grad.is_finite()) {
5631 return Err(EstimationError::InvalidInput(
5632 "Gaussian REML profiled search produced no finite candidate".to_string(),
5633 ));
5634 }
5635 Ok(ProfileSelection { rho: best_rho })
5636}
5637
5638fn enumerate_and_select_rho(
5639 eval: impl Fn(f64) -> ObjectiveEval,
5640 enclose: impl Fn(f64, f64) -> (Interval, Interval),
5641 init_rho: Option<f64>,
5642 visit: Option<&mut dyn FnMut(StationaryRoot, &ObjectiveEval)>,
5643) -> Result<ProfileSelection, EstimationError> {
5644 enumerate_and_select_rho_with_controls(
5645 eval,
5646 enclose,
5647 init_rho,
5648 ProfileSearchControls::PRODUCTION,
5649 visit,
5650 )
5651}
5652
5653fn optimize_rho(
5655 prepared: &GaussianRemlPrepared,
5656 init_rho: Option<f64>,
5657) -> Result<f64, EstimationError> {
5658 validate_reml_profile_residuals(
5659 &prepared.cache,
5660 prepared.ywy.view(),
5661 prepared.projected_rhs_squared.view(),
5662 RHO_LOWER,
5663 )?;
5664 if prepared.cache.penalty_rank == 0 {
5665 return Ok(init_rho.unwrap_or(0.0).clamp(RHO_LOWER, RHO_UPPER));
5666 }
5667 let eval = |rho: f64| prepared.evaluate(rho);
5668 let enclose = |a: f64, b: f64| {
5669 reml_deriv_enclosure(
5670 &prepared.cache,
5671 prepared.ywy.view(),
5672 prepared.projected_rhs_squared.view(),
5673 prepared.n_effective,
5674 prepared.n_outputs,
5675 a,
5676 b,
5677 )
5678 };
5679 Ok(enumerate_and_select_rho(eval, enclose, init_rho, None)?.rho)
5680}
5681
5682fn evaluate_reml_parts(
5683 cache: &GaussianRemlEigenCache,
5684 ywy: ArrayView1<'_, f64>,
5685 projected_rhs_squared: ArrayView2<'_, f64>,
5686 n_effective: usize,
5687 n_outputs: usize,
5688 rho: f64,
5689) -> ObjectiveEval {
5690 evaluate_reml_profile(
5691 cache,
5692 ywy,
5693 projected_rhs_squared,
5694 n_outputs,
5695 n_effective as f64 - cache.nullity as f64,
5696 rho,
5697 )
5698}
5699
5700fn evaluate_reml_profile(
5704 cache: &GaussianRemlEigenCache,
5705 ywy: ArrayView1<'_, f64>,
5706 projected_rhs_squared: ArrayView2<'_, f64>,
5707 logdet_output_count: usize,
5708 dispersion_dof: f64,
5709 rho: f64,
5710) -> ObjectiveEval {
5711 let d = logdet_output_count as f64;
5712
5713 let (logdet_term, edf) = gaussian_reml_logdet_term(cache, rho, d);
5716 let mut eval = ObjectiveEval {
5717 cost: 0.0,
5718 grad: 0.0,
5719 hess: 0.0,
5720 edf,
5721 cost_roundoff: 0.0,
5722 };
5723 eval += logdet_term;
5724 for output in 0..ywy.len() {
5725 eval += gaussian_reml_dispersion_term(
5726 cache,
5727 ywy,
5728 projected_rhs_squared,
5729 output,
5730 dispersion_dof,
5731 rho,
5732 );
5733 }
5734 eval
5735}
5736
5737fn invert_lower_triangular(lower: &Array2<f64>) -> Result<Array2<f64>, EstimationError> {
5738 let n = lower.nrows();
5739 if lower.ncols() != n {
5740 crate::bail_invalid_estim!("lower-triangular solve requires a square matrix");
5741 }
5742 let eye = Array2::eye(n);
5743 solve_lower_triangular_matrix(lower, &eye)
5744}
5745
5746fn solve_lower_triangular_matrix(
5747 lower: &Array2<f64>,
5748 rhs: &Array2<f64>,
5749) -> Result<Array2<f64>, EstimationError> {
5750 let n = lower.nrows();
5751 if lower.ncols() != n || rhs.nrows() != n {
5752 crate::bail_invalid_estim!("lower-triangular solve dimension mismatch");
5753 }
5754 if let Some(out) = gam_gpu::try_solve_lower_triangular_matrix(lower.view(), rhs.view()) {
5755 return Ok(out);
5756 }
5757 let mut out = Array2::<f64>::zeros(rhs.dim());
5758 for col in 0..rhs.ncols() {
5759 for i in 0..n {
5760 let mut value = rhs[[i, col]];
5761 for k in 0..i {
5762 value -= lower[[i, k]] * out[[k, col]];
5763 }
5764 let diag = lower[[i, i]];
5765 if !(diag.is_finite() && diag.abs() > 0.0) {
5766 return Err(EstimationError::ModelIsIllConditioned {
5767 condition_number: f64::INFINITY,
5768 });
5769 }
5770 out[[i, col]] = value / diag;
5771 }
5772 }
5773 Ok(out)
5774}
5775
5776fn solve_spd_from_lower_factor(
5780 lower: &Array2<f64>,
5781 rhs: &Array2<f64>,
5782) -> Result<Array2<f64>, EstimationError> {
5783 let forward = solve_lower_triangular_matrix(lower, rhs)?;
5784 solve_upper_triangular_matrix(&lower.t().to_owned(), &forward)
5785}
5786
5787fn solve_upper_triangular_matrix(
5788 upper: &Array2<f64>,
5789 rhs: &Array2<f64>,
5790) -> Result<Array2<f64>, EstimationError> {
5791 let n = upper.nrows();
5792 if upper.ncols() != n || rhs.nrows() != n {
5793 crate::bail_invalid_estim!("upper-triangular solve dimension mismatch");
5794 }
5795 if let Some(out) = gam_gpu::try_solve_upper_triangular_matrix(upper.view(), rhs.view()) {
5796 return Ok(out);
5797 }
5798 let mut out = Array2::<f64>::zeros(rhs.dim());
5799 for col in 0..rhs.ncols() {
5800 for i_rev in 0..n {
5801 let i = n - 1 - i_rev;
5802 let mut value = rhs[[i, col]];
5803 for k in (i + 1)..n {
5804 value -= upper[[i, k]] * out[[k, col]];
5805 }
5806 let diag = upper[[i, i]];
5807 if !(diag.is_finite() && diag.abs() > 0.0) {
5808 return Err(EstimationError::ModelIsIllConditioned {
5809 condition_number: f64::INFINITY,
5810 });
5811 }
5812 out[[i, col]] = value / diag;
5813 }
5814 }
5815 Ok(out)
5816}
5817
5818#[cfg(test)]
5819mod tests {
5820 use super::*;
5821 use ndarray::array;
5822
5823 #[test]
5871 fn point_enclosure_must_sign_an_order_one_derivative_2694_2703() {
5872 const ORDER_ONE_DERIVATIVE: f64 = 0.25;
5880
5881 fn point_check(
5882 tag: &str,
5883 prepared: &GaussianRemlPrepared,
5884 rho: f64,
5885 failures: &mut Vec<String>,
5886 ) -> bool {
5887 let exact = prepared.evaluate(rho);
5888 let (dv, _) = reml_deriv_enclosure(
5889 &prepared.cache,
5890 prepared.ywy.view(),
5891 prepared.projected_rhs_squared.view(),
5892 prepared.n_effective,
5893 prepared.n_outputs,
5894 rho,
5895 rho,
5896 );
5897 if !(dv.lo <= exact.grad && exact.grad <= dv.hi) {
5901 failures.push(format!(
5902 "{tag} rho={rho}: SOUNDNESS — the zero-width enclosure \
5903 [{:.9e}, {:.9e}] does not contain the evaluator's own \
5904 V'={:.9e}",
5905 dv.lo, dv.hi, exact.grad
5906 ));
5907 return false;
5908 }
5909 if exact.grad.abs() < ORDER_ONE_DERIVATIVE {
5910 return false;
5911 }
5912 if dv.hi - dv.lo > exact.grad.abs() {
5913 failures.push(format!(
5914 "{tag} rho={rho}: SHARPNESS — V'={:.9e} but the ZERO-WIDTH \
5915 enclosure is [{:.9e}, {:.9e}], width {:.9e}. A width larger \
5916 than the value it encloses cannot sign that value, so no \
5917 cell containing this point can ever satisfy the \
5918 branch-and-bound's `dv.lo > 0 || dv.hi < 0` prune test, at \
5919 any subdivision depth.",
5920 exact.grad,
5921 dv.lo,
5922 dv.hi,
5923 dv.hi - dv.lo
5924 ));
5925 }
5926 true
5927 }
5928
5929 let mut failures: Vec<String> = Vec::new();
5930
5931 let control_x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0]];
5937 let control_y = array![[1.0], [3.0], [5.0], [7.0], [9.0]];
5938 let control_penalty = array![[0.0, 0.0], [0.0, 1.0]];
5939 let control = prepare_gaussian_reml(
5940 control_x.view(),
5941 control_y.view(),
5942 control_penalty.view(),
5943 None,
5944 None,
5945 None,
5946 )
5947 .expect("the control design is finite and full rank");
5948 let mut control_asserted = 0usize;
5949 for rho in [RHO_LOWER, -10.0, 0.0, 10.0] {
5950 if point_check("CONTROL", &control, rho, &mut failures) {
5951 control_asserted += 1;
5952 }
5953 }
5954 if control_asserted == 0 {
5955 failures.push(
5956 "CONTROL: no rho carried an order-one V', so the sharpness \
5957 property was never asserted on the passing side — the gate's \
5958 instrument did not engage"
5959 .to_string(),
5960 );
5961 }
5962
5963 let n = 12usize;
5971 let witness_x = Array2::<f64>::from_shape_fn((n, 3), |(row, col)| {
5972 let t = 2.0 * std::f64::consts::PI * (row as f64) / (n as f64);
5973 match col {
5974 0 => 1.0,
5975 1 => t.sin(),
5976 _ => t.cos(),
5977 }
5978 });
5979 let witness_y = Array2::<f64>::from_elem((n, 1), 0.7);
5980 let witness_penalty = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
5981 let witness = prepare_gaussian_reml(
5982 witness_x.view(),
5983 witness_y.view(),
5984 witness_penalty.view(),
5985 None,
5986 None,
5987 None,
5988 )
5989 .expect("the witness design is finite and full rank");
5990
5991 let DispersionResidualParts {
5996 unpenalized_residual,
5997 penalized_residual,
5998 ..
5999 } = dispersion_residual_parts(
6000 &witness.cache,
6001 witness.ywy.view(),
6002 witness.projected_rhs_squared.view(),
6003 0,
6004 RHO_LOWER,
6005 );
6006 let ywy = witness.ywy[0];
6007 if !(penalized_residual >= 0.0 && penalized_residual < 1.0e-25 * ywy) {
6008 failures.push(format!(
6009 "WITNESS regime: the rho-dependent deviance at rho={RHO_LOWER} is \
6010 {penalized_residual:.9e} against ywy={ywy:.9e}; this design does \
6011 not interpolate its response, so `ywy − Σc²` never cancels and \
6012 the fixture has drifted OUT of the regime under test — a pass \
6013 below would mean nothing"
6014 ));
6015 }
6016 let mut witness_asserted = 0usize;
6017 for rho in [RHO_LOWER, -25.0, -20.0] {
6018 if point_check("WITNESS", &witness, rho, &mut failures) {
6019 witness_asserted += 1;
6020 }
6021 }
6022 if witness_asserted == 0 {
6023 failures.push(
6024 "WITNESS: no rho carried an order-one V', so the sharpness \
6025 property was never asserted on the failing side — the gate's \
6026 instrument did not engage"
6027 .to_string(),
6028 );
6029 }
6030
6031 println!(
6035 "[2694-gate] CONTROL asserted at {control_asserted} rho, WITNESS \
6036 asserted at {witness_asserted} rho, failed clauses {}, witness \
6037 ywy={ywy:.9e} unpenalized_residual={unpenalized_residual:.9e} \
6038 penalized_residual={penalized_residual:.9e}",
6039 failures.len()
6040 );
6041
6042 assert!(
6043 failures.is_empty(),
6044 "#2703/#2694 REGRESSION — the profiled-REML derivative enclosure has \
6045 lost its sharpness on an exactly-interpolating design.\n\
6046 This gate was red when it landed and went green with the repair in \
6047 `reml_deriv_enclosure_profile`: `r0` is ρ-INDEPENDENT, so it enters \
6048 the enclosure as the single value the evaluator uses, not as a \
6049 bracket over the digits its cancellation destroyed. Bracketing it \
6050 put `dp_lo` on the ratio numerator's own scale, pinned \
6051 `num_hi/dp_lo` at `1.0` whatever the data, and gave a ZERO-WIDTH \
6052 enclosure of width `half_nu`. If you are seeing this, check that \
6053 change first.\n\
6054 witness ywy={ywy:.9e} unpenalized_residual={unpenalized_residual:.9e} \
6055 penalized_residual={penalized_residual:.9e}\n{}",
6056 failures.join("\n")
6057 );
6058 }
6059
6060 #[test]
6115 fn rho_enumeration_resolves_the_small_lambda_rail_and_still_refuses_unresolvable_structure_2703()
6116 {
6117 let mut failures: Vec<String> = Vec::new();
6118
6119 let n = 12usize;
6121 let witness_x = Array2::<f64>::from_shape_fn((n, 3), |(row, col)| {
6122 let t = 2.0 * std::f64::consts::PI * (row as f64) / (n as f64);
6123 match col {
6124 0 => 1.0,
6125 1 => t.sin(),
6126 _ => t.cos(),
6127 }
6128 });
6129 let witness_y = Array2::<f64>::from_elem((n, 1), 0.7);
6130 let witness_penalty = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
6131 let witness = prepare_gaussian_reml(
6132 witness_x.view(),
6133 witness_y.view(),
6134 witness_penalty.view(),
6135 None,
6136 None,
6137 None,
6138 )
6139 .expect("the witness design is finite and full rank");
6140
6141 let DispersionResidualParts {
6145 unpenalized_residual,
6146 penalized_residual,
6147 ..
6148 } = dispersion_residual_parts(
6149 &witness.cache,
6150 witness.ywy.view(),
6151 witness.projected_rhs_squared.view(),
6152 0,
6153 RHO_LOWER,
6154 );
6155 let ywy = witness.ywy[0];
6156 if !(penalized_residual >= 0.0 && penalized_residual < 1.0e-25 * ywy) {
6157 failures.push(format!(
6158 "WITNESS regime: the rho-dependent deviance at rho={RHO_LOWER} is \
6159 {penalized_residual:.9e} against ywy={ywy:.9e}; this design does not \
6160 interpolate its response, so `ywy − Σc²` never cancels and the \
6161 fixture has drifted OUT of the regime under test — a pass below \
6162 would mean nothing"
6163 ));
6164 }
6165
6166 let witness_eval = |rho: f64| witness.evaluate(rho);
6167 let witness_enclose = |a: f64, b: f64| {
6168 reml_deriv_enclosure(
6169 &witness.cache,
6170 witness.ywy.view(),
6171 witness.projected_rhs_squared.view(),
6172 witness.n_effective,
6173 witness.n_outputs,
6174 a,
6175 b,
6176 )
6177 };
6178 let mut witness_rho = f64::NAN;
6179 match enumerate_and_select_rho_with_controls(
6180 &witness_eval,
6181 &witness_enclose,
6182 None,
6183 ProfileSearchControls::PRODUCTION,
6184 None,
6185 ) {
6186 Ok(selection) => {
6187 witness_rho = selection.rho;
6188 let at_selected = witness_eval(selection.rho).cost;
6189 let at_lower = witness_eval(RHO_LOWER).cost;
6190 let at_upper = witness_eval(RHO_UPPER).cost;
6191 if !(selection.rho.is_finite()
6192 && selection.rho >= RHO_LOWER
6193 && selection.rho <= RHO_UPPER)
6194 {
6195 failures.push(format!(
6196 "WITNESS: the selected rho={} is not inside the search window \
6197 [{RHO_LOWER}, {RHO_UPPER}]",
6198 selection.rho
6199 ));
6200 }
6201 if !(at_selected <= at_lower && at_selected <= at_upper) {
6202 failures.push(format!(
6203 "WITNESS: the selection is not the best candidate the search \
6204 saw — cost {at_selected:.9e} at rho={} against {at_lower:.9e} \
6205 at the lower rail and {at_upper:.9e} at the upper rail",
6206 selection.rho
6207 ));
6208 }
6209 }
6210 Err(error) => failures.push(format!(
6211 "WITNESS: the production branch-and-bound REFUSED an interpolating \
6212 design — this is the #2703 symptom itself: {error}"
6213 )),
6214 }
6215
6216 const CENTRE: f64 = 0.5;
6225 let objective = |amplitude: f64, wavenumber: f64| {
6226 move |rho: f64| {
6227 let phase = wavenumber * rho;
6228 ObjectiveEval {
6229 cost: 0.5 * (rho - CENTRE) * (rho - CENTRE) + amplitude * phase.sin(),
6230 grad: (rho - CENTRE) + amplitude * wavenumber * phase.cos(),
6231 hess: 1.0 - amplitude * wavenumber * wavenumber * phase.sin(),
6232 edf: 0.0,
6233 cost_roundoff: 0.0,
6237 }
6238 }
6239 };
6240 let enclosure = |amplitude: f64, wavenumber: f64| {
6241 move |a: f64, b: f64| {
6242 let grad_swing = (amplitude * wavenumber).abs();
6243 let hess_swing = (amplitude * wavenumber * wavenumber).abs();
6244 (
6245 Interval {
6246 lo: (a - CENTRE) - grad_swing,
6247 hi: (b - CENTRE) + grad_swing,
6248 },
6249 Interval {
6250 lo: 1.0 - hess_swing,
6251 hi: 1.0 + hess_swing,
6252 },
6253 )
6254 }
6255 };
6256
6257 let unresolvable_wavenumber = std::f64::consts::TAU / (0.25 * RHO_BRACKET_RESOLUTION);
6265 let unresolvable_amplitude = 1.0 / unresolvable_wavenumber;
6266 let mut positive_control_verdict = String::new();
6267 match enumerate_and_select_rho_with_controls(
6268 objective(unresolvable_amplitude, unresolvable_wavenumber),
6269 enclosure(unresolvable_amplitude, unresolvable_wavenumber),
6270 None,
6271 ProfileSearchControls::PRODUCTION,
6272 None,
6273 ) {
6274 Ok(selection) => failures.push(format!(
6275 "POSITIVE CONTROL: the enumerator MINTED rho={} on an objective whose \
6276 stationary points are spaced below its own bracket resolution \
6277 ({RHO_BRACKET_RESOLUTION:e}). The unresolvable-structure refusal can \
6278 no longer fire, which is a worse defect than the one #2703 reported",
6279 selection.rho
6280 )),
6281 Err(error) => {
6282 positive_control_verdict = error.to_string();
6283 if !positive_control_verdict.contains("remained non-monotone") {
6284 failures.push(format!(
6285 "POSITIVE CONTROL: refused, but not with the \
6286 unresolvable-structure verdict: {positive_control_verdict}"
6287 ));
6288 }
6289 }
6290 }
6291
6292 let mut negative_control_rho = f64::NAN;
6294 match enumerate_and_select_rho_with_controls(
6295 objective(0.0, unresolvable_wavenumber),
6296 enclosure(0.0, unresolvable_wavenumber),
6297 None,
6298 ProfileSearchControls::PRODUCTION,
6299 None,
6300 ) {
6301 Ok(selection) => {
6302 negative_control_rho = selection.rho;
6303 let scale =
6309 1.0 + selection.rho.abs().max(CENTRE.abs()) + RHO_BRACKET_RESOLUTION;
6310 let admissible = RHO_BRACKET_RESOLUTION * scale;
6311 if (selection.rho - CENTRE).abs() > admissible {
6312 failures.push(format!(
6313 "NEGATIVE CONTROL: selected rho={} against the analytic root \
6314 {CENTRE}, off by {:e} which exceeds the search's own bracket \
6315 acceptance {admissible:e}",
6316 selection.rho,
6317 (selection.rho - CENTRE).abs()
6318 ));
6319 }
6320 }
6321 Err(error) => failures.push(format!(
6322 "NEGATIVE CONTROL: the enumerator refused a smooth unimodal objective \
6323 with an interior root at {CENTRE}, so the positive control's refusal \
6324 above is attributable to the harness rather than to the structure: \
6325 {error}"
6326 )),
6327 }
6328
6329 println!(
6333 "[2703-gate] WITNESS selected rho={witness_rho:.9e} \
6334 (ywy={ywy:.9e} unpenalized_residual={unpenalized_residual:.9e} \
6335 penalized_residual={penalized_residual:.9e}), \
6336 POSITIVE CONTROL refusal={:?}, NEGATIVE CONTROL rho={negative_control_rho:.9e}, \
6337 failed clauses {}",
6338 positive_control_verdict
6339 .split(':')
6340 .next_back()
6341 .unwrap_or("")
6342 .trim(),
6343 failures.len()
6344 );
6345
6346 assert!(
6347 failures.is_empty(),
6348 "#2703 REGRESSION — the 1-D REML rho enumerator no longer resolves a \
6349 small-lambda-flat objective, or no longer refuses one it cannot \
6350 resolve.\n\
6351 The six `gam-sae inference::` failures this gate stands for were ONE \
6352 cause: `r0` entered the derivative enclosure as a BRACKET over the \
6353 digits its cancellation destroys, the width was rho-INDEPENDENT and so \
6354 survived every bisection, and no cell could ever be pruned. If the \
6355 WITNESS clause is red, check `reml_deriv_enclosure_profile`'s `r0` \
6356 first. If a CONTROL clause is red, the guard's ability to fire has \
6357 moved, which is the more serious direction.\n{}",
6358 failures.join("\n")
6359 );
6360 }
6361
6362 #[test]
6363 fn edf_does_not_double_count_penalty_nullspace() {
6364 let x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0],];
6365 let y = array![[0.0], [1.0], [1.8], [3.2], [4.1]];
6366 let penalty = array![[0.0, 0.0], [0.0, 1.0]];
6367 let result =
6368 gaussian_reml_multi_closed_form(x.view(), y.view(), penalty.view(), None, Some(0.0))
6369 .expect("small full-rank Gaussian REML fit");
6370
6371 assert!(result.edf >= result.cache.nullity as f64);
6372 assert!(result.edf <= x.ncols() as f64 + 1.0e-10);
6373 }
6374
6375 #[test]
6383 fn profiled_gaussian_reml_penalized_scale_selects_analytic_lambda_one_2496() {
6384 let x = array![[1.0], [0.0], [0.0]];
6385 let y = array![1.0, 1.0, 0.0];
6386 let penalty = array![[1.0]];
6387 let fit = gaussian_reml_closed_form_with_nullspace_dim(
6388 x.view(),
6389 y.view(),
6390 penalty.view(),
6391 Some(0),
6392 None,
6393 None,
6394 )
6395 .expect("analytic one-mode Gaussian REML profile");
6396
6397 eprintln!(
6398 "[#2496] analytic profile: lambda={:.12e} rho={:.12e} sigma2={:.12e}",
6399 fit.lambda, fit.rho, fit.sigma2,
6400 );
6401 assert!(
6402 fit.rho.abs() <= 1.0e-9,
6403 "analytic optimum is rho=log(lambda)=0, got {}",
6404 fit.rho,
6405 );
6406 assert!((fit.lambda - 1.0).abs() <= 1.0e-9);
6407 assert!((fit.coefficients[0] - 0.5).abs() <= 1.0e-9);
6408 assert!((fit.sigma2 - 0.5).abs() <= 1.0e-9);
6409 }
6410
6411 #[test]
6430 fn saturated_design_keeps_a_finite_small_lambda_enclosure_2585() {
6431 let n = 8usize;
6432 let mut x = Array2::<f64>::zeros((n, n));
6433 for i in 0..n {
6434 x[[i, i]] = 1.0;
6435 }
6436 let y =
6437 Array2::from_shape_vec((n, 1), vec![0.7, -1.3, 2.1, 0.4, -0.9, 1.6, -0.2, 1.1])
6438 .expect("saturated response");
6439 let mut penalty = Array2::<f64>::zeros((n, n));
6442 for i in 0..n - 1 {
6443 penalty[[i, i]] = 1.0e-3;
6444 }
6445
6446 let prepared =
6447 prepare_gaussian_reml(x.view(), y.view(), penalty.view(), None, None, None)
6448 .expect("saturated design must still prepare");
6449
6450 let a = RHO_LOWER;
6451 let b = RHO_LOWER + 1.0e-3;
6452 let (dv, dvv) = reml_deriv_enclosure(
6453 &prepared.cache,
6454 prepared.ywy.view(),
6455 prepared.projected_rhs_squared.view(),
6456 prepared.n_effective,
6457 prepared.n_outputs,
6458 a,
6459 b,
6460 );
6461 assert!(
6462 dv.lo.is_finite() && dv.hi.is_finite(),
6463 "saturated small-lambda cell produced an unbounded V' enclosure [{}, {}]",
6464 dv.lo,
6465 dv.hi
6466 );
6467 assert!(
6468 dvv.lo.is_finite() && dvv.hi.is_finite(),
6469 "saturated small-lambda cell produced an unbounded V'' enclosure [{}, {}]",
6470 dvv.lo,
6471 dvv.hi
6472 );
6473 for rho in [a, b] {
6474 let jet = prepared.evaluate(rho);
6475 assert!(
6476 interval_contains(dv, jet.grad),
6477 "V' enclosure [{}, {}] missed the endpoint gradient {} at rho={rho}",
6478 dv.lo,
6479 dv.hi,
6480 jet.grad
6481 );
6482 assert!(
6483 interval_contains(dvv, jet.hess),
6484 "V'' enclosure [{}, {}] missed the endpoint curvature {} at rho={rho}",
6485 dvv.lo,
6486 dvv.hi,
6487 jet.hess
6488 );
6489 }
6490
6491 let edge = prepared.evaluate(RHO_LOWER);
6497 assert!(
6498 edge.cost.is_finite() && edge.grad.is_finite() && edge.hess.is_finite(),
6499 "saturated small-lambda jet is not finite: cost={} grad={} hess={}",
6500 edge.cost,
6501 edge.grad,
6502 edge.hess
6503 );
6504 let sigma2 = prepared.sigma2(RHO_LOWER);
6505 assert!(
6506 sigma2.iter().all(|v| v.is_finite() && *v > 0.0),
6507 "saturated profiled dispersion collapsed to {sigma2:?}"
6508 );
6509
6510 }
6520
6521 #[test]
6528 fn profiled_gaussian_reml_is_penalty_scale_and_coefficient_chart_invariant_2496() {
6529 let x = array![[1.0], [0.0], [0.0]];
6530 let y = array![1.0, 1.0, 0.0];
6531 let penalty = array![[1.0]];
6532 let baseline = gaussian_reml_closed_form_with_nullspace_dim(
6533 x.view(),
6534 y.view(),
6535 penalty.view(),
6536 Some(0),
6537 None,
6538 None,
6539 )
6540 .expect("baseline analytic Gaussian REML profile");
6541
6542 for alpha in [1.0e-3_f64, 37.0, 1.0e4] {
6543 let scaled_penalty = penalty.mapv(|value| alpha * value);
6544 let scaled = gaussian_reml_closed_form_with_nullspace_dim(
6545 x.view(),
6546 y.view(),
6547 scaled_penalty.view(),
6548 Some(0),
6549 None,
6550 Some(baseline.rho - alpha.ln()),
6551 )
6552 .expect("penalty-scaled Gaussian REML profile");
6553 let score_tolerance = 1.0e-9 * (1.0 + baseline.reml_score.abs());
6554 assert!(
6555 (scaled.reml_score - baseline.reml_score).abs() <= score_tolerance,
6556 "S -> alpha S changed profiled evidence at alpha={alpha}: baseline={}, scaled={}",
6557 baseline.reml_score,
6558 scaled.reml_score,
6559 );
6560 assert!(
6561 (scaled.rho - (baseline.rho - alpha.ln())).abs() <= 1.0e-9,
6562 "S -> alpha S must shift rho by -log(alpha) at alpha={alpha}: baseline={}, scaled={}",
6563 baseline.rho,
6564 scaled.rho,
6565 );
6566 assert!(
6567 (alpha * scaled.lambda - baseline.lambda).abs() <= 1.0e-9,
6568 "physical lambda*S changed at alpha={alpha}",
6569 );
6570 for row in 0..x.nrows() {
6571 assert!((scaled.fitted[row] - baseline.fitted[row]).abs() <= 1.0e-9);
6572 }
6573 }
6574
6575 let coefficient_scale = 7.0_f64;
6576 let reparameterized_x = x.mapv(|value| value / coefficient_scale);
6577 let reparameterized_penalty =
6578 penalty.mapv(|value| value / coefficient_scale.powi(2));
6579 let reparameterized = gaussian_reml_closed_form_with_nullspace_dim(
6580 reparameterized_x.view(),
6581 y.view(),
6582 reparameterized_penalty.view(),
6583 Some(0),
6584 None,
6585 Some(baseline.rho),
6586 )
6587 .expect("coefficient-reparameterized Gaussian REML profile");
6588 let score_tolerance = 1.0e-9 * (1.0 + baseline.reml_score.abs());
6589 assert!(
6590 (reparameterized.reml_score - baseline.reml_score).abs() <= score_tolerance
6591 );
6592 assert!((reparameterized.rho - baseline.rho).abs() <= 1.0e-9);
6593 assert!(
6594 (reparameterized.coefficients[0]
6595 - coefficient_scale * baseline.coefficients[0])
6596 .abs()
6597 <= 1.0e-9
6598 );
6599 for row in 0..x.nrows() {
6600 assert!((reparameterized.fitted[row] - baseline.fitted[row]).abs() <= 1.0e-9);
6601 }
6602 eprintln!(
6603 "[#2496] gauges: base_rho={:.12e}, chart_rho={:.12e}, score={:.12e}",
6604 baseline.rho, reparameterized.rho, baseline.reml_score,
6605 );
6606 }
6607
6608 #[test]
6609 fn shared_dispersion_pools_projection_exact_and_missed_outputs() {
6610 let n = 12usize;
6611 let mut x = Array2::<f64>::zeros((n, 2));
6612 let mut y = Array2::<f64>::zeros((n, 2));
6613 for row in 0..n {
6614 let t = row as f64 - 5.5;
6615 x[[row, 0]] = 1.0;
6616 x[[row, 1]] = t;
6617 y[[row, 0]] = t;
6620 y[[row, 1]] = if row % 2 == 0 { -2.0 } else { 3.0 };
6622 }
6623 let penalty = Array2::<f64>::zeros((2, 2));
6624 let fit = gaussian_reml_multi_shared_dispersion_closed_form(
6625 x.view(),
6626 y.view(),
6627 penalty.view(),
6628 None,
6629 None,
6630 )
6631 .expect("shared-dispersion vector REML fit");
6632
6633 assert_eq!(fit.sigma2[0].to_bits(), fit.sigma2[1].to_bits());
6634 let mut pooled_rss = 0.0_f64;
6635 for row in 0..n {
6636 for output in 0..2 {
6637 let residual = y[[row, output]] - fit.fitted[[row, output]];
6638 pooled_rss += residual * residual;
6639 }
6640 }
6641 let shared_nu = (2 * (n - fit.cache.nullity)) as f64;
6642 let expected_sigma2 = pooled_rss / shared_nu;
6643 assert!(expected_sigma2 > 0.0);
6644 assert!(
6645 (fit.sigma2[0] - expected_sigma2).abs()
6646 <= f64::EPSILON.sqrt() * expected_sigma2.max(1.0),
6647 "shared sigma2 {} must equal pooled vector deviance / shared dof {}",
6648 fit.sigma2[0],
6649 expected_sigma2
6650 );
6651 }
6652
6653 #[test]
6654 fn shared_dispersion_penalty_envelope_gradient_matches_refitted_direction() {
6655 let n = 24usize;
6656 let mut x = Array2::<f64>::zeros((n, 3));
6657 let mut y = Array2::<f64>::zeros((n, 2));
6658 for row in 0..n {
6659 let t = -1.0 + 2.0 * row as f64 / (n - 1) as f64;
6660 x[[row, 0]] = 1.0;
6661 x[[row, 1]] = t;
6662 x[[row, 2]] = t * t;
6663 y[[row, 0]] = 0.3 + 1.2 * t - 0.8 * t * t + 0.04 * (7.0 * t).sin();
6664 y[[row, 1]] = -0.2 + 0.5 * t + 0.4 * t * t + 0.03 * (5.0 * t).cos();
6665 }
6666 let penalty = array![[0.0, 0.0, 0.0], [0.0, 0.7, 0.1], [0.0, 0.1, 1.4]];
6667 let direction = array![[0.0, 0.0, 0.0], [0.0, 0.3, -0.08], [0.0, -0.08, 0.6]];
6668 let fit = gaussian_reml_multi_shared_dispersion_closed_form(
6669 x.view(),
6670 y.view(),
6671 penalty.view(),
6672 None,
6673 None,
6674 )
6675 .unwrap();
6676 let gradient = gaussian_reml_multi_shared_dispersion_penalty_gradient_from_fit(
6677 x.view(),
6678 y.view(),
6679 penalty.view(),
6680 None,
6681 &fit,
6682 )
6683 .unwrap();
6684 let analytic = gradient
6685 .iter()
6686 .zip(direction.iter())
6687 .map(|(gradient, direction)| gradient * direction)
6688 .sum::<f64>();
6689
6690 let step = f64::EPSILON.cbrt();
6691 let plus_penalty = &penalty + &(direction.mapv(|value| step * value));
6692 let minus_penalty = &penalty - &(direction.mapv(|value| step * value));
6693 let plus = gaussian_reml_multi_shared_dispersion_closed_form(
6694 x.view(),
6695 y.view(),
6696 plus_penalty.view(),
6697 None,
6698 Some(fit.rho),
6699 )
6700 .unwrap();
6701 let minus = gaussian_reml_multi_shared_dispersion_closed_form(
6702 x.view(),
6703 y.view(),
6704 minus_penalty.view(),
6705 None,
6706 Some(fit.rho),
6707 )
6708 .unwrap();
6709 let numerical = (plus.reml_score - minus.reml_score) / (2.0 * step);
6710 let scale = analytic.abs().max(numerical.abs()).max(1.0);
6711 assert!(
6712 (analytic - numerical).abs() <= 2.0e-5 * scale,
6713 "shared-dispersion penalty envelope derivative mismatch: analytic={analytic}, refitted={numerical}"
6714 );
6715 }
6716
6717 #[test]
6718 fn block_orthogonal_score_matches_the_objective_derivative() {
6719 let gram = array![[3.0, 0.4], [0.4, 2.0]];
6720 let rhs = array![[1.2, -0.3], [0.6, 0.9]];
6721 let penalty = array![[1.0, 0.2], [0.2, 0.8]];
6722 let scale = array![1.3, 0.8];
6723 let rho = 0.37;
6724 let step = 1.0e-6;
6725 let eval = block_orthogonal_eval(&gram, &rhs, &penalty, rho).unwrap();
6726 let analytic = block_orthogonal_scale_objective(&eval, rho, scale.view(), 2).grad;
6727 let value_at = |candidate_rho: f64| {
6728 let candidate = block_orthogonal_eval(&gram, &rhs, &penalty, candidate_rho).unwrap();
6729 block_orthogonal_scale_objective(&candidate, candidate_rho, scale.view(), 2).value
6730 };
6731 let numerical = (value_at(rho + step) - value_at(rho - step)) / (2.0 * step);
6732 assert!(
6733 (analytic - numerical).abs() <= 1.0e-7 * analytic.abs().max(1.0),
6734 "analytic score {analytic:.12e} != objective derivative {numerical:.12e}"
6735 );
6736 }
6737
6738 #[test]
6739 fn block_orthogonal_profile_hessian_matches_the_profiled_objective() {
6740 let grams = [
6741 array![[3.0, 0.4], [0.4, 2.0]],
6742 array![[2.5, -0.2], [-0.2, 1.8]],
6743 ];
6744 let rhs = [
6745 array![[1.2, -0.3], [0.6, 0.9]],
6746 array![[0.5, 0.8], [-0.4, 0.7]],
6747 ];
6748 let penalties = [
6749 array![[1.0, 0.2], [0.2, 0.8]],
6750 array![[0.9, -0.1], [-0.1, 1.1]],
6751 ];
6752 let ranks = [2_usize, 2_usize];
6753 let ywy = array![8.0, 9.0];
6754 let nu = 7.0;
6755 let rhos = array![0.37, -0.21];
6756 let profile_value = |candidate_rhos: ArrayView1<'_, f64>| {
6757 let evals = (0..2)
6758 .map(|block| {
6759 block_orthogonal_eval(
6760 &grams[block],
6761 &rhs[block],
6762 &penalties[block],
6763 candidate_rhos[block],
6764 )
6765 .unwrap()
6766 })
6767 .collect::<Vec<_>>();
6768 let mut q = ywy.clone();
6769 for eval in &evals {
6770 q -= &eval.fitted_energy;
6771 }
6772 let determinant_term = evals
6773 .iter()
6774 .enumerate()
6775 .map(|(block, eval)| eval.logdet - ranks[block] as f64 * candidate_rhos[block])
6776 .sum::<f64>();
6777 0.5 * 2.0 * determinant_term + 0.5 * nu * q.iter().map(|value| value.ln()).sum::<f64>()
6778 };
6779 let evals = (0..2)
6780 .map(|block| {
6781 block_orthogonal_eval(&grams[block], &rhs[block], &penalties[block], rhos[block])
6782 .unwrap()
6783 })
6784 .collect::<Vec<_>>();
6785 let scale = block_orthogonal_conditional_scale(&evals, ywy.view(), nu).unwrap();
6786 let analytic =
6787 block_orthogonal_profile_hessian(&evals, rhos.view(), scale.view(), &ranks, nu)
6788 .unwrap();
6789 let step = 1.0e-4;
6790 let center = profile_value(rhos.view());
6791 let mut numerical = Array2::<f64>::zeros((2, 2));
6792 for coordinate in 0..2 {
6793 let mut plus = rhos.clone();
6794 let mut minus = rhos.clone();
6795 plus[coordinate] += step;
6796 minus[coordinate] -= step;
6797 numerical[[coordinate, coordinate]] = (profile_value(plus.view()) - 2.0 * center
6798 + profile_value(minus.view()))
6799 / (step * step);
6800 }
6801 let mut plus_plus = rhos.clone();
6802 let mut plus_minus = rhos.clone();
6803 let mut minus_plus = rhos.clone();
6804 let mut minus_minus = rhos.clone();
6805 plus_plus[0] += step;
6806 plus_plus[1] += step;
6807 plus_minus[0] += step;
6808 plus_minus[1] -= step;
6809 minus_plus[0] -= step;
6810 minus_plus[1] += step;
6811 minus_minus[0] -= step;
6812 minus_minus[1] -= step;
6813 let cross = (profile_value(plus_plus.view())
6814 - profile_value(plus_minus.view())
6815 - profile_value(minus_plus.view())
6816 + profile_value(minus_minus.view()))
6817 / (4.0 * step * step);
6818 numerical[[0, 1]] = cross;
6819 numerical[[1, 0]] = cross;
6820 for row in 0..2 {
6821 for col in 0..2 {
6822 assert!(
6823 (analytic[[row, col]] - numerical[[row, col]]).abs()
6824 <= 2.0e-6 * analytic[[row, col]].abs().max(1.0),
6825 "profile Hessian ({row}, {col}) analytic {:.12e} != numerical {:.12e}",
6826 analytic[[row, col]],
6827 numerical[[row, col]]
6828 );
6829 }
6830 }
6831 }
6832
6833 #[test]
6834 fn block_orthogonal_shared_scale_fit_carries_a_score_certificate() {
6835 let c0 = [1.0_f64; 8];
6841 let c1 = [1.0, 1.0, 1.0, 1.0, -1.0, -1.0, -1.0, -1.0];
6842 let c2 = [1.0, 1.0, -1.0, -1.0, 1.0, 1.0, -1.0, -1.0];
6843 let c3 = [1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0];
6844 let mut d1 = Array2::<f64>::zeros((8, 2));
6845 let mut d2 = Array2::<f64>::zeros((8, 2));
6846 for i in 0..8 {
6847 d1[[i, 0]] = c0[i];
6848 d1[[i, 1]] = c1[i];
6849 d2[[i, 0]] = c2[i];
6850 d2[[i, 1]] = c3[i];
6851 }
6852 let penalties = vec![Array2::<f64>::eye(2), Array2::<f64>::eye(2)];
6853 let bumps = [0.03, -0.05, 0.02, 0.01, -0.02, 0.04, -0.01, -0.02];
6854 let mut y = Array2::<f64>::zeros((8, 1));
6855 for i in 0..8 {
6856 y[[i, 0]] = c0[i] + 0.5 * c1[i] + 0.25 * c2[i] + bumps[i];
6857 }
6858
6859 let result = gaussian_reml_blocks_orthogonal_shared_scale(
6860 &[d1.clone(), d2.clone()],
6861 &penalties,
6862 y.view(),
6863 None,
6864 None,
6865 )
6866 .expect("well-posed orthogonal-block fit must certify and mint");
6867
6868 let weight = Array1::<f64>::ones(8);
6869 let ywy = (0..8).map(|i| y[[i, 0]] * y[[i, 0]]).sum::<f64>();
6870 let nu = 8.0_f64;
6872 let mut evals = Vec::new();
6873 for (block, design) in [&d1, &d2].into_iter().enumerate() {
6874 let gram = canonicalize_penalty(dense_xt_diag_x(design.view(), weight.view()).view());
6875 let rhs = dense_xt_diag_y(design.view(), weight.view(), y.view());
6876 let pen = canonicalize_penalty(penalties[block].view());
6877 evals.push(
6878 block_orthogonal_eval(&gram, &rhs, &pen, result.log_lambdas[block])
6879 .expect("block eval at the minted rho"),
6880 );
6881 }
6882 let explained: f64 = evals.iter().map(|eval| eval.fitted_energy[0]).sum();
6883 let q = ywy - explained;
6884 assert!(q > 0.0);
6885 let scale = Array1::from_vec(vec![nu / q]);
6886 for (block, eval) in evals.iter().enumerate() {
6887 let derivs =
6888 block_orthogonal_scale_objective(eval, result.log_lambdas[block], scale.view(), 2);
6889 let residual = derivs.grad.abs() / 2.0;
6890 assert!(
6891 residual <= BLOCK_ORTHOGONAL_SCORE_TOL,
6892 "block {block} score residual {residual:.3e} exceeds the certificate tolerance"
6893 );
6894 }
6895 let curvature = block_orthogonal_profile_spectrum(
6896 &block_orthogonal_profile_hessian(
6897 &evals,
6898 result.log_lambdas.view(),
6899 scale.view(),
6900 &[2, 2],
6901 nu,
6902 )
6903 .unwrap(),
6904 )
6905 .unwrap()
6906 .curvature;
6907 assert!(
6908 curvature.min_eigenvalue >= -curvature.roundoff,
6909 "minted fit has negative profiled curvature {:.6e} beyond roundoff {:.3e}",
6910 curvature.min_eigenvalue,
6911 curvature.roundoff
6912 );
6913
6914 let err = gaussian_reml_blocks_orthogonal_shared_scale_with_controls(
6915 &[d1, d2],
6916 &penalties,
6917 y.view(),
6918 None,
6919 None,
6920 BlockOrthogonalControls {
6921 max_outer_passes: 0,
6922 ..BlockOrthogonalControls::default()
6923 },
6924 )
6925 .unwrap_err();
6926 match err {
6927 EstimationError::BlockOrthogonalRemlDidNotConverge {
6928 iterations,
6929 max_score_residual,
6930 rho_checkpoint,
6931 ..
6932 } => {
6933 assert_eq!(iterations, 0);
6934 assert!(max_score_residual.is_infinite());
6935 assert_eq!(rho_checkpoint, vec![0.0, 0.0]);
6936 }
6937 other => panic!("expected typed block-orthogonal exhaustion, got {other}"),
6938 }
6939 }
6940
6941 #[test]
6942 fn block_orthogonal_solver_rejects_cross_block_signal() {
6943 let first = array![[1.0], [1.0], [1.0], [1.0], [1.0], [1.0]];
6944 let second = array![[0.0], [1.0], [2.0], [3.0], [4.0], [5.0]];
6945 let penalties = vec![Array2::<f64>::eye(1), Array2::<f64>::eye(1)];
6946 let y = array![[0.2], [0.8], [1.7], [3.1], [3.9], [5.2]];
6947 let err = gaussian_reml_blocks_orthogonal_shared_scale(
6948 &[first, second],
6949 &penalties,
6950 y.view(),
6951 None,
6952 None,
6953 )
6954 .unwrap_err();
6955 assert!(
6956 matches!(&err, EstimationError::InvalidInput(_)),
6957 "nonorthogonal blocks must fail the decomposed-objective contract: {err}"
6958 );
6959 assert!(err.to_string().contains("weighted cross-product"));
6960 }
6961
6962 #[test]
6963 fn multi_output_duplicate_columns_match_scalar_fit() {
6964 let x = array![
6965 [1.0, -1.0],
6966 [1.0, -0.5],
6967 [1.0, 0.0],
6968 [1.0, 0.5],
6969 [1.0, 1.0],
6970 [1.0, 1.5],
6971 ];
6972 let y1 = array![0.5, 0.2, 0.0, 0.3, 1.1, 2.0];
6973 let y = Array2::from_shape_fn(
6974 (y1.len(), 2),
6975 |(i, j)| if j == 0 { y1[i] } else { 2.0 * y1[i] },
6976 );
6977 let penalty = array![[0.0, 0.0], [0.0, 1.0]];
6978
6979 let scalar =
6980 gaussian_reml_closed_form(x.view(), y1.view(), penalty.view(), None, Some(0.0))
6981 .expect("scalar Gaussian REML fit");
6982 let multi =
6983 gaussian_reml_multi_closed_form(x.view(), y.view(), penalty.view(), None, Some(0.0))
6984 .expect("multi-output Gaussian REML fit");
6985
6986 assert!((multi.rho - scalar.rho).abs() <= 1.0e-8);
6987 for i in 0..x.ncols() {
6988 assert!((multi.coefficients[[i, 0]] - scalar.coefficients[i]).abs() <= 1.0e-8);
6989 assert!((multi.coefficients[[i, 1]] - 2.0 * scalar.coefficients[i]).abs() <= 1.0e-8);
6990 }
6991 }
6992
6993 #[derive(Clone, Copy, Debug)]
6994 enum ForwardScalar {
6995 Lambda,
6996 RemlScore,
6997 Coefficient(usize, usize),
6998 Fitted(usize, usize),
6999 Edf,
7000 }
7001
7002 fn finite_difference_design() -> Array2<f64> {
7003 Array2::from_shape_fn((20, 5), |(row, col)| {
7004 let t = (row as f64 - 9.5) / 10.0;
7005 match col {
7006 0 => 1.0,
7007 1 => t,
7008 2 => 0.5 * (3.0 * t * t - 1.0),
7009 3 => 0.5 * (5.0 * t * t * t - 3.0 * t),
7010 4 => (35.0 * t.powi(4) - 30.0 * t * t + 3.0) / 8.0,
7011 _ => unreachable!(),
7012 }
7013 })
7014 }
7015
7016 fn finite_difference_response(outputs: usize) -> Array2<f64> {
7017 Array2::from_shape_fn((20, outputs), |(row, output)| {
7028 let t = (row as f64 - 9.5) / 10.0;
7029 let phase = output as f64 + 1.0;
7030 0.2 + 0.25 * phase * t - 0.12 * t * t
7031 + (0.08 + 0.03 * phase) * (1.1 * t + 0.3 * phase).sin()
7032 + 0.05 * (7.0 * t + 0.5 * phase).sin()
7033 })
7034 }
7035
7036 fn finite_difference_penalty() -> Array2<f64> {
7037 Array2::from_diag(&array![0.0, 0.8, 1.2, 1.7, 2.3])
7038 }
7039
7040 fn finite_difference_weights() -> Array1<f64> {
7041 Array1::from_shape_fn(20, |row| {
7042 let t = (row as f64 - 9.5) / 10.0;
7043 1.0 + 0.025 * (1.1 * t).sin() + 0.01 * t
7044 })
7045 }
7046
7047 fn one_hot_objective_try(
7054 x: ArrayView2<'_, f64>,
7055 y: ArrayView2<'_, f64>,
7056 penalty: ArrayView2<'_, f64>,
7057 weights: ArrayView1<'_, f64>,
7058 target: ForwardScalar,
7059 ) -> Option<f64> {
7060 let fit = gaussian_reml_multi_closed_form_with_cache(
7061 x,
7062 y,
7063 penalty,
7064 Some(weights),
7065 Some(0.85),
7066 None,
7067 )
7068 .ok()?;
7069 Some(match target {
7070 ForwardScalar::Lambda => fit.lambda,
7071 ForwardScalar::RemlScore => fit.reml_score,
7072 ForwardScalar::Coefficient(row, col) => fit.coefficients[[row, col]],
7073 ForwardScalar::Fitted(row, col) => fit.fitted[[row, col]],
7074 ForwardScalar::Edf => fit.edf,
7075 })
7076 }
7077
7078 fn one_hot_objective(
7079 x: ArrayView2<'_, f64>,
7080 y: ArrayView2<'_, f64>,
7081 penalty: ArrayView2<'_, f64>,
7082 weights: ArrayView1<'_, f64>,
7083 target: ForwardScalar,
7084 ) -> f64 {
7085 one_hot_objective_try(x, y, penalty, weights, target)
7086 .expect("finite-difference forward fit")
7087 }
7088
7089 fn one_hot_backward(
7090 x: ArrayView2<'_, f64>,
7091 y: ArrayView2<'_, f64>,
7092 penalty: ArrayView2<'_, f64>,
7093 weights: ArrayView1<'_, f64>,
7094 target: ForwardScalar,
7095 ) -> GaussianRemlBackwardResult {
7096 let mut grad_coefficients = Array2::<f64>::zeros((x.ncols(), y.ncols()));
7097 let mut grad_fitted = Array2::<f64>::zeros(y.dim());
7098 let (grad_lambda, grad_score, grad_edf, coefficient_upstream, fitted_upstream) =
7099 match target {
7100 ForwardScalar::Lambda => (1.0, 0.0, 0.0, None, None),
7101 ForwardScalar::RemlScore => (0.0, 1.0, 0.0, None, None),
7102 ForwardScalar::Coefficient(row, col) => {
7103 grad_coefficients[[row, col]] = 1.0;
7104 (0.0, 0.0, 0.0, Some(grad_coefficients.view()), None)
7105 }
7106 ForwardScalar::Fitted(row, col) => {
7107 grad_fitted[[row, col]] = 1.0;
7108 (0.0, 0.0, 0.0, None, Some(grad_fitted.view()))
7109 }
7110 ForwardScalar::Edf => (0.0, 0.0, 1.0, None, None),
7111 };
7112 gaussian_reml_multi_closed_form_backward(
7113 x,
7114 y,
7115 penalty,
7116 Some(weights),
7117 Some(0.85),
7118 grad_lambda,
7119 coefficient_upstream,
7120 fitted_upstream,
7121 grad_score,
7122 grad_edf,
7123 )
7124 .expect("analytic backward VJP")
7125 }
7126
7127 fn assert_fd_close(label: &str, analytic: f64, finite_difference: f64) {
7128 let rel_tol = 1.0e-6_f64;
7129 let abs_tol = 1.0e-6_f64;
7130 let tol = abs_tol.max(rel_tol * analytic.abs().max(finite_difference.abs()));
7131 let diff = (analytic - finite_difference).abs();
7132 assert!(
7133 diff <= tol,
7134 "{label}: analytic={analytic:.12e}, finite_difference={finite_difference:.12e}, diff={diff:.3e}, tol={tol:.3e}"
7135 );
7136 }
7137
7138 fn adaptive_central_difference(mut eval: impl FnMut(f64) -> f64) -> f64 {
7139 let steps: [f64; 5] = [1.0e-3, 5.0e-4, 2.5e-4, 1.25e-4, 6.25e-5];
7140 let mut best = f64::NAN;
7141 let mut best_delta = f64::INFINITY;
7142 let mut previous: Option<f64> = None;
7143 for h in steps {
7144 let d1 = (eval(h) - eval(-h)) / (2.0 * h);
7145 let half_h = 0.5 * h;
7146 let d2 = (eval(half_h) - eval(-half_h)) / (2.0 * half_h);
7147 let estimate: f64 = d2 + (d2 - d1) / 3.0;
7148 if let Some(prev) = previous {
7149 let delta = (estimate - prev).abs();
7150 if delta < best_delta {
7151 best_delta = delta;
7152 best = estimate;
7153 }
7154 } else {
7155 best = estimate;
7156 }
7157 previous = Some(estimate);
7158 }
7159 best
7160 }
7161
7162 fn assert_backward_matches_forward_finite_difference(outputs: usize) {
7163 let x = finite_difference_design();
7164 let y = finite_difference_response(outputs);
7165 let penalty = finite_difference_penalty();
7166 let weights = finite_difference_weights();
7167 let targets = [
7168 ForwardScalar::Lambda,
7169 ForwardScalar::RemlScore,
7170 ForwardScalar::Coefficient(3, outputs - 1),
7171 ForwardScalar::Fitted(12, outputs - 1),
7172 ForwardScalar::Edf,
7173 ];
7174 for target in targets {
7175 let backward =
7176 one_hot_backward(x.view(), y.view(), penalty.view(), weights.view(), target);
7177
7178 for row in 0..x.nrows() {
7179 for col in 0..x.ncols() {
7180 let eval = |delta: f64| {
7181 let mut candidate = x.clone();
7182 candidate[[row, col]] += delta;
7183 one_hot_objective(
7184 candidate.view(),
7185 y.view(),
7186 penalty.view(),
7187 weights.view(),
7188 target,
7189 )
7190 };
7191 let fd = adaptive_central_difference(eval);
7192 assert_fd_close(
7193 &format!("target={target:?} x[{row},{col}]"),
7194 backward.grad_x[[row, col]],
7195 fd,
7196 );
7197 }
7198 }
7199
7200 for row in 0..y.nrows() {
7201 for col in 0..y.ncols() {
7202 let eval = |delta: f64| {
7203 let mut candidate = y.clone();
7204 candidate[[row, col]] += delta;
7205 one_hot_objective(
7206 x.view(),
7207 candidate.view(),
7208 penalty.view(),
7209 weights.view(),
7210 target,
7211 )
7212 };
7213 let fd = adaptive_central_difference(eval);
7214 assert_fd_close(
7215 &format!("target={target:?} y[{row},{col}]"),
7216 backward.grad_y[[row, col]],
7217 fd,
7218 );
7219 }
7220 }
7221
7222 for row in 0..weights.len() {
7223 let eval = |delta: f64| {
7224 let mut candidate = weights.clone();
7225 candidate[row] += delta;
7226 one_hot_objective(x.view(), y.view(), penalty.view(), candidate.view(), target)
7227 };
7228 let fd = adaptive_central_difference(eval);
7229 assert_fd_close(
7230 &format!("target={target:?} weights[{row}]"),
7231 backward.grad_weights[row],
7232 fd,
7233 );
7234 }
7235
7236 let null_index = 0usize; let probe_h = 1.0e-3_f64; for r in 0..penalty.nrows() {
7260 for c in 0..penalty.ncols() {
7261 if r == null_index || c == null_index {
7262 continue;
7263 }
7264 let eval = |delta: f64| {
7265 let mut candidate = penalty.clone();
7266 candidate[[r, c]] += delta;
7267 one_hot_objective(
7268 x.view(),
7269 y.view(),
7270 candidate.view(),
7271 weights.view(),
7272 target,
7273 )
7274 };
7275 let cone_safe = {
7276 let mut s_plus = penalty.clone();
7277 let mut s_minus = penalty.clone();
7278 s_plus[[r, c]] += probe_h;
7279 s_minus[[r, c]] -= probe_h;
7280 one_hot_objective_try(
7281 x.view(),
7282 y.view(),
7283 s_plus.view(),
7284 weights.view(),
7285 target,
7286 )
7287 .is_some()
7288 && one_hot_objective_try(
7289 x.view(),
7290 y.view(),
7291 s_minus.view(),
7292 weights.view(),
7293 target,
7294 )
7295 .is_some()
7296 };
7297 if !cone_safe {
7298 continue;
7299 }
7300 let fd = adaptive_central_difference(eval);
7301 assert_fd_close(
7302 &format!("target={target:?} penalty[{r},{c}]"),
7303 backward.grad_penalty[[r, c]],
7304 fd,
7305 );
7306 }
7307 }
7308 }
7309 }
7310
7311 #[test]
7312 fn scalar_backward_matches_forward_finite_difference_for_all_x_y_and_weight_entries() {
7313 assert_backward_matches_forward_finite_difference(1);
7314 }
7315
7316 #[test]
7317 fn multi_output_backward_matches_forward_finite_difference_for_all_x_y_and_weight_entries() {
7318 assert_backward_matches_forward_finite_difference(3);
7319 }
7320
7321 #[test]
7322 fn backward_vjp_matches_finite_difference() {
7323 let x = array![
7324 [1.0, -1.0, 0.2],
7325 [1.0, -0.3, -0.1],
7326 [1.0, 0.2, 0.4],
7327 [1.0, 0.8, 0.1],
7328 [1.0, 1.4, 0.5],
7329 [1.0, 2.0, 0.9],
7330 ];
7331 let y = array![
7332 [0.1, -0.2],
7333 [0.2, 0.1],
7334 [0.7, 0.0],
7335 [1.1, 0.3],
7336 [1.8, 0.9],
7337 [2.4, 1.4],
7338 ];
7339 let weights = array![1.0, 0.9, 1.1, 1.2, 0.8, 1.3];
7340 let penalty = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.2], [0.0, 0.2, 1.7]];
7341 let upstream_coefficients = array![[0.2, -0.1], [0.05, 0.03], [-0.04, 0.07]];
7342 let upstream_fitted = array![
7343 [0.01, -0.02],
7344 [0.03, 0.01],
7345 [-0.01, 0.02],
7346 [0.04, -0.03],
7347 [0.02, 0.05],
7348 [-0.02, 0.01],
7349 ];
7350 let upstream_lambda = 0.17;
7351 let upstream_score = -0.11;
7352
7353 let backward = gaussian_reml_multi_closed_form_backward(
7354 x.view(),
7355 y.view(),
7356 penalty.view(),
7357 Some(weights.view()),
7358 Some(0.8),
7359 upstream_lambda,
7360 Some(upstream_coefficients.view()),
7361 Some(upstream_fitted.view()),
7362 upstream_score,
7363 0.0,
7364 )
7365 .expect("backward VJP");
7366
7367 let objective = |x_eval: &Array2<f64>, y_eval: &Array2<f64>, w_eval: &Array1<f64>| {
7368 let fit = gaussian_reml_multi_closed_form_with_cache(
7369 x_eval.view(),
7370 y_eval.view(),
7371 penalty.view(),
7372 Some(w_eval.view()),
7373 Some(0.8),
7374 None,
7375 )
7376 .expect("fit for objective");
7377 upstream_lambda * fit.lambda
7378 + upstream_score * fit.reml_score
7379 + (&fit.coefficients * &upstream_coefficients).sum()
7380 + (&fit.fitted * &upstream_fitted).sum()
7381 };
7382 let eps = 1.0e-6;
7383 assert!(objective(&x, &y, &weights).is_finite());
7384
7385 let mut x_plus = x.clone();
7386 let mut x_minus = x.clone();
7387 x_plus[[3, 2]] += eps;
7388 x_minus[[3, 2]] -= eps;
7389 let fd_x =
7390 (objective(&x_plus, &y, &weights) - objective(&x_minus, &y, &weights)) / (2.0 * eps);
7391 assert!(
7392 (fd_x - backward.grad_x[[3, 2]]).abs() <= 2.0e-4,
7393 "grad_x mismatch: analytic={} fd={}",
7394 backward.grad_x[[3, 2]],
7395 fd_x
7396 );
7397
7398 let mut y_plus = y.clone();
7399 let mut y_minus = y.clone();
7400 y_plus[[4, 1]] += eps;
7401 y_minus[[4, 1]] -= eps;
7402 let fd_y =
7403 (objective(&x, &y_plus, &weights) - objective(&x, &y_minus, &weights)) / (2.0 * eps);
7404 assert!(
7405 (fd_y - backward.grad_y[[4, 1]]).abs() <= 2.0e-4,
7406 "grad_y mismatch: analytic={} fd={}",
7407 backward.grad_y[[4, 1]],
7408 fd_y
7409 );
7410
7411 let mut w_plus = weights.clone();
7412 let mut w_minus = weights.clone();
7413 w_plus[2] += eps;
7414 w_minus[2] -= eps;
7415 let fd_w = (objective(&x, &y, &w_plus) - objective(&x, &y, &w_minus)) / (2.0 * eps);
7416 assert!(
7417 (fd_w - backward.grad_weights[2]).abs() <= 2.0e-4,
7418 "grad_weight mismatch: analytic={} fd={}",
7419 backward.grad_weights[2],
7420 fd_w
7421 );
7422
7423 let objective_s = |s_eval: &Array2<f64>| {
7435 let fit = gaussian_reml_multi_closed_form_with_cache(
7436 x.view(),
7437 y.view(),
7438 s_eval.view(),
7439 Some(weights.view()),
7440 Some(0.8),
7441 None,
7442 )
7443 .expect("fit for penalty objective");
7444 upstream_lambda * fit.lambda
7445 + upstream_score * fit.reml_score
7446 + (&fit.coefficients * &upstream_coefficients).sum()
7447 + (&fit.fitted * &upstream_fitted).sum()
7448 };
7449 for (r, c) in [(1usize, 1usize), (1, 2), (2, 2)] {
7453 let mut s_plus = penalty.clone();
7454 let mut s_minus = penalty.clone();
7455 s_plus[[r, c]] += eps;
7456 s_minus[[r, c]] -= eps;
7457 let fd_s = (objective_s(&s_plus) - objective_s(&s_minus)) / (2.0 * eps);
7458 assert!(
7459 (fd_s - backward.grad_penalty[[r, c]]).abs() <= 2.0e-4,
7460 "grad_penalty[{r},{c}] mismatch: analytic={} fd={}",
7461 backward.grad_penalty[[r, c]],
7462 fd_s
7463 );
7464 }
7465 }
7466
7467 #[test]
7468 fn batched_eigen_cache_matches_per_fit_build() {
7469 let xtwx_a = array![[4.0, 1.0], [1.0, 3.0]];
7475 let xtwx_b = array![[2.5, -0.5], [-0.5, 1.7]];
7476 let xtwx_c = array![[7.2, 0.3], [0.3, 5.1]];
7477 let penalty = array![[0.0, 0.0], [0.0, 1.0]];
7478
7479 let batched = build_gaussian_reml_eigen_cache_batched(
7480 vec![xtwx_a.clone(), xtwx_b.clone(), xtwx_c.clone()],
7481 penalty.view(),
7482 None,
7483 );
7484 assert_eq!(batched.len(), 3);
7485
7486 for (xtwx, batched_cache) in [&xtwx_a, &xtwx_b, &xtwx_c].into_iter().zip(batched.iter()) {
7487 let single = gaussian_reml_eigen_cache_from_xtwx(xtwx.clone(), penalty.view(), None)
7488 .expect("per-fit cache");
7489 let batched_cache = batched_cache.as_ref().expect("batched cache");
7490 assert_eq!(batched_cache.penalty_rank, single.penalty_rank);
7491 assert_eq!(batched_cache.nullity, single.nullity);
7492 assert_eq!(batched_cache.xtwx_fingerprint, single.xtwx_fingerprint);
7493 assert_eq!(
7494 batched_cache.penalty_fingerprint,
7495 single.penalty_fingerprint
7496 );
7497 assert!((batched_cache.logdet_xtwx - single.logdet_xtwx).abs() <= 1.0e-12);
7498 assert!(
7499 (batched_cache.logdet_penalty_positive - single.logdet_penalty_positive).abs()
7500 <= 1.0e-12
7501 );
7502 for (a, b) in batched_cache
7503 .penalty_eigenvalues
7504 .iter()
7505 .zip(single.penalty_eigenvalues.iter())
7506 {
7507 assert!((a - b).abs() <= 1.0e-12);
7508 }
7509 for ((a, b), _) in batched_cache
7510 .coefficient_basis
7511 .iter()
7512 .zip(single.coefficient_basis.iter())
7513 .zip(0..)
7514 {
7515 assert!((a - b).abs() <= 1.0e-12);
7516 }
7517 }
7518 }
7519
7520 struct Lcg(u64);
7523 impl Lcg {
7524 fn new(seed: u64) -> Self {
7525 Lcg(seed)
7526 }
7527 fn next_u64(&mut self) -> u64 {
7528 self.0 = self
7529 .0
7530 .wrapping_mul(6364136223846793005)
7531 .wrapping_add(1442695040888963407);
7532 self.0
7533 }
7534 fn unit(&mut self) -> f64 {
7535 (self.next_u64() >> 11) as f64 / (1u64 << 53) as f64
7536 }
7537 fn range(&mut self, lo: f64, hi: f64) -> f64 {
7538 lo + (hi - lo) * self.unit()
7539 }
7540 }
7541
7542 fn synthetic_cache(eigs: &[f64]) -> GaussianRemlEigenCache {
7547 let n = eigs.len();
7548 let tolerance = penalty_range_tolerance(ArrayView1::from(eigs));
7553 let rank = eigs.iter().filter(|&&delta| delta > tolerance).count();
7554 GaussianRemlEigenCache {
7555 penalty_eigenvalues: Array1::from(eigs.to_vec()),
7556 eigenvectors: Array2::eye(n),
7557 coefficient_basis: Array2::eye(n),
7558 xtwx_fingerprint: 0,
7559 penalty_fingerprint: 0,
7560 logdet_xtwx: 0.0,
7561 logdet_penalty_positive: 0.0,
7562 penalty_rank: rank,
7563 nullity: n - rank,
7564 }
7565 }
7566
7567 #[test]
7576 fn profiled_one_mode_certificate_matches_analytic_root_and_ignores_seed_as_candidate() {
7577 let delta = 4.0;
7578 let q = 2.0;
7579 let irreducible_residual = 3.0;
7580 let n_effective = 10usize;
7581 let cache = synthetic_cache(&[delta]);
7582 let ywy = array![q + irreducible_residual];
7583 let projected = array![[q]];
7584 let eval = |rho: f64| {
7585 evaluate_reml_parts(&cache, ywy.view(), projected.view(), n_effective, 1, rho)
7586 };
7587 let enclose = |a: f64, b: f64| {
7588 reml_deriv_enclosure(&cache, ywy.view(), projected.view(), n_effective, 1, a, b)
7589 };
7590 let expected_t =
7591 irreducible_residual / (((n_effective - 1) as f64) * q - irreducible_residual);
7592 let expected_rho = (expected_t / delta).ln();
7593 let mut roots = Vec::new();
7594 let selection = {
7597 let mut collect_root = |root: StationaryRoot, _: &ObjectiveEval| roots.push(root);
7598 enumerate_and_select_rho(&eval, &enclose, Some(-20.0), Some(&mut collect_root))
7599 .expect("profile certificate")
7600 };
7601
7602 assert_eq!(roots.len(), 1, "unexpected stationary set");
7603 assert!(
7604 roots[0].bracket[0] <= expected_rho && expected_rho <= roots[0].bracket[1],
7605 "analytic root {expected_rho} outside certified bracket {:?}",
7606 roots[0].bracket
7607 );
7608 assert!(
7609 (selection.rho - expected_rho).abs()
7610 <= RHO_BRACKET_RESOLUTION * (1.0 + expected_rho.abs()),
7611 "selected rho {} differs from analytic profiled root {expected_rho}",
7612 selection.rho
7613 );
7614 assert_ne!(
7615 selection.rho.to_bits(),
7616 (-20.0_f64).to_bits(),
7617 "a nonstationary warm hint must never enter the objective argmin"
7618 );
7619 }
7620
7621 #[test]
7622 fn unresolved_stationary_structure_is_a_typed_refusal() {
7623 let eval = |rho: f64| ObjectiveEval {
7624 cost: rho * rho,
7625 grad: 2.0 * rho,
7626 hess: 2.0,
7627 edf: 0.0,
7628 cost_roundoff: 0.0,
7630 };
7631 let enclose = |_: f64, _: f64| (Interval::entire(), Interval::entire());
7634 let error = enumerate_and_select_rho_with_controls(
7635 eval,
7636 enclose,
7637 None,
7638 ProfileSearchControls {
7639 lower: -1.0,
7640 upper: 1.0,
7641 resolution: 0.25,
7642 max_depth: 0,
7643 },
7644 None,
7645 )
7646 .expect_err("ambiguous stationary structure must refuse");
7647 assert!(matches!(error, EstimationError::RemlDidNotConverge { .. }));
7648 }
7649
7650 #[test]
7651 fn profiled_modal_evaluation_is_finite_beyond_exp_range() {
7652 let cache = synthetic_cache(&[4.0]);
7653 let ywy = array![5.0];
7654 let projected = array![[2.0]];
7655 for rho in [-1_000.0, 1_000.0] {
7656 let mode = modal_kernels(rho, 4.0);
7657 assert!(mode.log_one_plus_t.is_finite());
7658 assert!(mode.u.is_finite());
7659 assert!(mode.v.is_finite());
7660 assert!(mode.w.is_finite());
7661 assert!(mode.k.is_finite());
7662 let value = evaluate_reml_parts(&cache, ywy.view(), projected.view(), 10, 1, rho);
7663 assert!(value.cost.is_finite(), "non-finite cost at rho={rho}");
7664 assert!(value.grad.is_finite(), "non-finite gradient at rho={rho}");
7665 assert!(value.hess.is_finite(), "non-finite Hessian at rho={rho}");
7666 }
7667 }
7668
7669 #[test]
7672 fn selected_rho_beats_every_certified_profile_candidate() {
7673 let mut rng = Lcg::new(0x9911_7733_5522_0044);
7674 for _case in 0..40 {
7675 let n_eig = 2 + (rng.next_u64() % 4) as usize;
7676 let eigs: Vec<f64> = (0..n_eig).map(|_| rng.range(-5.0, 6.0).exp()).collect();
7677 let cache = synthetic_cache(&eigs);
7678 let c2: Vec<f64> = (0..n_eig)
7679 .map(|_| {
7680 let v = rng.range(0.0, 2.5);
7681 v * v
7682 })
7683 .collect();
7684 let sum_c2: f64 = c2.iter().sum();
7685 let prs = Array2::from_shape_vec((n_eig, 1), c2).unwrap();
7686 let ywy = Array1::from(vec![sum_c2 + rng.range(0.05, 2.0)]);
7687 let n_eff = 80usize;
7688 let n_out = 1usize;
7689
7690 let eval =
7691 |rho: f64| evaluate_reml_parts(&cache, ywy.view(), prs.view(), n_eff, n_out, rho);
7692 let enclose = |a: f64, b: f64| {
7693 reml_deriv_enclosure(&cache, ywy.view(), prs.view(), n_eff, n_out, a, b)
7694 };
7695 let mut roots = Vec::new();
7696 let selection = {
7697 let mut collect_rho = |root: StationaryRoot, _: &ObjectiveEval| roots.push(root.rho);
7698 enumerate_and_select_rho(&eval, &enclose, None, Some(&mut collect_rho)).unwrap()
7699 };
7700 let selected = selection.rho;
7701 let selected_cost = eval(selected).cost;
7702 let tol = 1.0e-8 * (1.0 + selected_cost.abs());
7703
7704 for &r in &roots {
7705 assert!(selected_cost <= eval(r).cost + tol);
7706 }
7707 assert!(selected_cost <= eval(RHO_LOWER).cost + tol);
7708 assert!(selected_cost <= eval(RHO_UPPER).cost + tol);
7709 }
7710 }
7711
7712 #[test]
7713 fn backward_from_fit_matches_backward_with_refit() {
7714 let x = array![[1.0, -0.9], [1.0, -0.4], [1.0, 0.1], [1.0, 0.6], [1.0, 1.1],];
7719 let y = array![[0.2, -0.1], [0.4, 0.1], [0.7, 0.3], [1.0, 0.5], [1.5, 0.8]];
7720 let penalty = array![[0.0, 0.0], [0.0, 1.5]];
7721 let weights = array![1.05, 0.95, 1.01, 0.99, 1.03];
7722
7723 let refit = gaussian_reml_multi_closed_form_backward(
7724 x.view(),
7725 y.view(),
7726 penalty.view(),
7727 Some(weights.view()),
7728 Some(0.85),
7729 0.2,
7730 None,
7731 None,
7732 -0.1,
7733 0.0,
7734 )
7735 .expect("refit backward");
7736
7737 let fit = gaussian_reml_multi_closed_form_with_cache(
7738 x.view(),
7739 y.view(),
7740 penalty.view(),
7741 Some(weights.view()),
7742 Some(0.85),
7743 None,
7744 )
7745 .expect("forward fit");
7746 let from_fit = gaussian_reml_multi_closed_form_backward_from_fit(
7747 x.view(),
7748 y.view(),
7749 penalty.view(),
7750 Some(weights.view()),
7751 &fit,
7752 0.2,
7753 None,
7754 None,
7755 -0.1,
7756 0.0,
7757 )
7758 .expect("from_fit backward");
7759
7760 for (a, b) in refit.grad_x.iter().zip(from_fit.grad_x.iter()) {
7761 assert!((a - b).abs() <= 1.0e-12);
7762 }
7763 for (a, b) in refit.grad_y.iter().zip(from_fit.grad_y.iter()) {
7764 assert!((a - b).abs() <= 1.0e-12);
7765 }
7766 for (a, b) in refit.grad_weights.iter().zip(from_fit.grad_weights.iter()) {
7767 assert!((a - b).abs() <= 1.0e-12);
7768 }
7769 }
7770
7771 #[test]
7781 fn backward_degrades_gracefully_when_k_is_near_singular() {
7782 let x = array![
7786 [1.0, -1.0, 0.5],
7787 [1.0, -0.5, 0.2],
7788 [1.0, 0.0, -0.1],
7789 [1.0, 0.5, 0.3],
7790 [1.0, 1.0, 0.8],
7791 [1.0, 1.5, 1.1],
7792 [1.0, 2.0, 1.5],
7793 [1.0, 2.5, 2.0],
7794 [1.0, 3.0, 2.6],
7795 [1.0, 3.5, 3.1],
7796 ];
7797 let y = array![
7798 [0.1],
7799 [0.3],
7800 [0.4],
7801 [0.7],
7802 [1.0],
7803 [1.5],
7804 [2.0],
7805 [2.7],
7806 [3.3],
7807 [4.0]
7808 ];
7809 let penalty = array![[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
7811
7812 let mut fit =
7813 gaussian_reml_multi_closed_form(x.view(), y.view(), penalty.view(), None, Some(0.0))
7814 .expect("forward fit must succeed for well-posed input");
7815 fit.reml_hess_rho = 0.0;
7819
7820 let result = gaussian_reml_multi_closed_form_backward_from_fit(
7821 x.view(),
7822 y.view(),
7823 penalty.view(),
7824 None,
7825 &fit,
7826 1.0,
7829 None,
7830 None,
7831 1.0,
7832 1.0,
7833 )
7834 .expect("backward must NOT error on near-singular K");
7835
7836 assert_eq!(result.grad_x.dim(), (x.nrows(), x.ncols()));
7837 assert_eq!(result.grad_y.dim(), (y.nrows(), y.ncols()));
7838 assert_eq!(result.grad_penalty.dim(), (x.ncols(), x.ncols()));
7839 assert_eq!(result.grad_weights.dim(), x.nrows());
7840 for v in result.grad_x.iter() {
7841 assert!(v.is_finite(), "grad_x must be finite, got {v}");
7842 }
7843 for v in result.grad_y.iter() {
7844 assert!(v.is_finite(), "grad_y must be finite, got {v}");
7845 }
7846 for v in result.grad_penalty.iter() {
7847 assert!(v.is_finite(), "grad_penalty must be finite, got {v}");
7848 }
7849 for v in result.grad_weights.iter() {
7850 assert!(v.is_finite(), "grad_weights must be finite, got {v}");
7851 }
7852 }
7853}
7854
7855pub struct GaussianRemlBlocksBackwardAnalytic {
7859 pub grad_designs: Vec<Array2<f64>>,
7860 pub grad_penalties: Vec<Array2<f64>>,
7861 pub grad_y: Array2<f64>,
7862 pub grad_weights: Array1<f64>,
7865}
7866
7867pub fn gaussian_reml_fit_blocks_backward_analytic(
7876 designs: &[Array2<f64>],
7877 penalties_raw: &[Array2<f64>],
7878 y: ArrayView1<'_, f64>,
7879 weights: ArrayView1<'_, f64>,
7880 rhos: &[f64],
7881 grad_coefficients: Option<ArrayView2<'_, f64>>,
7882 grad_fitted: Option<ArrayView2<'_, f64>>,
7883 grad_lambdas: Option<ArrayView1<'_, f64>>,
7884 grad_log_lambdas: Option<ArrayView1<'_, f64>>,
7885 grad_reml_score: f64,
7886 grad_edf: Option<ArrayView1<'_, f64>>,
7887) -> Result<GaussianRemlBlocksBackwardAnalytic, EstimationError> {
7888 let n = y.len();
7889 let f_blocks = designs.len();
7890 if f_blocks == 0 || penalties_raw.len() != f_blocks {
7891 return Err(EstimationError::InvalidInput(format!(
7892 "gaussian_reml_fit_blocks_backward requires equal non-zero design and penalty \
7893 block counts; got designs={}, penalties={}",
7894 f_blocks,
7895 penalties_raw.len()
7896 )));
7897 }
7898 let mut offsets = Vec::with_capacity(f_blocks + 1);
7899 let mut cursor = 0_usize;
7900 offsets.push(cursor);
7901 for (block, design) in designs.iter().enumerate() {
7902 if design.nrows() != n {
7903 return Err(EstimationError::InvalidInput(format!(
7904 "designs[{block}] has {} rows, expected {n}",
7905 design.nrows()
7906 )));
7907 }
7908 if penalties_raw[block].dim() != (design.ncols(), design.ncols()) {
7909 return Err(EstimationError::InvalidInput(format!(
7910 "penalties[{block}] has shape {}x{}, expected {}x{}",
7911 penalties_raw[block].nrows(),
7912 penalties_raw[block].ncols(),
7913 design.ncols(),
7914 design.ncols()
7915 )));
7916 }
7917 cursor += design.ncols();
7918 offsets.push(cursor);
7919 }
7920 let p_total = cursor;
7923 if n == 0 || p_total == 0 {
7924 return Err(EstimationError::InvalidInput(
7925 "gaussian_reml_fit_blocks_backward requires non-empty rows and at least one coefficient column"
7926 .to_string(),
7927 ));
7928 }
7929
7930 if rhos.len() != f_blocks {
7931 return Err(EstimationError::InvalidInput(format!(
7932 "log_lambdas length mismatch: expected {f_blocks}, got {}",
7933 rhos.len()
7934 )));
7935 }
7936 if let Some(gc) = grad_coefficients {
7937 if gc.dim() != (p_total, 1) {
7938 return Err(EstimationError::InvalidInput(format!(
7939 "grad_coefficients shape mismatch: expected {}x1, got {}x{}",
7940 p_total,
7941 gc.nrows(),
7942 gc.ncols()
7943 )));
7944 }
7945 }
7946 if let Some(gf) = grad_fitted {
7947 if gf.dim() != (n, 1) {
7948 return Err(EstimationError::InvalidInput(format!(
7949 "grad_fitted shape mismatch: expected {}x1, got {}x{}",
7950 n,
7951 gf.nrows(),
7952 gf.ncols()
7953 )));
7954 }
7955 }
7956 if !grad_reml_score.is_finite() {
7957 return Err(EstimationError::InvalidInput(format!(
7958 "grad_reml_score must be finite; got {grad_reml_score}"
7959 )));
7960 }
7961 if let Some(vec) = grad_lambdas {
7962 if vec.len() != f_blocks {
7963 return Err(EstimationError::InvalidInput(format!(
7964 "grad_lambdas length mismatch: expected {f_blocks}, got {}",
7965 vec.len()
7966 )));
7967 }
7968 }
7969 if let Some(vec) = grad_log_lambdas {
7970 if vec.len() != f_blocks {
7971 return Err(EstimationError::InvalidInput(format!(
7972 "grad_log_lambdas length mismatch: expected {f_blocks}, got {}",
7973 vec.len()
7974 )));
7975 }
7976 }
7977 if let Some(vec) = grad_edf {
7978 if vec.len() != f_blocks {
7979 return Err(EstimationError::InvalidInput(format!(
7980 "grad_edf length mismatch: expected {f_blocks}, got {}",
7981 vec.len()
7982 )));
7983 }
7984 }
7985 if let Some(gc) = grad_coefficients {
7986 if let Some(((row, col), value)) = gc.indexed_iter().find(|(_, value)| !value.is_finite()) {
7987 return Err(EstimationError::InvalidInput(format!(
7988 "grad_coefficients[{row},{col}] must be finite; got {value}"
7989 )));
7990 }
7991 }
7992 if let Some(gf) = grad_fitted {
7993 if let Some(((row, col), value)) = gf.indexed_iter().find(|(_, value)| !value.is_finite()) {
7994 return Err(EstimationError::InvalidInput(format!(
7995 "grad_fitted[{row},{col}] must be finite; got {value}"
7996 )));
7997 }
7998 }
7999 if let Some(vec) = grad_lambdas {
8000 if let Some((block, value)) = vec.iter().enumerate().find(|(_, value)| !value.is_finite()) {
8001 return Err(EstimationError::InvalidInput(format!(
8002 "grad_lambdas[{block}] must be finite; got {value}"
8003 )));
8004 }
8005 }
8006 if let Some(vec) = grad_log_lambdas {
8007 if let Some((block, value)) = vec.iter().enumerate().find(|(_, value)| !value.is_finite()) {
8008 return Err(EstimationError::InvalidInput(format!(
8009 "grad_log_lambdas[{block}] must be finite; got {value}"
8010 )));
8011 }
8012 }
8013 if let Some(vec) = grad_edf {
8014 if let Some((block, value)) = vec.iter().enumerate().find(|(_, value)| !value.is_finite()) {
8015 return Err(EstimationError::InvalidInput(format!(
8016 "grad_edf[{block}] must be finite; got {value}"
8017 )));
8018 }
8019 }
8020 for (block, design) in designs.iter().enumerate() {
8021 if let Some(((row, col), value)) =
8022 design.indexed_iter().find(|(_, value)| !value.is_finite())
8023 {
8024 return Err(EstimationError::InvalidInput(format!(
8025 "designs[{block}][{row},{col}] must be finite; got {value}"
8026 )));
8027 }
8028 }
8029 for (block, penalty) in penalties_raw.iter().enumerate() {
8030 if let Some(((row, col), value)) =
8031 penalty.indexed_iter().find(|(_, value)| !value.is_finite())
8032 {
8033 return Err(EstimationError::InvalidInput(format!(
8034 "penalties[{block}][{row},{col}] must be finite; got {value}"
8035 )));
8036 }
8037 }
8038 if let Some((row, value)) = y.iter().enumerate().find(|(_, value)| !value.is_finite()) {
8039 return Err(EstimationError::InvalidInput(format!(
8040 "y[{row}] must be finite; got {value}"
8041 )));
8042 }
8043 if let Some((row, value)) = weights
8044 .iter()
8045 .enumerate()
8046 .find(|(_, value)| !value.is_finite() || **value < 0.0)
8047 {
8048 return Err(EstimationError::InvalidInput(format!(
8049 "weights[{row}] must be finite and non-negative; got {value}"
8050 )));
8051 }
8052
8053 let mut z = Array2::<f64>::zeros((n, p_total));
8054 for k in 0..f_blocks {
8055 z.slice_mut(s![.., offsets[k]..offsets[k + 1]])
8056 .assign(&designs[k]);
8057 }
8058
8059 let blockwise_penalties: Vec<BlockwisePenalty> = penalties_raw
8060 .iter()
8061 .enumerate()
8062 .map(|(block, penalty)| {
8063 BlockwisePenalty::new(offsets[block]..offsets[block + 1], penalty.clone())
8064 })
8065 .collect();
8066 let domain = GaussianRemlBlocksDomain::from_blockwise_penalties(p_total, &blockwise_penalties)?;
8067 let lambdas = Array1::from_vec(
8068 gam_problem::checked_exp_log_strengths(rhos.iter().copied())
8069 .map_err(|error| EstimationError::InvalidInput(error.to_string()))?,
8070 );
8071 let k_matrix = domain.certify_joint_coefficient_map(z.view(), weights, lambdas.view())?;
8072
8073 if f_blocks == 1 {
8080 let mut upstream_lambda = grad_lambdas.map_or(0.0, |gradient| gradient[0]);
8081 if let Some(gradient) = grad_log_lambdas {
8082 upstream_lambda += gradient[0] / lambdas[0];
8083 }
8084 let y_owned = y.to_owned().insert_axis(Axis(1));
8085 let weights_owned = weights.to_owned();
8086 let fit = gaussian_reml_multi_closed_form_with_cache(
8087 z.view(),
8088 y_owned.view(),
8089 penalties_raw[0].view(),
8090 Some(weights_owned.view()),
8091 Some(lambdas[0]),
8092 None,
8093 )?;
8094 let backward = gaussian_reml_multi_closed_form_backward_from_fit(
8095 z.view(),
8096 y_owned.view(),
8097 penalties_raw[0].view(),
8098 Some(weights_owned.view()),
8099 &fit,
8100 upstream_lambda,
8101 grad_coefficients,
8102 grad_fitted,
8103 grad_reml_score,
8104 grad_edf.map_or(0.0, |gradient| gradient[0]),
8105 )?;
8106 return Ok(GaussianRemlBlocksBackwardAnalytic {
8107 grad_designs: vec![backward.grad_x],
8108 grad_penalties: vec![backward.grad_penalty],
8109 grad_y: backward.grad_y,
8110 grad_weights: backward.grad_weights,
8111 });
8112 }
8113
8114 let penalties = domain.local_penalties();
8115 let pinvs = domain.penalty_pseudoinverses()?;
8116 let r = gam_linalg::utils::certified_spd_inverse(
8117 &k_matrix,
8118 "block Gaussian REML penalized normal matrix",
8119 )
8120 .map(gam_linalg::utils::CertifiedSpdInverse::into_inverse)
8121 .map_err(|error| {
8122 EstimationError::InvalidInput(format!(
8123 "block Gaussian REML requires an exact SPD penalized normal matrix: {error}"
8124 ))
8125 })?;
8126
8127 let mut xtwy = Array1::<f64>::zeros(p_total);
8128 for row in 0..n {
8129 let wy = weights[row] * y[row];
8130 for col in 0..p_total {
8131 xtwy[col] += z[[row, col]] * wy;
8132 }
8133 }
8134 let beta = r.dot(&xtwy);
8135 let fitted = z.dot(&beta);
8136 if let Some((col, value)) = beta
8137 .iter()
8138 .enumerate()
8139 .find(|(_, value)| !value.is_finite())
8140 {
8141 return Err(EstimationError::InvalidInput(format!(
8142 "solved coefficient {col} is non-finite: {value}"
8143 )));
8144 }
8145 let residual = &y.to_owned() - &fitted;
8146 let weighted_residual = &residual * &weights.to_owned();
8147 let mut q = residual
8148 .iter()
8149 .zip(weights.iter())
8150 .map(|(&value, &weight)| weight * value * value)
8151 .sum::<f64>();
8152 for block in 0..f_blocks {
8153 let start = offsets[block];
8154 let end = offsets[block + 1];
8155 let beta_block = beta.slice(s![start..end]);
8156 q += lambdas[block] * beta_block.dot(&penalties[block].dot(&beta_block));
8157 }
8158 if !q.is_finite() || q <= 0.0 {
8159 return Err(EstimationError::InvalidInput(format!(
8160 "Gaussian REML residual quadratic form must be finite and positive; got {q}"
8161 )));
8162 }
8163 let nullity = domain.nullspace_dims().iter().sum::<usize>();
8164 let nu = effective_observation_count(weights) as f64 - nullity as f64;
8167 if !(nu.is_finite() && nu > 0.0) {
8168 return Err(EstimationError::InvalidInput(format!(
8169 "Gaussian REML residual degrees of freedom must be positive; got {nu}"
8170 )));
8171 }
8172 let tau = nu / q;
8173 let tau_q = -nu / (q * q);
8174 if !(tau.is_finite() && tau_q.is_finite()) {
8175 return Err(EstimationError::InvalidInput(format!(
8176 "Gaussian REML scale derivatives are non-finite: tau={tau}, tau_q={tau_q}"
8177 )));
8178 }
8179
8180 let mut grad_z = Array2::<f64>::zeros((n, p_total));
8181 let mut g_kernel = Array2::<f64>::zeros((p_total, p_total));
8182 let mut h_kernel = Array1::<f64>::zeros(p_total);
8183 let mut q_kernel = 0.0_f64;
8184 let mut j_blocks: Vec<Array2<f64>> = penalties
8185 .iter()
8186 .map(|p| Array2::<f64>::zeros(p.dim()))
8187 .collect();
8188
8189 let mut beta_tilde = Array1::<f64>::zeros(p_total);
8190 if let Some(gc) = grad_coefficients {
8191 beta_tilde += &gc.column(0).to_owned();
8192 }
8193 if let Some(gf) = grad_fitted {
8194 let gf_col = gf.column(0).to_owned();
8195 beta_tilde += &z.t().dot(&gf_col);
8196 for row in 0..n {
8197 for col in 0..p_total {
8198 grad_z[[row, col]] += gf_col[row] * beta[col];
8199 }
8200 }
8201 }
8202
8203 let u = r.dot(&beta_tilde);
8208 h_kernel += &u;
8209 for i in 0..p_total {
8210 for j in 0..p_total {
8211 g_kernel[[i, j]] -= 0.5 * (beta[i] * u[j] + u[i] * beta[j]);
8212 }
8213 }
8214
8215 let mut alpha = Array1::<f64>::zeros(f_blocks);
8216 if let Some(gl) = grad_lambdas {
8217 for block in 0..f_blocks {
8218 alpha[block] += gl[block] * lambdas[block];
8219 }
8220 }
8221 if let Some(grho) = grad_log_lambdas {
8222 alpha += &grho.to_owned();
8223 }
8224
8225 let mut p_betas = Vec::with_capacity(f_blocks);
8226 let mut m_vectors = Vec::with_capacity(f_blocks);
8227 let mut rp_matrices = Vec::with_capacity(f_blocks);
8228 let mut rpr_matrices = Vec::with_capacity(f_blocks);
8229 let mut b_values = Array1::<f64>::zeros(f_blocks);
8230 let mut t_values = Array1::<f64>::zeros(f_blocks);
8231
8232 for block in 0..f_blocks {
8233 let start = offsets[block];
8234 let end = offsets[block + 1];
8235 let beta_k = beta.slice(s![start..end]).to_owned();
8236 let s_beta = penalties[block].dot(&beta_k);
8237 let lambda = lambdas[block];
8238 let lambda_s_beta = s_beta.mapv(|value| lambda * value);
8239 let mut p_beta = Array1::<f64>::zeros(p_total);
8240 for local_i in 0..(end - start) {
8241 p_beta[start + local_i] = lambda_s_beta[local_i];
8242 }
8243 let weighted_penalty = penalties[block].mapv(|value| lambda * value);
8244 let rp_block = r.slice(s![.., start..end]).dot(&weighted_penalty);
8245 let mut rp = Array2::<f64>::zeros((p_total, p_total));
8246 rp.slice_mut(s![.., start..end]).assign(&rp_block);
8247 let rpr = rp_block.dot(&r.slice(s![start..end, ..]));
8248 let m = r.slice(s![.., start..end]).dot(&lambda_s_beta);
8249 b_values[block] = beta.dot(&p_beta);
8250 t_values[block] = (0..(end - start))
8251 .map(|local_i| rp_block[[start + local_i, local_i]])
8252 .sum::<f64>();
8253 alpha[block] -= u.dot(&p_beta);
8254 p_betas.push(p_beta);
8255 m_vectors.push(m);
8256 rp_matrices.push(rp);
8257 rpr_matrices.push(rpr);
8258 }
8259
8260 if grad_reml_score != 0.0 {
8261 q_kernel += 0.5 * grad_reml_score * tau;
8262 g_kernel += &(r.clone() * (0.5 * grad_reml_score));
8263 for block in 0..f_blocks {
8264 j_blocks[block] -= &(pinvs[block].clone() * (0.5 * grad_reml_score / lambdas[block]));
8265 }
8266 }
8267
8268 let mut trace_pairs = Array2::<f64>::zeros((f_blocks, f_blocks));
8269 for i in 0..f_blocks {
8270 for j in 0..f_blocks {
8271 trace_pairs[[i, j]] =
8272 gam_linalg::utils::trace_of_product(rp_matrices[i].view(), rp_matrices[j].view());
8273 }
8274 }
8275
8276 if let Some(ge) = grad_edf {
8277 for edf_block in 0..f_blocks {
8278 let scale = ge[edf_block];
8279 if scale == 0.0 {
8280 continue;
8281 }
8282 let start = offsets[edf_block];
8283 let end = offsets[edf_block + 1];
8284 g_kernel += &(rpr_matrices[edf_block].clone() * scale);
8285 j_blocks[edf_block] -= &(r.slice(s![start..end, start..end]).to_owned() * scale);
8286 for rho_block in 0..f_blocks {
8287 alpha[rho_block] += scale * trace_pairs[[edf_block, rho_block]];
8288 if rho_block == edf_block {
8289 alpha[rho_block] -= scale * t_values[edf_block];
8290 }
8291 }
8292 }
8293 }
8294
8295 if let Some((block, value)) = alpha
8296 .iter()
8297 .enumerate()
8298 .find(|(_, value)| !value.is_finite())
8299 {
8300 return Err(EstimationError::InvalidInput(format!(
8301 "rho adjoint seed for block {block} is non-finite: {value}"
8302 )));
8303 }
8304
8305 if alpha.iter().any(|value| *value != 0.0) {
8306 let mut outer_h = Array2::<f64>::zeros((f_blocks, f_blocks));
8307 for k in 0..f_blocks {
8308 for j in 0..f_blocks {
8309 let beta_pk_r_pj_beta = p_betas[k].dot(&m_vectors[j]);
8310 outer_h[[k, j]] = 0.5 * trace_pairs[[k, j]] + tau * beta_pk_r_pj_beta
8311 - if k == j {
8312 0.5 * (t_values[k] + tau * b_values[k])
8313 } else {
8314 0.0
8315 }
8316 - 0.5 * tau_q * b_values[k] * b_values[j];
8317 }
8318 }
8319 gam_linalg::matrix::symmetrize_in_place(&mut outer_h);
8324 if let Some(((row, col), value)) =
8325 outer_h.indexed_iter().find(|(_, value)| !value.is_finite())
8326 {
8327 return Err(EstimationError::InvalidInput(format!(
8328 "outer rho curvature entry ({row},{col}) is non-finite: {value}"
8329 )));
8330 }
8331 let rho_adj = gam_linalg::utils::certified_symmetric_solve(
8332 &outer_h,
8333 &alpha,
8334 "block Gaussian REML outer-rho adjoint",
8335 )
8336 .map(gam_linalg::utils::CertifiedSymmetricSolution::into_solution)
8337 .map_err(|error| {
8338 EstimationError::InvalidInput(format!(
8339 "block Gaussian REML outer-rho adjoint is not exactly solvable: {error}"
8340 ))
8341 })?;
8342 if let Some((block, value)) = rho_adj
8343 .iter()
8344 .enumerate()
8345 .find(|(_, value)| !value.is_finite())
8346 {
8347 return Err(EstimationError::InvalidInput(format!(
8348 "outer rho adjoint for block {block} is non-finite: {value}"
8349 )));
8350 }
8351 let weighted_b_sum = rho_adj
8352 .iter()
8353 .zip(b_values.iter())
8354 .map(|(&zk, &bk)| zk * bk)
8355 .sum::<f64>();
8356 q_kernel += 0.5 * tau_q * weighted_b_sum;
8357 for block in 0..f_blocks {
8358 let zk = rho_adj[block];
8359 if zk == 0.0 {
8360 continue;
8361 }
8362 g_kernel -= &(rpr_matrices[block].clone() * (0.5 * zk));
8363 let m = &m_vectors[block];
8364 for i in 0..p_total {
8365 h_kernel[i] += tau * zk * m[i];
8366 for j in 0..p_total {
8367 g_kernel[[i, j]] -= 0.5 * tau * zk * (beta[i] * m[j] + m[i] * beta[j]);
8368 }
8369 }
8370 let start = offsets[block];
8371 let end = offsets[block + 1];
8372 j_blocks[block] += &(r.slice(s![start..end, start..end]).to_owned() * (0.5 * zk));
8373 for i in 0..(end - start) {
8374 for j in 0..(end - start) {
8375 j_blocks[block][[i, j]] += 0.5 * tau * zk * beta[start + i] * beta[start + j];
8376 }
8377 }
8378 }
8379 }
8380
8381 for row in 0..n {
8382 for col in 0..p_total {
8383 grad_z[[row, col]] += -2.0 * q_kernel * weighted_residual[row] * beta[col];
8384 }
8385 }
8386 let zg = z.dot(&g_kernel);
8387 for row in 0..n {
8388 for col in 0..p_total {
8389 grad_z[[row, col]] += 2.0 * weights[row] * zg[[row, col]];
8390 }
8391 }
8392 let wy = y.to_owned() * &weights.to_owned();
8393 for row in 0..n {
8394 for col in 0..p_total {
8395 grad_z[[row, col]] += wy[row] * h_kernel[col];
8396 }
8397 }
8398
8399 let mut grad_y = Array2::<f64>::zeros((n, 1));
8400 let zh = z.dot(&h_kernel);
8401 for row in 0..n {
8402 grad_y[[row, 0]] = 2.0 * q_kernel * weighted_residual[row] + weights[row] * zh[row];
8403 }
8404
8405 let mut grad_weights = Array1::<f64>::zeros(n);
8406 for row in 0..n {
8407 let diag_zgz = (0..p_total)
8408 .map(|col| z[[row, col]] * zg[[row, col]])
8409 .sum::<f64>();
8410 grad_weights[row] = q_kernel * residual[row] * residual[row] + diag_zgz + y[row] * zh[row];
8411 }
8412 finish_gaussian_reml_weight_vjp(weights, 1, grad_reml_score, &mut grad_weights);
8413
8414 let mut grad_penalties = Vec::with_capacity(f_blocks);
8415 for block in 0..f_blocks {
8416 let start = offsets[block];
8417 let end = offsets[block + 1];
8418 let mut local = g_kernel.slice(s![start..end, start..end]).to_owned();
8419 for i in 0..(end - start) {
8420 for j in 0..(end - start) {
8421 local[[i, j]] += q_kernel * beta[start + i] * beta[start + j];
8422 }
8423 }
8424 local += &j_blocks[block];
8425 local *= lambdas[block];
8426 gam_linalg::matrix::symmetrize_in_place(&mut local);
8427 grad_penalties.push(local);
8428 }
8429
8430 let mut grad_designs = Vec::with_capacity(f_blocks);
8431 for block in 0..f_blocks {
8432 grad_designs.push(
8433 grad_z
8434 .slice(s![.., offsets[block]..offsets[block + 1]])
8435 .to_owned(),
8436 );
8437 }
8438
8439 Ok(GaussianRemlBlocksBackwardAnalytic {
8440 grad_designs,
8441 grad_penalties,
8442 grad_y,
8443 grad_weights,
8444 })
8445}
8446
8447pub struct DenseFisherGaussianFit {
8451 pub coefficients: Array2<f64>,
8452 pub fitted: Array2<f64>,
8453 pub sigma2: Array1<f64>,
8454 pub objective: f64,
8455}
8456
8457pub fn add_block_diagonal_penalty(
8460 hessian: &mut Array2<f64>,
8461 penalty: ArrayView2<'_, f64>,
8462 lambda: f64,
8463 n_outputs: usize,
8464) -> Result<(), EstimationError> {
8465 let k = penalty.ncols();
8466 if penalty.nrows() != k {
8467 return Err(EstimationError::InvalidInput(format!(
8468 "penalty must be square for dense Fisher fit; got {}x{}",
8469 penalty.nrows(),
8470 penalty.ncols()
8471 )));
8472 }
8473 if hessian.dim() != (k * n_outputs, k * n_outputs) {
8474 return Err(EstimationError::InvalidInput(
8475 "dense Fisher Hessian shape mismatch while adding penalty".to_string(),
8476 ));
8477 }
8478 for output in 0..n_outputs {
8479 let offset = output * k;
8480 for row in 0..k {
8481 for col in 0..k {
8482 let s_sym = 0.5 * (penalty[[row, col]] + penalty[[col, row]]);
8483 hessian[[offset + row, offset + col]] += lambda * s_sym;
8484 }
8485 }
8486 }
8487 Ok(())
8488}
8489
8490pub fn dense_fisher_gaussian_fit(
8497 design: ArrayView2<'_, f64>,
8498 y: ArrayView2<'_, f64>,
8499 penalty: ArrayView2<'_, f64>,
8500 row_weights: ArrayView1<'_, f64>,
8501 fisher_w: ArrayView3<'_, f64>,
8502 lambda: f64,
8503 latent_prior_score: f64,
8504) -> Result<DenseFisherGaussianFit, EstimationError> {
8505 let n_obs = design.nrows();
8506 let k = design.ncols();
8507 let n_outputs = y.ncols();
8508 let mut hessian = crate::pirls::dense_block_xtwx(design, fisher_w, Some(row_weights))?;
8509 add_block_diagonal_penalty(&mut hessian, penalty, lambda, n_outputs)?;
8510 let rhs = crate::pirls::dense_block_xtwy(design, fisher_w, y, Some(row_weights))?;
8511 let beta_vec =
8512 gam_linalg::utils::solve_dense_block_system(&hessian, &rhs, "dense Fisher Gaussian")
8513 .map_err(EstimationError::InvalidInput)?;
8514 let mut coefficients = Array2::<f64>::zeros((k, n_outputs));
8515 for output in 0..n_outputs {
8516 for col in 0..k {
8517 coefficients[[col, output]] = beta_vec[output * k + col];
8518 }
8519 }
8520 let fitted = design.dot(&coefficients);
8521 let mut sigma2 = Array1::<f64>::zeros(n_outputs);
8522 let mut objective = latent_prior_score;
8523 for row in 0..n_obs {
8524 for a in 0..n_outputs {
8525 let ra = y[[row, a]] - fitted[[row, a]];
8526 sigma2[a] += row_weights[row] * ra * ra;
8527 for b in 0..n_outputs {
8528 objective += 0.5
8529 * row_weights[row]
8530 * ra
8531 * fisher_w[[row, a, b]]
8532 * (y[[row, b]] - fitted[[row, b]]);
8533 }
8534 }
8535 }
8536 for output in 0..n_outputs {
8537 sigma2[output] /= (n_obs.saturating_sub(k).max(1)) as f64;
8538 let beta_col = coefficients.column(output);
8539 let s_beta = penalty.dot(&beta_col);
8540 objective += 0.5 * lambda * beta_col.dot(&s_beta);
8541 }
8542 Ok(DenseFisherGaussianFit {
8543 coefficients,
8544 fitted,
8545 sigma2,
8546 objective,
8547 })
8548}
8549
8550#[cfg(test)]
8571mod perfect_fit_refusal_tests {
8572 use super::*;
8573 use ndarray::array;
8574
8575 fn zero_residual_designs() -> Vec<(&'static str, Array2<f64>, Array2<f64>, Array2<f64>)> {
8581 let n = 12usize;
8585 let a_x = Array2::<f64>::from_shape_fn((n, 3), |(row, col)| {
8586 let t = 2.0 * std::f64::consts::PI * (row as f64) / (n as f64);
8587 match col {
8588 0 => 1.0,
8589 1 => t.sin(),
8590 _ => t.cos(),
8591 }
8592 });
8593 let a_y = Array2::<f64>::from_elem((n, 1), 0.7);
8594 let a_penalty = array![[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]];
8595
8596 let b_x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0]];
8601 let b_y = array![[1.0], [3.0], [5.0], [7.0], [9.0]];
8602 let b_penalty = array![[0.0, 0.0], [0.0, 1.0]];
8603
8604 let c_x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0]];
8608 let c_y = array![[1.0], [1.0], [1.0], [1.0], [1.0]];
8609 let c_penalty = array![[0.0, 0.0], [0.0, 1.0]];
8610
8611 let d_x = c_x.clone();
8613 let d_y = Array2::<f64>::zeros((5, 1));
8614 let d_penalty = c_penalty.clone();
8615
8616 vec![
8617 ("A irrational basis, constant response", a_x, a_y, a_penalty),
8618 ("B integer basis, penalized mass present", b_x, b_y, b_penalty),
8619 ("C integer basis, all mass in null(S)", c_x, c_y, c_penalty),
8620 ("D identically zero response", d_x, d_y, d_penalty),
8621 ]
8622 }
8623
8624 #[test]
8629 fn every_zero_residual_design_is_refused_alike() {
8630 let mut failures: Vec<String> = Vec::new();
8631 let mut verdicts: Vec<(&'static str, bool)> = Vec::new();
8632
8633 for (name, x, y, penalty) in zero_residual_designs() {
8634 let prepared =
8635 prepare_gaussian_reml(x.view(), y.view(), penalty.view(), None, None, None)
8636 .unwrap_or_else(|error| panic!("{name}: preparation failed: {error}"));
8637
8638 let DispersionResidualParts {
8642 unpenalized_residual,
8643 penalized_residual,
8644 ..
8645 } = dispersion_residual_parts(
8646 &prepared.cache,
8647 prepared.ywy.view(),
8648 prepared.projected_rhs_squared.view(),
8649 0,
8650 RHO_LOWER,
8651 );
8652 let residual = unpenalized_residual + penalized_residual;
8653 let ywy = prepared.ywy[0];
8654 let resolution = profile_residual_resolution(&prepared.cache, ywy);
8655 if !(residual <= f64::EPSILON.sqrt() * ywy.max(1.0)) {
8656 failures.push(format!(
8657 "{name}: REGIME — residual {residual:.6e} against ywy {ywy:.6e} is far above \
8658 the cancellation scale, so this design does NOT interpolate its response and \
8659 the fixture has drifted out of the regime under test"
8660 ));
8661 }
8662
8663 let verdict = validate_reml_profile_residuals(
8664 &prepared.cache,
8665 prepared.ywy.view(),
8666 prepared.projected_rhs_squared.view(),
8667 RHO_LOWER,
8668 );
8669 verdicts.push((name, verdict.is_ok()));
8670 if verdict.is_ok() {
8671 failures.push(format!(
8672 "{name}: ACCEPTED a residual of {residual:.6e} (ywy {ywy:.6e}, resolution \
8673 {resolution:.6e}) whose true value is exactly zero; the profiled dispersion \
8674 it carries is pure roundoff"
8675 ));
8676 }
8677 }
8678
8679 let accepted: Vec<&str> = verdicts
8680 .iter()
8681 .filter(|(_, ok)| *ok)
8682 .map(|(name, _)| *name)
8683 .collect();
8684 assert!(
8685 failures.is_empty(),
8686 "#2723: the four exactly-zero-residual designs did not agree on refusal. \
8687 Accepted: {accepted:?}. Details:\n - {}",
8688 failures.join("\n - ")
8689 );
8690 }
8691
8692 #[test]
8696 fn a_genuine_residual_is_accepted_at_every_scale() {
8697 let x = array![[1.0, 0.0], [1.0, 1.0], [1.0, 2.0], [1.0, 3.0], [1.0, 4.0]];
8698 let base_y = array![[1.1], [2.9], [5.2], [6.8], [9.1]];
8699 let penalty = array![[0.0, 0.0], [0.0, 1.0]];
8700
8701 for scale in [1.0e-6, 1.0, 1.0e6] {
8702 let y = base_y.mapv(|value| value * scale);
8703 let prepared =
8704 prepare_gaussian_reml(x.view(), y.view(), penalty.view(), None, None, None)
8705 .expect("the control design is finite and full rank");
8706 let DispersionResidualParts {
8707 unpenalized_residual,
8708 penalized_residual,
8709 ..
8710 } = dispersion_residual_parts(
8711 &prepared.cache,
8712 prepared.ywy.view(),
8713 prepared.projected_rhs_squared.view(),
8714 0,
8715 RHO_LOWER,
8716 );
8717 let residual = unpenalized_residual + penalized_residual;
8718 let ywy = prepared.ywy[0];
8719 assert!(
8720 residual > f64::EPSILON.sqrt() * ywy,
8721 "control at scale {scale:e}: residual {residual:.6e} is at cancellation scale \
8722 against ywy {ywy:.6e}, so this is not a genuine-residual control"
8723 );
8724 let verdict = validate_reml_profile_residuals(
8725 &prepared.cache,
8726 prepared.ywy.view(),
8727 prepared.projected_rhs_squared.view(),
8728 RHO_LOWER,
8729 );
8730 assert!(
8731 verdict.is_ok(),
8732 "control at scale {scale:e}: a genuine residual {residual:.6e} (ywy {ywy:.6e}) was \
8733 refused: {:?}",
8734 verdict.err()
8735 );
8736 }
8737 }
8738
8739 #[test]
8742 fn the_refusal_is_invariant_to_the_response_scale() {
8743 for (name, x, y, penalty) in zero_residual_designs() {
8744 for scale in [1.0e-8, 1.0, 1.0e8] {
8745 let scaled = y.mapv(|value| value * scale);
8746 let prepared =
8747 prepare_gaussian_reml(x.view(), scaled.view(), penalty.view(), None, None, None)
8748 .unwrap_or_else(|error| panic!("{name} at {scale:e}: {error}"));
8749 let verdict = validate_reml_profile_residuals(
8750 &prepared.cache,
8751 prepared.ywy.view(),
8752 prepared.projected_rhs_squared.view(),
8753 RHO_LOWER,
8754 );
8755 assert!(
8756 verdict.is_err(),
8757 "{name}: rescaling the response by {scale:e} flipped the perfect-fit verdict \
8758 to ACCEPTED; the bar is not scale-invariant"
8759 );
8760 }
8761 }
8762 }
8763}
8764
8765#[cfg(test)]
8782mod eigenvalue_range_predicate_agreement_2740_tests {
8783 use super::*;
8784 use ndarray::array;
8785
8786 const LARGEST: f64 = 4.0;
8787
8788 fn disputed_band_cache() -> GaussianRemlEigenCache {
8797 let eigenvalues = array![LARGEST, 1.0, 5.0e-11, 3.2e-18, 0.0];
8798 let p = eigenvalues.len();
8799 GaussianRemlEigenCache {
8800 penalty_eigenvalues: eigenvalues,
8801 eigenvectors: Array2::eye(p),
8802 coefficient_basis: Array2::eye(p),
8803 xtwx_fingerprint: 0,
8804 penalty_fingerprint: 0,
8805 logdet_xtwx: 0.0,
8806 logdet_penalty_positive: LARGEST.ln() + 1.0_f64.ln(),
8807 penalty_rank: 2,
8808 nullity: 3,
8809 }
8810 }
8811
8812 #[test]
8814 fn the_range_count_the_null_count_and_penalty_rank_are_one_predicate() {
8815 let cache = disputed_band_cache();
8816 let spectrum = PenaltyRangeSpectrum::of(&cache);
8817
8818 assert_eq!(
8821 spectrum.tolerance,
8822 LARGEST * EIGEN_REL_TOL,
8823 "the range threshold must be the relative one that defines penalty_rank"
8824 );
8825
8826 let absolute_positive = cache
8830 .penalty_eigenvalues
8831 .iter()
8832 .filter(|delta| **delta > 0.0)
8833 .count();
8834 assert!(
8835 absolute_positive > cache.penalty_rank,
8836 "precondition unmet: the fixture carries no positive-but-null direction \
8837 (penalty_rank={}, eigenvalues passing `> 0.0`={absolute_positive})",
8838 cache.penalty_rank
8839 );
8840
8841 assert_eq!(
8842 spectrum.rank(),
8843 cache.penalty_rank,
8844 "the classified range count must be the rank the cache reports"
8845 );
8846 assert_eq!(
8847 spectrum.iter().filter(|delta| *delta > 0.0).count(),
8848 cache.penalty_rank,
8849 "a `δ > 0.0` read of the CLASSIFIED spectrum must select exactly the \
8850 directions penalty_rank counted"
8851 );
8852 assert_eq!(
8853 spectrum.iter().filter(|delta| *delta == 0.0).count(),
8854 cache.nullity,
8855 "the null set must be the exact complement of the range set"
8856 );
8857 }
8858
8859 #[test]
8865 fn the_large_rho_logdet_gradient_vanishes_because_sum_and_offset_share_a_population() {
8866 let cache = disputed_band_cache();
8867 let spectrum = PenaltyRangeSpectrum::of(&cache);
8868 let lambda = RHO_UPPER.exp();
8869 let n_outputs = 1.0_f64;
8870
8871 let (term, _edf) = gaussian_reml_logdet_term(&cache, RHO_UPPER, n_outputs);
8872
8873 let residual: f64 = spectrum
8876 .iter()
8877 .filter(|delta| *delta > 0.0)
8878 .map(|delta| 1.0 / (1.0 + lambda * delta))
8879 .sum();
8880 let bound = 0.5 * n_outputs * residual
8881 + ((spectrum.len() + 4) as f64) * f64::EPSILON * (cache.penalty_rank as f64);
8882 assert!(
8883 term.grad.abs() <= bound,
8884 "the large-λ log-determinant gradient is {} but the range-populated \
8885 residual bounds it by {bound:e}",
8886 term.grad
8887 );
8888
8889 let disputed_trace: f64 = (0..spectrum.len())
8894 .filter(|index| spectrum.get(*index) == 0.0)
8895 .map(|index| {
8896 let t = lambda * cache.penalty_eigenvalues[index];
8897 t / (1.0 + t)
8898 })
8899 .sum();
8900 assert!(
8901 disputed_trace > 0.5,
8902 "precondition unmet: the disputed band contributes only {disputed_trace} to \
8903 the trace at rho={RHO_UPPER}, so an absolute `> 0.0` sum would barely differ \
8904 from the classified one and this test would be mute"
8905 );
8906 assert!(
8907 0.5 * n_outputs * disputed_trace > bound,
8908 "the disputed band's contribution {disputed_trace} does not clear the bound \
8909 {bound:e}; the assertion above cannot distinguish the two predicates"
8910 );
8911 }
8912}