Skip to main content

fdars_core/
peer.rs

1//! Structured-penalty scalar-on-function regression via PEER.
2//!
3//! PEER (Partially Empirical Eigenvectors for Regression) estimates the
4//! coefficient function β(t) via penalized normal equations
5//! `(W_c'W_c + λQ)β = W_c'y_c`, where `W_c` is the centered, integration-
6//! weighted design matrix and `Q` is a penalty chosen from three families.
7//!
8//! # Penalty families
9//!
10//! - [`PeerPenalty::Ridge`] — identity penalty (uniform shrinkage).
11//! - [`PeerPenalty::Difference`] — second-difference roughness (D'D, order 2).
12//! - [`PeerPenalty::Decree`] — caller-supplied structured penalty matrix.
13//!
14//! # Quick start
15//!
16//! ```
17//! use fdars_core::matrix::FdMatrix;
18//! use fdars_core::peer::{peer, PeerConfig, PeerPenalty, LambdaChoice};
19//! use fdars_core::helpers::simpsons_weights;
20//!
21//! // Tiny synthetic dataset: n=10 observations, m=5 evaluation points.
22//! let (n, m) = (10_usize, 5_usize);
23//! let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
24//! let mut data = FdMatrix::zeros(n, m);
25//! let mut y = vec![0.0_f64; n];
26//! let true_beta: Vec<f64> = argvals.iter()
27//!     .map(|&t| (std::f64::consts::PI * t).sin()).collect();
28//! let w = simpsons_weights(&argvals);
29//! for i in 0..n {
30//!     for j in 0..m {
31//!         let xi = ((i * m + j) as f64 * 0.3).sin();
32//!         data[(i, j)] = xi;
33//!         y[i] += xi * true_beta[j] * w[j];
34//!     }
35//! }
36//! // Fit PEER with Ridge penalty and a fixed λ.
37//! let config = PeerConfig {
38//!     penalty: PeerPenalty::Ridge,
39//!     lambda: LambdaChoice::Fixed(1e-3),
40//! };
41//! let fit = peer(&data, &y, &argvals, &config).unwrap();
42//! assert_eq!(fit.beta.len(), m);
43//! assert!(fit.fitted_values.iter().all(|v| v.is_finite()));
44//! // Predict on training data — must reproduce fitted_values (self-consistency).
45//! let preds = fit.predict(&data, &argvals).unwrap();
46//! for (p, f) in preds.iter().zip(&fit.fitted_values) {
47//!     assert!((p - f).abs() < 1e-9);
48//! }
49//! ```
50
51use 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// ---------------------------------------------------------------------------
59// Public types
60// ---------------------------------------------------------------------------
61
62/// Penalty family for the PEER estimator.
63///
64/// Selects the penalty matrix Q (m×m) entering the penalized normal equations
65/// `(W_c'W_c + λQ)β = W_c'y_c`.
66#[derive(Debug, Clone, PartialEq)]
67#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
68pub enum PeerPenalty {
69    /// Identity penalty Q = I_m (uniform ridge shrinkage).
70    Ridge,
71    /// Second-difference roughness penalty Q = D'D (only `order = 2` is
72    /// supported in this release; other orders return
73    /// [`FdarError::InvalidParameter`]).
74    Difference { order: usize },
75    /// Caller-supplied penalty matrix.
76    ///
77    /// The tuple holds `(q_flat, p)` where `q_flat` is a flat row-major m×m
78    /// symmetric PSD matrix (length `p*p`) and `p` is its dimension. `p` must
79    /// equal the number of argvals grid points `m`; a mismatch returns
80    /// [`FdarError::InvalidDimension`].
81    ///
82    /// Q must be symmetric (row-major and column-major are equivalent for
83    /// symmetric matrices) and positive semi-definite for Cholesky stability.
84    Decree(Vec<f64>, usize),
85}
86
87impl Default for PeerPenalty {
88    fn default() -> Self {
89        PeerPenalty::Difference { order: 2 }
90    }
91}
92
93/// How to choose the smoothing parameter λ for the PEER estimator.
94///
95/// The default is [`LambdaChoice::Gcv`] which automatically selects λ by
96/// minimising the GCV score over a fixed internal log-spaced grid.
97#[derive(Debug, Clone, Default, PartialEq)]
98#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
99pub enum LambdaChoice {
100    /// Use this value verbatim; no grid search or EM is performed.
101    Fixed(f64),
102    /// Select λ by minimising the GCV score on a fixed 40-point log-spaced
103    /// internal grid over [1e-6, 1e4].  Fully deterministic; ties resolve to
104    /// the smaller grid index.
105    #[default]
106    Gcv,
107    /// Estimate λ via a self-contained REML EM using the eigendecomposition of
108    /// the penalty matrix Q (null space → fixed effect; range space → random
109    /// effect b ∼ N(0, σ²_u I)).  Returns λ = σ²_e / σ²_u.  Fully
110    /// deterministic; fixed initialisation and 100-iteration cap.
111    Reml,
112}
113
114/// Records which λ-selection path actually ran.
115#[derive(Debug, Clone, PartialEq)]
116#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
117pub enum LambdaMethod {
118    /// A fixed value was supplied via [`LambdaChoice::Fixed`]; no search ran.
119    Fixed,
120    /// λ was selected by GCV grid search.
121    Gcv,
122    /// λ was estimated by the self-contained REML EM.
123    Reml,
124}
125
126/// Configuration for the [`peer`] estimator.
127#[derive(Debug, Clone, Default, PartialEq)]
128#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
129pub struct PeerConfig {
130    /// Penalty family selecting the structured Q matrix.
131    pub penalty: PeerPenalty,
132    /// How to choose (or fix) the smoothing parameter λ.
133    ///
134    /// - [`LambdaChoice::Fixed(v)`](LambdaChoice::Fixed) — use `v` verbatim.
135    /// - [`LambdaChoice::Gcv`] — automatic GCV grid search (default).
136    /// - [`LambdaChoice::Reml`] — automatic REML EM estimation.
137    pub lambda: LambdaChoice,
138}
139
140/// Result of the [`peer`] estimator.
141///
142/// Carries the estimated coefficient function β(t), model diagnostics, and
143/// the penalty configuration used — all on the `argvals` grid.
144#[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    /// Estimated coefficient function β(t), length m (on the argvals grid).
150    pub beta: Vec<f64>,
151    /// Intercept (= ȳ, the mean of the response vector).
152    ///
153    /// This is the centered-response mean, not the out-of-sample prediction
154    /// intercept. Prediction on a new curve x* uses
155    /// `ȳ + Σ_j (x*[j]·w[j] − w_bar[j])·β[j] = (ȳ − w_bar·β) + Σ_j x*[j]·w[j]·β[j]`,
156    /// so a predictor must combine `intercept` with [`w_bar`](Self::w_bar) and
157    /// `beta`. Storing `w_bar` here keeps that reconstruction exact (consumed by
158    /// out-of-sample prediction in a later phase).
159    pub intercept: f64,
160    /// Column means of the Simpson-weighted design `W[i,j] = data[(i,j)]·w[j]`,
161    /// length m. Retained so out-of-sample prediction can reproduce the same
162    /// centering the fit used (see [`intercept`](Self::intercept)).
163    pub w_bar: Vec<f64>,
164    /// Fitted values ŷ_i, length n.
165    pub fitted_values: Vec<f64>,
166    /// Effective degrees of freedom tr(H) = tr((W_c'W_c + λQ)^{-1} W_c'W_c).
167    pub effective_df: f64,
168    /// Smoothing parameter λ that was used.
169    pub lambda: f64,
170    /// Penalty family that was used.
171    pub penalty_type: PeerPenalty,
172    /// GCV score at the selected λ.  `Some(score)` when [`LambdaMethod::Gcv`]
173    /// ran; `None` when [`LambdaMethod::Fixed`] or [`LambdaMethod::Reml`].
174    pub gcv: Option<f64>,
175    /// Which λ-selection path ran.
176    pub lambda_method: LambdaMethod,
177}
178
179/// Result of the [`lpeer`] longitudinal PEER estimator.
180///
181/// Carries the estimated coefficient function β(t), subject-level variance
182/// components, and the penalty configuration used — all on the `argvals` grid.
183///
184/// Both variance components are non-negative (clamped to a positive floor by
185/// `famm::fit_scalar_mixed_model`).
186#[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    /// Estimated coefficient function β(t), length m (on the argvals grid).
192    pub beta: Vec<f64>,
193    /// Intercept (= ȳ, the mean of the response vector). Matches `peer()` convention.
194    pub intercept: f64,
195    /// Column means of the Simpson-weighted design `W[i,j] = data[(i,j)]·w[j]`,
196    /// length m. Computed identically to `peer()` so [`predict`](Self::predict)
197    /// reproduces training `fitted_values` exactly.
198    pub w_bar: Vec<f64>,
199    /// Fitted values ŷ_i (marginal, fixed-effect only), length n.
200    pub fitted_values: Vec<f64>,
201    /// Between-subject variance σ²_u (≥ 0). Estimated by REML EM inside
202    /// `famm::fit_scalar_mixed_model`; clamped to a positive floor.
203    pub sigma2_subject: f64,
204    /// Residual variance σ²_ε (≥ 0). Estimated by REML EM; clamped to a positive floor.
205    pub sigma2_resid: f64,
206    /// Number of unique subjects derived from `subject_map`.
207    pub n_subjects: usize,
208    /// Smoothing parameter λ that was used (from `PeerConfig`).
209    pub lambda: f64,
210    /// Penalty family that was used.
211    pub penalty_type: PeerPenalty,
212    /// GCV score at the selected λ. `Some(score)` when [`LambdaMethod::Gcv`]
213    /// ran; `None` for [`LambdaMethod::Fixed`] or [`LambdaMethod::Reml`].
214    pub gcv: Option<f64>,
215    /// Which λ-selection path ran.
216    pub lambda_method: LambdaMethod,
217}
218
219// ---------------------------------------------------------------------------
220// Public entry point
221// ---------------------------------------------------------------------------
222
223/// Fit the PEER scalar-on-function regression model.
224///
225/// Estimates the coefficient function β(t) on the `argvals` grid by solving
226/// the penalized normal equations `(W_c'W_c + λQ)β = W_c'y_c`, where
227/// `W_c[i,j] = (data[(i,j)] · w[j]) − column_mean` and `w` are Simpson's
228/// integration weights.
229///
230/// # Arguments
231///
232/// * `data`    — n×m functional predictor matrix (rows = observations,
233///   columns = evaluation points; column-major [`FdMatrix`]).
234/// * `y`       — scalar response vector, length n.
235/// * `argvals` — evaluation grid, length m (must equal `data.ncols()`).
236/// * `config`  — penalty family and λ selection.
237///
238/// # Errors
239///
240/// Returns [`FdarError::InvalidDimension`] when dimensions are inconsistent
241/// or a Decree Q has wrong size, [`FdarError::InvalidParameter`] for
242/// unsupported `Difference` orders, and [`FdarError::ComputationFailed`]
243/// when the penalized system is numerically singular.
244pub 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    // --- Entry validation ---
253    if n < 2 {
254        // Centering collapses a single observation to β = 0 (all-zero design);
255        // require at least 2 so the fit is not silently degenerate.
256        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    // Non-finite inputs would propagate into finite-but-wrong normal equations,
284    // bypassing the post-solve β NaN guard — reject them explicitly.
285    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    // Simpson's weights assume a strictly increasing grid; a non-monotone grid
298    // yields negative weights that silently corrupt the design integral.
299    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    // 1. Integration weights w[j]
307    let w = simpsons_weights(argvals);
308
309    // 2. Weighted design: wmat[i,j] = data[i,j] * w[j]
310    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    // 3. Center response and design
318    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    // 4. Build penalty Q (m×m, row-major)
333    let q = build_q(m, &config.penalty)?;
334
335    // 5. WtW (m×m, row-major, symmetric) and wty (length m)
336    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    // 6. λ selection — dispatch based on config.lambda
349    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    // 7. A = WtW + λQ; solve A β = wty via Cholesky
362    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    // NaN guard: non-finite β should not be returned silently
369    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    // 8. Effective degrees of freedom: tr(H) = tr(A^{-1} WtW)
377    let effective_df = compute_peer_trace_hat(&wtw, &q, lambda, m, n);
378
379    // 9. Fitted values: ŷ[i] = ȳ + Σ_j wc[i,j] · β[j]
380    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
397/// Fit the longitudinal PEER scalar-on-function regression model with subject
398/// random effects.
399///
400/// Extends [`peer`] to grouped/repeated-measures data by reducing the
401/// functional predictor to FPC scores (via `fdata_to_pc_1d`) and calling
402/// `famm::fit_scalar_mixed_model` for subject-level random intercepts.
403/// The PEER penalty (via `config`) regularises β(t) through the same λ
404/// dispatch as `peer()`; the mixed model then replaces the OLS second pass
405/// with a GLS+REML-EM pass over the FPC scores.
406///
407/// **ncomp cap:** `min(n − 1, m, 10)` — documented; prevents over-parameterisation
408/// at small n.
409///
410/// # Arguments
411///
412/// * `data`        — n×m functional predictor matrix.
413/// * `y`           — scalar response vector, length n.
414/// * `argvals`     — evaluation grid, length m.
415/// * `subject_map` — subject index per observation, length n. Non-contiguous
416///   IDs are re-indexed internally via `famm::build_subject_map`.
417/// * `config`      — penalty family and λ selection (same as `peer()`).
418///
419/// # Errors
420///
421/// Returns [`FdarError::InvalidDimension`] on dimension mismatches,
422/// [`FdarError::InvalidParameter`] when `n_subjects < 2` or argvals is
423/// non-monotone/non-finite, and [`FdarError::ComputationFailed`] when the
424/// mixed-model produces non-finite coefficients.
425///
426/// # Example
427///
428/// ```
429/// use fdars_core::matrix::FdMatrix;
430/// use fdars_core::peer::{lpeer, PeerConfig, PeerPenalty, LambdaChoice};
431///
432/// let (n, m) = (12_usize, 5_usize);
433/// let argvals: Vec<f64> = (0..m).map(|i| i as f64 / (m - 1) as f64).collect();
434/// let mut data = FdMatrix::zeros(n, m);
435/// let mut y = vec![0.0_f64; n];
436/// // 3 subjects × 4 observations each
437/// let subject_map: Vec<usize> = (0..n).map(|i| i / 4).collect();
438/// for i in 0..n {
439///     for j in 0..m {
440///         let xi = ((i * m + j) as f64 * 0.3).sin();
441///         data[(i, j)] = xi;
442///         y[i] += xi * 0.5;
443///     }
444/// }
445/// let config = PeerConfig {
446///     penalty: PeerPenalty::Ridge,
447///     lambda: LambdaChoice::Fixed(1e-2),
448/// };
449/// let fit = lpeer(&data, &y, &argvals, &subject_map, &config).unwrap();
450/// assert_eq!(fit.beta.len(), m);
451/// assert!(fit.sigma2_subject >= 0.0);
452/// assert!(fit.sigma2_resid >= 0.0);
453/// let preds = fit.predict(&data, &argvals).unwrap();
454/// assert_eq!(preds.len(), n);
455/// ```
456pub 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    // --- Entry validation (mirrors peer()) ---
466    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    // Build dense 0-indexed subject map (handles non-contiguous IDs)
521    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    // 1. Integration weights + weighted design (identical to peer())
530    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    // w_bar computed identically to peer() so predict formula is identical
542    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    // 2. λ selection (identical dispatch to peer())
554    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    // 3. FPC score reduction — cap at min(n-1, m, 10); raw scores, no h.sqrt() rescaling
581    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; // n×ncomp FdMatrix (raw FPC scores)
584
585    // 4. Mixed-model fit over FPC scores (pass yc, not y — intercept = y_bar).
586    //    This estimates the subject-level random effect: σ²_subject (between-subject)
587    //    and σ²_resid, plus a baseline (unpenalized) fixed-effect γ.
588    let result =
589        crate::famm::fit_scalar_mixed_model(&yc, &sm_dense, n_subjects, Some(scores), ncomp);
590
591    // NaN guard on gamma
592    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    // 5. Apply the PEER structured penalty λQ in the FPC-score space via a
600    //    penalized-GLS re-solve, reusing the mixed model's variance components.
601    //    The random-intercept marginal covariance is Σ = σ²_e·I + σ²_u·ZZ'; for a
602    //    subject g of size n_g, (Σ⁻¹v)_i = (v_i − s_g·σ²_u/(σ²_e+n_g·σ²_u))/σ²_e
603    //    with s_g the subject sum of v. In score space the β-penalty λ·β'Qβ becomes
604    //    λ·γ'(Φ'QΦ)γ, so γ solves (SᵀΣ⁻¹S + λ·Φ'QΦ)γ = SᵀΣ⁻¹y_c. λ (Fixed/GCV/REML)
605    //    now genuinely regularizes β(t); it is not a mere label.
606    let su = result.sigma2_u;
607    let se = result.sigma2_eps;
608    // Per-subject observation counts (for the random-intercept shrinkage).
609    let mut gcnt = vec![0usize; n_subjects];
610    for &g in &sm_dense {
611        gcnt[g] += 1;
612    }
613    // Σ⁻¹ applied to a length-n vector.
614    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    // Σ⁻¹ S  (n×ncomp)
628    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    // A = Sᵀ Σ⁻¹ S  (ncomp×ncomp) and b = Sᵀ Σ⁻¹ y_c  (ncomp)
637    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    // P = Φ'QΦ  (ncomp×ncomp): penalty λ·β'Qβ expressed on the FPC coefficients.
647    // QΦ  (m×ncomp)
648    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    // A_pen = A + λ·Φ'QΦ
657    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    // Penalized-GLS coefficients; fall back to the unpenalized mixed-model γ on failure.
666    let gamma = cholesky_solve(&a_mat, &b_vec, ncomp).unwrap_or_else(|_| result.gamma.clone());
667
668    // 6. Back-project γ → β(t): beta[j] = Σ_k gamma[k] * rotation[(j, k)]
669    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    // 7. Fitted values (same formula as peer())
685    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    /// Predict scalar responses for new functional observations (marginal prediction).
707    ///
708    /// Applies the fitted coefficient function β(t) to `new_data` via the same
709    /// formula as [`PeerResult::predict`]: new-subject random effect = 0 (marginal /
710    /// fixed-effect-only prediction).
711    ///
712    /// **Self-consistency:** re-passing the training data and argvals reproduces
713    /// the training [`fitted_values`](Self::fitted_values) to within 1e-9.
714    ///
715    /// # Errors
716    ///
717    /// Returns [`FdarError::InvalidDimension`] when `new_data.ncols()` or
718    /// `argvals.len()` does not equal the number of training grid points `m`.
719    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
724// ---------------------------------------------------------------------------
725// Private helpers
726// ---------------------------------------------------------------------------
727
728/// Build the m×m penalty matrix Q (flat row-major) for the given family.
729fn 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
758/// Compute effective df = tr(H) = tr(A^{-1} WtW) where A = WtW + λQ.
759///
760/// Uses the column-solve trick: for each column j of WtW, solve A z = WtW[:,j]
761/// and accumulate z[j].  Falls back to `m as f64` on Cholesky failure;
762/// clamps result to `n as f64`.
763fn 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; // fallback
770    };
771    let mut trace = 0.0_f64;
772    for j in 0..m {
773        // Extract column j of WtW (stored row-major: WtW[row, col] = wtw[row*m + col])
774        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
781/// Build a 40-point log-spaced grid on [1e-6, 1e4].
782///
783/// Grid point i is `10^(-6 + 10·i/39)` for i in 0..40.  The grid is purely
784/// a function of constants — no RNG — so GCV selection is fully deterministic.
785fn 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
791/// Select λ by minimising the GCV score over the fixed 40-point grid.
792///
793/// GCV score: `GCV(λ) = n·RSS(λ) / (n − tr(H(λ)))²`.
794///
795/// Tie-break: when two grid points share the minimum GCV (within floating-point
796/// equality), the smaller grid index (smaller λ) is retained.
797///
798/// Guard: if `n − tr(H) ≤ 0` for a candidate (degenerate, over-smoothed), that
799/// grid point is skipped.  If no grid point yields a finite score, the smallest
800/// grid λ is returned with score `f64::INFINITY`.
801///
802/// Returns `(best_lambda, gcv_at_best)`.
803fn 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        // Build A = WtW + lam*Q
818        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        // Solve for beta
823        let Ok(beta) = cholesky_solve(&a, wty, m) else {
824            continue;
825        };
826        // RSS = ||y_c - W_c beta||^2
827        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        // tr(H) reuses compute_peer_trace_hat
834        let trh = compute_peer_trace_hat(wtw, q, lam, m, n);
835        // Guard: skip if denominator <= 0
836        let denom = n as f64 - trh;
837        if denom <= 0.0 {
838            continue;
839        }
840        let gcv = n as f64 * rss / (denom * denom);
841        // Strict < keeps first (smaller-λ) grid index on ties
842        if gcv < best_gcv {
843            best_gcv = gcv;
844            best_lam = lam;
845        }
846    }
847    (best_lam, best_gcv)
848}
849
850/// Select λ via a self-contained REML EM using the eigendecomposition of Q.
851///
852/// The PEER-as-mixed-model equivalence partitions the m-dimensional coefficient
853/// space into:
854/// - Null space of Q (s dimensions) → fixed (unpenalised) effect α.
855/// - Range space of Q (r dimensions) → random effect b ∼ N(0, σ²_u I_r).
856///
857/// An EM loop estimates σ²_u and σ²_e; the returned λ = σ²_e / σ²_u.
858///
859/// Edge cases:
860/// - Ridge (Q = I_m, s = 0): GLS α-update is skipped; r_alpha = y_c.
861/// - Zero-Q Decree (r = 0): returns the documented fallback λ = 1e-4.
862/// - σ²_u → 0: clamped to 1e-12 (keeps λ finite).
863///
864/// Deterministic: fixed initialisation + 100-iteration cap, no RNG.
865fn select_lambda_reml_peer(wc: &FdMatrix, yc: &[f64], q: &[f64], m: usize, n: usize) -> f64 {
866    // --- STEP 1: eigendecompose Q (ascending eigenvalue order, null space first) ---
867    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    // --- STEP 2: edge — zero-Q Decree (range space empty) ---
897    if r == 0 {
898        return 1e-4; // documented fallback: no random effect to estimate
899    }
900
901    // --- STEP 3: project designs into eigenbasis ---
902    // z_null[i*s + col]  = Σ_j wc[(i,j)] * V_null[(j, col)]
903    // z_range[i*r + col] = Σ_j wc[(i,j)] * V_range[(j, col)]
904    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    // Pre-compute ZtZ_range (r×r, row-major, symmetric) — built once, reused each iter
924    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    // --- STEP 4: deterministic init (no RNG) ---
936    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    // OLS init of alpha (if s > 0): solve (Z_null' Z_null + 1e-10 I_s) alpha = Z_null' yc
942    let mut alpha = vec![0.0_f64; s];
943    if s > 0 {
944        // ZtZ_null (s×s) + ridge
945        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        // Z_null' yc
957        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    // --- STEP 5: EM loop, fixed cap 100 iterations ---
966    for _iter in 0..100 {
967        let su_old = sigma2_u;
968        let se_old = sigma2_e;
969
970        // r_alpha = yc - Z_null * alpha  (residual after fixed effects)
971        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        // E-step: M = ZtZ_range/sigma2_e + I_r/sigma2_u  (r×r)
983        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        // Cholesky of the E-step matrix M. On failure the EM cannot advance, so
992        // restore the last stable variance components and stop — the function then
993        // returns λ = σ²_e/σ²_u from those last-good values (defined graceful
994        // fallback, never a panic or NaN; WR-03).
995        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        // Sigma_b = M^{-1}: solve column by column (r solves)
1005        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        // Z_range' r_alpha  (length r)
1018        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        // b_hat = Sigma_b * ztr / sigma2_e  (length r)
1023        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        // M-step: sigma2_u_new = (b_hat'b_hat + tr(Sigma_b)) / r
1033        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        // resid = r_alpha - Z_range * b_hat  (length n)
1037        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        // tr(Z_range Sigma_b Z_range') = tr(Sigma_b ZtZ_range) = Σ_{a,j} sigma_b[a*r+j]*ztz_range[j*r+a]
1045        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        // Clamp both variance components to a positive floor BEFORE they are used
1057        // as divisors in the GLS block below — an unclamped σ²_e ≈ 0 would divide
1058        // to Inf/NaN in the Woodbury solve (WR-02).
1059        let sigma2_u_c = sigma2_u_new.max(1e-12);
1060        let sigma2_e_c = sigma2_e_new.max(1e-12);
1061
1062        // GLS null-space update (only when s > 0) using Woodbury identity
1063        // Sigma^{-1} = (1/sigma2_e)(I_n - Z_range * K^{-1} * Z_range')
1064        // where K = sigma2_e/sigma2_u * I_r + ZtZ_range  (r×r)
1065        if s > 0 {
1066            // Build K = ZtZ_range + (sigma2_e/sigma2_u)*I_r
1067            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            // Cholesky of K for the Woodbury inverse
1073            if let Ok(l_k) = cholesky_factor(&k_mat, r) {
1074                // Compute Sigma^{-1} Z_null and Z_null' Sigma^{-1} Z_null for the s×s GLS system
1075                // Sigma^{-1} z_null_col = (1/sigma2_e)(z_null_col - Z_range * K^{-1} Z_range' z_null_col)
1076                let mut gls_lhs = vec![0.0_f64; s * s]; // Z_null' Sigma^{-1} Z_null  (s×s)
1077                let mut gls_rhs = vec![0.0_f64; s]; // Z_null' Sigma^{-1} y_c  (length s)
1078
1079                // Precompute Z_range' y_c (length r)
1080                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                // K^{-1} Z_range' y_c  (length r)
1084                let kinv_zr_yc = cholesky_forward_back(&l_k, &zr_yc, r);
1085                // Sigma^{-1} y_c  (length n): (1/sigma2_e)(y_c - Z_range * K^{-1} Z_range' y_c)
1086                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                    // Z_range' z_null_col  (length r)
1097                    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                    // K^{-1} Z_range' z_null_col
1105                    let kinv_zr_zn = cholesky_forward_back(&l_k, &zr_zn, r);
1106                    // Sigma^{-1} z_null_col  (length n)
1107                    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                    // Z_null' Sigma^{-1} z_null_col → column col_null of gls_lhs
1117                    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                    // Z_null' Sigma^{-1} y_c: entry col_null of gls_rhs
1122                    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                // Add ridge to GLS LHS for numerical stability
1127                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        // Commit the clamped variance components for this iteration.
1138        sigma2_u = sigma2_u_c;
1139        sigma2_e = sigma2_e_c;
1140
1141        // Measure convergence on the PRE-clamp M-step updates so a component
1142        // resting on the 1e-12 floor two iterations running is not mistaken for
1143        // genuine convergence (WR-01) — that would return λ = σ²_e/floor
1144        // (extreme over-smoothing) instead of the REML optimum.
1145        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    // --- STEP 6: return λ = σ²_e / σ²_u ---
1152    (sigma2_e / sigma2_u).max(1e-15)
1153}
1154
1155// ---------------------------------------------------------------------------
1156// Shared prediction helper
1157// ---------------------------------------------------------------------------
1158
1159/// Shared prediction core for both `PeerResult` and `LpeerResult`.
1160///
1161/// Computes out-of-sample predictions via
1162/// `ŷ*[i] = (intercept − w_bar·β) + Σ_j x*[i,j] · w[j] · β[j]`
1163/// where `w = simpsons_weights(argvals)`.
1164///
1165/// Self-consistency: re-passing the training `data` and `argvals` reproduces
1166/// the training `fitted_values` to within floating-point rounding (≤ 1e-9).
1167fn 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    // Simpson's weights assume a strictly increasing grid; a non-monotone grid
1191    // yields negative weights that would silently corrupt the prediction.
1192    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    // w_bar always has length m for a well-formed fit result; assert in debug to
1199    // catch a corrupted result before it silently truncates the base offset.
1200    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    /// Predict scalar responses for new functional observations.
1215    ///
1216    /// Applies the fitted PEER coefficient function β(t) to `new_data` via
1217    /// `ŷ*[i] = (intercept − w_bar·β) + Σ_j x*[i,j] · w[j] · β[j]`
1218    /// where `w = simpsons_weights(argvals)`.
1219    ///
1220    /// **Self-consistency:** re-passing the training data and training `argvals`
1221    /// reproduces the training [`fitted_values`](Self::fitted_values) to within
1222    /// floating-point rounding (≤ 1e-9 in absolute error per observation).
1223    ///
1224    /// # Errors
1225    ///
1226    /// Returns [`FdarError::InvalidDimension`] when `new_data.ncols()` or
1227    /// `argvals.len()` does not equal the number of training grid points `m`.
1228    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// ---------------------------------------------------------------------------
1234// Tests
1235// ---------------------------------------------------------------------------
1236
1237#[cfg(test)]
1238mod tests {
1239    use super::*;
1240    use crate::test_helpers::uniform_grid;
1241
1242    /// Stateless deterministic pseudo-random unit in [-1, 1] (splitmix64 finalizer).
1243    /// Used to build spanning design curves without an RNG dependency.
1244    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    /// Build the shared synthetic fixture: n=200 curves, m=40 grid, true β(t)=sin(π·t).
1254    /// Generates y_i = Σ_j X_i(t_j)·β(t_j)·w_j + small_noise.
1255    ///
1256    /// The design curves are deterministic pseudo-random (spanning the full m-dim
1257    /// grid space, so W has full column rank and β(t) is identifiable) — a smooth
1258    /// low-rank family (e.g. phase-shifted single-frequency sinusoids) would live
1259    /// in a 2-D subspace and leave β(t) unrecoverable regardless of the estimator.
1260    /// n ≫ m keeps W'W well-conditioned and makes the fixed λ negligible relative to
1261    /// the signal, so the penalty bias stays well under the recovery tolerance.
1262    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            // Small deterministic noise
1280            y[i] += 0.005 * hash_unit(1_000_000 + i as u64);
1281        }
1282        (data, y, t, true_beta)
1283    }
1284
1285    // -------------------------------------------------------------------------
1286    // Phase 66 migrated tests: Difference{2} tracer
1287    // -------------------------------------------------------------------------
1288
1289    #[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    // -------------------------------------------------------------------------
1356    // Phase 66 migrated tests: all three penalty families
1357    // -------------------------------------------------------------------------
1358
1359    #[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        // Use the same penalty_matrix as a caller-supplied Decree matrix
1384        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    // -------------------------------------------------------------------------
1416    // Phase 66 migrated: Decree partition-structured Q yields β(t) distinct from Difference{2}
1417    // -------------------------------------------------------------------------
1418
1419    #[test]
1420    fn test_peer_decree_distinct_from_roughness() {
1421        let (data, y, t, _) = make_fixture();
1422        let m = t.len();
1423        // Use a moderate lambda so the penalty visibly shapes β(t)
1424        let lambda = LambdaChoice::Fixed(1.0);
1425
1426        // Fit (a) with plain Difference{2}
1427        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        // Build a partition-aware Q: a second-difference operator whose stencil
1434        // is NOT applied across the mid-grid boundary b = m/2.
1435        // This leaves the two halves independently penalized.
1436        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            // Skip stencil rows that straddle the boundary
1440            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        // Add a small ridge to each diagonal to ensure PSD for Cholesky stability
1451        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        // Both fits must be all-finite
1463        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        // The two β(t) must differ meaningfully (> 1e-3 max pointwise)
1473        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    // -------------------------------------------------------------------------
1487    // Phase 66 migrated: Error/NaN surface — wrong-dim Q, dim mismatch, no NaN
1488    // -------------------------------------------------------------------------
1489
1490    #[test]
1491    fn test_peer_decree_wrong_dim() {
1492        let (data, y, t, _) = make_fixture();
1493        let m = t.len();
1494        // Q sized (m-1)×(m-1) ≠ m×m
1495        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        // argvals has length m+1 instead of m
1511        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        // Ridge
1528        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        // Difference{2}
1547        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        // Decree using penalty_matrix(m) as a valid caller-supplied Q
1566        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    // -------------------------------------------------------------------------
1587    // Phase 66 migrated: Prediction-contract + hardened-validation tests
1588    // -------------------------------------------------------------------------
1589
1590    #[test]
1591    fn test_peer_stores_w_bar_for_prediction() {
1592        // The stored w_bar must let an out-of-sample predictor reproduce the
1593        // training fitted values exactly via ŷ = (ȳ − w_bar·β) + Σ_j x[j]·w[j]·β[j].
1594        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); // first pair now decreasing
1630        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    // -------------------------------------------------------------------------
1666    // Task 2 tests: GCV grid-search selector
1667    // -------------------------------------------------------------------------
1668
1669    #[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        // Non-degenerate lambda: not at extreme ends of the grid
1697        assert!(
1698            result.lambda > 1e-10 && result.lambda < 1e6,
1699            "GCV lambda should be non-degenerate, got {}",
1700            result.lambda
1701        );
1702        // GCV score recorded
1703        assert!(result.gcv.is_some(), "gcv field must be Some for Gcv");
1704        assert_eq!(result.lambda_method, LambdaMethod::Gcv);
1705
1706        // β recovery within tolerance
1707        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    // -------------------------------------------------------------------------
1717    // Task 3 tests: REML EM selector
1718    // -------------------------------------------------------------------------
1719
1720    #[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        // Use the high-SNR fixture (noise=0.005) for the agreement test.
1759        // REML on this problem correctly identifies a small λ (the range-space
1760        // component of sin(πt) is large, so σ²_u is large, giving λ = σ²_e/σ²_u
1761        // near zero — valid REML behavior for this signal shape).  The resulting
1762        // beta still recovers the truth within 0.15 because at this SNR even
1763        // near-OLS estimates are close.  GCV picks a larger but also reasonable λ.
1764        let (data, y, t, true_beta) = make_fixture(); // noise=0.005, high SNR
1765
1766        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        // REML beta(t) recovers true beta within 0.15
1779        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        // GCV and REML beta agree within documented tolerance (0.2 max abs diff)
1791        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    // -------------------------------------------------------------------------
1806    // Task 1 (Phase 68): PeerResult::predict — self-consistency, dim validation,
1807    // and all-finite on fresh curves.
1808    // -------------------------------------------------------------------------
1809
1810    #[test]
1811    fn test_peer_predict_self_consistent() {
1812        // Re-pass training curves → predictions must equal fitted_values within 1e-9.
1813        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        // new_data with wrong column count must yield FdarError::InvalidDimension.
1834        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        // Build a new_data matrix with m-1 columns (wrong ncols)
1843        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        // Predict on fresh curves (offset by a large constant) — all results finite.
1861        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        // Build genuinely new curves: large constant shift from training data
1870        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        // Ridge penalty (Q = I_m, null space empty, s=0) — must not panic
1891        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        // Zero-Q Decree (all-zero matrix, r=0) — must return documented fallback λ=1e-4
1904        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    // -------------------------------------------------------------------------
1919    // Task 2 (Phase 68): lpeer() — longitudinal PEER with subject random effects
1920    // -------------------------------------------------------------------------
1921
1922    /// Longitudinal test fixture: n_subjects=20, obs_per=10 (n=200), m=10.
1923    /// True β(t)=sin(πt); injects deterministic between-subject effects at
1924    /// known variance σ²_u_true=1.0.  Returns (data, y, t, subject_map, sigma2_u_true).
1925    ///
1926    /// Uses m=10 so that ncomp = min(n-1, m, 10) = 10 = m, giving full FPC coverage
1927    /// for accurate β(t) back-projection.  n=200 provides reliable REML EM convergence.
1928    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            // Deterministic subject random effect at scale sqrt(sigma2_u_true)
1945            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; // inject between-subject effect
1955                             // small within-subject noise
1956                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        // λ must genuinely regularize β(t): a heavy structured penalty produces a
2013        // materially different (smoother) coefficient function than a near-zero one.
2014        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        // REML EM is consistent but not exact at n=50, 10 subjects — allow factor ~3
2048        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        // subject_map of wrong length (n-1 instead of n)
2060        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        // All observations from the same subject → n_subjects = 1 (degenerate)
2078        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        // LpeerResult::predict on training data must reproduce fitted_values within 1e-9.
2094        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    // -------------------------------------------------------------------------
2114    // Task 4 (Phase 68): Crate-root + prelude export reachability compile checks
2115    // -------------------------------------------------------------------------
2116
2117    #[test]
2118    fn test_crate_root_exports_compile() {
2119        // Verify crate-root pub use block brings all 8 PEER symbols into scope.
2120        // In-crate tests use `crate::` paths; the re-export block in lib.rs makes
2121        // them available at the root so external crates can import `fdars_core::{peer, ...}`.
2122        // We test the crate-root by referencing each symbol through the module path
2123        // that lib.rs re-exports from.
2124        use crate::error::FdarError;
2125
2126        // Verify function signatures match expected types (compile-only checks).
2127        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        // Verify enum/struct variants compile
2138        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        // Verify LpeerResult and PeerResult are constructible via the module path
2146        let _ = std::mem::size_of::<PeerResult>();
2147        let _ = std::mem::size_of::<LpeerResult>();
2148    }
2149
2150    #[test]
2151    fn test_prelude_exports_compile() {
2152        // Verify the prelude re-exports all 8 PEER symbols.
2153        // `use crate::prelude::*` brings them into this scope.
2154        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}