1use crate::error::FdarError;
52use crate::function_on_scalar::penalty_matrix;
53use crate::helpers::simpsons_weights;
54use crate::linalg::{cholesky_factor, cholesky_forward_back, cholesky_solve};
55use crate::matrix::FdMatrix;
56use nalgebra::DMatrix;
57
58#[derive(Debug, Clone, PartialEq)]
67#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
68pub enum PeerPenalty {
69 Ridge,
71 Difference { order: usize },
75 Decree(Vec<f64>, usize),
85}
86
87impl Default for PeerPenalty {
88 fn default() -> Self {
89 PeerPenalty::Difference { order: 2 }
90 }
91}
92
93#[derive(Debug, Clone, Default, PartialEq)]
98#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
99pub enum LambdaChoice {
100 Fixed(f64),
102 #[default]
106 Gcv,
107 Reml,
112}
113
114#[derive(Debug, Clone, PartialEq)]
116#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
117pub enum LambdaMethod {
118 Fixed,
120 Gcv,
122 Reml,
124}
125
126#[derive(Debug, Clone, Default, PartialEq)]
128#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
129pub struct PeerConfig {
130 pub penalty: PeerPenalty,
132 pub lambda: LambdaChoice,
138}
139
140#[derive(Debug, Clone, PartialEq)]
145#[non_exhaustive]
146#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
147#[must_use = "expensive computation whose result should not be discarded"]
148pub struct PeerResult {
149 pub beta: Vec<f64>,
151 pub intercept: f64,
160 pub w_bar: Vec<f64>,
164 pub fitted_values: Vec<f64>,
166 pub effective_df: f64,
168 pub lambda: f64,
170 pub penalty_type: PeerPenalty,
172 pub gcv: Option<f64>,
175 pub lambda_method: LambdaMethod,
177}
178
179#[derive(Debug, Clone, PartialEq)]
187#[non_exhaustive]
188#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
189#[must_use = "expensive computation whose result should not be discarded"]
190pub struct LpeerResult {
191 pub beta: Vec<f64>,
193 pub intercept: f64,
195 pub w_bar: Vec<f64>,
199 pub fitted_values: Vec<f64>,
201 pub sigma2_subject: f64,
204 pub sigma2_resid: f64,
206 pub n_subjects: usize,
208 pub lambda: f64,
210 pub penalty_type: PeerPenalty,
212 pub gcv: Option<f64>,
215 pub lambda_method: LambdaMethod,
217}
218
219pub fn peer(
245 data: &FdMatrix,
246 y: &[f64],
247 argvals: &[f64],
248 config: &PeerConfig,
249) -> Result<PeerResult, FdarError> {
250 let (n, m) = data.shape();
251
252 if n < 2 {
254 return Err(FdarError::InvalidDimension {
257 parameter: "data",
258 expected: "at least 2 observations".to_string(),
259 actual: format!("{n} rows"),
260 });
261 }
262 if m < 3 {
263 return Err(FdarError::InvalidDimension {
264 parameter: "data",
265 expected: "at least 3 evaluation points (m >= 3)".to_string(),
266 actual: format!("{m} columns"),
267 });
268 }
269 if argvals.len() != m {
270 return Err(FdarError::InvalidDimension {
271 parameter: "argvals",
272 expected: format!("{m}"),
273 actual: format!("{}", argvals.len()),
274 });
275 }
276 if y.len() != n {
277 return Err(FdarError::InvalidDimension {
278 parameter: "y",
279 expected: format!("{n}"),
280 actual: format!("{}", y.len()),
281 });
282 }
283 if y.iter().any(|v| !v.is_finite()) {
286 return Err(FdarError::InvalidParameter {
287 parameter: "y",
288 message: "response contains non-finite values (NaN/Inf)".to_string(),
289 });
290 }
291 if argvals.iter().any(|v| !v.is_finite()) {
292 return Err(FdarError::InvalidParameter {
293 parameter: "argvals",
294 message: "argvals contains non-finite values (NaN/Inf)".to_string(),
295 });
296 }
297 if argvals.windows(2).any(|w| w[1] <= w[0]) {
300 return Err(FdarError::InvalidParameter {
301 parameter: "argvals",
302 message: "argvals must be strictly increasing".to_string(),
303 });
304 }
305
306 let w = simpsons_weights(argvals);
308
309 let mut wmat = FdMatrix::zeros(n, m);
311 for i in 0..n {
312 for j in 0..m {
313 wmat[(i, j)] = data[(i, j)] * w[j];
314 }
315 }
316
317 let y_bar: f64 = y.iter().sum::<f64>() / n as f64;
319 let yc: Vec<f64> = y.iter().map(|&yi| yi - y_bar).collect();
320
321 let w_bar: Vec<f64> = (0..m)
322 .map(|j| (0..n).map(|i| wmat[(i, j)]).sum::<f64>() / n as f64)
323 .collect();
324
325 let mut wc = FdMatrix::zeros(n, m);
326 for i in 0..n {
327 for j in 0..m {
328 wc[(i, j)] = wmat[(i, j)] - w_bar[j];
329 }
330 }
331
332 let q = build_q(m, &config.penalty)?;
334
335 let mut wtw = vec![0.0_f64; m * m];
337 for j in 0..m {
338 for k in j..m {
339 let s: f64 = (0..n).map(|i| wc[(i, j)] * wc[(i, k)]).sum();
340 wtw[j * m + k] = s;
341 wtw[k * m + j] = s;
342 }
343 }
344 let wty: Vec<f64> = (0..m)
345 .map(|j| (0..n).map(|i| wc[(i, j)] * yc[i]).sum())
346 .collect();
347
348 let (lambda, gcv_score, lambda_method) = match &config.lambda {
350 LambdaChoice::Fixed(lam) => (*lam, None, LambdaMethod::Fixed),
351 LambdaChoice::Gcv => {
352 let (lam, g) = select_lambda_gcv_peer(&wc, &yc, &wtw, &wty, &q, m, n);
353 (lam, Some(g), LambdaMethod::Gcv)
354 }
355 LambdaChoice::Reml => {
356 let lam = select_lambda_reml_peer(&wc, &yc, &q, m, n);
357 (lam, None, LambdaMethod::Reml)
358 }
359 };
360
361 let mut a = vec![0.0_f64; m * m];
363 for i in 0..m * m {
364 a[i] = wtw[i] + lambda * q[i];
365 }
366 let beta = cholesky_solve(&a, &wty, m)?;
367
368 if beta.iter().any(|v| !v.is_finite()) {
370 return Err(FdarError::ComputationFailed {
371 operation: "peer",
372 detail: "non-finite coefficient (singular penalized system)".into(),
373 });
374 }
375
376 let effective_df = compute_peer_trace_hat(&wtw, &q, lambda, m, n);
378
379 let fitted_values: Vec<f64> = (0..n)
381 .map(|i| y_bar + (0..m).map(|j| wc[(i, j)] * beta[j]).sum::<f64>())
382 .collect();
383
384 Ok(PeerResult {
385 beta,
386 intercept: y_bar,
387 w_bar,
388 fitted_values,
389 effective_df,
390 lambda,
391 penalty_type: config.penalty.clone(),
392 gcv: gcv_score,
393 lambda_method,
394 })
395}
396
397pub fn lpeer(
457 data: &FdMatrix,
458 y: &[f64],
459 argvals: &[f64],
460 subject_map: &[usize],
461 config: &PeerConfig,
462) -> Result<LpeerResult, FdarError> {
463 let (n, m) = data.shape();
464
465 if n < 2 {
467 return Err(FdarError::InvalidDimension {
468 parameter: "data",
469 expected: "at least 2 observations".to_string(),
470 actual: format!("{n} rows"),
471 });
472 }
473 if m < 3 {
474 return Err(FdarError::InvalidDimension {
475 parameter: "data",
476 expected: "at least 3 evaluation points (m >= 3)".to_string(),
477 actual: format!("{m} columns"),
478 });
479 }
480 if argvals.len() != m {
481 return Err(FdarError::InvalidDimension {
482 parameter: "argvals",
483 expected: format!("{m}"),
484 actual: format!("{}", argvals.len()),
485 });
486 }
487 if y.len() != n {
488 return Err(FdarError::InvalidDimension {
489 parameter: "y",
490 expected: format!("{n}"),
491 actual: format!("{}", y.len()),
492 });
493 }
494 if subject_map.len() != n {
495 return Err(FdarError::InvalidDimension {
496 parameter: "subject_map",
497 expected: format!("{n}"),
498 actual: format!("{}", subject_map.len()),
499 });
500 }
501 if y.iter().any(|v| !v.is_finite()) {
502 return Err(FdarError::InvalidParameter {
503 parameter: "y",
504 message: "response contains non-finite values (NaN/Inf)".to_string(),
505 });
506 }
507 if argvals.iter().any(|v| !v.is_finite()) {
508 return Err(FdarError::InvalidParameter {
509 parameter: "argvals",
510 message: "argvals contains non-finite values (NaN/Inf)".to_string(),
511 });
512 }
513 if argvals.windows(2).any(|w| w[1] <= w[0]) {
514 return Err(FdarError::InvalidParameter {
515 parameter: "argvals",
516 message: "argvals must be strictly increasing".to_string(),
517 });
518 }
519
520 let (sm_dense, n_subjects) = crate::famm::build_subject_map(subject_map);
522 if n_subjects < 2 {
523 return Err(FdarError::InvalidParameter {
524 parameter: "subject_map",
525 message: "at least 2 distinct subjects required for lpeer".to_string(),
526 });
527 }
528
529 let w = simpsons_weights(argvals);
531 let mut wmat = FdMatrix::zeros(n, m);
532 for i in 0..n {
533 for j in 0..m {
534 wmat[(i, j)] = data[(i, j)] * w[j];
535 }
536 }
537
538 let y_bar: f64 = y.iter().sum::<f64>() / n as f64;
539 let yc: Vec<f64> = y.iter().map(|&yi| yi - y_bar).collect();
540
541 let w_bar: Vec<f64> = (0..m)
543 .map(|j| (0..n).map(|i| wmat[(i, j)]).sum::<f64>() / n as f64)
544 .collect();
545
546 let mut wc = FdMatrix::zeros(n, m);
547 for i in 0..n {
548 for j in 0..m {
549 wc[(i, j)] = wmat[(i, j)] - w_bar[j];
550 }
551 }
552
553 let q = build_q(m, &config.penalty)?;
555
556 let mut wtw = vec![0.0_f64; m * m];
557 for j in 0..m {
558 for k in j..m {
559 let s: f64 = (0..n).map(|i| wc[(i, j)] * wc[(i, k)]).sum();
560 wtw[j * m + k] = s;
561 wtw[k * m + j] = s;
562 }
563 }
564 let wty: Vec<f64> = (0..m)
565 .map(|j| (0..n).map(|i| wc[(i, j)] * yc[i]).sum())
566 .collect();
567
568 let (lambda, gcv_score, lambda_method) = match &config.lambda {
569 LambdaChoice::Fixed(lam) => (*lam, None, LambdaMethod::Fixed),
570 LambdaChoice::Gcv => {
571 let (lam, g) = select_lambda_gcv_peer(&wc, &yc, &wtw, &wty, &q, m, n);
572 (lam, Some(g), LambdaMethod::Gcv)
573 }
574 LambdaChoice::Reml => {
575 let lam = select_lambda_reml_peer(&wc, &yc, &q, m, n);
576 (lam, None, LambdaMethod::Reml)
577 }
578 };
579
580 let ncomp = (n - 1).min(m).min(10);
582 let fpca = crate::regression::fdata_to_pc_1d(data, ncomp, argvals)?;
583 let scores = &fpca.scores; let result =
589 crate::famm::fit_scalar_mixed_model(&yc, &sm_dense, n_subjects, Some(scores), ncomp);
590
591 if result.gamma.iter().any(|v| !v.is_finite()) {
593 return Err(FdarError::ComputationFailed {
594 operation: "lpeer",
595 detail: "non-finite gamma from mixed model".into(),
596 });
597 }
598
599 let su = result.sigma2_u;
607 let se = result.sigma2_eps;
608 let mut gcnt = vec![0usize; n_subjects];
610 for &g in &sm_dense {
611 gcnt[g] += 1;
612 }
613 let sigma_inv = |v: &[f64]| -> Vec<f64> {
615 let mut gsum = vec![0.0_f64; n_subjects];
616 for i in 0..n {
617 gsum[sm_dense[i]] += v[i];
618 }
619 (0..n)
620 .map(|i| {
621 let g = sm_dense[i];
622 let shrink = su / (se + gcnt[g] as f64 * su);
623 (v[i] - shrink * gsum[g]) / se
624 })
625 .collect::<Vec<f64>>()
626 };
627 let mut sinv_s = vec![0.0_f64; n * ncomp];
629 for k in 0..ncomp {
630 let col: Vec<f64> = (0..n).map(|i| scores[(i, k)]).collect();
631 let sc = sigma_inv(&col);
632 for i in 0..n {
633 sinv_s[i * ncomp + k] = sc[i];
634 }
635 }
636 let sinv_yc = sigma_inv(&yc);
638 let mut a_mat = vec![0.0_f64; ncomp * ncomp];
639 let mut b_vec = vec![0.0_f64; ncomp];
640 for k in 0..ncomp {
641 for l in 0..ncomp {
642 a_mat[k * ncomp + l] = (0..n).map(|i| scores[(i, k)] * sinv_s[i * ncomp + l]).sum();
643 }
644 b_vec[k] = (0..n).map(|i| scores[(i, k)] * sinv_yc[i]).sum();
645 }
646 let mut q_phi = vec![0.0_f64; m * ncomp];
649 for i in 0..m {
650 for l in 0..ncomp {
651 q_phi[i * ncomp + l] = (0..m)
652 .map(|jj| q[i * m + jj] * fpca.rotation[(jj, l)])
653 .sum();
654 }
655 }
656 for k in 0..ncomp {
658 for l in 0..ncomp {
659 let p_kl: f64 = (0..m)
660 .map(|i| fpca.rotation[(i, k)] * q_phi[i * ncomp + l])
661 .sum();
662 a_mat[k * ncomp + l] += lambda * p_kl;
663 }
664 }
665 let gamma = cholesky_solve(&a_mat, &b_vec, ncomp).unwrap_or_else(|_| result.gamma.clone());
667
668 let beta: Vec<f64> = (0..m)
670 .map(|j| {
671 (0..ncomp.min(gamma.len()))
672 .map(|k| gamma[k] * fpca.rotation[(j, k)])
673 .sum()
674 })
675 .collect();
676
677 if beta.iter().any(|v| !v.is_finite()) {
678 return Err(FdarError::ComputationFailed {
679 operation: "lpeer",
680 detail: "non-finite beta after back-projection".into(),
681 });
682 }
683
684 let base = y_bar - w_bar.iter().zip(&beta).map(|(wb, b)| wb * b).sum::<f64>();
686 let fitted_values: Vec<f64> = (0..n)
687 .map(|i| base + (0..m).map(|j| data[(i, j)] * w[j] * beta[j]).sum::<f64>())
688 .collect();
689
690 Ok(LpeerResult {
691 beta,
692 intercept: y_bar,
693 w_bar,
694 fitted_values,
695 sigma2_subject: result.sigma2_u,
696 sigma2_resid: result.sigma2_eps,
697 n_subjects,
698 lambda,
699 penalty_type: config.penalty.clone(),
700 gcv: gcv_score,
701 lambda_method,
702 })
703}
704
705impl LpeerResult {
706 pub fn predict(&self, new_data: &FdMatrix, argvals: &[f64]) -> Result<Vec<f64>, FdarError> {
720 peer_predict_core(&self.beta, self.intercept, &self.w_bar, new_data, argvals)
721 }
722}
723
724fn build_q(m: usize, penalty: &PeerPenalty) -> Result<Vec<f64>, FdarError> {
730 match penalty {
731 PeerPenalty::Ridge => {
732 let mut q = vec![0.0_f64; m * m];
733 for i in 0..m {
734 q[i * m + i] = 1.0;
735 }
736 Ok(q)
737 }
738 PeerPenalty::Difference { order: 2 } => Ok(penalty_matrix(m)),
739 PeerPenalty::Difference { order } => Err(FdarError::InvalidParameter {
740 parameter: "penalty",
741 message: format!(
742 "Difference order {order} unsupported; only order 2 is available in this release"
743 ),
744 }),
745 PeerPenalty::Decree(q_raw, p_q) => {
746 if *p_q != m || q_raw.len() != m * m {
747 return Err(FdarError::InvalidDimension {
748 parameter: "penalty",
749 expected: format!("{m}x{m} ({} elems)", m * m),
750 actual: format!("{p_q}x{p_q} ({} elems)", q_raw.len()),
751 });
752 }
753 Ok(q_raw.clone())
754 }
755 }
756}
757
758fn compute_peer_trace_hat(wtw: &[f64], q: &[f64], lambda: f64, m: usize, n: usize) -> f64 {
764 let mut a = vec![0.0_f64; m * m];
765 for i in 0..m * m {
766 a[i] = wtw[i] + lambda * q[i];
767 }
768 let Ok(l) = cholesky_factor(&a, m) else {
769 return m as f64; };
771 let mut trace = 0.0_f64;
772 for j in 0..m {
773 let col: Vec<f64> = (0..m).map(|i| wtw[i * m + j]).collect();
775 let z = cholesky_forward_back(&l, &col, m);
776 trace += z[j];
777 }
778 trace.min(n as f64)
779}
780
781fn gcv_lambda_grid() -> Vec<f64> {
786 (0..40)
787 .map(|i| 10.0_f64.powf(-6.0 + 10.0 * i as f64 / 39.0))
788 .collect()
789}
790
791fn select_lambda_gcv_peer(
804 wc: &FdMatrix,
805 yc: &[f64],
806 wtw: &[f64],
807 wty: &[f64],
808 q: &[f64],
809 m: usize,
810 n: usize,
811) -> (f64, f64) {
812 let grid = gcv_lambda_grid();
813 let mut best_lam = grid[0];
814 let mut best_gcv = f64::INFINITY;
815
816 for &lam in &grid {
817 let mut a = vec![0.0_f64; m * m];
819 for i in 0..m * m {
820 a[i] = wtw[i] + lam * q[i];
821 }
822 let Ok(beta) = cholesky_solve(&a, wty, m) else {
824 continue;
825 };
826 let rss: f64 = (0..n)
828 .map(|i| {
829 let yhat: f64 = (0..m).map(|j| wc[(i, j)] * beta[j]).sum();
830 (yc[i] - yhat).powi(2)
831 })
832 .sum();
833 let trh = compute_peer_trace_hat(wtw, q, lam, m, n);
835 let denom = n as f64 - trh;
837 if denom <= 0.0 {
838 continue;
839 }
840 let gcv = n as f64 * rss / (denom * denom);
841 if gcv < best_gcv {
843 best_gcv = gcv;
844 best_lam = lam;
845 }
846 }
847 (best_lam, best_gcv)
848}
849
850fn select_lambda_reml_peer(wc: &FdMatrix, yc: &[f64], q: &[f64], m: usize, n: usize) -> f64 {
866 let q_mat = DMatrix::from_row_slice(m, m, q);
868 let eigen = q_mat.symmetric_eigen();
869
870 let mut idx_sorted: Vec<usize> = (0..m).collect();
871 idx_sorted.sort_by(|&a, &b| {
872 eigen.eigenvalues[a]
873 .partial_cmp(&eigen.eigenvalues[b])
874 .unwrap_or(std::cmp::Ordering::Equal)
875 });
876
877 let max_ev = idx_sorted
878 .iter()
879 .map(|&i| eigen.eigenvalues[i].abs())
880 .fold(0.0_f64, f64::max);
881 let tol = 1e-8 * max_ev.max(1.0);
882
883 let null_idx: Vec<usize> = idx_sorted
884 .iter()
885 .copied()
886 .filter(|&i| eigen.eigenvalues[i].abs() < tol)
887 .collect();
888 let range_idx: Vec<usize> = idx_sorted
889 .iter()
890 .copied()
891 .filter(|&i| eigen.eigenvalues[i].abs() >= tol)
892 .collect();
893 let s = null_idx.len();
894 let r = range_idx.len();
895
896 if r == 0 {
898 return 1e-4; }
900
901 let mut z_null = vec![0.0_f64; n * s];
905 let mut z_range = vec![0.0_f64; n * r];
906 for i in 0..n {
907 for (col, &ev_idx) in null_idx.iter().enumerate() {
908 let mut val = 0.0;
909 for j in 0..m {
910 val += wc[(i, j)] * eigen.eigenvectors[(j, ev_idx)];
911 }
912 z_null[i * s + col] = val;
913 }
914 for (col, &ev_idx) in range_idx.iter().enumerate() {
915 let mut val = 0.0;
916 for j in 0..m {
917 val += wc[(i, j)] * eigen.eigenvectors[(j, ev_idx)];
918 }
919 z_range[i * r + col] = val;
920 }
921 }
922
923 let mut ztz_range = vec![0.0_f64; r * r];
925 for a in 0..r {
926 for b in a..r {
927 let s_val: f64 = (0..n)
928 .map(|i| z_range[i * r + a] * z_range[i * r + b])
929 .sum();
930 ztz_range[a * r + b] = s_val;
931 ztz_range[b * r + a] = s_val;
932 }
933 }
934
935 let y_mean = yc.iter().sum::<f64>() / n as f64;
937 let y_var = yc.iter().map(|&v| (v - y_mean).powi(2)).sum::<f64>() / (n - 1).max(1) as f64;
938 let mut sigma2_e = y_var.max(1e-12);
939 let mut sigma2_u = (sigma2_e * 0.1).max(1e-12);
940
941 let mut alpha = vec![0.0_f64; s];
943 if s > 0 {
944 let mut ztz_null = vec![0.0_f64; s * s];
946 for a in 0..s {
947 for b in a..s {
948 let sv: f64 = (0..n).map(|i| z_null[i * s + a] * z_null[i * s + b]).sum();
949 ztz_null[a * s + b] = sv;
950 ztz_null[b * s + a] = sv;
951 }
952 }
953 for diag in 0..s {
954 ztz_null[diag * s + diag] += 1e-10;
955 }
956 let zty_null: Vec<f64> = (0..s)
958 .map(|col| (0..n).map(|i| z_null[i * s + col] * yc[i]).sum())
959 .collect();
960 if let Ok(a_init) = cholesky_solve(&ztz_null, &zty_null, s) {
961 alpha = a_init;
962 }
963 }
964
965 for _iter in 0..100 {
967 let su_old = sigma2_u;
968 let se_old = sigma2_e;
969
970 let r_alpha: Vec<f64> = if s > 0 {
972 (0..n)
973 .map(|i| {
974 let za: f64 = (0..s).map(|col| z_null[i * s + col] * alpha[col]).sum();
975 yc[i] - za
976 })
977 .collect()
978 } else {
979 yc.to_vec()
980 };
981
982 let mut big_m = vec![0.0_f64; r * r];
984 for idx in 0..r * r {
985 big_m[idx] = ztz_range[idx] / sigma2_e;
986 }
987 for diag in 0..r {
988 big_m[diag * r + diag] += 1.0 / sigma2_u;
989 }
990
991 let l_m = match cholesky_factor(&big_m, r) {
996 Ok(l) => l,
997 Err(_) => {
998 sigma2_u = su_old;
999 sigma2_e = se_old;
1000 break;
1001 }
1002 };
1003
1004 let mut sigma_b = vec![0.0_f64; r * r];
1006 let mut trace_sigma_b = 0.0_f64;
1007 for col in 0..r {
1008 let mut e_col = vec![0.0_f64; r];
1009 e_col[col] = 1.0;
1010 let sol = cholesky_forward_back(&l_m, &e_col, r);
1011 for row in 0..r {
1012 sigma_b[row * r + col] = sol[row];
1013 }
1014 trace_sigma_b += sol[col];
1015 }
1016
1017 let ztr: Vec<f64> = (0..r)
1019 .map(|col| (0..n).map(|i| z_range[i * r + col] * r_alpha[i]).sum())
1020 .collect();
1021
1022 let b_hat: Vec<f64> = (0..r)
1024 .map(|row| {
1025 (0..r)
1026 .map(|col| sigma_b[row * r + col] * ztr[col])
1027 .sum::<f64>()
1028 / sigma2_e
1029 })
1030 .collect();
1031
1032 let b_sq: f64 = b_hat.iter().map(|&v| v * v).sum();
1034 let sigma2_u_new = (b_sq + trace_sigma_b) / r as f64;
1035
1036 let resid: Vec<f64> = (0..n)
1038 .map(|i| {
1039 let zb: f64 = (0..r).map(|col| z_range[i * r + col] * b_hat[col]).sum();
1040 r_alpha[i] - zb
1041 })
1042 .collect();
1043
1044 let tr_zsz: f64 = (0..r)
1046 .map(|a| {
1047 (0..r)
1048 .map(|j| sigma_b[a * r + j] * ztz_range[j * r + a])
1049 .sum::<f64>()
1050 })
1051 .sum();
1052
1053 let resid_sq: f64 = resid.iter().map(|&v| v * v).sum();
1054 let sigma2_e_new = (resid_sq + tr_zsz) / n as f64;
1055
1056 let sigma2_u_c = sigma2_u_new.max(1e-12);
1060 let sigma2_e_c = sigma2_e_new.max(1e-12);
1061
1062 if s > 0 {
1066 let ratio = sigma2_e_c / sigma2_u_c;
1068 let mut k_mat = ztz_range.clone();
1069 for diag in 0..r {
1070 k_mat[diag * r + diag] += ratio;
1071 }
1072 if let Ok(l_k) = cholesky_factor(&k_mat, r) {
1074 let mut gls_lhs = vec![0.0_f64; s * s]; let mut gls_rhs = vec![0.0_f64; s]; let zr_yc: Vec<f64> = (0..r)
1081 .map(|col| (0..n).map(|i| z_range[i * r + col] * yc[i]).sum())
1082 .collect();
1083 let kinv_zr_yc = cholesky_forward_back(&l_k, &zr_yc, r);
1085 let sinv_yc: Vec<f64> = (0..n)
1087 .map(|i| {
1088 let zk: f64 = (0..r)
1089 .map(|col| z_range[i * r + col] * kinv_zr_yc[col])
1090 .sum();
1091 (yc[i] - zk) / sigma2_e_c
1092 })
1093 .collect();
1094
1095 for col_null in 0..s {
1096 let zr_zn: Vec<f64> = (0..r)
1098 .map(|col| {
1099 (0..n)
1100 .map(|i| z_range[i * r + col] * z_null[i * s + col_null])
1101 .sum()
1102 })
1103 .collect();
1104 let kinv_zr_zn = cholesky_forward_back(&l_k, &zr_zn, r);
1106 let sinv_zn: Vec<f64> = (0..n)
1108 .map(|i| {
1109 let zk: f64 = (0..r)
1110 .map(|col| z_range[i * r + col] * kinv_zr_zn[col])
1111 .sum();
1112 (z_null[i * s + col_null] - zk) / sigma2_e_c
1113 })
1114 .collect();
1115
1116 for row_null in 0..s {
1118 let v: f64 = (0..n).map(|i| z_null[i * s + row_null] * sinv_zn[i]).sum();
1119 gls_lhs[row_null * s + col_null] = v;
1120 }
1121 let rhs_val: f64 = (0..n).map(|i| z_null[i * s + col_null] * sinv_yc[i]).sum();
1123 gls_rhs[col_null] = rhs_val;
1124 }
1125
1126 for diag in 0..s {
1128 gls_lhs[diag * s + diag] += 1e-10;
1129 }
1130
1131 if let Ok(alpha_new) = cholesky_solve(&gls_lhs, &gls_rhs, s) {
1132 alpha = alpha_new;
1133 }
1134 }
1135 }
1136
1137 sigma2_u = sigma2_u_c;
1139 sigma2_e = sigma2_e_c;
1140
1141 let delta = (sigma2_u_new - su_old).abs() + (sigma2_e_new - se_old).abs();
1146 if delta < 1e-8 * (su_old + se_old) {
1147 break;
1148 }
1149 }
1150
1151 (sigma2_e / sigma2_u).max(1e-15)
1153}
1154
1155fn peer_predict_core(
1168 beta: &[f64],
1169 intercept: f64,
1170 w_bar: &[f64],
1171 new_data: &FdMatrix,
1172 argvals: &[f64],
1173) -> Result<Vec<f64>, FdarError> {
1174 let (n_new, m_new) = new_data.shape();
1175 let m = beta.len();
1176 if m_new != m {
1177 return Err(FdarError::InvalidDimension {
1178 parameter: "new_data",
1179 expected: format!("{m} columns (training grid length)"),
1180 actual: format!("{m_new}"),
1181 });
1182 }
1183 if argvals.len() != m {
1184 return Err(FdarError::InvalidDimension {
1185 parameter: "argvals",
1186 expected: format!("{m}"),
1187 actual: format!("{}", argvals.len()),
1188 });
1189 }
1190 if argvals.windows(2).any(|w| w[1] <= w[0]) {
1193 return Err(FdarError::InvalidParameter {
1194 parameter: "argvals",
1195 message: "argvals must be strictly increasing".to_string(),
1196 });
1197 }
1198 debug_assert_eq!(w_bar.len(), m, "w_bar length must equal beta length");
1201 let w = simpsons_weights(argvals);
1202 let base: f64 = intercept - w_bar.iter().zip(beta).map(|(wb, b)| wb * b).sum::<f64>();
1203 let preds: Vec<f64> = (0..n_new)
1204 .map(|i| {
1205 base + (0..m)
1206 .map(|j| new_data[(i, j)] * w[j] * beta[j])
1207 .sum::<f64>()
1208 })
1209 .collect();
1210 Ok(preds)
1211}
1212
1213impl PeerResult {
1214 pub fn predict(&self, new_data: &FdMatrix, argvals: &[f64]) -> Result<Vec<f64>, FdarError> {
1229 peer_predict_core(&self.beta, self.intercept, &self.w_bar, new_data, argvals)
1230 }
1231}
1232
1233#[cfg(test)]
1238mod tests {
1239 use super::*;
1240 use crate::test_helpers::uniform_grid;
1241
1242 fn hash_unit(k: u64) -> f64 {
1245 let mut z = k
1246 .wrapping_add(0x9E37_79B9_7F4A_7C15)
1247 .wrapping_mul(0xBF58_476D_1CE4_E5B9);
1248 z = (z ^ (z >> 30)).wrapping_mul(0x94D0_49BB_1331_11EB);
1249 z ^= z >> 31;
1250 ((z >> 11) as f64 / (1_u64 << 53) as f64) * 2.0 - 1.0
1251 }
1252
1253 fn make_fixture() -> (FdMatrix, Vec<f64>, Vec<f64>, Vec<f64>) {
1263 let (n, m) = (200_usize, 40_usize);
1264 let t = uniform_grid(m);
1265 let true_beta: Vec<f64> = t
1266 .iter()
1267 .map(|&ti| (std::f64::consts::PI * ti).sin())
1268 .collect();
1269 let w = simpsons_weights(&t);
1270
1271 let mut data = FdMatrix::zeros(n, m);
1272 let mut y = vec![0.0_f64; n];
1273 for i in 0..n {
1274 for j in 0..m {
1275 let xi = hash_unit((i * m + j) as u64);
1276 data[(i, j)] = xi;
1277 y[i] += xi * true_beta[j] * w[j];
1278 }
1279 y[i] += 0.005 * hash_unit(1_000_000 + i as u64);
1281 }
1282 (data, y, t, true_beta)
1283 }
1284
1285 #[test]
1290 fn test_peer_difference_beta_recovery() {
1291 let (data, y, t, true_beta) = make_fixture();
1292 let config = PeerConfig {
1293 penalty: PeerPenalty::Difference { order: 2 },
1294 lambda: LambdaChoice::Fixed(1e-4),
1295 };
1296 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1297
1298 let max_err = result
1299 .beta
1300 .iter()
1301 .zip(true_beta.iter())
1302 .map(|(a, b)| (a - b).abs())
1303 .fold(0.0_f64, f64::max);
1304 assert!(
1305 max_err < 0.1,
1306 "Difference beta recovery error: {max_err} >= 0.1"
1307 );
1308 assert!(
1309 result.fitted_values.iter().all(|v| v.is_finite()),
1310 "fitted_values contain non-finite values"
1311 );
1312 }
1313
1314 #[test]
1315 fn test_peer_result_shape() {
1316 let (data, y, t, _) = make_fixture();
1317 let m = t.len();
1318 let n = y.len();
1319 let config = PeerConfig {
1320 penalty: PeerPenalty::Difference { order: 2 },
1321 lambda: LambdaChoice::Fixed(1e-4),
1322 };
1323 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1324
1325 assert_eq!(result.beta.len(), m, "beta length should be m={m}");
1326 assert_eq!(
1327 result.fitted_values.len(),
1328 n,
1329 "fitted_values length should be n={n}"
1330 );
1331 assert!(
1332 result.effective_df.is_finite() && result.effective_df > 0.0,
1333 "effective_df={} must be finite and positive",
1334 result.effective_df
1335 );
1336 assert!(
1337 (result.lambda - 1e-4).abs() < 1e-15,
1338 "lambda should be 1e-4, got {}",
1339 result.lambda
1340 );
1341 assert_eq!(
1342 result.penalty_type,
1343 PeerPenalty::Difference { order: 2 },
1344 "penalty_type should be Difference{{order:2}}"
1345 );
1346
1347 let y_bar = y.iter().sum::<f64>() / n as f64;
1348 assert!(
1349 (result.intercept - y_bar).abs() < 1e-9,
1350 "intercept={} should equal mean(y)={y_bar} within 1e-9",
1351 result.intercept
1352 );
1353 }
1354
1355 #[test]
1360 fn test_peer_ridge_fits() {
1361 let (data, y, t, _) = make_fixture();
1362 let m = t.len();
1363 let config = PeerConfig {
1364 penalty: PeerPenalty::Ridge,
1365 lambda: LambdaChoice::Fixed(1e-4),
1366 };
1367 let result = peer(&data, &y, &t, &config).expect("peer() with Ridge should succeed");
1368 assert_eq!(result.beta.len(), m);
1369 assert!(
1370 result.beta.iter().all(|v| v.is_finite()),
1371 "Ridge beta contains non-finite values"
1372 );
1373 assert!(
1374 result.fitted_values.iter().all(|v| v.is_finite()),
1375 "Ridge fitted_values contain non-finite values"
1376 );
1377 }
1378
1379 #[test]
1380 fn test_peer_decree_fits() {
1381 let (data, y, t, _) = make_fixture();
1382 let m = t.len();
1383 let q_flat = penalty_matrix(m);
1385 let config = PeerConfig {
1386 penalty: PeerPenalty::Decree(q_flat, m),
1387 lambda: LambdaChoice::Fixed(1e-4),
1388 };
1389 let result = peer(&data, &y, &t, &config).expect("peer() with Decree should succeed");
1390 assert_eq!(result.beta.len(), m);
1391 assert!(
1392 result.beta.iter().all(|v| v.is_finite()),
1393 "Decree beta contains non-finite values"
1394 );
1395 assert!(
1396 result.fitted_values.iter().all(|v| v.is_finite()),
1397 "Decree fitted_values contain non-finite values"
1398 );
1399 }
1400
1401 #[test]
1402 fn test_peer_difference_order_rejected() {
1403 let (data, y, t, _) = make_fixture();
1404 let config = PeerConfig {
1405 penalty: PeerPenalty::Difference { order: 3 },
1406 lambda: LambdaChoice::Fixed(1e-4),
1407 };
1408 let err = peer(&data, &y, &t, &config).expect_err("order 3 should be rejected");
1409 assert!(
1410 matches!(err, FdarError::InvalidParameter { .. }),
1411 "expected InvalidParameter, got {err:?}"
1412 );
1413 }
1414
1415 #[test]
1420 fn test_peer_decree_distinct_from_roughness() {
1421 let (data, y, t, _) = make_fixture();
1422 let m = t.len();
1423 let lambda = LambdaChoice::Fixed(1.0);
1425
1426 let config_diff = PeerConfig {
1428 penalty: PeerPenalty::Difference { order: 2 },
1429 lambda: lambda.clone(),
1430 };
1431 let res_diff = peer(&data, &y, &t, &config_diff).expect("Difference{2} fit should succeed");
1432
1433 let b = m / 2;
1437 let mut q_partition = vec![0.0_f64; m * m];
1438 for i in 0..m.saturating_sub(2) {
1439 if i + 1 == b || i + 2 == b || i == b {
1441 continue;
1442 }
1443 let coeffs = [(i, 1.0_f64), (i + 1, -2.0), (i + 2, 1.0)];
1444 for &(r, cr) in &coeffs {
1445 for &(c, cc) in &coeffs {
1446 q_partition[r * m + c] += cr * cc;
1447 }
1448 }
1449 }
1450 for i in 0..m {
1452 q_partition[i * m + i] += 1e-6;
1453 }
1454
1455 let config_dec = PeerConfig {
1456 penalty: PeerPenalty::Decree(q_partition, m),
1457 lambda,
1458 };
1459 let res_dec =
1460 peer(&data, &y, &t, &config_dec).expect("Decree partition fit should succeed");
1461
1462 assert!(
1464 res_diff.beta.iter().all(|v| v.is_finite()),
1465 "Difference beta contains non-finite"
1466 );
1467 assert!(
1468 res_dec.beta.iter().all(|v| v.is_finite()),
1469 "Decree beta contains non-finite"
1470 );
1471
1472 let max_diff = res_diff
1474 .beta
1475 .iter()
1476 .zip(res_dec.beta.iter())
1477 .map(|(a, b)| (a - b).abs())
1478 .fold(0.0_f64, f64::max);
1479 assert!(
1480 max_diff > 1e-3,
1481 "Decree partition Q should yield a β(t) distinct from Difference{{2}} \
1482 by > 1e-3; got max_diff={max_diff}"
1483 );
1484 }
1485
1486 #[test]
1491 fn test_peer_decree_wrong_dim() {
1492 let (data, y, t, _) = make_fixture();
1493 let m = t.len();
1494 let config = PeerConfig {
1496 penalty: PeerPenalty::Decree(vec![0.0; (m - 1) * (m - 1)], m - 1),
1497 lambda: LambdaChoice::Fixed(1.0),
1498 };
1499 let err = peer(&data, &y, &t, &config).expect_err("wrong-dim Decree should return Err");
1500 assert!(
1501 matches!(err, FdarError::InvalidDimension { .. }),
1502 "expected InvalidDimension, got {err:?}"
1503 );
1504 }
1505
1506 #[test]
1507 fn test_peer_argvals_mismatch() {
1508 let (data, y, t, _) = make_fixture();
1509 let m = t.len();
1510 let argvals_bad: Vec<f64> = (0..=m).map(|i| i as f64 / m as f64).collect();
1512 let config = PeerConfig::default();
1513 let err = peer(&data, &y, &argvals_bad, &config)
1514 .expect_err("mismatched argvals should return Err");
1515 assert!(
1516 matches!(err, FdarError::InvalidDimension { .. }),
1517 "expected InvalidDimension, got {err:?}"
1518 );
1519 }
1520
1521 #[test]
1522 fn test_peer_no_nan_all_families() {
1523 let (data, y, t, _) = make_fixture();
1524 let m = t.len();
1525 let lambda = LambdaChoice::Fixed(1.0);
1526
1527 let config_ridge = PeerConfig {
1529 penalty: PeerPenalty::Ridge,
1530 lambda: lambda.clone(),
1531 };
1532 let res_ridge = peer(&data, &y, &t, &config_ridge).expect("Ridge should succeed");
1533 assert!(
1534 res_ridge.beta.iter().all(|v| v.is_finite()),
1535 "Ridge beta contains NaN/Inf"
1536 );
1537 assert!(
1538 res_ridge.fitted_values.iter().all(|v| v.is_finite()),
1539 "Ridge fitted_values contains NaN/Inf"
1540 );
1541 assert!(
1542 res_ridge.effective_df.is_finite(),
1543 "Ridge effective_df is not finite"
1544 );
1545
1546 let config_diff = PeerConfig {
1548 penalty: PeerPenalty::Difference { order: 2 },
1549 lambda: lambda.clone(),
1550 };
1551 let res_diff = peer(&data, &y, &t, &config_diff).expect("Difference{2} should succeed");
1552 assert!(
1553 res_diff.beta.iter().all(|v| v.is_finite()),
1554 "Difference beta contains NaN/Inf"
1555 );
1556 assert!(
1557 res_diff.fitted_values.iter().all(|v| v.is_finite()),
1558 "Difference fitted_values contains NaN/Inf"
1559 );
1560 assert!(
1561 res_diff.effective_df.is_finite(),
1562 "Difference effective_df is not finite"
1563 );
1564
1565 let q_flat = penalty_matrix(m);
1567 let config_dec = PeerConfig {
1568 penalty: PeerPenalty::Decree(q_flat, m),
1569 lambda,
1570 };
1571 let res_dec = peer(&data, &y, &t, &config_dec).expect("Decree should succeed");
1572 assert!(
1573 res_dec.beta.iter().all(|v| v.is_finite()),
1574 "Decree beta contains NaN/Inf"
1575 );
1576 assert!(
1577 res_dec.fitted_values.iter().all(|v| v.is_finite()),
1578 "Decree fitted_values contains NaN/Inf"
1579 );
1580 assert!(
1581 res_dec.effective_df.is_finite(),
1582 "Decree effective_df is not finite"
1583 );
1584 }
1585
1586 #[test]
1591 fn test_peer_stores_w_bar_for_prediction() {
1592 let (data, y, t, _) = make_fixture();
1595 let (n, m) = data.shape();
1596 let w = simpsons_weights(&t);
1597 let config = PeerConfig {
1598 penalty: PeerPenalty::Difference { order: 2 },
1599 lambda: LambdaChoice::Fixed(1e-4),
1600 };
1601 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1602
1603 assert_eq!(result.w_bar.len(), m, "w_bar length must equal m");
1604
1605 let base = result.intercept
1606 - result
1607 .w_bar
1608 .iter()
1609 .zip(&result.beta)
1610 .map(|(wb, b)| wb * b)
1611 .sum::<f64>();
1612 for i in 0..n {
1613 let pred = base
1614 + (0..m)
1615 .map(|j| data[(i, j)] * w[j] * result.beta[j])
1616 .sum::<f64>();
1617 assert!(
1618 (pred - result.fitted_values[i]).abs() < 1e-9,
1619 "prediction reconstruction mismatch at row {i}: {pred} vs {}",
1620 result.fitted_values[i]
1621 );
1622 }
1623 }
1624
1625 #[test]
1626 fn test_peer_rejects_non_monotonic_argvals() {
1627 let (data, y, t, _) = make_fixture();
1628 let mut bad = t.clone();
1629 bad.swap(0, 1); let config = PeerConfig::default();
1631 let res = peer(&data, &y, &bad, &config);
1632 assert!(
1633 matches!(res, Err(FdarError::InvalidParameter { .. })),
1634 "non-monotonic argvals must be rejected, got {res:?}"
1635 );
1636 }
1637
1638 #[test]
1639 fn test_peer_rejects_single_observation() {
1640 let m = 5;
1641 let t = uniform_grid(m);
1642 let mut data = FdMatrix::zeros(1, m);
1643 for j in 0..m {
1644 data[(0, j)] = 1.0 + j as f64;
1645 }
1646 let y = vec![1.0];
1647 let res = peer(&data, &y, &t, &PeerConfig::default());
1648 assert!(
1649 matches!(res, Err(FdarError::InvalidDimension { .. })),
1650 "single observation must be rejected, got {res:?}"
1651 );
1652 }
1653
1654 #[test]
1655 fn test_peer_rejects_non_finite_y() {
1656 let (data, mut y, t, _) = make_fixture();
1657 y[0] = f64::NAN;
1658 let res = peer(&data, &y, &t, &PeerConfig::default());
1659 assert!(
1660 matches!(res, Err(FdarError::InvalidParameter { .. })),
1661 "non-finite y must be rejected, got {res:?}"
1662 );
1663 }
1664
1665 #[test]
1670 fn test_peer_gcv_deterministic() {
1671 let (data, y, t, _) = make_fixture();
1672 let config = PeerConfig {
1673 penalty: PeerPenalty::Difference { order: 2 },
1674 lambda: LambdaChoice::Gcv,
1675 };
1676 let r1 = peer(&data, &y, &t, &config).expect("first GCV call should succeed");
1677 let r2 = peer(&data, &y, &t, &config).expect("second GCV call should succeed");
1678 assert_eq!(
1679 r1.lambda, r2.lambda,
1680 "GCV lambda must be bit-exact across two runs: {} vs {}",
1681 r1.lambda, r2.lambda
1682 );
1683 assert_eq!(r1.lambda_method, LambdaMethod::Gcv);
1684 assert!(r1.gcv.is_some(), "GCV result must have gcv score");
1685 }
1686
1687 #[test]
1688 fn test_peer_gcv_recovers_beta() {
1689 let (data, y, t, true_beta) = make_fixture();
1690 let config = PeerConfig {
1691 penalty: PeerPenalty::Difference { order: 2 },
1692 lambda: LambdaChoice::Gcv,
1693 };
1694 let result = peer(&data, &y, &t, &config).expect("GCV peer() should succeed");
1695
1696 assert!(
1698 result.lambda > 1e-10 && result.lambda < 1e6,
1699 "GCV lambda should be non-degenerate, got {}",
1700 result.lambda
1701 );
1702 assert!(result.gcv.is_some(), "gcv field must be Some for Gcv");
1704 assert_eq!(result.lambda_method, LambdaMethod::Gcv);
1705
1706 let max_err = result
1708 .beta
1709 .iter()
1710 .zip(true_beta.iter())
1711 .map(|(a, b)| (a - b).abs())
1712 .fold(0.0_f64, f64::max);
1713 assert!(max_err < 0.15, "GCV beta recovery error: {max_err} >= 0.15");
1714 }
1715
1716 #[test]
1721 fn test_peer_reml_deterministic() {
1722 let (data, y, t, _) = make_fixture();
1723 let config = PeerConfig {
1724 penalty: PeerPenalty::Difference { order: 2 },
1725 lambda: LambdaChoice::Reml,
1726 };
1727 let r1 = peer(&data, &y, &t, &config).expect("first REML call should succeed");
1728 let r2 = peer(&data, &y, &t, &config).expect("second REML call should succeed");
1729 assert_eq!(
1730 r1.lambda, r2.lambda,
1731 "REML lambda must be bit-exact across two runs: {} vs {}",
1732 r1.lambda, r2.lambda
1733 );
1734 assert_eq!(r1.lambda_method, LambdaMethod::Reml);
1735 assert!(r1.gcv.is_none(), "REML result must have gcv == None");
1736 }
1737
1738 #[test]
1739 fn test_peer_reml_lambda_positive() {
1740 let (data, y, t, _) = make_fixture();
1741 let config = PeerConfig {
1742 penalty: PeerPenalty::Difference { order: 2 },
1743 lambda: LambdaChoice::Reml,
1744 };
1745 let result = peer(&data, &y, &t, &config).expect("REML peer() should succeed");
1746
1747 assert!(
1748 result.lambda > 0.0 && result.lambda.is_finite(),
1749 "REML lambda must be positive and finite, got {}",
1750 result.lambda
1751 );
1752 assert!(result.gcv.is_none(), "gcv must be None for Reml");
1753 assert_eq!(result.lambda_method, LambdaMethod::Reml);
1754 }
1755
1756 #[test]
1757 fn test_peer_reml_gcv_beta_agreement() {
1758 let (data, y, t, true_beta) = make_fixture(); let config_gcv = PeerConfig {
1767 penalty: PeerPenalty::Difference { order: 2 },
1768 lambda: LambdaChoice::Gcv,
1769 };
1770 let config_reml = PeerConfig {
1771 penalty: PeerPenalty::Difference { order: 2 },
1772 lambda: LambdaChoice::Reml,
1773 };
1774
1775 let res_gcv = peer(&data, &y, &t, &config_gcv).expect("GCV fit should succeed");
1776 let res_reml = peer(&data, &y, &t, &config_reml).expect("REML fit should succeed");
1777
1778 let reml_err = res_reml
1780 .beta
1781 .iter()
1782 .zip(true_beta.iter())
1783 .map(|(a, b)| (a - b).abs())
1784 .fold(0.0_f64, f64::max);
1785 assert!(
1786 reml_err < 0.15,
1787 "REML beta recovery error: {reml_err} >= 0.15"
1788 );
1789
1790 let beta_diff = res_gcv
1792 .beta
1793 .iter()
1794 .zip(res_reml.beta.iter())
1795 .map(|(a, b)| (a - b).abs())
1796 .fold(0.0_f64, f64::max);
1797 assert!(
1798 beta_diff < 0.2,
1799 "REML vs GCV beta disagreement: {beta_diff} >= 0.2 (lambdas: gcv={}, reml={})",
1800 res_gcv.lambda,
1801 res_reml.lambda
1802 );
1803 }
1804
1805 #[test]
1811 fn test_peer_predict_self_consistent() {
1812 let (data, y, t, _) = make_fixture();
1814 let config = PeerConfig {
1815 penalty: PeerPenalty::Difference { order: 2 },
1816 lambda: LambdaChoice::Fixed(1e-4),
1817 };
1818 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1819 let preds = result
1820 .predict(&data, &t)
1821 .expect("predict on training data should succeed");
1822 assert_eq!(preds.len(), result.fitted_values.len());
1823 for (i, (p, f)) in preds.iter().zip(&result.fitted_values).enumerate() {
1824 assert!(
1825 (p - f).abs() < 1e-9,
1826 "predict vs fitted_values mismatch at row {i}: {p} vs {f}"
1827 );
1828 }
1829 }
1830
1831 #[test]
1832 fn test_predict_wrong_ncols() {
1833 let (data, y, t, _) = make_fixture();
1835 let (_, m) = data.shape();
1836 let config = PeerConfig {
1837 penalty: PeerPenalty::Ridge,
1838 lambda: LambdaChoice::Fixed(1e-4),
1839 };
1840 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1841
1842 let mut bad_data = FdMatrix::zeros(5, m - 1);
1844 for i in 0..5 {
1845 for j in 0..(m - 1) {
1846 bad_data[(i, j)] = hash_unit((i * m + j) as u64);
1847 }
1848 }
1849 let err = result
1850 .predict(&bad_data, &t[..m - 1])
1851 .expect_err("wrong ncols should return Err");
1852 assert!(
1853 matches!(err, FdarError::InvalidDimension { .. }),
1854 "expected InvalidDimension, got {err:?}"
1855 );
1856 }
1857
1858 #[test]
1859 fn test_predict_new_curves_finite() {
1860 let (data, y, t, _) = make_fixture();
1862 let (n, m) = data.shape();
1863 let config = PeerConfig {
1864 penalty: PeerPenalty::Ridge,
1865 lambda: LambdaChoice::Fixed(1e-4),
1866 };
1867 let result = peer(&data, &y, &t, &config).expect("peer() should succeed");
1868
1869 let mut new_data = FdMatrix::zeros(n, m);
1871 for i in 0..n {
1872 for j in 0..m {
1873 new_data[(i, j)] = hash_unit((100_000 + i * m + j) as u64) + 10.0;
1874 }
1875 }
1876 let preds = result
1877 .predict(&new_data, &t)
1878 .expect("predict on new curves should succeed");
1879 assert!(
1880 preds.iter().all(|v| v.is_finite()),
1881 "predictions on fresh curves contain non-finite values"
1882 );
1883 }
1884
1885 #[test]
1886 fn test_peer_reml_ridge_and_zeroq_edges() {
1887 let (data, y, t, _) = make_fixture();
1888 let m = t.len();
1889
1890 let config_ridge = PeerConfig {
1892 penalty: PeerPenalty::Ridge,
1893 lambda: LambdaChoice::Reml,
1894 };
1895 let res_ridge =
1896 peer(&data, &y, &t, &config_ridge).expect("REML with Ridge (s=0) should not panic");
1897 assert!(
1898 res_ridge.lambda > 0.0 && res_ridge.lambda.is_finite(),
1899 "Ridge REML lambda must be positive finite, got {}",
1900 res_ridge.lambda
1901 );
1902
1903 let zero_q = vec![0.0_f64; m * m];
1905 let config_zero = PeerConfig {
1906 penalty: PeerPenalty::Decree(zero_q, m),
1907 lambda: LambdaChoice::Reml,
1908 };
1909 let res_zero =
1910 peer(&data, &y, &t, &config_zero).expect("REML with zero-Q should not panic");
1911 assert!(
1912 (res_zero.lambda - 1e-4).abs() < 1e-15,
1913 "zero-Q REML fallback lambda should be 1e-4, got {}",
1914 res_zero.lambda
1915 );
1916 }
1917
1918 fn make_lpeer_fixture() -> (FdMatrix, Vec<f64>, Vec<f64>, Vec<usize>, f64) {
1929 let (n_subjects, obs_per, m) = (20_usize, 10_usize, 10_usize);
1930 let n = n_subjects * obs_per;
1931 let t = uniform_grid(m);
1932 let true_beta: Vec<f64> = t
1933 .iter()
1934 .map(|&ti| (std::f64::consts::PI * ti).sin())
1935 .collect();
1936 let w = simpsons_weights(&t);
1937 let sigma2_u_true = 1.0_f64;
1938
1939 let mut data = FdMatrix::zeros(n, m);
1940 let mut y = vec![0.0_f64; n];
1941 let mut subject_map = vec![0_usize; n];
1942
1943 for s in 0..n_subjects {
1944 let u_s = hash_unit(s as u64) * sigma2_u_true.sqrt();
1946 for obs in 0..obs_per {
1947 let i = s * obs_per + obs;
1948 subject_map[i] = s;
1949 for j in 0..m {
1950 let xi = hash_unit((i * m + j) as u64);
1951 data[(i, j)] = xi;
1952 y[i] += xi * true_beta[j] * w[j];
1953 }
1954 y[i] += u_s; y[i] += 0.02 * hash_unit(1_000_000 + i as u64);
1957 }
1958 }
1959 (data, y, t, subject_map, sigma2_u_true)
1960 }
1961
1962 #[test]
1963 fn test_lpeer_variance_non_negative() {
1964 let (data, y, t, subject_map, _) = make_lpeer_fixture();
1965 let config = PeerConfig {
1966 penalty: PeerPenalty::Ridge,
1967 lambda: LambdaChoice::Fixed(1e-3),
1968 };
1969 let fit = lpeer(&data, &y, &t, &subject_map, &config)
1970 .expect("lpeer() should succeed on longitudinal fixture");
1971 assert!(
1972 fit.sigma2_subject >= 0.0,
1973 "sigma2_subject must be non-negative, got {}",
1974 fit.sigma2_subject
1975 );
1976 assert!(
1977 fit.sigma2_resid >= 0.0,
1978 "sigma2_resid must be non-negative, got {}",
1979 fit.sigma2_resid
1980 );
1981 }
1982
1983 #[test]
1984 fn test_lpeer_beta_recovery() {
1985 let (data, y, t, subject_map, _) = make_lpeer_fixture();
1986 let m = t.len();
1987 let true_beta: Vec<f64> = t
1988 .iter()
1989 .map(|&ti| (std::f64::consts::PI * ti).sin())
1990 .collect();
1991 let config = PeerConfig {
1992 penalty: PeerPenalty::Ridge,
1993 lambda: LambdaChoice::Fixed(1e-3),
1994 };
1995 let fit = lpeer(&data, &y, &t, &subject_map, &config)
1996 .expect("lpeer() should succeed on longitudinal fixture");
1997 assert_eq!(fit.beta.len(), m);
1998 let max_err = fit
1999 .beta
2000 .iter()
2001 .zip(true_beta.iter())
2002 .map(|(a, b)| (a - b).abs())
2003 .fold(0.0_f64, f64::max);
2004 assert!(
2005 max_err < 0.5,
2006 "lpeer β(t) recovery error vs sin(πt): {max_err} >= 0.5"
2007 );
2008 }
2009
2010 #[test]
2011 fn test_lpeer_lambda_regularizes() {
2012 let (data, y, t, subject_map, _) = make_lpeer_fixture();
2015 let cfg_light = PeerConfig {
2016 penalty: PeerPenalty::Difference { order: 2 },
2017 lambda: LambdaChoice::Fixed(0.0),
2018 };
2019 let cfg_heavy = PeerConfig {
2020 penalty: PeerPenalty::Difference { order: 2 },
2021 lambda: LambdaChoice::Fixed(1e6),
2022 };
2023 let light = lpeer(&data, &y, &t, &subject_map, &cfg_light).expect("light λ fit");
2024 let heavy = lpeer(&data, &y, &t, &subject_map, &cfg_heavy).expect("heavy λ fit");
2025 let max_diff = light
2026 .beta
2027 .iter()
2028 .zip(heavy.beta.iter())
2029 .map(|(a, b)| (a - b).abs())
2030 .fold(0.0_f64, f64::max);
2031 assert!(
2032 max_diff > 1e-3,
2033 "λ does not regularize β(t): max|β(λ=0) − β(λ=1e6)| = {max_diff} <= 1e-3"
2034 );
2035 assert!(heavy.beta.iter().all(|v| v.is_finite()));
2036 }
2037
2038 #[test]
2039 fn test_lpeer_sigma2_tracks_injection() {
2040 let (data, y, t, subject_map, sigma2_u_true) = make_lpeer_fixture();
2041 let config = PeerConfig {
2042 penalty: PeerPenalty::Ridge,
2043 lambda: LambdaChoice::Fixed(1e-3),
2044 };
2045 let fit = lpeer(&data, &y, &t, &subject_map, &config)
2046 .expect("lpeer() should succeed on longitudinal fixture");
2047 assert!(
2049 fit.sigma2_subject > 0.1 && fit.sigma2_subject < 5.0,
2050 "lpeer sigma2_subject={} should track injected {sigma2_u_true} within (0.1, 5.0)",
2051 fit.sigma2_subject
2052 );
2053 }
2054
2055 #[test]
2056 fn test_lpeer_invalid_subject_map() {
2057 let (data, y, t, _, _) = make_lpeer_fixture();
2058 let n = y.len();
2059 let bad_map: Vec<usize> = (0..n - 1).map(|i| i / 5).collect();
2061 let config = PeerConfig {
2062 penalty: PeerPenalty::Ridge,
2063 lambda: LambdaChoice::Fixed(1e-3),
2064 };
2065 let err = lpeer(&data, &y, &t, &bad_map, &config)
2066 .expect_err("wrong-length subject_map should return Err");
2067 assert!(
2068 matches!(err, FdarError::InvalidDimension { .. }),
2069 "expected InvalidDimension, got {err:?}"
2070 );
2071 }
2072
2073 #[test]
2074 fn test_lpeer_single_subject_rejected() {
2075 let (data, y, t, _, _) = make_lpeer_fixture();
2076 let n = y.len();
2077 let single_map = vec![0_usize; n];
2079 let config = PeerConfig {
2080 penalty: PeerPenalty::Ridge,
2081 lambda: LambdaChoice::Fixed(1e-3),
2082 };
2083 let err = lpeer(&data, &y, &t, &single_map, &config)
2084 .expect_err("single-subject map should be rejected");
2085 assert!(
2086 matches!(err, FdarError::InvalidParameter { .. }),
2087 "expected InvalidParameter, got {err:?}"
2088 );
2089 }
2090
2091 #[test]
2092 fn test_lpeer_predict_self_consistent() {
2093 let (data, y, t, subject_map, _) = make_lpeer_fixture();
2095 let config = PeerConfig {
2096 penalty: PeerPenalty::Ridge,
2097 lambda: LambdaChoice::Fixed(1e-3),
2098 };
2099 let fit = lpeer(&data, &y, &t, &subject_map, &config)
2100 .expect("lpeer() should succeed on longitudinal fixture");
2101 let preds = fit
2102 .predict(&data, &t)
2103 .expect("LpeerResult::predict on training data should succeed");
2104 assert_eq!(preds.len(), fit.fitted_values.len());
2105 for (i, (p, f)) in preds.iter().zip(&fit.fitted_values).enumerate() {
2106 assert!(
2107 (p - f).abs() < 1e-9,
2108 "LpeerResult::predict vs fitted_values mismatch at row {i}: {p} vs {f}"
2109 );
2110 }
2111 }
2112
2113 #[test]
2118 fn test_crate_root_exports_compile() {
2119 use crate::error::FdarError;
2125
2126 let _peer_fn: fn(&FdMatrix, &[f64], &[f64], &PeerConfig) -> Result<PeerResult, FdarError> =
2128 crate::peer::peer;
2129 let _lpeer_fn: fn(
2130 &FdMatrix,
2131 &[f64],
2132 &[f64],
2133 &[usize],
2134 &PeerConfig,
2135 ) -> Result<LpeerResult, FdarError> = crate::peer::lpeer;
2136
2137 let _lc = LambdaChoice::Fixed(1.0);
2139 let _lm = LambdaMethod::Fixed;
2140 let _pp = PeerPenalty::Ridge;
2141 let _ = PeerConfig {
2142 penalty: _pp,
2143 lambda: _lc,
2144 };
2145 let _ = std::mem::size_of::<PeerResult>();
2147 let _ = std::mem::size_of::<LpeerResult>();
2148 }
2149
2150 #[test]
2151 fn test_prelude_exports_compile() {
2152 use crate::error::FdarError;
2155 use crate::prelude::*;
2156
2157 let _peer_fn: fn(&FdMatrix, &[f64], &[f64], &PeerConfig) -> Result<PeerResult, FdarError> =
2158 peer;
2159 let _lpeer_fn: fn(
2160 &FdMatrix,
2161 &[f64],
2162 &[f64],
2163 &[usize],
2164 &PeerConfig,
2165 ) -> Result<LpeerResult, FdarError> = lpeer;
2166
2167 let _lc = LambdaChoice::Gcv;
2168 let _lm = LambdaMethod::Gcv;
2169 let _pp = PeerPenalty::Ridge;
2170 let _ = PeerConfig {
2171 penalty: _pp,
2172 lambda: _lc,
2173 };
2174 let _ = std::mem::size_of::<PeerResult>();
2175 let _ = std::mem::size_of::<LpeerResult>();
2176 }
2177}