1use std::collections::HashMap;
94use std::sync::Arc;
95
96use faer::Side;
97use gam_linalg::faer_ndarray::FaerEigh;
98use gam_math::score_opt::{ClosedInterval, DerivativeEnclosure, ScoreJet, maximize_score_1d};
99use gam_terms::grid_spline_2d::{chol_solve, cholesky_logdet};
100use ndarray::Array2;
101
102const OVERLAP: f64 = 2.0;
106const H0_FRACTION: f64 = 0.5;
108const INITIAL_LEVELS: usize = 3;
110const MAX_LEVELS: usize = 16;
112const MAX_CENTERS: usize = 200_000;
114const REFINE_TOL: f64 = 1e-3;
117
118const DENSE_GRAM_MAX: usize = 1536;
121
122const CG_RTOL: f64 = 1e-9;
128const CG_MAX_ITERS: usize = 4000;
129
130const COARSE_DOMINANCE: f64 = 4.0;
157const COARSE_SPACE_MAX: usize = 4096;
169const MIN_COARSE_LEVELS: usize = 2;
170
171const QUASI_UNIFORMITY_MAX_ASPECT: f64 = 1.0e3;
195
196const SLQ_PROBES: usize = 24;
200const SLQ_LANCZOS_STEPS: usize = 48;
201
202const RNG_SEED: u64 = 0x1032_CA5C_ADE0_5EED;
204
205const EIG_FLOOR: f64 = 1e-300;
207
208struct SplitMix64(u64);
212
213impl SplitMix64 {
214 fn new(seed: u64) -> Self {
215 SplitMix64(seed)
216 }
217
218 fn next_u64(&mut self) -> u64 {
219 gam_linalg::utils::splitmix64(&mut self.0)
220 }
221
222 fn next_unit(&mut self) -> f64 {
224 ((self.next_u64() >> 11) as f64 + 0.5) / 9_007_199_254_740_992.0
225 }
226
227 fn next_normal(&mut self) -> f64 {
229 let u1 = self.next_unit();
230 let u2 = self.next_unit();
231 (-2.0 * u1.ln()).sqrt() * (std::f64::consts::TAU * u2).cos()
232 }
233
234 fn next_sign(&mut self) -> f64 {
236 if self.next_u64() & 1 == 0 { 1.0 } else { -1.0 }
237 }
238}
239
240#[inline]
245fn cell_of(z: &[f64; 3], dim: usize, width: f64) -> (i32, i32, i32) {
246 let mut c = [0_i32; 3];
247 for a in 0..dim {
248 c[a] = (z[a] / width).floor() as i32;
249 }
250 (c[0], c[1], c[2])
251}
252
253struct HashGrid {
256 width: f64,
257 dim: usize,
258 cells: HashMap<(i32, i32, i32), Vec<u32>>,
259}
260
261impl HashGrid {
262 fn new(width: f64, dim: usize) -> Self {
263 HashGrid {
264 width,
265 dim,
266 cells: HashMap::new(),
267 }
268 }
269
270 fn insert(&mut self, idx: u32, z: &[f64; 3]) {
271 let key = cell_of(z, self.dim, self.width);
272 self.cells.entry(key).or_default().push(idx);
273 }
274
275 fn for_neighbors(&self, z: &[f64; 3], mut visit: impl FnMut(u32)) {
278 let (c0, c1, c2) = cell_of(z, self.dim, self.width);
279 let d2 = if self.dim > 2 { 1 } else { 0 };
280 let d1 = if self.dim > 1 { 1 } else { 0 };
281 for i0 in -1..=1_i32 {
282 for i1 in -d1..=d1 {
283 for i2 in -d2..=d2 {
284 if let Some(bucket) = self.cells.get(&(c0 + i0, c1 + i1, c2 + i2)) {
285 for &idx in bucket {
286 visit(idx);
287 }
288 }
289 }
290 }
291 }
292 }
293}
294
295#[inline]
296fn dist2(a: &[f64; 3], b: &[f64; 3], dim: usize) -> f64 {
297 let mut s = 0.0;
298 for k in 0..dim {
299 let d = a[k] - b[k];
300 s += d * d;
301 }
302 s
303}
304
305#[inline]
308fn wendland(r: f64) -> f64 {
309 if r >= 1.0 {
310 return 0.0;
311 }
312 let v = 1.0 - r;
313 let v2 = v * v;
314 v2 * v2 * (4.0 * r + 1.0)
315}
316
317struct Level {
323 h: f64,
324 delta: f64,
325 weight: f64,
327 centers: Vec<[f64; 3]>,
328 col_offset: usize,
330 grid: HashGrid,
331}
332
333struct Core {
335 dim: usize,
336 metric: [f64; 3],
337 z_lo: [f64; 3],
340 z_range: [f64; 3],
341 sobolev_s: f64,
342 levels: Vec<Level>,
343 net: Vec<[f64; 3]>,
346 m: usize,
348 row_ptr: Vec<usize>,
350 col_idx: Vec<u32>,
351 vals: Vec<f64>,
352 w: Vec<f64>,
354 y: Vec<f64>,
355 z: Vec<[f64; 3]>,
357 rhs: Vec<f64>,
359 ytwy: f64,
360 gram_diag: Vec<f64>,
361 pen_diag: Vec<f64>,
363 pen_logdet_const: f64,
366 dense_gram: Option<Vec<f64>>,
369 predict_chol: Option<Vec<f64>>,
376}
377
378#[derive(Clone, Copy, Debug, PartialEq, Eq)]
380pub enum LogdetMethod {
381 DenseExact,
383 Slq,
386}
387
388#[derive(Clone, Copy, Debug)]
390pub struct CascadeCertificate {
391 pub solve_rel_residual: f64,
394 pub solve_iters: usize,
397 pub logdet_method: LogdetMethod,
399}
400
401#[derive(Clone, Copy, Debug)]
404pub struct RefinementCertificate {
405 pub next_level_gain_bound: f64,
407 pub tolerance: f64,
409}
410
411#[derive(Clone, Copy, Debug, PartialEq, Eq)]
416pub enum RefinementObstruction {
417 LevelCapacity {
419 levels: usize,
420 maximum_levels: usize,
421 },
422 CenterCapacity {
424 centers: usize,
425 maximum_centers: usize,
426 },
427}
428
429impl std::fmt::Display for RefinementObstruction {
430 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
431 match *self {
432 Self::LevelCapacity {
433 levels,
434 maximum_levels,
435 } => write!(
436 f,
437 "level capacity reached ({levels} of {maximum_levels} levels)"
438 ),
439 Self::CenterCapacity {
440 centers,
441 maximum_centers,
442 } => write!(
443 f,
444 "center capacity exceeded ({centers} centers for capacity {maximum_centers})"
445 ),
446 }
447 }
448}
449
450#[derive(Clone, Copy, Debug, PartialEq)]
454pub enum NextLevelAssessment {
455 EmptyNet,
457 GainBound(f64),
459 CapacityExceeded {
463 obstruction: RefinementObstruction,
464 gain_bound: f64,
465 },
466}
467
468pub struct ResidualCascadeDesign {
472 core: Arc<Core>,
473}
474
475pub struct ResidualCascadeFit {
477 core: Arc<Core>,
478 predict_chol: Option<Vec<f64>>,
482 pub coeff: Vec<f64>,
484 log_lambda: f64,
486 pub sigma2: f64,
488 pub restricted_loglik: f64,
492 pub rss_pen: f64,
494 pub certificate: CascadeCertificate,
496 pub refinement: Option<RefinementCertificate>,
498}
499
500pub struct ResidualCascadeCheckpoint {
507 iterate: ResidualCascadeFit,
508}
509
510impl ResidualCascadeCheckpoint {
511 fn new(iterate: ResidualCascadeFit) -> Self {
512 Self { iterate }
513 }
514
515 pub fn num_levels(&self) -> usize {
517 self.iterate.num_levels()
518 }
519
520 pub fn num_centers(&self) -> usize {
522 self.iterate.num_centers()
523 }
524
525 pub fn log_lambda(&self) -> f64 {
527 self.iterate.log_lambda
528 }
529
530 pub fn rss_pen(&self) -> f64 {
532 self.iterate.rss_pen
533 }
534
535 pub fn certificate(&self) -> CascadeCertificate {
537 self.iterate.certificate
538 }
539}
540
541impl std::fmt::Debug for ResidualCascadeCheckpoint {
542 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
543 f.debug_struct("ResidualCascadeCheckpoint")
544 .field("num_levels", &self.num_levels())
545 .field("num_centers", &self.num_centers())
546 .field("log_lambda", &self.log_lambda())
547 .field("rss_pen", &self.rss_pen())
548 .field("certificate", &self.certificate())
549 .finish_non_exhaustive()
550 }
551}
552
553#[derive(Debug)]
555pub enum ResidualCascadeError {
556 Computation(String),
558 Underresolved {
562 checkpoint: ResidualCascadeCheckpoint,
563 gain_bound: f64,
564 requested_tolerance: f64,
565 obstruction: RefinementObstruction,
566 },
567}
568
569impl From<String> for ResidualCascadeError {
570 fn from(reason: String) -> Self {
571 Self::Computation(reason)
572 }
573}
574
575impl std::fmt::Display for ResidualCascadeError {
576 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
577 match self {
578 Self::Computation(reason) => f.write_str(reason),
579 Self::Underresolved {
580 checkpoint,
581 gain_bound,
582 requested_tolerance,
583 obstruction,
584 } => write!(
585 f,
586 "residual cascade underresolved after {} levels: next-level gain bound \
587 {gain_bound:.6e} exceeds requested tolerance {requested_tolerance:.6e}; \
588 {obstruction}",
589 checkpoint.num_levels()
590 ),
591 }
592 }
593}
594
595impl std::error::Error for ResidualCascadeError {}
596
597#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
601pub struct LevelState {
602 pub h: f64,
603 pub delta: f64,
604 pub weight: f64,
605 pub col_offset: u64,
606 pub centers: Vec<f64>,
608}
609
610#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
627pub struct ResidualCascadeState {
628 pub dim: u64,
629 pub metric: [f64; 3],
631 pub z_lo: [f64; 3],
632 pub z_range: [f64; 3],
633 pub sobolev_s: f64,
634 pub levels: Vec<LevelState>,
635 pub m: u64,
637 pub pen_logdet_const: f64,
640 pub coeff: Vec<f64>,
642 pub log_lambda: f64,
643 pub sigma2: f64,
644 pub restricted_loglik: f64,
645 pub rss_pen: f64,
646 pub predict_chol: Vec<f64>,
649}
650
651fn forward_sub_into(l: &[f64], p: usize, b: &[f64], out: &mut [f64]) {
653 for i in 0..p {
654 let mut s = b[i];
655 for t in 0..i {
656 s -= l[i * p + t] * out[t];
657 }
658 out[i] = s / l[i * p + i];
659 }
660}
661
662fn back_sub_into(l: &[f64], p: usize, y: &[f64], out: &mut [f64]) {
664 for i in (0..p).rev() {
665 let mut s = y[i];
666 for t in i + 1..p {
667 s -= l[t * p + i] * out[t];
668 }
669 out[i] = s / l[i * p + i];
670 }
671}
672
673struct Preconditioner {
686 ncoarse: usize,
688 coarse_chol: Vec<f64>,
690 coarse_logdet: f64,
692 inv_fine: Vec<f64>,
694 inv_sqrt_fine: Vec<f64>,
696 fine_logdet: f64,
698}
699
700impl Preconditioner {
701 fn solve(&self, r: &[f64], out: &mut [f64]) {
703 let nc = self.ncoarse;
704 let zc = chol_solve(&self.coarse_chol, nc, &r[..nc]);
705 out[..nc].copy_from_slice(&zc);
706 for (k, o) in out[nc..].iter_mut().enumerate() {
707 *o = r[nc + k] * self.inv_fine[k];
708 }
709 }
710
711 fn apply_r_inv_t(&self, v: &[f64], out: &mut [f64]) {
713 let nc = self.ncoarse;
714 back_sub_into(&self.coarse_chol, nc, &v[..nc], &mut out[..nc]);
715 for (k, o) in out[nc..].iter_mut().enumerate() {
716 *o = v[nc + k] * self.inv_sqrt_fine[k];
717 }
718 }
719
720 fn apply_r_inv(&self, v: &[f64], out: &mut [f64]) {
722 let nc = self.ncoarse;
723 forward_sub_into(&self.coarse_chol, nc, &v[..nc], &mut out[..nc]);
724 for (k, o) in out[nc..].iter_mut().enumerate() {
725 *o = v[nc + k] * self.inv_sqrt_fine[k];
726 }
727 }
728
729 fn logdet(&self) -> f64 {
731 self.coarse_logdet + self.fine_logdet
732 }
733}
734
735#[derive(Clone, Copy)]
741struct CascadeSpectralMode {
742 eigenvalue: f64,
743 weight: f64,
744}
745
746struct CascadeRemlProfile<'a> {
758 core: &'a Core,
759 null_logdet: f64,
760 modes: Vec<CascadeSpectralMode>,
761}
762
763struct CascadeScoreEvaluation {
764 jet: ScoreJet,
765 normalized_logdet: f64,
767}
768
769impl CascadeRemlProfile<'_> {
770 fn log_lambda_domain(&self) -> Result<(f64, f64), String> {
776 let mut smallest = f64::INFINITY;
777 let mut largest = 0.0_f64;
778 for mode in &self.modes {
779 if mode.weight > 0.0 && mode.eigenvalue > 0.0 {
780 smallest = smallest.min(mode.eigenvalue);
781 largest = largest.max(mode.eigenvalue);
782 }
783 }
784 if !(smallest.is_finite() && smallest > 0.0 && largest.is_finite() && largest > 0.0) {
785 return Err(
786 "residual cascade: the data identify no positive penalized Schur mode; log lambda is not estimable"
787 .into(),
788 );
789 }
790 let log_relative_resolution = f64::EPSILON.sqrt().ln();
791 let lo = (smallest.ln() + log_relative_resolution).max(f64::MIN_POSITIVE.ln());
792 let hi = (largest.ln() - log_relative_resolution).min(f64::MAX.ln());
793 if !(lo.is_finite() && hi.is_finite() && lo < hi) {
794 return Err(format!(
795 "residual cascade: invalid spectrum-derived log-lambda domain [{lo}, {hi}]"
796 ));
797 }
798 Ok((lo, hi))
799 }
800
801 fn evaluate(&self, log_lambda: f64) -> Result<CascadeScoreEvaluation, String> {
802 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
803 .map_err(|error| format!("residual cascade: {error}"))?;
804
805 let core = self.core;
806 let (coeff, _, _) = core.solve_coeff(lambda, &core.rhs, None)?;
807 let rss = core.rss_pen(&coeff);
808 if !(rss.is_finite() && rss > 0.0) {
809 return Err(format!(
810 "residual cascade: degenerate penalized residual {rss}"
811 ));
812 }
813
814 let dc: Vec<f64> = coeff
821 .iter()
822 .zip(core.pen_diag.iter())
823 .map(|(&c, &d)| d * c)
824 .collect();
825 let penalty_energy = coeff
826 .iter()
827 .zip(dc.iter())
828 .map(|(&c, &v)| c * v)
829 .sum::<f64>();
830 let (u, _, _) = core.solve_coeff(lambda, &dc, None)?;
831 let inverse_penalty_energy = dc.iter().zip(u.iter()).map(|(&a, &b)| a * b).sum::<f64>();
832 let third_energy = u
833 .iter()
834 .zip(core.pen_diag.iter())
835 .map(|(&v, &d)| d * v * v)
836 .sum::<f64>();
837 let rss_d1 = lambda * penalty_energy;
838 let lambda2 = lambda * lambda;
839 let rss_d2 = rss_d1 - 2.0 * lambda2 * inverse_penalty_energy;
840 let rss_d3 =
841 rss_d1 - 6.0 * lambda2 * inverse_penalty_energy + 6.0 * lambda2 * lambda * third_energy;
842
843 let mut normalized_logdet = self.null_logdet;
844 let mut determinant_d1 = 0.0;
845 let mut determinant_d2 = 0.0;
846 for mode in &self.modes {
847 let theta = mode.eigenvalue;
848 let weight = mode.weight;
849 if theta == 0.0 || weight == 0.0 {
850 continue;
851 }
852 let log_theta = theta.ln();
855 normalized_logdet += weight
856 * if log_theta > log_lambda {
857 (log_theta - log_lambda) + (log_lambda - log_theta).exp().ln_1p()
858 } else {
859 (log_theta - log_lambda).exp().ln_1p()
860 };
861 let t = if theta > lambda {
862 1.0 / (1.0 + lambda / theta)
863 } else {
864 theta / (lambda + theta)
865 };
866 determinant_d1 -= weight * t;
867 determinant_d2 += weight * t * (1.0 - t);
868 }
869
870 let dof = (core.y.len() - core.nullity()) as f64;
871 let rss_log_d1 = rss_d1 / rss;
872 let rss_log_d2 = rss_d2 / rss - rss_log_d1 * rss_log_d1;
873 let rss_log_d3 = rss_d3 / rss - 3.0 * rss_d1 * rss_d2 / (rss * rss)
874 + 2.0 * rss_log_d1 * rss_log_d1 * rss_log_d1;
875 if !(rss_log_d3.is_finite()) {
876 return Err(format!(
877 "residual cascade: non-finite analytic residual derivative at log lambda {log_lambda}"
878 ));
879 }
880 let jet = ScoreJet {
881 value: -0.5 * (normalized_logdet + dof * (rss / dof).ln()),
882 derivative: -0.5 * (determinant_d1 + dof * rss_log_d1),
883 curvature: -0.5 * (determinant_d2 + dof * rss_log_d2),
884 };
885 if !(jet.value.is_finite() && jet.derivative.is_finite() && jet.curvature.is_finite()) {
886 return Err(format!(
887 "residual cascade: non-finite REML jet at log lambda {log_lambda}: value {}, derivative {}, curvature {}",
888 jet.value, jet.derivative, jet.curvature
889 ));
890 }
891 Ok(CascadeScoreEvaluation {
892 jet,
893 normalized_logdet,
894 })
895 }
896
897 fn enclose(&self, lo: f64, hi: f64) -> Result<DerivativeEnclosure, String> {
906 if !(lo.is_finite() && hi.is_finite() && lo <= hi) {
907 return Err(format!(
908 "residual cascade: invalid score-enclosure interval [{lo}, {hi}]"
909 ));
910 }
911 let left = self.evaluate(lo)?.jet;
912 let right = self.evaluate(hi)?.jet;
913 let width = hi - lo;
914 let rank = (self.core.m - self.core.nullity()) as f64;
915 let dof = (self.core.y.len() - self.core.nullity()) as f64;
916 let curvature_abs_bound = 0.5 * (0.25 * rank + 2.0 * dof);
917 let third_abs_bound = 0.5 * (0.25 * rank + 6.0 * dof);
918 let derivative_radius = curvature_abs_bound * width;
919 let curvature_radius = third_abs_bound * width;
920 Ok(DerivativeEnclosure {
921 derivative: ClosedInterval::outward(
922 (left.derivative - derivative_radius).min(right.derivative - derivative_radius),
923 (left.derivative + derivative_radius).max(right.derivative + derivative_radius),
924 ),
925 curvature: ClosedInterval::outward(
926 (left.curvature - curvature_radius).min(right.curvature - curvature_radius),
927 (left.curvature + curvature_radius).max(right.curvature + curvature_radius),
928 ),
929 })
930 }
931}
932
933impl Core {
934 #[inline]
935 fn dense_gram_entry(&self, row: usize, col: usize) -> Option<f64> {
936 let gram = self.dense_gram.as_ref()?;
937 let (i, j) = if row <= col { (row, col) } else { (col, row) };
938 Some(gram[i * self.m + j])
939 }
940
941 fn null_gram_factor(&self) -> Result<(Vec<f64>, f64), String> {
944 let q = self.nullity();
945 let mut gram = vec![0.0; q * q];
946 if self.dense_gram.is_some() {
947 for i in 0..q {
948 for j in i..q {
949 let value = self.dense_gram_entry(i, j).expect("dense Gram exists");
950 gram[i * q + j] = value;
951 gram[j * q + i] = value;
952 }
953 }
954 } else {
955 for row in 0..self.w.len() {
956 let lo = self.row_ptr[row];
957 let hi = self.row_ptr[row + 1];
958 for ea in lo..hi {
959 let ca = self.col_idx[ea] as usize;
960 if ca >= q {
961 break;
962 }
963 let weighted = self.w[row] * self.vals[ea];
964 for eb in ea..hi {
965 let cb = self.col_idx[eb] as usize;
966 if cb >= q {
967 break;
968 }
969 gram[ca * q + cb] += weighted * self.vals[eb];
970 }
971 }
972 }
973 for i in 0..q {
974 for j in i + 1..q {
975 gram[j * q + i] = gram[i * q + j];
976 }
977 }
978 }
979 let logdet = cholesky_logdet(&mut gram, q).map_err(|error| {
980 format!("residual cascade: polynomial null-space factorization failed: {error}")
981 })?;
982 Ok((gram, logdet))
983 }
984
985 fn schur_whitened_matvec(
989 &self,
990 null_chol: &[f64],
991 input: &[f64],
992 output: &mut [f64],
993 full: &mut [f64],
994 gram_full: &mut [f64],
995 projected_null: &mut [f64],
996 ) {
997 let q = self.nullity();
998 full.fill(0.0);
999 for (i, &value) in input.iter().enumerate() {
1000 full[q + i] = value / self.pen_diag[q + i].sqrt();
1001 }
1002 self.matvec(0.0, full, gram_full);
1003 let null_coeff = chol_solve(null_chol, q, &gram_full[..q]);
1004 full.fill(0.0);
1005 full[..q].copy_from_slice(&null_coeff);
1006 self.matvec(0.0, full, projected_null);
1007 for i in 0..output.len() {
1008 output[i] = (gram_full[q + i] - projected_null[q + i]) / self.pen_diag[q + i].sqrt();
1009 }
1010 }
1011
1012 fn dense_cascade_spectrum(
1013 &self,
1014 null_chol: &[f64],
1015 ) -> Result<Vec<CascadeSpectralMode>, String> {
1016 let q = self.nullity();
1017 let rank = self.m - q;
1018 let mut schur = Array2::<f64>::zeros((rank, rank));
1019 let mut cross = vec![0.0; q];
1020 for j in 0..rank {
1021 for (k, value) in cross.iter_mut().enumerate() {
1022 *value = self.dense_gram_entry(k, q + j).expect("dense Gram exists");
1023 }
1024 let projected = chol_solve(null_chol, q, &cross);
1025 for i in 0..=j {
1026 let mut value = self
1027 .dense_gram_entry(q + i, q + j)
1028 .expect("dense Gram exists");
1029 for (k, &coefficient) in projected.iter().enumerate() {
1030 value -=
1031 self.dense_gram_entry(q + i, k).expect("dense Gram exists") * coefficient;
1032 }
1033 value /= (self.pen_diag[q + i] * self.pen_diag[q + j]).sqrt();
1034 schur[(i, j)] = value;
1035 schur[(j, i)] = value;
1036 }
1037 }
1038 let (eigenvalues, _) = schur.eigh(Side::Lower).map_err(|error| {
1039 format!("residual cascade: Schur-complement eigendecomposition failed: {error}")
1040 })?;
1041 let scale = eigenvalues
1042 .iter()
1043 .copied()
1044 .map(f64::abs)
1045 .fold(0.0, f64::max);
1046 let roundoff = f64::EPSILON * rank.max(1) as f64 * scale.max(f64::MIN_POSITIVE);
1047 eigenvalues
1048 .iter()
1049 .copied()
1050 .enumerate()
1051 .map(|(index, eigenvalue)| {
1052 if !eigenvalue.is_finite() || eigenvalue < -roundoff {
1053 Err(format!(
1054 "residual cascade: penalty-whitened Schur mode {index} is not positive semidefinite ({eigenvalue})"
1055 ))
1056 } else {
1057 Ok(CascadeSpectralMode {
1058 eigenvalue: eigenvalue.max(0.0),
1059 weight: 1.0,
1060 })
1061 }
1062 })
1063 .collect()
1064 }
1065
1066 fn iterative_cascade_spectrum(
1072 &self,
1073 null_chol: &[f64],
1074 ) -> Result<Vec<CascadeSpectralMode>, String> {
1075 let q0 = self.nullity();
1076 let rank = self.m - q0;
1077 let steps = SLQ_LANCZOS_STEPS.min(rank);
1078 let mut modes = Vec::with_capacity(SLQ_PROBES * steps);
1079 let mut full = vec![0.0; self.m];
1080 let mut gram_full = vec![0.0; self.m];
1081 let mut projected_null = vec![0.0; self.m];
1082 let mut matvec = vec![0.0; rank];
1083 let mut basis: Vec<Vec<f64>> = Vec::with_capacity(steps);
1084
1085 for probe in 0..SLQ_PROBES {
1086 let mut rng =
1087 SplitMix64::new(RNG_SEED ^ (probe as u64).wrapping_mul(0xD134_2543_DE82_EF95));
1088 let inv_norm = 1.0 / (rank as f64).sqrt();
1089 let mut q = (0..rank)
1090 .map(|_| rng.next_sign() * inv_norm)
1091 .collect::<Vec<_>>();
1092 let mut q_previous: Option<Vec<f64>> = None;
1093 let mut alpha = Vec::with_capacity(steps);
1094 let mut beta = Vec::with_capacity(steps.saturating_sub(1));
1095 basis.clear();
1096
1097 for _ in 0..steps {
1098 self.schur_whitened_matvec(
1099 null_chol,
1100 &q,
1101 &mut matvec,
1102 &mut full,
1103 &mut gram_full,
1104 &mut projected_null,
1105 );
1106 let diagonal = matvec
1107 .iter()
1108 .zip(q.iter())
1109 .map(|(&a, &b)| a * b)
1110 .sum::<f64>();
1111 alpha.push(diagonal);
1112 let mut residual = matvec.clone();
1113 for i in 0..rank {
1114 residual[i] -= diagonal * q[i];
1115 }
1116 if let Some(previous) = &q_previous {
1117 let previous_beta = beta.last().copied().unwrap_or(0.0);
1118 for i in 0..rank {
1119 residual[i] -= previous_beta * previous[i];
1120 }
1121 }
1122 basis.push(q.clone());
1123 for direction in &basis {
1124 let projection = residual
1125 .iter()
1126 .zip(direction.iter())
1127 .map(|(&a, &b)| a * b)
1128 .sum::<f64>();
1129 for i in 0..rank {
1130 residual[i] -= projection * direction[i];
1131 }
1132 }
1133 let norm = residual
1134 .iter()
1135 .map(|value| value * value)
1136 .sum::<f64>()
1137 .sqrt();
1138 if !norm.is_finite() {
1139 return Err(
1140 "residual cascade: Schur-spectrum Lanczos produced a non-finite norm"
1141 .into(),
1142 );
1143 }
1144 let rounding_floor =
1145 f64::EPSILON * rank.max(1) as f64 * diagonal.abs().max(f64::MIN_POSITIVE);
1146 if norm <= rounding_floor {
1147 break;
1148 }
1149 beta.push(norm);
1150 q_previous = Some(std::mem::replace(&mut q, residual));
1151 for value in &mut q {
1152 *value /= norm;
1153 }
1154 }
1155
1156 beta.truncate(alpha.len().saturating_sub(1));
1157 let (eigenvalues, first_components) = symmetric_tridiagonal_eigen(&alpha, &beta)?;
1158 let scale = eigenvalues
1159 .iter()
1160 .copied()
1161 .map(f64::abs)
1162 .fold(0.0, f64::max);
1163 let roundoff = f64::EPSILON * alpha.len().max(1) as f64 * scale.max(f64::MIN_POSITIVE);
1164 for (index, (&eigenvalue, &first)) in
1165 eigenvalues.iter().zip(first_components.iter()).enumerate()
1166 {
1167 if !eigenvalue.is_finite() || eigenvalue < -roundoff {
1168 return Err(format!(
1169 "residual cascade: Schur-spectrum Ritz value {index} is not positive semidefinite ({eigenvalue})"
1170 ));
1171 }
1172 let weight = rank as f64 * first * first / SLQ_PROBES as f64;
1173 if !(weight.is_finite() && weight >= 0.0) {
1174 return Err(format!(
1175 "residual cascade: invalid Schur-spectrum quadrature weight {weight}"
1176 ));
1177 }
1178 modes.push(CascadeSpectralMode {
1179 eigenvalue: eigenvalue.max(0.0),
1180 weight,
1181 });
1182 }
1183 }
1184 Ok(modes)
1185 }
1186
1187 fn reml_profile(&self) -> Result<CascadeRemlProfile<'_>, String> {
1188 let (null_chol, null_logdet) = self.null_gram_factor()?;
1189 let modes = if self.dense_gram.is_some() {
1190 self.dense_cascade_spectrum(&null_chol)?
1191 } else {
1192 self.iterative_cascade_spectrum(&null_chol)?
1193 };
1194 Ok(CascadeRemlProfile {
1195 core: self,
1196 null_logdet,
1197 modes,
1198 })
1199 }
1200
1201 fn scale_point(&self, x: &[f64]) -> [f64; 3] {
1203 let mut z = [0.0_f64; 3];
1204 for a in 0..self.dim {
1205 z[a] = self.metric[a] * x[a] - self.z_lo[a];
1206 }
1207 z
1208 }
1209
1210 fn basis_row_scaled(&self, z: &[f64; 3]) -> Vec<(usize, f64)> {
1213 let mut row = Vec::with_capacity(self.dim + 1 + self.levels.len() * 8);
1214 row.push((0, 1.0));
1215 for a in 0..self.dim {
1216 row.push((a + 1, 2.0 * z[a] / self.z_range[a] - 1.0));
1217 }
1218 for level in &self.levels {
1219 let start = row.len();
1220 level.grid.for_neighbors(z, |j| {
1221 let c = &level.centers[j as usize];
1222 let r = dist2(z, c, self.dim).sqrt() / level.delta;
1223 let v = wendland(r);
1224 if v > 0.0 {
1225 row.push((level.col_offset + j as usize, v));
1226 }
1227 });
1228 row[start..].sort_unstable_by_key(|&(col, _)| col);
1229 }
1230 row
1231 }
1232
1233 fn matvec(&self, lambda: f64, v: &[f64], out: &mut [f64]) {
1235 for (o, (&d, &x)) in out.iter_mut().zip(self.pen_diag.iter().zip(v.iter())) {
1236 *o = lambda * d * x;
1237 }
1238 for i in 0..self.w.len() {
1239 let lo = self.row_ptr[i];
1240 let hi = self.row_ptr[i + 1];
1241 let mut t = 0.0;
1242 for e in lo..hi {
1243 t += self.vals[e] * v[self.col_idx[e] as usize];
1244 }
1245 t *= self.w[i];
1246 for e in lo..hi {
1247 out[self.col_idx[e] as usize] += self.vals[e] * t;
1248 }
1249 }
1250 }
1251
1252 fn coarse_space_cols(&self, lambda: f64) -> usize {
1263 let mut ncoarse = self.nullity();
1264 let mut buf: Vec<f64> = Vec::new();
1265 for (li, level) in self.levels.iter().enumerate() {
1266 let a = level.col_offset;
1267 let b = a + level.centers.len();
1268 if b <= a {
1269 continue;
1270 }
1271 if b > COARSE_SPACE_MAX {
1272 break;
1273 }
1274 let dominated = if li < MIN_COARSE_LEVELS {
1275 true
1276 } else {
1277 buf.clear();
1278 buf.extend_from_slice(&self.gram_diag[a..b]);
1279 buf.sort_unstable_by(|x, y| x.partial_cmp(y).unwrap());
1280 let gram_median = buf[buf.len() / 2];
1281 lambda * level.weight < COARSE_DOMINANCE * gram_median
1282 };
1283 if dominated {
1284 ncoarse = b;
1285 } else {
1286 break;
1287 }
1288 }
1289 let ncoarse = ncoarse.min(self.m);
1293 if log::log_enabled!(log::Level::Debug) {
1298 let mut s = String::new();
1299 for (li, level) in self.levels.iter().enumerate() {
1300 let a = level.col_offset;
1301 let b = a + level.centers.len();
1302 let mut buf: Vec<f64> = self.gram_diag[a..b].to_vec();
1303 buf.sort_unstable_by(|x, y| x.partial_cmp(y).unwrap());
1304 let med = if buf.is_empty() {
1305 0.0
1306 } else {
1307 buf[buf.len() / 2]
1308 };
1309 let coarse = b <= ncoarse;
1310 s.push_str(&format!(
1311 " L{li}[{}c off{a} w={:.2e} λw={:.2e} med={:.2e} {}]",
1312 level.centers.len(),
1313 level.weight,
1314 lambda * level.weight,
1315 med,
1316 if coarse { "C" } else { "F" }
1317 ));
1318 }
1319 log::debug!(
1320 "[1032-COARSE] λ={lambda:.3e} m={} ncoarse={ncoarse} cap={COARSE_SPACE_MAX}{s}",
1321 self.m
1322 );
1323 }
1324 ncoarse
1325 }
1326
1327 fn build_preconditioner(&self, lambda: f64) -> Result<Preconditioner, String> {
1333 let m = self.m;
1334 let nc = self.coarse_space_cols(lambda);
1335 let mut acc = vec![0.0_f64; nc * nc];
1336 for i in 0..self.w.len() {
1337 let lo = self.row_ptr[i];
1338 let hi = self.row_ptr[i + 1];
1339 let mut end = lo;
1341 while end < hi && (self.col_idx[end] as usize) < nc {
1342 end += 1;
1343 }
1344 for ea in lo..end {
1345 let ca = self.col_idx[ea] as usize;
1346 let va = self.w[i] * self.vals[ea];
1347 for eb in ea..end {
1348 let cb = self.col_idx[eb] as usize;
1349 acc[ca * nc + cb] += va * self.vals[eb];
1350 }
1351 }
1352 }
1353 for i in 0..nc {
1354 for j in i + 1..nc {
1355 acc[j * nc + i] = acc[i * nc + j];
1356 }
1357 }
1358 for i in 0..nc {
1359 acc[i * nc + i] += lambda * self.pen_diag[i];
1360 }
1361 let coarse_logdet = cholesky_logdet(&mut acc, nc)?;
1362 let mut inv_fine = Vec::with_capacity(m - nc);
1363 let mut inv_sqrt_fine = Vec::with_capacity(m - nc);
1364 let mut fine_logdet = 0.0;
1365 for j in nc..m {
1366 let p = self.gram_diag[j] + lambda * self.pen_diag[j];
1367 if !(p.is_finite() && p > EIG_FLOOR) {
1368 return Err(format!(
1369 "residual cascade: non-positive preconditioner diagonal {p} at column {j}"
1370 ));
1371 }
1372 inv_fine.push(1.0 / p);
1373 inv_sqrt_fine.push(1.0 / p.sqrt());
1374 fine_logdet += p.ln();
1375 }
1376 Ok(Preconditioner {
1377 ncoarse: nc,
1378 coarse_chol: acc,
1379 coarse_logdet,
1380 inv_fine,
1381 inv_sqrt_fine,
1382 fine_logdet,
1383 })
1384 }
1385
1386 fn pcg(
1389 &self,
1390 lambda: f64,
1391 b: &[f64],
1392 warm: Option<&[f64]>,
1393 ) -> Result<(Vec<f64>, f64, usize), String> {
1394 let m = self.m;
1395 let prec = self.build_preconditioner(lambda)?;
1396 let b_norm = b.iter().map(|v| v * v).sum::<f64>().sqrt();
1397 if b_norm == 0.0 {
1398 return Ok((vec![0.0; m], 0.0, 0));
1399 }
1400 let mut zv = vec![0.0; m];
1401 let mut x = match warm {
1402 Some(x0) => {
1403 if x0.len() != m {
1404 return Err(format!(
1405 "residual cascade: warm-start length {} != system size {m}",
1406 x0.len()
1407 ));
1408 }
1409 x0.to_vec()
1410 }
1411 None => {
1412 prec.solve(b, &mut zv);
1413 zv.clone()
1414 }
1415 };
1416 let mut r = vec![0.0; m];
1417 self.matvec(lambda, &x, &mut r);
1418 for (ri, &bi) in r.iter_mut().zip(b.iter()) {
1419 *ri = bi - *ri;
1420 }
1421 prec.solve(&r, &mut zv);
1422 let mut p_dir = zv.clone();
1423 let mut rz: f64 = r.iter().zip(zv.iter()).map(|(&a, &c)| a * c).sum();
1424 let mut ap = vec![0.0; m];
1425 let max_iters = CG_MAX_ITERS;
1426 for iter in 0..max_iters {
1427 let r_norm = r.iter().map(|v| v * v).sum::<f64>().sqrt();
1428 if r_norm <= CG_RTOL * b_norm {
1429 return Ok((x, r_norm / b_norm, iter));
1430 }
1431 self.matvec(lambda, &p_dir, &mut ap);
1432 let pap: f64 = p_dir.iter().zip(ap.iter()).map(|(&a, &c)| a * c).sum();
1433 if !(pap.is_finite() && pap > 0.0) {
1434 return Err(format!(
1435 "residual cascade: CG curvature breakdown (p'Ap = {pap}) at iteration {iter}"
1436 ));
1437 }
1438 let alpha = rz / pap;
1439 for j in 0..m {
1440 x[j] += alpha * p_dir[j];
1441 r[j] -= alpha * ap[j];
1442 }
1443 prec.solve(&r, &mut zv);
1444 let rz_new: f64 = r.iter().zip(zv.iter()).map(|(&a, &c)| a * c).sum();
1445 let beta = rz_new / rz;
1446 rz = rz_new;
1447 for j in 0..m {
1448 p_dir[j] = zv[j] + beta * p_dir[j];
1449 }
1450 }
1451 Err(format!(
1452 "residual cascade: CG failed to reach relative residual {CG_RTOL} within \
1453 {CG_MAX_ITERS} iterations (the coarse-space additive-Schwarz preconditioner should \
1454 make this n-independent; this indicates a degenerate design)"
1455 ))
1456 }
1457
1458 fn dense_system(&self, lambda: f64) -> Option<Vec<f64>> {
1460 let gram = self.dense_gram.as_ref()?;
1461 let m = self.m;
1462 let mut a = vec![0.0; m * m];
1463 for i in 0..m {
1464 for j in i..m {
1465 let mut v = gram[i * m + j];
1466 if i == j {
1467 v += lambda * self.pen_diag[i];
1468 }
1469 a[i * m + j] = v;
1470 a[j * m + i] = v;
1471 }
1472 }
1473 Some(a)
1474 }
1475
1476 fn logdet_dense(&self, lambda: f64) -> Result<f64, String> {
1479 let mut a = self.dense_system(lambda).ok_or_else(|| {
1480 format!(
1481 "residual cascade: dense logdet requested past the sizing cap \
1482 (m = {} > {DENSE_GRAM_MAX})",
1483 self.m
1484 )
1485 })?;
1486 cholesky_logdet(&mut a, self.m)
1487 }
1488
1489 fn logdet_slq(&self, lambda: f64) -> Result<f64, String> {
1499 let m = self.m;
1500 let prec = self.build_preconditioner(lambda)?;
1501 let logdet = prec.logdet();
1502 let mut scratch_in = vec![0.0; m];
1504 let mut scratch_out = vec![0.0; m];
1505 let mut vbuf = vec![0.0; m];
1506 let mut trace_est = 0.0;
1507 let steps = SLQ_LANCZOS_STEPS.min(m);
1508 let mut basis: Vec<Vec<f64>> = Vec::with_capacity(steps);
1509 for probe in 0..SLQ_PROBES {
1510 let mut rng =
1511 SplitMix64::new(RNG_SEED ^ (probe as u64).wrapping_mul(0xD134_2543_DE82_EF95));
1512 let mut q = vec![0.0; m];
1513 for qj in q.iter_mut() {
1514 *qj = rng.next_sign();
1515 }
1516 let z_norm2 = m as f64;
1517 let inv_norm = 1.0 / (m as f64).sqrt();
1518 for qj in q.iter_mut() {
1519 *qj *= inv_norm;
1520 }
1521 basis.clear();
1523 let mut alpha = Vec::with_capacity(steps);
1524 let mut beta: Vec<f64> = Vec::with_capacity(steps);
1525 let mut q_prev: Option<Vec<f64>> = None;
1526 for _step in 0..steps {
1527 prec.apply_r_inv_t(&q, &mut scratch_in);
1529 self.matvec(lambda, &scratch_in, &mut scratch_out);
1530 prec.apply_r_inv(&scratch_out, &mut vbuf);
1531 let mut v: Vec<f64> = vbuf.clone();
1532 let a: f64 = v.iter().zip(q.iter()).map(|(&x, &y)| x * y).sum();
1533 alpha.push(a);
1534 for j in 0..m {
1535 v[j] -= a * q[j];
1536 }
1537 if let Some(prev) = &q_prev {
1538 let b_prev = beta.last().copied().unwrap_or(0.0);
1539 for j in 0..m {
1540 v[j] -= b_prev * prev[j];
1541 }
1542 }
1543 basis.push(q.clone());
1545 for qb in &basis {
1546 let proj: f64 = v.iter().zip(qb.iter()).map(|(&x, &y)| x * y).sum();
1547 for j in 0..m {
1548 v[j] -= proj * qb[j];
1549 }
1550 }
1551 let b: f64 = v.iter().map(|x| x * x).sum::<f64>().sqrt();
1552 if !(b.is_finite()) {
1553 return Err("residual cascade: Lanczos breakdown (non-finite norm)".into());
1554 }
1555 if b < 1e-13 {
1556 break;
1557 }
1558 beta.push(b);
1559 q_prev = Some(std::mem::replace(&mut q, v));
1560 for qj in q.iter_mut() {
1561 *qj /= b;
1562 }
1563 }
1564 beta.truncate(alpha.len().saturating_sub(1));
1565 let (theta, tau) = symmetric_tridiagonal_eigen(&alpha, &beta)?;
1566 let mut quad = 0.0;
1567 for (&t, &w0) in theta.iter().zip(tau.iter()) {
1568 if !(t.is_finite() && t > EIG_FLOOR) {
1569 return Err(format!(
1570 "residual cascade: non-positive Ritz value {t} in SLQ (system not PD)"
1571 ));
1572 }
1573 quad += w0 * w0 * t.ln();
1574 }
1575 trace_est += z_norm2 * quad;
1576 }
1577 Ok(logdet + trace_est / SLQ_PROBES as f64)
1578 }
1579
1580 fn logdet(&self, lambda: f64) -> Result<(f64, LogdetMethod), String> {
1582 if self.dense_gram.is_some() {
1583 Ok((self.logdet_dense(lambda)?, LogdetMethod::DenseExact))
1584 } else {
1585 Ok((self.logdet_slq(lambda)?, LogdetMethod::Slq))
1586 }
1587 }
1588
1589 fn solve_coeff(
1591 &self,
1592 lambda: f64,
1593 b: &[f64],
1594 warm: Option<&[f64]>,
1595 ) -> Result<(Vec<f64>, f64, usize), String> {
1596 if let Some(l) = &self.predict_chol {
1600 return Ok((chol_solve(l, self.m, b), 0.0, 0));
1601 }
1602 if let Some(mut a) = self.dense_system(lambda) {
1603 cholesky_logdet(&mut a, self.m)?;
1604 return Ok((chol_solve(&a, self.m, b), 0.0, 0));
1605 }
1606 self.pcg(lambda, b, warm)
1607 }
1608
1609 fn assemble_predict_factor(&self, lambda: f64) -> Result<Vec<f64>, String> {
1617 let m = self.m;
1618 let mut a = vec![0.0_f64; m * m];
1619 if let Some(gram) = &self.dense_gram {
1620 for i in 0..m {
1621 for j in i..m {
1622 let v = gram[i * m + j];
1623 a[i * m + j] = v;
1624 a[j * m + i] = v;
1625 }
1626 }
1627 } else {
1628 for i in 0..self.w.len() {
1629 let lo = self.row_ptr[i];
1630 let hi = self.row_ptr[i + 1];
1631 for ea in lo..hi {
1632 let ca = self.col_idx[ea] as usize;
1633 let va = self.w[i] * self.vals[ea];
1634 for eb in ea..hi {
1635 let cb = self.col_idx[eb] as usize;
1636 a[ca * m + cb] += va * self.vals[eb];
1637 }
1638 }
1639 }
1640 for i in 0..m {
1642 for j in i + 1..m {
1643 a[j * m + i] = a[i * m + j];
1644 }
1645 }
1646 }
1647 for (i, d) in self.pen_diag.iter().enumerate() {
1648 a[i * m + i] += lambda * d;
1649 }
1650 cholesky_logdet(&mut a, m)?;
1651 Ok(a)
1652 }
1653
1654 fn rss_pen(&self, coeff: &[f64]) -> f64 {
1656 let mut quad = 0.0;
1657 for (c, r) in coeff.iter().zip(self.rhs.iter()) {
1658 quad += c * r;
1659 }
1660 self.ytwy - quad
1661 }
1662
1663 fn nullity(&self) -> usize {
1665 self.dim + 1
1666 }
1667
1668 fn residuals(&self, coeff: &[f64]) -> Vec<f64> {
1670 let n = self.y.len();
1671 let mut r = Vec::with_capacity(n);
1672 for i in 0..n {
1673 let mut fit = 0.0;
1674 for e in self.row_ptr[i]..self.row_ptr[i + 1] {
1675 fit += self.vals[e] * coeff[self.col_idx[e] as usize];
1676 }
1677 r.push(self.y[i] - fit);
1678 }
1679 r
1680 }
1681}
1682
1683fn symmetric_tridiagonal_eigen(d: &[f64], e: &[f64]) -> Result<(Vec<f64>, Vec<f64>), String> {
1690 let n = d.len();
1691 if n == 0 {
1692 return Ok((Vec::new(), Vec::new()));
1693 }
1694 let mut diag = d.to_vec();
1695 let mut off = vec![0.0; n];
1696 off[..n - 1].copy_from_slice(&e[..n - 1]);
1697 let mut first = vec![0.0; n];
1698 first[0] = 1.0;
1699 for l in 0..n {
1700 let mut iter = 0;
1701 loop {
1702 let mut msplit = n - 1;
1704 for mm in l..n - 1 {
1705 let dd = diag[mm].abs() + diag[mm + 1].abs();
1706 if off[mm].abs() <= f64::EPSILON * dd {
1707 msplit = mm;
1708 break;
1709 }
1710 }
1711 if msplit == l {
1712 break;
1713 }
1714 iter += 1;
1715 if iter > 60 {
1716 return Err("residual cascade: tridiagonal QL failed to converge".into());
1717 }
1718 let mut g = (diag[l + 1] - diag[l]) / (2.0 * off[l]);
1719 let mut r = g.hypot(1.0);
1720 g = diag[msplit] - diag[l] + off[l] / (g + r.copysign(g));
1721 let (mut s, mut c) = (1.0, 1.0);
1722 let mut p = 0.0;
1723 let mut broke_early = false;
1724 for i in (l..msplit).rev() {
1725 let mut f = s * off[i];
1726 let b = c * off[i];
1727 r = f.hypot(g);
1728 off[i + 1] = r;
1729 if r == 0.0 {
1730 diag[i + 1] -= p;
1731 off[msplit] = 0.0;
1732 broke_early = true;
1733 break;
1734 }
1735 s = f / r;
1736 c = g / r;
1737 g = diag[i + 1] - p;
1738 r = (diag[i] - g) * s + 2.0 * c * b;
1739 p = s * r;
1740 diag[i + 1] = g + p;
1741 g = c * r - b;
1742 f = first[i + 1];
1744 first[i + 1] = s * first[i] + c * f;
1745 first[i] = c * first[i] - s * f;
1746 }
1747 if broke_early {
1748 continue;
1749 }
1750 diag[l] -= p;
1751 off[l] = g;
1752 off[msplit] = 0.0;
1753 }
1754 }
1755 Ok((diag, first))
1756}
1757
1758fn extend_net(
1779 net: &mut Vec<[f64; 3]>,
1780 points: &[[f64; 3]],
1781 dim: usize,
1782 h: f64,
1783 box_hi: &[f64; 3],
1784) -> Vec<[f64; 3]> {
1785 let mut grid = HashGrid::new(h, dim);
1786 for (idx, c) in net.iter().enumerate() {
1787 grid.insert(idx as u32, c);
1788 }
1789 let h2 = h * h;
1790 let mut new_centers = Vec::new();
1791 let try_add = |net: &mut Vec<[f64; 3]>,
1792 grid: &mut HashGrid,
1793 new_centers: &mut Vec<[f64; 3]>,
1794 p: &[f64; 3]| {
1795 let mut covered = false;
1796 grid.for_neighbors(p, |j| {
1797 if !covered && dist2(p, &net[j as usize], dim) <= h2 {
1798 covered = true;
1799 }
1800 });
1801 if !covered {
1802 let idx = net.len() as u32;
1803 net.push(*p);
1804 grid.insert(idx, p);
1805 new_centers.push(*p);
1806 }
1807 };
1808 for p in points {
1809 try_add(net, &mut grid, &mut new_centers, p);
1810 if net.len() > MAX_CENTERS {
1811 return new_centers;
1812 }
1813 }
1814 let mut cells = [1_i64; 3];
1828 for a in 0..dim {
1829 cells[a] = (box_hi[a] / h).ceil() as i64 + 1;
1830 }
1831 let mut c = [0.0_f64; 3];
1832 'fill: for i0 in 0..cells[0] {
1833 c[0] = (i0 as f64 + 0.5) * h;
1834 for i1 in 0..cells[1] {
1835 if dim > 1 {
1836 c[1] = (i1 as f64 + 0.5) * h;
1837 }
1838 for i2 in 0..cells[2] {
1839 if dim > 2 {
1840 c[2] = (i2 as f64 + 0.5) * h;
1841 }
1842 try_add(net, &mut grid, &mut new_centers, &c);
1843 if net.len() > MAX_CENTERS {
1844 break 'fill;
1845 }
1846 }
1847 }
1848 }
1849 new_centers
1850}
1851
1852impl ResidualCascadeDesign {
1853 pub fn build(
1862 xs: &[&[f64]],
1863 y: &[f64],
1864 w: &[f64],
1865 metric: &[f64],
1866 sobolev_s: f64,
1867 levels: usize,
1868 ) -> Result<Self, String> {
1869 let dim = xs.len();
1870 if !(dim == 2 || dim == 3) {
1871 return Err(format!(
1872 "residual cascade: built for scattered 2-3D smooths, got {dim} axes"
1873 ));
1874 }
1875 let n = y.len();
1876 if w.len() != n || xs.iter().any(|x| x.len() != n) {
1877 return Err(format!(
1878 "residual cascade: length mismatch (y={n}, w={}, axes={:?})",
1879 w.len(),
1880 xs.iter().map(|x| x.len()).collect::<Vec<_>>()
1881 ));
1882 }
1883 if n <= dim + 1 {
1884 return Err(format!(
1885 "residual cascade: needs more than {} rows for the profiled REML degrees of \
1886 freedom, got {n}",
1887 dim + 1
1888 ));
1889 }
1890 if metric.len() != dim || metric.iter().any(|&s| !(s.is_finite() && s > 0.0)) {
1891 return Err(format!(
1892 "residual cascade: metric must be {dim} finite positive scales, got {metric:?}"
1893 ));
1894 }
1895 if !(sobolev_s > dim as f64 / 2.0 && sobolev_s <= (dim as f64 + 3.0) / 2.0) {
1896 return Err(format!(
1897 "residual cascade: sobolev_s must lie in (d/2, (d+3)/2] = ({}, {}] for the \
1898 Wendland-(3,1) bump, got {sobolev_s}",
1899 dim as f64 / 2.0,
1900 (dim as f64 + 3.0) / 2.0
1901 ));
1902 }
1903 if levels == 0 || levels > MAX_LEVELS {
1904 return Err(format!(
1905 "residual cascade: levels must be in 1..={MAX_LEVELS}, got {levels}"
1906 ));
1907 }
1908 for i in 0..n {
1909 if !(y[i].is_finite() && w[i].is_finite() && w[i] > 0.0)
1910 || xs.iter().any(|x| !x[i].is_finite())
1911 {
1912 return Err(format!(
1913 "residual cascade: non-finite or non-positive input at row {i}"
1914 ));
1915 }
1916 }
1917 let mut z_lo = [f64::INFINITY; 3];
1919 let mut z_hi = [f64::NEG_INFINITY; 3];
1920 for a in 0..dim {
1921 for &v in xs[a] {
1922 let s = metric[a] * v;
1923 z_lo[a] = z_lo[a].min(s);
1924 z_hi[a] = z_hi[a].max(s);
1925 }
1926 }
1927 let mut z_range = [1.0_f64; 3];
1928 let mut max_range = 0.0_f64;
1929 for a in 0..dim {
1930 if !(z_hi[a] > z_lo[a]) {
1931 return Err(format!(
1932 "residual cascade: degenerate axis {a} bounding box [{}, {}]",
1933 z_lo[a], z_hi[a]
1934 ));
1935 }
1936 z_range[a] = z_hi[a] - z_lo[a];
1937 max_range = max_range.max(z_range[a]);
1938 }
1939 for a in dim..3 {
1940 z_lo[a] = 0.0;
1941 }
1942 let z: Vec<[f64; 3]> = (0..n)
1943 .map(|i| {
1944 let mut p = [0.0_f64; 3];
1945 for a in 0..dim {
1946 p[a] = metric[a] * xs[a][i] - z_lo[a];
1947 }
1948 p
1949 })
1950 .collect();
1951 let mut metric3 = [1.0_f64; 3];
1952 metric3[..dim].copy_from_slice(metric);
1953
1954 let h0 = H0_FRACTION * max_range;
1955 let mut net: Vec<[f64; 3]> = Vec::new();
1956 let mut level_specs = Vec::with_capacity(levels);
1957 let mut col = dim + 1;
1958 let mut pen_logdet_const = 0.0;
1959 for l in 0..levels {
1960 let h = h0 * 0.5_f64.powi(l as i32);
1961 let new_centers = extend_net(&mut net, &z, dim, h, &z_range);
1962 if net.len() > MAX_CENTERS {
1963 return Err(format!(
1964 "residual cascade: center cap {MAX_CENTERS} exceeded at level {l}"
1965 ));
1966 }
1967 let weight = level_weight(l, sobolev_s, dim);
1968 pen_logdet_const += new_centers.len() as f64 * weight.ln();
1969 let delta = OVERLAP * h;
1970 let mut grid = HashGrid::new(delta, dim);
1971 for (j, c) in new_centers.iter().enumerate() {
1972 grid.insert(j as u32, c);
1973 }
1974 let col_offset = col;
1975 col += new_centers.len();
1976 level_specs.push(Level {
1977 h,
1978 delta,
1979 weight,
1980 centers: new_centers,
1981 col_offset,
1982 grid,
1983 });
1984 }
1985 let m = col;
1986
1987 let mut row_ptr = Vec::with_capacity(n + 1);
1989 row_ptr.push(0_usize);
1990 let mut col_idx: Vec<u32> = Vec::new();
1991 let mut vals: Vec<f64> = Vec::new();
1992 let mut rhs = vec![0.0_f64; m];
1993 let mut gram_diag = vec![0.0_f64; m];
1994 let mut ytwy = 0.0_f64;
1995 let probe_core = CoreScaffold {
1996 dim,
1997 z_range,
1998 levels: &level_specs,
1999 };
2000 for i in 0..n {
2001 let row = probe_core.basis_row(&z[i]);
2002 for &(c, v) in &row {
2003 col_idx.push(c as u32);
2004 vals.push(v);
2005 rhs[c] += w[i] * y[i] * v;
2006 gram_diag[c] += w[i] * v * v;
2007 }
2008 ytwy += w[i] * y[i] * y[i];
2009 row_ptr.push(col_idx.len());
2010 }
2011 let mut pen_diag = vec![0.0_f64; m];
2012 for level in &level_specs {
2013 for j in 0..level.centers.len() {
2014 pen_diag[level.col_offset + j] = level.weight;
2015 }
2016 }
2017
2018 let dense_gram = if m <= DENSE_GRAM_MAX {
2021 let mut gram = vec![0.0_f64; m * m];
2022 for i in 0..n {
2023 let lo = row_ptr[i];
2024 let hi = row_ptr[i + 1];
2025 for ea in lo..hi {
2026 let ca = col_idx[ea] as usize;
2027 let va = w[i] * vals[ea];
2028 for eb in ea..hi {
2029 gram[ca * m + col_idx[eb] as usize] += va * vals[eb];
2030 }
2031 }
2032 }
2033 Some(gram)
2034 } else {
2035 None
2036 };
2037
2038 Ok(ResidualCascadeDesign {
2039 core: Arc::new(Core {
2040 dim,
2041 metric: metric3,
2042 z_lo,
2043 z_range,
2044 sobolev_s,
2045 levels: level_specs,
2046 net,
2047 m,
2048 row_ptr,
2049 col_idx,
2050 vals,
2051 w: w.to_vec(),
2052 y: y.to_vec(),
2053 z,
2054 rhs,
2055 ytwy,
2056 gram_diag,
2057 pen_diag,
2058 pen_logdet_const,
2059 dense_gram,
2060 predict_chol: None,
2061 }),
2062 })
2063 }
2064
2065 pub fn num_levels(&self) -> usize {
2067 self.core.levels.len()
2068 }
2069
2070 pub fn metric_scaled_aspect_ratio(&self) -> f64 {
2078 let dim = self.core.dim;
2079 let n = self.core.z.len();
2080 if dim == 0 || n == 0 {
2081 return 1.0;
2082 }
2083 let mut mean = [0.0_f64; 3];
2084 for p in &self.core.z {
2085 for a in 0..dim {
2086 mean[a] += p[a];
2087 }
2088 }
2089 for m in mean.iter_mut().take(dim) {
2090 *m /= n as f64;
2091 }
2092 let mut var = [0.0_f64; 3];
2093 for p in &self.core.z {
2094 for a in 0..dim {
2095 let d = p[a] - mean[a];
2096 var[a] += d * d;
2097 }
2098 }
2099 let mut sd_lo = f64::INFINITY;
2100 let mut sd_hi = 0.0_f64;
2101 for v in var.iter().take(dim) {
2102 let sd = (v / n as f64).sqrt();
2103 sd_lo = sd_lo.min(sd);
2104 sd_hi = sd_hi.max(sd);
2105 }
2106 if !(sd_lo > 0.0 && sd_lo.is_finite()) {
2107 return f64::INFINITY;
2109 }
2110 sd_hi / sd_lo
2111 }
2112
2113 pub fn quasi_uniformity_certified(&self) -> bool {
2121 self.metric_scaled_aspect_ratio() <= QUASI_UNIFORMITY_MAX_ASPECT
2122 }
2123
2124 pub fn coarse_space_cols(&self, log_lambda: f64) -> Result<usize, String> {
2132 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
2133 .map_err(|error| format!("residual cascade: {error}"))?;
2134 Ok(self.core.coarse_space_cols(lambda))
2135 }
2136
2137 pub fn num_coeffs(&self) -> usize {
2139 self.core.m
2140 }
2141
2142 pub fn num_nonzeros(&self) -> usize {
2150 self.core.col_idx.len()
2151 }
2152
2153 pub fn num_centers(&self) -> usize {
2155 self.core.m - self.core.nullity()
2156 }
2157
2158 pub fn centers(&self, level: usize) -> Vec<Vec<f64>> {
2160 let lv = &self.core.levels[level];
2161 lv.centers
2162 .iter()
2163 .map(|c| {
2164 (0..self.core.dim)
2165 .map(|a| (c[a] + self.core.z_lo[a]) / self.core.metric[a])
2166 .collect()
2167 })
2168 .collect()
2169 }
2170
2171 pub fn basis_row(&self, x: &[f64]) -> Result<Vec<(usize, f64)>, String> {
2175 self.check_point(x)?;
2176 Ok(self.core.basis_row_scaled(&self.core.scale_point(x)))
2177 }
2178
2179 fn check_point(&self, x: &[f64]) -> Result<(), String> {
2180 if x.len() != self.core.dim || x.iter().any(|v| !v.is_finite()) {
2181 return Err(format!(
2182 "residual cascade: point must be {} finite coordinates, got {x:?}",
2183 self.core.dim
2184 ));
2185 }
2186 Ok(())
2187 }
2188
2189 pub fn penalty_value(&self, coeff: &[f64]) -> Result<f64, String> {
2191 if coeff.len() != self.core.m {
2192 return Err(format!(
2193 "residual cascade: coefficient length {} != {}",
2194 coeff.len(),
2195 self.core.m
2196 ));
2197 }
2198 Ok(coeff
2199 .iter()
2200 .zip(self.core.pen_diag.iter())
2201 .map(|(&c, &d)| d * c * c)
2202 .sum())
2203 }
2204
2205 pub fn logdet_exact(&self, log_lambda: f64) -> Result<f64, String> {
2208 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
2209 .map_err(|error| format!("residual cascade: {error}"))?;
2210 self.core.logdet_dense(lambda)
2211 }
2212
2213 pub fn logdet_slq(&self, log_lambda: f64) -> Result<f64, String> {
2216 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
2217 .map_err(|error| format!("residual cascade: {error}"))?;
2218 self.core.logdet_slq(lambda)
2219 }
2220
2221 pub fn criterion(&self, log_lambda: f64) -> Result<f64, String> {
2225 Ok(self.core.reml_profile()?.evaluate(log_lambda)?.jet.value)
2226 }
2227
2228 pub fn fit_at(
2230 &self,
2231 log_lambda: f64,
2232 sigma2: Option<f64>,
2233 ) -> Result<ResidualCascadeFit, String> {
2234 self.fit_at_with_warm(log_lambda, sigma2, None, None)
2235 }
2236
2237 fn fit_at_with_warm(
2238 &self,
2239 log_lambda: f64,
2240 sigma2: Option<f64>,
2241 warm: Option<&[f64]>,
2242 profile_normalized_logdet: Option<f64>,
2243 ) -> Result<ResidualCascadeFit, String> {
2244 let core = &self.core;
2245 let lambda = gam_problem::checked_exp_log_strength(log_lambda)
2246 .map_err(|error| format!("residual cascade: {error}"))?;
2247 let (coeff, rel_res, iters) = core.solve_coeff(lambda, &core.rhs, warm)?;
2248 let rss_pen = core.rss_pen(&coeff);
2249 let dof = (core.y.len() - core.nullity()) as f64;
2250 let sigma2 = match sigma2 {
2251 Some(s) => {
2252 if !(s.is_finite() && s > 0.0) {
2253 return Err(format!("residual cascade: invalid sigma2 {s}"));
2254 }
2255 s
2256 }
2257 None => {
2258 if !(rss_pen > 0.0) {
2259 return Err(format!(
2260 "residual cascade: degenerate penalized residual {rss_pen}"
2261 ));
2262 }
2263 rss_pen / dof
2264 }
2265 };
2266 let r = (core.m - core.nullity()) as f64;
2267 let (logdet, logdet_method) = match profile_normalized_logdet {
2268 Some(normalized) => (
2269 normalized + r * log_lambda + core.pen_logdet_const,
2270 if core.dense_gram.is_some() {
2271 LogdetMethod::DenseExact
2272 } else {
2273 LogdetMethod::Slq
2274 },
2275 ),
2276 None => core.logdet(lambda)?,
2277 };
2278 let restricted_loglik = -0.5
2281 * (logdet - r * log_lambda - core.pen_logdet_const
2282 + dof * sigma2.ln()
2283 + rss_pen / sigma2);
2284 let predict_chol = if core.dense_gram.is_some() {
2285 Some(core.assemble_predict_factor(lambda)?)
2286 } else {
2287 None
2288 };
2289 Ok(ResidualCascadeFit {
2290 core: Arc::clone(&self.core),
2291 predict_chol,
2292 coeff,
2293 log_lambda,
2294 sigma2,
2295 restricted_loglik,
2296 rss_pen,
2297 certificate: CascadeCertificate {
2298 solve_rel_residual: rel_res,
2299 solve_iters: iters,
2300 logdet_method,
2301 },
2302 refinement: None,
2303 })
2304 }
2305
2306 pub fn fit_reml(&self) -> Result<ResidualCascadeFit, String> {
2313 let profile = self.core.reml_profile()?;
2314 let (log_lambda_lo, log_lambda_hi) = profile.log_lambda_domain()?;
2315 let search = maximize_score_1d(
2316 log_lambda_lo,
2317 log_lambda_hi,
2318 f64::EPSILON.sqrt(),
2319 |log_lambda| {
2320 profile
2321 .evaluate(log_lambda)
2322 .map(|evaluation| evaluation.jet)
2323 },
2324 |lo, hi| profile.enclose(lo, hi),
2325 )
2326 .map_err(|error| format!("residual cascade: REML stationary isolation failed: {error}"))?;
2327 let selected = profile.evaluate(search.optimum.x)?;
2328 self.fit_at_with_warm(
2329 search.optimum.x,
2330 None,
2331 None,
2332 Some(selected.normalized_logdet),
2333 )
2334 }
2335
2336 pub fn assess_next_level(
2343 &self,
2344 fit: &ResidualCascadeFit,
2345 ) -> Result<NextLevelAssessment, String> {
2346 let core = &self.core;
2347 if !Arc::ptr_eq(core, &fit.core) {
2348 return Err("residual cascade: fit does not belong to this design".into());
2349 }
2350 let next_l = core.levels.len();
2351 let h = core.levels[next_l - 1].h * 0.5;
2352 let mut net = core.net.clone();
2353 let candidates = extend_net(&mut net, &core.z, core.dim, h, &core.z_range);
2354 if candidates.is_empty() {
2355 return Ok(NextLevelAssessment::EmptyNet);
2356 }
2357 if net.len() > MAX_CENTERS {
2358 return Ok(NextLevelAssessment::CapacityExceeded {
2359 obstruction: RefinementObstruction::CenterCapacity {
2360 centers: net.len(),
2361 maximum_centers: MAX_CENTERS,
2362 },
2363 gain_bound: f64::INFINITY,
2368 });
2369 }
2370 let delta = OVERLAP * h;
2371 let mut grid = HashGrid::new(delta, core.dim);
2372 for (j, c) in candidates.iter().enumerate() {
2373 grid.insert(j as u32, c);
2374 }
2375 let r = core.residuals(&fit.coeff);
2376 let mut g = vec![0.0_f64; candidates.len()];
2377 for (i, zi) in core.z.iter().enumerate() {
2378 let wr = core.w[i] * r[i];
2379 grid.for_neighbors(zi, |j| {
2380 let rad = dist2(zi, &candidates[j as usize], core.dim).sqrt() / delta;
2381 g[j as usize] += wr * wendland(rad);
2382 });
2383 }
2384 let g2: f64 = g.iter().map(|v| v * v).sum();
2385 let d_next = level_weight(next_l, core.sobolev_s, core.dim);
2386 let lambda = gam_problem::checked_exp_log_strength(fit.log_lambda)
2387 .map_err(|error| format!("residual cascade refinement: {error}"))?;
2388 let gain_bound = g2 / (lambda * d_next);
2389 if next_l >= MAX_LEVELS {
2390 Ok(NextLevelAssessment::CapacityExceeded {
2391 obstruction: RefinementObstruction::LevelCapacity {
2392 levels: next_l,
2393 maximum_levels: MAX_LEVELS,
2394 },
2395 gain_bound,
2396 })
2397 } else {
2398 Ok(NextLevelAssessment::GainBound(gain_bound))
2399 }
2400 }
2401}
2402
2403fn level_weight(l: usize, sobolev_s: f64, dim: usize) -> f64 {
2405 (4.0_f64).powf(l as f64 * (sobolev_s - dim as f64 / 2.0))
2406}
2407
2408struct CoreScaffold<'a> {
2412 dim: usize,
2413 z_range: [f64; 3],
2414 levels: &'a [Level],
2415}
2416
2417impl CoreScaffold<'_> {
2418 fn basis_row(&self, z: &[f64; 3]) -> Vec<(usize, f64)> {
2419 let mut row = Vec::with_capacity(self.dim + 1 + self.levels.len() * 8);
2420 row.push((0, 1.0));
2421 for a in 0..self.dim {
2422 row.push((a + 1, 2.0 * z[a] / self.z_range[a] - 1.0));
2423 }
2424 for level in self.levels {
2425 let start = row.len();
2426 level.grid.for_neighbors(z, |j| {
2427 let c = &level.centers[j as usize];
2428 let r = dist2(z, c, self.dim).sqrt() / level.delta;
2429 let v = wendland(r);
2430 if v > 0.0 {
2431 row.push((level.col_offset + j as usize, v));
2432 }
2433 });
2434 row[start..].sort_unstable_by_key(|&(col, _)| col);
2435 }
2436 row
2437 }
2438}
2439
2440impl ResidualCascadeFit {
2441 pub fn log_lambda(&self) -> f64 {
2442 self.log_lambda
2443 }
2444
2445 pub fn lambda(&self) -> f64 {
2446 gam_problem::checked_exp_log_strength(self.log_lambda)
2447 .expect("ResidualCascadeFit construction validates its private log strength")
2448 }
2449
2450 pub fn predict(&self, x: &[f64]) -> Result<(f64, f64), String> {
2454 let core = &self.core;
2455 if x.len() != core.dim || x.iter().any(|v| !v.is_finite()) {
2456 return Err(format!(
2457 "residual cascade: prediction point must be {} finite coordinates, got {x:?}",
2458 core.dim
2459 ));
2460 }
2461 let row = core.basis_row_scaled(&core.scale_point(x));
2462 let mut mean = 0.0;
2463 let mut dense_row = vec![0.0_f64; core.m];
2464 for &(c, v) in &row {
2465 mean += v * self.coeff[c];
2466 dense_row[c] += v;
2467 }
2468 let lambda = gam_problem::checked_exp_log_strength(self.log_lambda)
2469 .map_err(|error| format!("residual cascade fit: {error}"))?;
2470 let zsol = if let Some(l) = &self.predict_chol {
2471 chol_solve(l, core.m, &dense_row)
2472 } else {
2473 core.solve_coeff(lambda, &dense_row, None)?.0
2474 };
2475 let mut quad = 0.0;
2476 for (a, b) in dense_row.iter().zip(zsol.iter()) {
2477 quad += a * b;
2478 }
2479 Ok((mean, self.sigma2 * quad))
2480 }
2481
2482 pub fn sample_coefficients(&self, n_samples: usize) -> Result<Vec<Vec<f64>>, String> {
2487 let core = &self.core;
2488 let lambda = gam_problem::checked_exp_log_strength(self.log_lambda)
2489 .map_err(|error| format!("residual cascade fit: {error}"))?;
2490 let sigma = self.sigma2.sqrt();
2491 let sqrt_lambda = lambda.sqrt();
2492 let n = core.y.len();
2493 let mut rng = SplitMix64::new(RNG_SEED ^ 0xA11C_E5A_u64);
2494 let mut samples = Vec::with_capacity(n_samples);
2495 for _ in 0..n_samples {
2496 let mut b = core.rhs.clone();
2497 for i in 0..n {
2499 let f = sigma * core.w[i].sqrt() * rng.next_normal();
2500 for e in core.row_ptr[i]..core.row_ptr[i + 1] {
2501 b[core.col_idx[e] as usize] += f * core.vals[e];
2502 }
2503 }
2504 for (bj, &dj) in b.iter_mut().zip(core.pen_diag.iter()) {
2506 if dj > 0.0 {
2507 *bj += sigma * sqrt_lambda * dj.sqrt() * rng.next_normal();
2508 }
2509 }
2510 let (c, _, _) = core.solve_coeff(lambda, &b, Some(&self.coeff))?;
2511 samples.push(c);
2512 }
2513 Ok(samples)
2514 }
2515
2516 pub fn num_levels(&self) -> usize {
2518 self.core.levels.len()
2519 }
2520
2521 pub fn num_coeffs(&self) -> usize {
2523 self.core.m
2524 }
2525
2526 pub fn num_centers(&self) -> usize {
2528 self.core.m - self.core.nullity()
2529 }
2530
2531 pub fn to_state(&self) -> Result<ResidualCascadeState, String> {
2537 let core = &self.core;
2538 let lambda = gam_problem::checked_exp_log_strength(self.log_lambda)
2539 .map_err(|error| format!("residual cascade fit: {error}"))?;
2540 let predict_chol = if let Some(l) = &self.predict_chol {
2541 l.clone()
2542 } else if let Some(l) = &core.predict_chol {
2543 l.clone()
2544 } else {
2545 core.assemble_predict_factor(lambda)?
2546 };
2547 let dim = core.dim;
2548 let levels = core
2549 .levels
2550 .iter()
2551 .map(|level| {
2552 let mut centers = Vec::with_capacity(level.centers.len() * dim);
2553 for c in &level.centers {
2554 centers.extend_from_slice(&c[..dim]);
2555 }
2556 LevelState {
2557 h: level.h,
2558 delta: level.delta,
2559 weight: level.weight,
2560 col_offset: level.col_offset as u64,
2561 centers,
2562 }
2563 })
2564 .collect();
2565 Ok(ResidualCascadeState {
2566 dim: dim as u64,
2567 metric: core.metric,
2568 z_lo: core.z_lo,
2569 z_range: core.z_range,
2570 sobolev_s: core.sobolev_s,
2571 levels,
2572 m: core.m as u64,
2573 pen_logdet_const: core.pen_logdet_const,
2574 coeff: self.coeff.clone(),
2575 log_lambda: self.log_lambda,
2576 sigma2: self.sigma2,
2577 restricted_loglik: self.restricted_loglik,
2578 rss_pen: self.rss_pen,
2579 predict_chol,
2580 })
2581 }
2582
2583 pub fn from_state(state: &ResidualCascadeState) -> Result<Self, String> {
2592 let dim = state.dim as usize;
2593 if !(dim == 2 || dim == 3) {
2594 return Err(format!(
2595 "residual cascade state: dim must be 2 or 3, got {dim}"
2596 ));
2597 }
2598 if !(state.sobolev_s > dim as f64 / 2.0 && state.sobolev_s <= (dim as f64 + 3.0) / 2.0) {
2599 return Err(format!(
2600 "residual cascade state: sobolev_s {} outside the Wendland window ({}, {}]",
2601 state.sobolev_s,
2602 dim as f64 / 2.0,
2603 (dim as f64 + 3.0) / 2.0
2604 ));
2605 }
2606 for a in 0..dim {
2607 if !(state.metric[a].is_finite() && state.metric[a] > 0.0) {
2608 return Err(format!(
2609 "residual cascade state: metric axis {a} must be finite positive, got {}",
2610 state.metric[a]
2611 ));
2612 }
2613 if !(state.z_range[a].is_finite()
2614 && state.z_range[a] > 0.0
2615 && state.z_lo[a].is_finite())
2616 {
2617 return Err(format!(
2618 "residual cascade state: degenerate box on axis {a} (lo={}, range={})",
2619 state.z_lo[a], state.z_range[a]
2620 ));
2621 }
2622 }
2623 let m = state.m as usize;
2624 let mut metric3 = [1.0_f64; 3];
2625 metric3[..dim].copy_from_slice(&state.metric[..dim]);
2626 let mut z_lo = [0.0_f64; 3];
2627 let mut z_range = [1.0_f64; 3];
2628 z_lo[..dim].copy_from_slice(&state.z_lo[..dim]);
2629 z_range[..dim].copy_from_slice(&state.z_range[..dim]);
2630
2631 let mut levels = Vec::with_capacity(state.levels.len());
2634 let mut net: Vec<[f64; 3]> = Vec::new();
2635 let mut pen_diag = vec![0.0_f64; m];
2636 let mut expected_offset = dim + 1;
2637 for (li, ls) in state.levels.iter().enumerate() {
2638 if !(ls.h.is_finite() && ls.h > 0.0 && ls.delta.is_finite() && ls.delta > 0.0) {
2639 return Err(format!(
2640 "residual cascade state: level {li} has non-positive h/delta ({}, {})",
2641 ls.h, ls.delta
2642 ));
2643 }
2644 if !(ls.weight.is_finite() && ls.weight > 0.0) {
2645 return Err(format!(
2646 "residual cascade state: level {li} has non-positive prior weight {}",
2647 ls.weight
2648 ));
2649 }
2650 if ls.centers.len() % dim != 0 {
2651 return Err(format!(
2652 "residual cascade state: level {li} centers length {} not a multiple of dim {dim}",
2653 ls.centers.len()
2654 ));
2655 }
2656 let n_centers = ls.centers.len() / dim;
2657 let col_offset = ls.col_offset as usize;
2658 if col_offset != expected_offset {
2659 return Err(format!(
2660 "residual cascade state: level {li} col_offset {col_offset} ≠ expected {expected_offset}"
2661 ));
2662 }
2663 let mut grid = HashGrid::new(ls.delta, dim);
2664 let mut centers = Vec::with_capacity(n_centers);
2665 for j in 0..n_centers {
2666 let mut c = [0.0_f64; 3];
2667 for a in 0..dim {
2668 let v = ls.centers[j * dim + a];
2669 if !v.is_finite() {
2670 return Err(format!(
2671 "residual cascade state: non-finite center coordinate at level {li}, center {j}"
2672 ));
2673 }
2674 c[a] = v;
2675 }
2676 grid.insert(j as u32, &c);
2677 centers.push(c);
2678 net.push(c);
2679 let col = col_offset + j;
2680 if col >= m {
2681 return Err(format!(
2682 "residual cascade state: level {li} column {col} exceeds m {m}"
2683 ));
2684 }
2685 pen_diag[col] = ls.weight;
2686 }
2687 expected_offset = col_offset + n_centers;
2688 levels.push(Level {
2689 h: ls.h,
2690 delta: ls.delta,
2691 weight: ls.weight,
2692 centers,
2693 col_offset,
2694 grid,
2695 });
2696 }
2697 if expected_offset != m {
2698 return Err(format!(
2699 "residual cascade state: column accounting mismatch (dim+1+Σcenters = {expected_offset} ≠ m {m})"
2700 ));
2701 }
2702 if state.coeff.len() != m {
2703 return Err(format!(
2704 "residual cascade state: coeff length {} ≠ m {m}",
2705 state.coeff.len()
2706 ));
2707 }
2708 if state.predict_chol.len() != m * m {
2709 return Err(format!(
2710 "residual cascade state: predict_chol must be m×m = {m}² = {}, got {}",
2711 m * m,
2712 state.predict_chol.len()
2713 ));
2714 }
2715 for (i, v) in state
2716 .coeff
2717 .iter()
2718 .chain(state.predict_chol.iter())
2719 .enumerate()
2720 {
2721 if !v.is_finite() {
2722 return Err(format!("residual cascade state: non-finite entry at {i}"));
2723 }
2724 }
2725 for g in 0..m {
2726 let piv = state.predict_chol[g * m + g];
2727 if !(piv.is_finite() && piv > 0.0) {
2728 return Err(format!(
2729 "residual cascade state: non-positive Cholesky pivot {piv} at index {g}"
2730 ));
2731 }
2732 }
2733 gam_problem::validate_log_strength(state.log_lambda)
2734 .map_err(|error| format!("residual cascade state: {error}"))?;
2735 if !(state.sigma2.is_finite()
2736 && state.sigma2 > 0.0
2737 && state.restricted_loglik.is_finite()
2738 && state.rss_pen.is_finite())
2739 {
2740 return Err(format!(
2741 "residual cascade state: invalid scalars (log_lambda={}, sigma2={}, restricted_loglik={}, rss_pen={})",
2742 state.log_lambda, state.sigma2, state.restricted_loglik, state.rss_pen
2743 ));
2744 }
2745 let core = Core {
2746 dim,
2747 metric: metric3,
2748 z_lo,
2749 z_range,
2750 sobolev_s: state.sobolev_s,
2751 levels,
2752 net,
2753 m,
2754 row_ptr: Vec::new(),
2755 col_idx: Vec::new(),
2756 vals: Vec::new(),
2757 w: Vec::new(),
2758 y: Vec::new(),
2759 z: Vec::new(),
2760 rhs: Vec::new(),
2761 ytwy: 0.0,
2762 gram_diag: Vec::new(),
2763 pen_diag,
2764 pen_logdet_const: state.pen_logdet_const,
2765 dense_gram: None,
2766 predict_chol: Some(state.predict_chol.clone()),
2767 };
2768 Ok(ResidualCascadeFit {
2769 core: Arc::new(core),
2770 predict_chol: None,
2771 coeff: state.coeff.clone(),
2772 log_lambda: state.log_lambda,
2773 sigma2: state.sigma2,
2774 restricted_loglik: state.restricted_loglik,
2775 rss_pen: state.rss_pen,
2776 certificate: CascadeCertificate {
2777 solve_rel_residual: 0.0,
2778 solve_iters: 0,
2779 logdet_method: LogdetMethod::DenseExact,
2780 },
2781 refinement: None,
2782 })
2783 }
2784}
2785
2786#[derive(Clone, Copy, Debug, PartialEq)]
2787enum RefinementDecision {
2788 Converged {
2789 gain_bound: f64,
2790 },
2791 Refine,
2792 Underresolved {
2793 gain_bound: f64,
2794 obstruction: RefinementObstruction,
2795 },
2796}
2797
2798fn decide_refinement(
2802 assessment: NextLevelAssessment,
2803 requested_tolerance: f64,
2804) -> RefinementDecision {
2805 match assessment {
2806 NextLevelAssessment::EmptyNet => RefinementDecision::Converged { gain_bound: 0.0 },
2807 NextLevelAssessment::GainBound(gain_bound) if gain_bound <= requested_tolerance => {
2808 RefinementDecision::Converged { gain_bound }
2809 }
2810 NextLevelAssessment::GainBound(_) => RefinementDecision::Refine,
2811 NextLevelAssessment::CapacityExceeded {
2812 gain_bound,
2813 obstruction: _,
2814 } if gain_bound <= requested_tolerance => RefinementDecision::Converged { gain_bound },
2815 NextLevelAssessment::CapacityExceeded {
2816 gain_bound,
2817 obstruction,
2818 } => RefinementDecision::Underresolved {
2819 gain_bound,
2820 obstruction,
2821 },
2822 }
2823}
2824
2825pub fn fit_residual_cascade(
2834 xs: &[&[f64]],
2835 y: &[f64],
2836 w: &[f64],
2837 metric: &[f64],
2838 sobolev_s: f64,
2839) -> Result<ResidualCascadeFit, ResidualCascadeError> {
2840 let mut levels = INITIAL_LEVELS;
2841 loop {
2842 let design = ResidualCascadeDesign::build(xs, y, w, metric, sobolev_s, levels)?;
2843 if levels == INITIAL_LEVELS && !design.quasi_uniformity_certified() {
2852 return Err(format!(
2853 "residual cascade: metric-scaled aspect ratio {:.3e} exceeds the \
2854 quasi-uniformity ceiling {QUASI_UNIFORMITY_MAX_ASPECT:.0e}; the BPX \
2855 iteration bound is not trustworthy on this (near-degenerate) metric — \
2856 fall back to the dense kernel path",
2857 design.metric_scaled_aspect_ratio()
2858 )
2859 .into());
2860 }
2861 let mut fit = design.fit_reml()?;
2862 let assessment = design.assess_next_level(&fit)?;
2873 let requested_tolerance = REFINE_TOL * fit.rss_pen;
2874 match decide_refinement(assessment, requested_tolerance) {
2875 RefinementDecision::Converged { gain_bound } => {
2876 fit.refinement = Some(RefinementCertificate {
2877 next_level_gain_bound: gain_bound,
2878 tolerance: requested_tolerance,
2879 });
2880 return Ok(fit);
2881 }
2882 RefinementDecision::Refine => {
2883 levels += 1;
2884 }
2885 RefinementDecision::Underresolved {
2886 gain_bound,
2887 obstruction,
2888 } => {
2889 return Err(ResidualCascadeError::Underresolved {
2890 checkpoint: ResidualCascadeCheckpoint::new(fit),
2891 gain_bound,
2892 requested_tolerance,
2893 obstruction,
2894 });
2895 }
2896 }
2897 }
2898}
2899
2900#[cfg(test)]
2901mod refinement_decision_tests {
2902 use super::*;
2903
2904 const TOLERANCE: f64 = 0.25;
2905
2906 #[test]
2907 fn only_empty_or_passing_bound_converges() {
2908 assert_eq!(
2909 decide_refinement(NextLevelAssessment::EmptyNet, TOLERANCE),
2910 RefinementDecision::Converged { gain_bound: 0.0 }
2911 );
2912 assert_eq!(
2913 decide_refinement(NextLevelAssessment::GainBound(0.2), TOLERANCE),
2914 RefinementDecision::Converged { gain_bound: 0.2 }
2915 );
2916 assert_eq!(
2917 decide_refinement(NextLevelAssessment::GainBound(0.3), TOLERANCE),
2918 RefinementDecision::Refine
2919 );
2920 }
2921
2922 #[test]
2923 fn capacity_above_tolerance_is_underresolved() {
2924 let obstruction = RefinementObstruction::LevelCapacity {
2925 levels: MAX_LEVELS,
2926 maximum_levels: MAX_LEVELS,
2927 };
2928 assert_eq!(
2929 decide_refinement(
2930 NextLevelAssessment::CapacityExceeded {
2931 obstruction,
2932 gain_bound: 0.3,
2933 },
2934 TOLERANCE,
2935 ),
2936 RefinementDecision::Underresolved {
2937 gain_bound: 0.3,
2938 obstruction,
2939 }
2940 );
2941
2942 let center_obstruction = RefinementObstruction::CenterCapacity {
2943 centers: MAX_CENTERS + 1,
2944 maximum_centers: MAX_CENTERS,
2945 };
2946 assert_eq!(
2947 decide_refinement(
2948 NextLevelAssessment::CapacityExceeded {
2949 obstruction: center_obstruction,
2950 gain_bound: f64::INFINITY,
2951 },
2952 TOLERANCE,
2953 ),
2954 RefinementDecision::Underresolved {
2955 gain_bound: f64::INFINITY,
2956 obstruction: center_obstruction,
2957 }
2958 );
2959 }
2960
2961 #[test]
2962 fn capacity_does_not_block_an_independently_passing_bound() {
2963 assert_eq!(
2964 decide_refinement(
2965 NextLevelAssessment::CapacityExceeded {
2966 obstruction: RefinementObstruction::LevelCapacity {
2967 levels: MAX_LEVELS,
2968 maximum_levels: MAX_LEVELS,
2969 },
2970 gain_bound: 0.2,
2971 },
2972 TOLERANCE,
2973 ),
2974 RefinementDecision::Converged { gain_bound: 0.2 }
2975 );
2976 }
2977
2978 #[test]
2979 fn dense_spectral_profile_matches_factorization_and_analytic_slope() {
2980 let side = 6usize;
2981 let mut x1 = Vec::with_capacity(side * side);
2982 let mut x2 = Vec::with_capacity(side * side);
2983 let mut y = Vec::with_capacity(side * side);
2984 for i in 0..side {
2985 for j in 0..side {
2986 let a = i as f64 / (side - 1) as f64;
2987 let b = j as f64 / (side - 1) as f64;
2988 x1.push(a);
2989 x2.push(b);
2990 y.push((2.3 * a).sin() + (1.7 * b).cos() + 0.07 * ((3 * i + 5 * j) % 7) as f64);
2991 }
2992 }
2993 let weights = vec![1.0; y.len()];
2994 let axes: [&[f64]; 2] = [&x1, &x2];
2995 let design = ResidualCascadeDesign::build(&axes, &y, &weights, &[1.0, 1.0], 2.0, 2)
2996 .expect("cascade design");
2997 assert!(design.core.dense_gram.is_some());
2998 let profile = design.core.reml_profile().expect("spectral profile");
2999 let rank = (design.core.m - design.core.nullity()) as f64;
3000 let dof = (design.core.y.len() - design.core.nullity()) as f64;
3001
3002 for log_lambda in [-4.0, 0.0, 3.0] {
3003 let evaluation = profile.evaluate(log_lambda).expect("analytic score");
3004 let lambda = log_lambda.exp();
3005 let logdet = design.core.logdet_dense(lambda).expect("dense logdet");
3006 let coefficients = design
3007 .core
3008 .solve_coeff(lambda, &design.core.rhs, None)
3009 .expect("dense solve")
3010 .0;
3011 let rss = design.core.rss_pen(&coefficients);
3012 let direct = -0.5
3013 * (logdet - rank * log_lambda - design.core.pen_logdet_const
3014 + dof * (rss / dof).ln());
3015 assert!(
3016 (evaluation.jet.value - direct).abs() <= f64::EPSILON.sqrt() * (1.0 + direct.abs()),
3017 "spectral/direct score mismatch at {log_lambda}: {} versus {direct}",
3018 evaluation.jet.value,
3019 );
3020
3021 let step = f64::EPSILON.cbrt();
3024 let right = profile
3025 .evaluate(log_lambda + step)
3026 .expect("right score")
3027 .jet
3028 .value;
3029 let left = profile
3030 .evaluate(log_lambda - step)
3031 .expect("left score")
3032 .jet
3033 .value;
3034 let numerical_slope = (right - left) / (2.0 * step);
3035 assert!(
3036 (evaluation.jet.derivative - numerical_slope).abs()
3037 <= f64::EPSILON.sqrt() * (1.0 + numerical_slope.abs()),
3038 "analytic slope mismatch at {log_lambda}: {} versus {numerical_slope}",
3039 evaluation.jet.derivative,
3040 );
3041 }
3042 }
3043}