use crate::error::FdarError;
use crate::linalg::{cholesky_factor, cholesky_forward_back, compute_xtx};
use crate::matrix::FdMatrix;
use crate::regression::{fdata_to_pc_1d, FpcaResult};
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct FofResult {
pub intercept: Vec<f64>,
pub beta_surface: FdMatrix,
pub fitted: FdMatrix,
pub residuals: FdMatrix,
pub r_squared_t: Vec<f64>,
pub r_squared: f64,
pub ncomp_x: usize,
pub ncomp_y: usize,
pub fpca_x: FpcaResult,
pub fpca_y: FpcaResult,
pub coef_matrix: FdMatrix,
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn fof_regression(
x_data: &FdMatrix,
y_data: &FdMatrix,
x_argvals: &[f64],
y_argvals: &[f64],
ncomp_x: usize,
ncomp_y: usize,
) -> Result<FofResult, FdarError> {
let (n_x, m_x) = x_data.shape();
let (n_y, m_y) = y_data.shape();
if n_x != n_y {
return Err(FdarError::InvalidDimension {
parameter: "y_data",
expected: format!("{n_x} rows (matching x_data)"),
actual: format!("{n_y} rows"),
});
}
let n = n_x;
if n < 3 {
return Err(FdarError::InvalidDimension {
parameter: "x_data",
expected: "at least 3 observations".to_string(),
actual: format!("{n}"),
});
}
if x_argvals.len() != m_x {
return Err(FdarError::InvalidDimension {
parameter: "x_argvals",
expected: format!("{m_x} elements"),
actual: format!("{} elements", x_argvals.len()),
});
}
if y_argvals.len() != m_y {
return Err(FdarError::InvalidDimension {
parameter: "y_argvals",
expected: format!("{m_y} elements"),
actual: format!("{} elements", y_argvals.len()),
});
}
if ncomp_x == 0 {
return Err(FdarError::InvalidParameter {
parameter: "ncomp_x",
message: "must be >= 1".to_string(),
});
}
if ncomp_y == 0 {
return Err(FdarError::InvalidParameter {
parameter: "ncomp_y",
message: "must be >= 1".to_string(),
});
}
let ncomp_x = ncomp_x.min(n - 1).min(m_x);
let ncomp_y = ncomp_y.min(n - 1).min(m_y);
let fpca_x = fdata_to_pc_1d(x_data, ncomp_x, x_argvals)?;
let fpca_y = fdata_to_pc_1d(y_data, ncomp_y, y_argvals)?;
let x_scores = fpca_x.project(x_data)?;
let y_scores = fpca_y.project(y_data)?;
let mut xtx = compute_xtx(&x_scores);
let ridge = 1e-8 * (0..ncomp_x).map(|k| xtx[k * ncomp_x + k]).sum::<f64>() / ncomp_x as f64;
for k in 0..ncomp_x {
xtx[k * ncomp_x + k] += ridge.max(1e-12);
}
let l = cholesky_factor(&xtx, ncomp_x)?;
let mut coef_matrix = FdMatrix::zeros(ncomp_x, ncomp_y);
for l_col in 0..ncomp_y {
let mut xty = vec![0.0; ncomp_x];
for k in 0..ncomp_x {
let mut s = 0.0;
for i in 0..n {
s += x_scores[(i, k)] * y_scores[(i, l_col)];
}
xty[k] = s;
}
let b_col = cholesky_forward_back(&l, &xty, ncomp_x);
for k in 0..ncomp_x {
coef_matrix[(k, l_col)] = b_col[k];
}
}
let mut beta_surface = FdMatrix::zeros(m_y, m_x);
for si in 0..m_y {
for tj in 0..m_x {
let mut val = 0.0;
for k in 0..ncomp_x {
for l_col in 0..ncomp_y {
val += coef_matrix[(k, l_col)]
* fpca_x.rotation[(tj, k)]
* fpca_y.rotation[(si, l_col)];
}
}
beta_surface[(si, tj)] = val;
}
}
let mut fitted_scores = FdMatrix::zeros(n, ncomp_y);
for i in 0..n {
for l_col in 0..ncomp_y {
let mut s = 0.0;
for k in 0..ncomp_x {
s += x_scores[(i, k)] * coef_matrix[(k, l_col)];
}
fitted_scores[(i, l_col)] = s;
}
}
let mut fitted = FdMatrix::zeros(n, m_y);
for i in 0..n {
for j in 0..m_y {
let mut val = fpca_y.mean[j];
for l_col in 0..ncomp_y {
val += fitted_scores[(i, l_col)] * fpca_y.rotation[(j, l_col)];
}
fitted[(i, j)] = val;
}
}
let mut residuals = FdMatrix::zeros(n, m_y);
for i in 0..n {
for j in 0..m_y {
residuals[(i, j)] = y_data[(i, j)] - fitted[(i, j)];
}
}
let intercept = fpca_y.mean.clone();
let mut r_squared_t = vec![0.0; m_y];
for j in 0..m_y {
let y_mean_j = fpca_y.mean[j];
let mut ss_tot = 0.0;
let mut ss_res = 0.0;
for i in 0..n {
ss_tot += (y_data[(i, j)] - y_mean_j).powi(2);
ss_res += residuals[(i, j)].powi(2);
}
r_squared_t[j] = if ss_tot > 0.0 {
1.0 - ss_res / ss_tot
} else {
0.0
};
}
let r_squared = r_squared_t.iter().sum::<f64>() / m_y as f64;
Ok(FofResult {
intercept,
beta_surface,
fitted,
residuals,
r_squared_t,
r_squared,
ncomp_x,
ncomp_y,
fpca_x,
fpca_y,
coef_matrix,
})
}
pub fn predict_fof(fit: &FofResult, new_x: &FdMatrix) -> Result<FdMatrix, FdarError> {
let (n_new, _m_x) = new_x.shape();
let x_scores = fit.fpca_x.project(new_x)?;
let ncomp_x = fit.ncomp_x;
let ncomp_y = fit.ncomp_y;
let m_y = fit.fpca_y.mean.len();
let mut pred_scores = FdMatrix::zeros(n_new, ncomp_y);
for i in 0..n_new {
for l_col in 0..ncomp_y {
let mut s = 0.0;
for k in 0..ncomp_x {
s += x_scores[(i, k)] * fit.coef_matrix[(k, l_col)];
}
pred_scores[(i, l_col)] = s;
}
}
let mut predicted = FdMatrix::zeros(n_new, m_y);
for i in 0..n_new {
for j in 0..m_y {
let mut val = fit.fpca_y.mean[j];
for l_col in 0..ncomp_y {
val += pred_scores[(i, l_col)] * fit.fpca_y.rotation[(j, l_col)];
}
predicted[(i, j)] = val;
}
}
Ok(predicted)
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct FofCvResult {
pub candidates: Vec<(usize, usize)>,
pub cv_errors: Vec<f64>,
pub optimal: (usize, usize),
pub min_cv_mse: f64,
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn fof_cv(
x_data: &FdMatrix,
y_data: &FdMatrix,
x_argvals: &[f64],
y_argvals: &[f64],
ncomp_x_max: usize,
ncomp_y_max: usize,
n_folds: usize,
seed: u64,
) -> Result<FofCvResult, FdarError> {
let n = x_data.nrows();
if n < n_folds {
return Err(FdarError::InvalidDimension {
parameter: "x_data",
expected: format!("at least {n_folds} rows"),
actual: format!("{n}"),
});
}
let folds = crate::cv::create_folds(n, n_folds, seed);
let ncomp_x_max = ncomp_x_max.min(n - 2);
let ncomp_y_max = ncomp_y_max.min(n - 2);
let m_y = y_data.ncols();
let y_weights = crate::helpers::simpsons_weights(y_argvals);
let mut candidates = Vec::new();
let mut cv_errors = Vec::new();
let mut best = (1, 1);
let mut best_mse = f64::INFINITY;
for ncx in 1..=ncomp_x_max {
for ncy in 1..=ncomp_y_max {
let mut total_imse = 0.0;
let mut count = 0;
for fold in 0..n_folds {
let train_idx: Vec<usize> = (0..n).filter(|&i| folds[i] != fold).collect();
let test_idx: Vec<usize> = (0..n).filter(|&i| folds[i] == fold).collect();
let n_test = test_idx.len();
if n_test == 0 || train_idx.len() < ncx.max(ncy) + 2 {
continue;
}
let train_x = x_data.select_rows(&train_idx);
let train_y = y_data.select_rows(&train_idx);
let test_x = x_data.select_rows(&test_idx);
let test_y = y_data.select_rows(&test_idx);
let Ok(fit) = fof_regression(&train_x, &train_y, x_argvals, y_argvals, ncx, ncy)
else {
continue;
};
let Ok(predicted) = predict_fof(&fit, &test_x) else {
continue;
};
for ti in 0..n_test {
let imse: f64 = (0..m_y)
.map(|j| (test_y[(ti, j)] - predicted[(ti, j)]).powi(2) * y_weights[j])
.sum();
total_imse += imse;
count += 1;
}
}
let mse = if count > 0 {
total_imse / count as f64
} else {
f64::INFINITY
};
candidates.push((ncx, ncy));
cv_errors.push(mse);
if mse < best_mse {
best_mse = mse;
best = (ncx, ncy);
}
}
}
if candidates.is_empty() {
return Err(FdarError::ComputationFailed {
operation: "fof_cv",
detail: "no valid (ncomp_x, ncomp_y) produced CV errors".into(),
});
}
Ok(FofCvResult {
candidates,
cv_errors,
optimal: best,
min_cv_mse: best_mse,
})
}
#[cfg(test)]
mod tests {
use super::*;
use std::f64::consts::PI;
fn make_fof_data(
n: usize,
mx: usize,
my: usize,
seed: u64,
) -> (FdMatrix, FdMatrix, Vec<f64>, Vec<f64>) {
let tx: Vec<f64> = (0..mx).map(|j| j as f64 / (mx - 1).max(1) as f64).collect();
let ty: Vec<f64> = (0..my).map(|j| j as f64 / (my - 1).max(1) as f64).collect();
let mut x = FdMatrix::zeros(n, mx);
let mut y = FdMatrix::zeros(n, my);
for i in 0..n {
let a =
((seed.wrapping_mul(17).wrapping_add(i as u64 * 31) % 1000) as f64 / 500.0) - 1.0;
let b =
((seed.wrapping_mul(7).wrapping_add(i as u64 * 53) % 1000) as f64 / 500.0) - 1.0;
let c =
((seed.wrapping_mul(3).wrapping_add(i as u64 * 79) % 1000) as f64 / 500.0) - 1.0;
for j in 0..mx {
x[(i, j)] = a * (2.0 * PI * tx[j]).sin() + b * (4.0 * PI * tx[j]).cos() + c * tx[j];
}
for j in 0..my {
y[(i, j)] = 1.5 * a * (2.0 * PI * ty[j]).cos() - 0.8 * b * (3.0 * PI * ty[j]).sin()
+ 0.5 * c * ty[j].powi(2)
+ 0.01 * (seed.wrapping_add(i as u64 + j as u64) % 10) as f64;
}
}
(x, y, tx, ty)
}
#[test]
fn test_fof_regression_dimensions() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let fit = fof_regression(&x, &y, &tx, &ty, 3, 3).unwrap();
assert_eq!(fit.fitted.shape(), (30, 25));
assert_eq!(fit.residuals.shape(), (30, 25));
assert_eq!(fit.beta_surface.shape(), (25, 40));
assert_eq!(fit.intercept.len(), 25);
assert_eq!(fit.r_squared_t.len(), 25);
assert_eq!(fit.coef_matrix.shape(), (3, 3));
assert_eq!(fit.ncomp_x, 3);
assert_eq!(fit.ncomp_y, 3);
}
#[test]
fn test_fof_regression_r_squared_positive() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let fit = fof_regression(&x, &y, &tx, &ty, 3, 3).unwrap();
assert!(
fit.r_squared > 0.0,
"R² should be positive for correlated data, got {}",
fit.r_squared
);
}
#[test]
fn test_predict_fof_training_matches_fitted() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let fit = fof_regression(&x, &y, &tx, &ty, 3, 3).unwrap();
let predicted = predict_fof(&fit, &x).unwrap();
assert_eq!(predicted.shape(), fit.fitted.shape());
let (n, my) = predicted.shape();
for i in 0..n {
for j in 0..my {
assert!(
(predicted[(i, j)] - fit.fitted[(i, j)]).abs() < 1e-6,
"predicted should match fitted at ({i}, {j}): {} vs {}",
predicted[(i, j)],
fit.fitted[(i, j)]
);
}
}
}
#[test]
fn test_predict_fof_new_data_finite() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let fit = fof_regression(&x, &y, &tx, &ty, 3, 3).unwrap();
let n_new = 10;
let mx = 40;
let mut new_x = FdMatrix::zeros(n_new, mx);
for i in 0..n_new {
let p = (i as f64 + 0.5) * PI / n_new as f64;
for j in 0..mx {
new_x[(i, j)] = (2.0 * PI * tx[j] + p).cos();
}
}
let predicted = predict_fof(&fit, &new_x).unwrap();
assert_eq!(predicted.shape(), (n_new, 25));
for i in 0..n_new {
for j in 0..25 {
assert!(
predicted[(i, j)].is_finite(),
"prediction should be finite at ({i}, {j})"
);
}
}
}
#[test]
fn test_fof_regression_mismatched_n() {
let (x, _y, tx, ty) = make_fof_data(30, 40, 25, 42);
let y_bad = FdMatrix::zeros(20, 25);
let result = fof_regression(&x, &y_bad, &tx, &ty, 3, 3);
assert!(result.is_err());
}
#[test]
fn test_fof_regression_bad_argvals() {
let (x, y, _tx, ty) = make_fof_data(30, 40, 25, 42);
let bad_tx: Vec<f64> = (0..10).map(|j| j as f64).collect(); let result = fof_regression(&x, &y, &bad_tx, &ty, 3, 3);
assert!(result.is_err());
}
#[test]
fn test_fof_regression_zero_ncomp() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
assert!(fof_regression(&x, &y, &tx, &ty, 0, 3).is_err());
assert!(fof_regression(&x, &y, &tx, &ty, 3, 0).is_err());
}
#[test]
fn test_fof_cv() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let cv = fof_cv(&x, &y, &tx, &ty, 4, 4, 5, 42).unwrap();
assert!(!cv.candidates.is_empty());
assert!(cv.optimal.0 >= 1);
assert!(cv.optimal.1 >= 1);
assert!(cv.min_cv_mse.is_finite());
}
#[test]
fn test_fof_regression_residuals_consistent() {
let (x, y, tx, ty) = make_fof_data(30, 40, 25, 42);
let fit = fof_regression(&x, &y, &tx, &ty, 3, 3).unwrap();
let (n, my) = y.shape();
for i in 0..n {
for j in 0..my {
let expected_resid = y[(i, j)] - fit.fitted[(i, j)];
assert!(
(fit.residuals[(i, j)] - expected_resid).abs() < 1e-10,
"residual mismatch at ({i}, {j})"
);
}
}
}
}