use super::nonparametric::{compute_pairwise_distances, gaussian_kernel, select_bandwidth_loo};
use crate::error::FdarError;
use crate::matrix::FdMatrix;
use crate::regression::{fdata_to_pc_1d, FpcaResult};
use crate::smoothing::{nadaraya_watson, optim_bandwidth, CvCriterion};
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FamConfig {
pub ncomp: usize,
pub bandwidth: f64,
pub kernel: String,
pub n_grid_bandwidth: usize,
}
impl Default for FamConfig {
fn default() -> Self {
Self {
ncomp: 0,
bandwidth: 0.0,
kernel: "gaussian".to_string(),
n_grid_bandwidth: 20,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct GkamConfig {
pub bandwidth: f64,
pub kernel: String,
pub max_iter: usize,
pub epsilon: f64,
}
impl Default for GkamConfig {
fn default() -> Self {
Self {
bandwidth: 0.0,
kernel: "gaussian".to_string(),
max_iter: 50,
epsilon: 1e-6,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct GsamConfig {
pub ncomp: usize,
pub bandwidth: f64,
pub kernel: String,
pub n_grid_bandwidth: usize,
}
impl Default for GsamConfig {
fn default() -> Self {
Self {
ncomp: 0,
bandwidth: 0.0,
kernel: "gaussian".to_string(),
n_grid_bandwidth: 20,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct FamResult {
pub fitted_values: Vec<f64>,
pub residuals: Vec<f64>,
pub component_fits: Vec<Vec<f64>>,
pub intercept: f64,
pub bandwidths: Vec<f64>,
pub ncomp: usize,
pub r_squared: f64,
pub fpca: FpcaResult,
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct GkamResult {
pub fitted_values: Vec<f64>,
pub residuals: Vec<f64>,
pub component_fits: Vec<Vec<f64>>,
pub intercept: f64,
pub bandwidths: Vec<f64>,
pub iterations: usize,
pub converged: bool,
pub r_squared: f64,
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct GsamResult {
pub fitted_values: Vec<f64>,
pub residuals: Vec<f64>,
pub component_fits: Vec<Vec<f64>>,
pub intercept: f64,
pub bandwidths: Vec<f64>,
pub ncomp: usize,
pub r_squared: f64,
pub fpca: FpcaResult,
}
fn resolve_ncomp_additive(
ncomp: usize,
n: usize,
m: usize,
data: &FdMatrix,
y: &[f64],
argvals: &[f64],
kernel: &str,
n_grid: usize,
) -> Result<usize, FdarError> {
let max_ncomp = n.min(m);
if ncomp == 0 {
let cap = max_ncomp.clamp(1, 10);
let fpca_full = fdata_to_pc_1d(data, cap, argvals)?;
let mu_y = y.iter().sum::<f64>() / n as f64;
let mut best_ncomp = 1usize;
let mut best_gcv = f64::INFINITY;
let mut component_fits_acc: Vec<Vec<f64>> = Vec::with_capacity(cap);
for j in 0..cap {
let xi_j: Vec<f64> = (0..n).map(|i| fpca_full.scores[(i, j)]).collect();
let partial: Vec<f64> = (0..n)
.map(|i| {
let prior_sum: f64 = component_fits_acc.iter().map(|cf| cf[i]).sum();
y[i] - mu_y - prior_sum
})
.collect();
let bw_result =
optim_bandwidth(&xi_j, &partial, None, CvCriterion::Gcv, kernel, n_grid);
let gcv_j = bw_result.value;
let fit_j = nadaraya_watson(&xi_j, &partial, &xi_j, bw_result.h_opt, kernel)
.unwrap_or_else(|_| vec![0.0; n]);
component_fits_acc.push(fit_j);
if gcv_j < best_gcv {
best_gcv = gcv_j;
best_ncomp = j + 1; }
}
Ok(best_ncomp)
} else if ncomp > max_ncomp {
Err(FdarError::InvalidParameter {
parameter: "config.ncomp",
message: format!(
"ncomp ({ncomp}) exceeds min(n, m) = {max_ncomp}; reduce ncomp or provide more data"
),
})
} else {
Ok(ncomp)
}
}
#[allow(clippy::too_many_arguments)]
fn fpc_additive_smooth(
fpca: &FpcaResult,
y: &[f64],
n: usize,
ncomp: usize,
bandwidth: f64,
kernel: &str,
n_grid: usize,
scalar_covariates: Option<&FdMatrix>,
) -> Result<(Vec<Vec<f64>>, Vec<f64>, f64, Vec<f64>, Vec<f64>, f64), FdarError> {
let mu_y = y.iter().sum::<f64>() / n as f64;
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let total_comp = ncomp + p_scalar;
let mut all_scores: Vec<Vec<f64>> = Vec::with_capacity(total_comp);
for k in 0..ncomp {
all_scores.push((0..n).map(|i| fpca.scores[(i, k)]).collect());
}
if let Some(sc) = scalar_covariates {
for j in 0..p_scalar {
all_scores.push((0..n).map(|i| sc[(i, j)]).collect());
}
}
let mut component_fits: Vec<Vec<f64>> = vec![vec![0.0; n]; total_comp];
let mut bandwidths = vec![0.0_f64; total_comp];
for k in 0..total_comp {
let partial: Vec<f64> = (0..n)
.map(|i| {
let others: f64 = (0..total_comp)
.filter(|&j| j != k)
.map(|j| component_fits[j][i])
.sum();
y[i] - mu_y - others
})
.collect();
let xi_k = &all_scores[k];
let h = if bandwidth > 0.0 {
bandwidth
} else {
optim_bandwidth(xi_k, &partial, None, CvCriterion::Gcv, kernel, n_grid).h_opt
};
bandwidths[k] = h;
component_fits[k] = nadaraya_watson(xi_k, &partial, xi_k, h, kernel)?;
}
let fitted_values: Vec<f64> = (0..n)
.map(|i| mu_y + (0..total_comp).map(|k| component_fits[k][i]).sum::<f64>())
.collect();
let residuals: Vec<f64> = y
.iter()
.zip(&fitted_values)
.map(|(&yi, &yh)| yi - yh)
.collect();
let (r_squared, _) = super::compute_r_squared(y, &residuals, total_comp);
Ok((
component_fits,
bandwidths,
mu_y,
fitted_values,
residuals,
r_squared,
))
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn fam(
data: &FdMatrix,
y: &[f64],
argvals: &[f64],
scalar_covariates: Option<&FdMatrix>,
config: &FamConfig,
) -> Result<FamResult, FdarError> {
let (n, m) = data.shape();
if n == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 row".to_string(),
actual: "0".to_string(),
});
}
if m == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 column".to_string(),
actual: "0".to_string(),
});
}
if y.len() != n {
return Err(FdarError::InvalidDimension {
parameter: "y",
expected: format!("{n}"),
actual: format!("{}", y.len()),
});
}
if argvals.len() != m {
return Err(FdarError::InvalidDimension {
parameter: "argvals",
expected: format!("{m}"),
actual: format!("{}", argvals.len()),
});
}
if let Some(sc) = scalar_covariates {
if sc.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "scalar_covariates",
expected: format!("{n} rows"),
actual: format!("{} rows", sc.nrows()),
});
}
}
let ncomp = resolve_ncomp_additive(
config.ncomp,
n,
m,
data,
y,
argvals,
&config.kernel,
config.n_grid_bandwidth,
)?;
let fpca = fdata_to_pc_1d(data, ncomp, argvals)?;
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let total_comp = ncomp + p_scalar;
let (component_fits_all, bandwidths_all, intercept, fitted_values, residuals, r_squared) =
fpc_additive_smooth(
&fpca,
y,
n,
ncomp,
config.bandwidth,
&config.kernel,
config.n_grid_bandwidth,
scalar_covariates,
)?;
let component_fits: Vec<Vec<f64>> = component_fits_all.into_iter().take(total_comp).collect();
let bandwidths: Vec<f64> = bandwidths_all.into_iter().take(total_comp).collect();
Ok(FamResult {
fitted_values,
residuals,
component_fits,
intercept,
bandwidths,
ncomp,
r_squared,
fpca,
})
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn fregre_gkam(
predictors: &[&FdMatrix],
y: &[f64],
argvals_list: &[&[f64]],
scalar_covariates: Option<&FdMatrix>,
config: &GkamConfig,
) -> Result<GkamResult, FdarError> {
let n = y.len();
if n == 0 {
return Err(FdarError::InvalidDimension {
parameter: "y",
expected: "at least 1 observation".to_string(),
actual: "0".to_string(),
});
}
if predictors.is_empty() {
return Err(FdarError::InvalidDimension {
parameter: "predictors",
expected: "at least 1 functional predictor".to_string(),
actual: "0".to_string(),
});
}
if predictors.len() != argvals_list.len() {
return Err(FdarError::InvalidDimension {
parameter: "argvals_list",
expected: format!("{} (matching predictors.len())", predictors.len()),
actual: format!("{}", argvals_list.len()),
});
}
for (k, pred) in predictors.iter().enumerate() {
if pred.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "predictors[k].nrows()",
expected: format!("{n} (y.len())"),
actual: format!("{} for predictor {k}", pred.nrows()),
});
}
if argvals_list[k].len() != pred.ncols() {
return Err(FdarError::InvalidDimension {
parameter: "argvals_list[k]",
expected: format!("{} (predictors[k].ncols())", pred.ncols()),
actual: format!("{} for predictor {k}", argvals_list[k].len()),
});
}
}
if let Some(sc) = scalar_covariates {
if sc.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "scalar_covariates",
expected: format!("{n} rows"),
actual: format!("{} rows", sc.nrows()),
});
}
}
let q = predictors.len();
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let total_comp = q + p_scalar;
let mu_y = y.iter().sum::<f64>() / n as f64;
let dist_matrices: Vec<Vec<f64>> = predictors
.iter()
.zip(argvals_list.iter())
.map(|(pred, argvals)| compute_pairwise_distances(pred, argvals))
.collect();
let bandwidths_func: Vec<f64> = if config.bandwidth > 0.0 {
vec![config.bandwidth; q]
} else {
dist_matrices
.iter()
.map(|dists| select_bandwidth_loo(dists, y, n, None))
.collect()
};
let scalar_dists: Vec<Vec<f64>> = if let Some(sc) = scalar_covariates {
(0..p_scalar)
.map(|j| {
let mut d = vec![0.0_f64; n * n];
for i in 0..n {
for jj in (i + 1)..n {
let diff = sc[(i, j)] - sc[(jj, j)];
let dist = diff.abs();
d[i * n + jj] = dist;
d[jj * n + i] = dist;
}
}
d
})
.collect()
} else {
Vec::new()
};
let scalar_bandwidths: Vec<f64> = if p_scalar > 0 {
if config.bandwidth > 0.0 {
vec![config.bandwidth; p_scalar]
} else {
scalar_dists
.iter()
.map(|dists| select_bandwidth_loo(dists, y, n, None))
.collect()
}
} else {
Vec::new()
};
let mut all_bandwidths = bandwidths_func.clone();
all_bandwidths.extend_from_slice(&scalar_bandwidths);
let mut component_fits = vec![vec![0.0_f64; n]; total_comp];
let mut converged = false;
let mut iterations = 0;
for iter in 0..config.max_iter {
let mut max_delta = 0.0_f64;
for k in 0..q {
let h_k = all_bandwidths[k];
let dists_k = &dist_matrices[k];
let adjusted: Vec<f64> = (0..n)
.map(|i| {
let others: f64 = (0..total_comp)
.filter(|&j| j != k)
.map(|j| component_fits[j][i])
.sum();
y[i] - mu_y - others
})
.collect();
let new_fk: Vec<f64> = (0..n)
.map(|i| {
let mut num = 0.0_f64;
let mut den = 0.0_f64;
for j in 0..n {
let w = gaussian_kernel(dists_k[i * n + j], h_k);
num += w * adjusted[j];
den += w;
}
if den > 1e-15 {
num / den
} else {
adjusted[i]
}
})
.collect();
let delta = component_fits[k]
.iter()
.zip(&new_fk)
.map(|(old, &new)| (old - new).abs())
.fold(0.0_f64, f64::max);
max_delta = max_delta.max(delta);
component_fits[k] = new_fk;
}
for s_idx in 0..p_scalar {
let k = q + s_idx;
let h_k = all_bandwidths[k];
let dists_k = &scalar_dists[s_idx];
let adjusted: Vec<f64> = (0..n)
.map(|i| {
let others: f64 = (0..total_comp)
.filter(|&j| j != k)
.map(|j| component_fits[j][i])
.sum();
y[i] - mu_y - others
})
.collect();
let new_fk: Vec<f64> = (0..n)
.map(|i| {
let mut num = 0.0_f64;
let mut den = 0.0_f64;
for j in 0..n {
let w = gaussian_kernel(dists_k[i * n + j], h_k);
num += w * adjusted[j];
den += w;
}
if den > 1e-15 {
num / den
} else {
adjusted[i]
}
})
.collect();
let delta = component_fits[k]
.iter()
.zip(&new_fk)
.map(|(old, &new)| (old - new).abs())
.fold(0.0_f64, f64::max);
max_delta = max_delta.max(delta);
component_fits[k] = new_fk;
}
iterations = iter + 1;
if max_delta < config.epsilon {
converged = true;
break;
}
}
let fitted_values: Vec<f64> = (0..n)
.map(|i| mu_y + (0..total_comp).map(|k| component_fits[k][i]).sum::<f64>())
.collect();
let residuals: Vec<f64> = y
.iter()
.zip(&fitted_values)
.map(|(&yi, &yh)| yi - yh)
.collect();
let (r_squared, _) = super::compute_r_squared(y, &residuals, total_comp);
Ok(GkamResult {
fitted_values,
residuals,
component_fits,
intercept: mu_y,
bandwidths: all_bandwidths,
iterations,
converged,
r_squared,
})
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn fregre_gsam(
data: &FdMatrix,
y: &[f64],
argvals: &[f64],
scalar_covariates: Option<&FdMatrix>,
config: &GsamConfig,
) -> Result<GsamResult, FdarError> {
let (n, m) = data.shape();
if n == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 row".to_string(),
actual: "0".to_string(),
});
}
if m == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 column".to_string(),
actual: "0".to_string(),
});
}
if y.len() != n {
return Err(FdarError::InvalidDimension {
parameter: "y",
expected: format!("{n}"),
actual: format!("{}", y.len()),
});
}
if argvals.len() != m {
return Err(FdarError::InvalidDimension {
parameter: "argvals",
expected: format!("{m}"),
actual: format!("{}", argvals.len()),
});
}
if let Some(sc) = scalar_covariates {
if sc.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "scalar_covariates",
expected: format!("{n} rows"),
actual: format!("{} rows", sc.nrows()),
});
}
}
let ncomp = resolve_ncomp_additive(
config.ncomp,
n,
m,
data,
y,
argvals,
&config.kernel,
config.n_grid_bandwidth,
)?;
let fpca = fdata_to_pc_1d(data, ncomp, argvals)?;
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let total_comp = ncomp + p_scalar;
let (component_fits_all, bandwidths_all, intercept, fitted_values, residuals, r_squared) =
fpc_additive_smooth(
&fpca,
y,
n,
ncomp,
config.bandwidth,
&config.kernel,
config.n_grid_bandwidth,
scalar_covariates,
)?;
let component_fits: Vec<Vec<f64>> = component_fits_all.into_iter().take(total_comp).collect();
let bandwidths: Vec<f64> = bandwidths_all.into_iter().take(total_comp).collect();
Ok(GsamResult {
fitted_values,
residuals,
component_fits,
intercept,
bandwidths,
ncomp,
r_squared,
fpca,
})
}
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum VarSelectPenalty {
GroupLasso,
GroupMcp,
GroupScad,
Ls,
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct VarSelectConfig {
pub ncomp: usize,
pub penalty: VarSelectPenalty,
pub lambda: f64,
pub max_iter: usize,
pub epsilon: f64,
pub lambda_n_grid: usize,
}
impl Default for VarSelectConfig {
fn default() -> Self {
Self {
ncomp: 3,
penalty: VarSelectPenalty::GroupLasso,
lambda: 0.0,
max_iter: 100,
epsilon: 1e-5,
lambda_n_grid: 20,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct VarSelectResult {
pub active_predictors: Vec<bool>,
pub coefficients: Vec<Vec<f64>>,
pub fitted_values: Vec<f64>,
pub residuals: Vec<f64>,
pub intercept: f64,
pub lambda: f64,
pub r_squared: f64,
pub iterations: usize,
pub converged: bool,
pub fpcas: Vec<FpcaResult>,
}
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub enum PermTestStatistic {
R2,
FittedNorm,
ComponentNorm,
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct PermTestConfig {
pub n_perm: usize,
pub seed: u64,
pub statistic: PermTestStatistic,
}
impl Default for PermTestConfig {
fn default() -> Self {
Self {
n_perm: 999,
seed: 42,
statistic: PermTestStatistic::R2,
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
pub struct PermTestResult {
pub p_value: f64,
pub observed_statistic: f64,
pub null_statistics: Vec<f64>,
pub n_perm_success: usize,
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct HistoryIndexConfig {
pub window: f64,
pub n_lags: usize,
pub bandwidth: f64,
pub kernel: String,
}
impl Default for HistoryIndexConfig {
fn default() -> Self {
Self {
window: 1.0,
n_lags: 20,
bandwidth: 0.0,
kernel: "gaussian".to_string(),
}
}
}
#[derive(Debug, Clone, PartialEq)]
#[non_exhaustive]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct HistoryIndexResult {
pub fitted_values: Vec<f64>,
pub residuals: Vec<f64>,
pub intercept: f64,
pub slope: f64,
pub gamma: Vec<f64>,
pub lag_grid: Vec<f64>,
pub history_scores: Vec<f64>,
pub r_squared: f64,
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn variable_selection(
predictors: &[&FdMatrix],
y: &[f64],
argvals_list: &[&[f64]],
scalar_covariates: Option<&FdMatrix>,
config: &VarSelectConfig,
) -> Result<VarSelectResult, FdarError> {
match config.penalty {
VarSelectPenalty::GroupMcp | VarSelectPenalty::GroupScad => {
return Err(FdarError::InvalidParameter {
parameter: "config.penalty",
message: "GroupMcp and GroupScad are not yet implemented; use GroupLasso"
.to_string(),
});
}
VarSelectPenalty::GroupLasso | VarSelectPenalty::Ls => {}
}
let n = y.len();
if predictors.is_empty() {
return Err(FdarError::InvalidDimension {
parameter: "predictors",
expected: "at least 1 functional predictor".to_string(),
actual: "0".to_string(),
});
}
if predictors.len() != argvals_list.len() {
return Err(FdarError::InvalidDimension {
parameter: "argvals_list",
expected: format!("{} (matching predictors.len())", predictors.len()),
actual: format!("{}", argvals_list.len()),
});
}
for (p, pred) in predictors.iter().enumerate() {
if pred.nrows() != n {
return Err(FdarError::InvalidDimension {
parameter: "predictors[p].nrows()",
expected: format!("{n} (y.len())"),
actual: format!("{} for predictor {p}", pred.nrows()),
});
}
}
let big_p = predictors.len();
let mu_y = y.iter().sum::<f64>() / n as f64;
let ncomp_per = if config.ncomp == 0 { 3 } else { config.ncomp };
let mut fpcas: Vec<FpcaResult> = Vec::with_capacity(big_p);
let mut score_groups: Vec<Vec<Vec<f64>>> = Vec::with_capacity(big_p);
for p in 0..big_p {
let pred = predictors[p];
let argvals = argvals_list[p];
let (np, mp) = pred.shape();
let k_p = ncomp_per.min(np.min(mp).saturating_sub(1).max(1));
let fpca_p = fdata_to_pc_1d(pred, k_p, argvals)?;
let k_actual = fpca_p.scores.ncols();
let group_scores: Vec<Vec<f64>> = (0..k_actual)
.map(|k| (0..n).map(|i| fpca_p.scores[(i, k)]).collect())
.collect();
score_groups.push(group_scores);
fpcas.push(fpca_p);
}
if config.penalty == VarSelectPenalty::Ls {
return variable_selection_ls(y, n, mu_y, big_p, fpcas, score_groups, scalar_covariates);
}
let k_sizes: Vec<usize> = score_groups.iter().map(|g| g.len()).collect();
let y_centered: Vec<f64> = y.iter().map(|&yi| yi - mu_y).collect();
let lambda_max = k_sizes
.iter()
.zip(score_groups.iter())
.map(|(&k_g, group)| {
let norm_sq: f64 = group
.iter()
.map(|col| {
let xgty: f64 = col.iter().zip(&y_centered).map(|(&x, &yc)| x * yc).sum();
xgty * xgty
})
.sum::<f64>();
norm_sq.sqrt() / (k_g as f64).sqrt()
})
.fold(0.0_f64, f64::max)
.max(1e-10);
let lambda = if config.lambda > 0.0 {
config.lambda
} else {
select_group_lasso_lambda(
y,
&y_centered,
mu_y,
n,
&score_groups,
&k_sizes,
lambda_max,
config.lambda_n_grid,
config.max_iter,
config.epsilon,
scalar_covariates,
)
};
let (coefficients, iterations, converged) = group_lasso_cd(
y,
&y_centered,
mu_y,
n,
&score_groups,
&k_sizes,
lambda,
config.max_iter,
config.epsilon,
scalar_covariates,
)?;
let fitted_values: Vec<f64> = compute_varselect_fitted(
n,
mu_y,
&score_groups,
&coefficients,
scalar_covariates,
big_p,
);
let residuals: Vec<f64> = y
.iter()
.zip(&fitted_values)
.map(|(&yi, &yh)| yi - yh)
.collect();
let (r_squared, _) = super::compute_r_squared(y, &residuals, k_sizes.iter().sum::<usize>());
let active_predictors: Vec<bool> = coefficients[..big_p]
.iter()
.map(|beta_g| {
let norm: f64 = beta_g.iter().map(|&b| b * b).sum::<f64>();
norm.sqrt() > config.epsilon
})
.collect();
Ok(VarSelectResult {
active_predictors,
coefficients,
fitted_values,
residuals,
intercept: mu_y,
lambda,
r_squared,
iterations,
converged,
fpcas,
})
}
fn variable_selection_ls(
y: &[f64],
n: usize,
mu_y: f64,
big_p: usize,
fpcas: Vec<FpcaResult>,
score_groups: Vec<Vec<Vec<f64>>>,
scalar_covariates: Option<&FdMatrix>,
) -> Result<VarSelectResult, FdarError> {
let k_sizes: Vec<usize> = score_groups.iter().map(|g| g.len()).collect();
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let total_cols = k_sizes.iter().sum::<usize>() + p_scalar;
let mut x_flat = vec![0.0_f64; n * total_cols];
let mut col_offset = 0;
for grp in &score_groups {
for col in grp {
for (i, &v) in col.iter().enumerate() {
x_flat[col_offset * n + i] = v;
}
col_offset += 1;
}
}
if let Some(sc) = scalar_covariates {
for j in 0..p_scalar {
for i in 0..n {
x_flat[col_offset * n + i] = sc[(i, j)];
}
col_offset += 1;
}
}
let x_mat = FdMatrix::from_column_major(x_flat, n, total_cols).map_err(|e| {
FdarError::ComputationFailed {
operation: "variable_selection_ls design matrix",
detail: e.to_string(),
}
})?;
let y_centered: Vec<f64> = y.iter().map(|&yi| yi - mu_y).collect();
let xtx = super::compute_xtx(&x_mat);
let xty: Vec<f64> = (0..total_cols)
.map(|k| {
x_mat
.column(k)
.iter()
.zip(&y_centered)
.map(|(&xv, &yv)| xv * yv)
.sum::<f64>()
})
.collect();
let l = super::cholesky_factor(&xtx, total_cols).map_err(|_| FdarError::ComputationFailed {
operation: "variable_selection_ls cholesky",
detail: "design matrix is singular".to_string(),
})?;
let flat_coeffs = super::cholesky_forward_back(&l, &xty, total_cols);
let mut coefficients: Vec<Vec<f64>> = Vec::with_capacity(big_p + 1);
let mut offset = 0;
for &k_g in &k_sizes {
coefficients.push(flat_coeffs[offset..offset + k_g].to_vec());
offset += k_g;
}
coefficients.push(flat_coeffs[offset..offset + p_scalar].to_vec());
let fitted_values = compute_varselect_fitted(
n,
mu_y,
&score_groups,
&coefficients,
scalar_covariates,
big_p,
);
let residuals: Vec<f64> = y
.iter()
.zip(&fitted_values)
.map(|(&yi, &yh)| yi - yh)
.collect();
let (r_squared, _) = super::compute_r_squared(y, &residuals, total_cols);
let active_predictors = vec![true; big_p];
Ok(VarSelectResult {
active_predictors,
coefficients,
fitted_values,
residuals,
intercept: mu_y,
lambda: 0.0,
r_squared,
iterations: 1,
converged: true,
fpcas,
})
}
fn compute_varselect_fitted(
n: usize,
mu_y: f64,
score_groups: &[Vec<Vec<f64>>],
coefficients: &[Vec<f64>],
scalar_covariates: Option<&FdMatrix>,
big_p: usize,
) -> Vec<f64> {
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
(0..n)
.map(|i| {
let mut yhat = mu_y;
for p in 0..big_p {
for (k, col) in score_groups[p].iter().enumerate() {
yhat += coefficients[p][k] * col[i];
}
}
if let Some(sc) = scalar_covariates {
for j in 0..p_scalar {
yhat += coefficients[big_p][j] * sc[(i, j)];
}
}
yhat
})
.collect()
}
#[allow(clippy::too_many_arguments)]
fn select_group_lasso_lambda(
y: &[f64],
_y_centered: &[f64],
_mu_y: f64,
n: usize,
score_groups: &[Vec<Vec<f64>>],
_k_sizes: &[usize],
lambda_max: f64,
n_grid: usize,
max_iter: usize,
epsilon: f64,
scalar_covariates: Option<&FdMatrix>,
) -> f64 {
let grid_size = n_grid.max(2);
let big_p = score_groups.len();
let n_folds = 5_usize.min(n).max(2);
let fold_of: Vec<usize> = (0..n).map(|i| i % n_folds).collect();
let mut best_lambda = lambda_max * 0.1;
let mut best_cv_err = f64::INFINITY;
for gi in 0..grid_size {
let frac = (gi as f64 + 1.0) / grid_size as f64;
let lam = lambda_max * (0.01_f64.powf(1.0 - frac));
let mut cv_sq_err = 0.0_f64;
let mut cv_count = 0usize;
for fold in 0..n_folds {
let train_idx: Vec<usize> = (0..n).filter(|&i| fold_of[i] != fold).collect();
let val_idx: Vec<usize> = (0..n).filter(|&i| fold_of[i] == fold).collect();
if train_idx.is_empty() || val_idx.is_empty() {
continue;
}
let n_tr = train_idx.len();
let mu_tr = train_idx.iter().map(|&i| y[i]).sum::<f64>() / n_tr as f64;
let y_tr_centered: Vec<f64> = train_idx.iter().map(|&i| y[i] - mu_tr).collect();
let y_tr: Vec<f64> = train_idx.iter().map(|&i| y[i]).collect();
let sg_tr: Vec<Vec<Vec<f64>>> = score_groups
.iter()
.map(|grp| {
grp.iter()
.map(|col| train_idx.iter().map(|&i| col[i]).collect())
.collect()
})
.collect();
let sc_tr_mat: Option<FdMatrix> = scalar_covariates.and_then(|sc| {
let p_sc = sc.ncols();
let mut cm = vec![0.0_f64; n_tr * p_sc];
for (row, &orig_i) in train_idx.iter().enumerate() {
for j in 0..p_sc {
cm[j * n_tr + row] = sc[(orig_i, j)];
}
}
FdMatrix::from_column_major(cm, n_tr, p_sc).ok()
});
let k_sizes_tr: Vec<usize> = sg_tr.iter().map(|g| g.len()).collect();
let fit_result = group_lasso_cd(
&y_tr,
&y_tr_centered,
mu_tr,
n_tr,
&sg_tr,
&k_sizes_tr,
lam,
max_iter,
epsilon,
sc_tr_mat.as_ref(),
);
if let Ok((coeffs_tr, _, _)) = fit_result {
for &i in &val_idx {
let mut yhat = mu_tr;
for p in 0..big_p {
for (k, col) in score_groups[p].iter().enumerate() {
yhat += coeffs_tr[p][k] * col[i];
}
}
if let Some(sc) = scalar_covariates {
let p_sc = sc.ncols();
for j in 0..p_sc {
yhat += coeffs_tr[big_p][j] * sc[(i, j)];
}
}
let err = y[i] - yhat;
cv_sq_err += err * err;
cv_count += 1;
}
}
}
if cv_count > 0 {
let cv_mse = cv_sq_err / cv_count as f64;
if cv_mse < best_cv_err {
best_cv_err = cv_mse;
best_lambda = lam;
}
}
}
best_lambda
}
#[allow(clippy::too_many_arguments)]
fn group_lasso_cd(
_y: &[f64],
y_centered: &[f64],
_mu_y: f64,
n: usize,
score_groups: &[Vec<Vec<f64>>],
k_sizes: &[usize],
lambda: f64,
max_iter: usize,
epsilon: f64,
scalar_covariates: Option<&FdMatrix>,
) -> Result<(Vec<Vec<f64>>, usize, bool), FdarError> {
let big_p = score_groups.len();
let p_scalar = scalar_covariates.map_or(0, FdMatrix::ncols);
let mut beta_groups: Vec<Vec<f64>> = score_groups
.iter()
.map(|grp| vec![0.0_f64; grp.len()])
.collect();
let mut beta_scalar: Vec<f64> = vec![0.0_f64; p_scalar];
let mut converged = false;
let mut iterations = 0;
for _iter in 0..max_iter {
let mut max_delta = 0.0_f64;
for p in 0..big_p {
let k_g = k_sizes[p];
let group = &score_groups[p];
let partial: Vec<f64> = (0..n)
.map(|i| {
let mut res = y_centered[i];
for q in 0..big_p {
if q != p {
for (k, col) in score_groups[q].iter().enumerate() {
res -= beta_groups[q][k] * col[i];
}
}
}
if let Some(sc) = scalar_covariates {
for j in 0..p_scalar {
res -= beta_scalar[j] * sc[(i, j)];
}
}
res
})
.collect();
let mut xtx_g = vec![0.0_f64; k_g * k_g];
let mut xty_g = vec![0.0_f64; k_g];
for a in 0..k_g {
for b in 0..k_g {
let dot: f64 = group[a]
.iter()
.zip(&group[b])
.map(|(&xa, &xb)| xa * xb)
.sum();
xtx_g[a * k_g + b] = dot;
}
xty_g[a] = group[a]
.iter()
.zip(&partial)
.map(|(&xa, &pa)| xa * pa)
.sum();
}
let beta_ols =
crate::linalg::cholesky_solve(&xtx_g, &xty_g, k_g).unwrap_or_else(|_| {
let diag_sum: f64 = (0..k_g).map(|d| xtx_g[d * k_g + d].abs()).sum();
let delta = (1e-6 * diag_sum / k_g as f64).max(1e-8);
let mut xtx_ridge = xtx_g.clone();
for d in 0..k_g {
xtx_ridge[d * k_g + d] += delta;
}
crate::linalg::cholesky_solve(&xtx_ridge, &xty_g, k_g)
.unwrap_or_else(|_| vec![0.0; k_g]) });
let norm_ols: f64 = beta_ols.iter().map(|&b| b * b).sum::<f64>().sqrt();
let threshold = lambda * (k_g as f64).sqrt();
let scale = if norm_ols > 1e-15 {
(1.0 - threshold / norm_ols).max(0.0)
} else {
0.0
};
let new_beta: Vec<f64> = beta_ols.iter().map(|&b| b * scale).collect();
let delta = new_beta
.iter()
.zip(&beta_groups[p])
.map(|(&nb, &ob)| (nb - ob).abs())
.fold(0.0_f64, f64::max);
max_delta = max_delta.max(delta);
beta_groups[p] = new_beta;
}
if let Some(sc) = scalar_covariates {
for j in 0..p_scalar {
let partial_j: Vec<f64> = (0..n)
.map(|i| {
let mut res = y_centered[i];
for p in 0..big_p {
for (k, col) in score_groups[p].iter().enumerate() {
res -= beta_groups[p][k] * col[i];
}
}
for jj in 0..p_scalar {
if jj != j {
res -= beta_scalar[jj] * sc[(i, jj)];
}
}
res
})
.collect();
let col_j: Vec<f64> = (0..n).map(|i| sc[(i, j)]).collect();
let xjxj: f64 = col_j.iter().map(|&v| v * v).sum();
let xjy: f64 = col_j.iter().zip(&partial_j).map(|(&x, &p)| x * p).sum();
let new_bj = if xjxj > 1e-15 { xjy / xjxj } else { 0.0 };
let delta = (new_bj - beta_scalar[j]).abs();
max_delta = max_delta.max(delta);
beta_scalar[j] = new_bj;
}
}
iterations = _iter + 1;
if max_delta < epsilon {
converged = true;
break;
}
}
let mut coefficients: Vec<Vec<f64>> = beta_groups;
coefficients.push(beta_scalar);
Ok((coefficients, iterations, converged))
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn permutation_test_fam(
data: &FdMatrix,
y: &[f64],
argvals: &[f64],
scalar_covariates: Option<&FdMatrix>,
config: &FamConfig,
perm_config: &PermTestConfig,
) -> Result<PermTestResult, FdarError> {
if perm_config.n_perm == 0 {
return Err(FdarError::InvalidParameter {
parameter: "perm_config.n_perm",
message: "n_perm must be >= 1 for a meaningful permutation test".to_string(),
});
}
let original_fit = fam(data, y, argvals, scalar_covariates, config)?;
let observed_statistic = extract_perm_stat(&original_fit, perm_config.statistic);
use rand::prelude::*;
let mut rng = StdRng::seed_from_u64(perm_config.seed);
let n_perm = perm_config.n_perm;
let mut null_statistics: Vec<f64> = Vec::with_capacity(n_perm);
let mut n_ge = 0usize;
let mut n_perm_success = 0usize;
let mut y_perm: Vec<f64> = y.to_vec();
for _ in 0..n_perm {
y_perm.copy_from_slice(y);
y_perm.shuffle(&mut rng);
match fam(data, &y_perm, argvals, scalar_covariates, config) {
Ok(perm_fit) => {
let t_perm = extract_perm_stat(&perm_fit, perm_config.statistic);
null_statistics.push(t_perm);
if t_perm >= observed_statistic {
n_ge += 1;
}
n_perm_success += 1;
}
Err(_) => {
}
}
}
let p_value = (n_ge + 1) as f64 / (n_perm_success + 1) as f64;
Ok(PermTestResult {
p_value,
observed_statistic,
null_statistics,
n_perm_success,
})
}
fn extract_perm_stat(fit: &FamResult, stat: PermTestStatistic) -> f64 {
match stat {
PermTestStatistic::R2 => fit.r_squared,
PermTestStatistic::FittedNorm => {
fit.fitted_values.iter().map(|&v| v * v).sum::<f64>().sqrt()
}
PermTestStatistic::ComponentNorm => fit
.component_fits
.iter()
.map(|cf| cf.iter().map(|&v| v * v).sum::<f64>().sqrt())
.sum::<f64>(),
}
}
#[must_use = "expensive computation whose result should not be discarded"]
pub fn history_index(
data: &FdMatrix,
y: &[f64],
argvals: &[f64],
config: &HistoryIndexConfig,
) -> Result<HistoryIndexResult, FdarError> {
let (n, m) = data.shape();
if n == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 row".to_string(),
actual: "0".to_string(),
});
}
if m == 0 {
return Err(FdarError::InvalidDimension {
parameter: "data",
expected: "at least 1 column".to_string(),
actual: "0".to_string(),
});
}
if y.len() != n {
return Err(FdarError::InvalidDimension {
parameter: "y",
expected: format!("{n}"),
actual: format!("{}", y.len()),
});
}
if argvals.len() != m {
return Err(FdarError::InvalidDimension {
parameter: "argvals",
expected: format!("{m}"),
actual: format!("{}", argvals.len()),
});
}
let argvals_min = argvals.first().copied().unwrap_or(0.0);
let argvals_max = argvals.last().copied().unwrap_or(0.0);
let argvals_range = argvals_max - argvals_min;
if config.window <= 0.0 || config.window > argvals_range {
return Err(FdarError::InvalidParameter {
parameter: "config.window",
message: format!(
"window ({:.6}) must be positive and <= argvals range ({:.6})",
config.window, argvals_range
),
});
}
let n_lags = config.n_lags.max(1);
let delta_u = config.window / n_lags as f64;
let big_t = argvals_max;
let lag_grid: Vec<f64> = (0..n_lags).map(|l| l as f64 * delta_u).collect();
let x_lag: Vec<Vec<f64>> = (0..n)
.map(|i| {
lag_grid
.iter()
.map(|&u_l| {
let t_target = big_t - u_l;
let j = argvals
.partition_point(|&v| v < t_target)
.saturating_sub(1)
.min(m - 1);
data[(i, j)]
})
.collect()
})
.collect();
let mu_y = y.iter().sum::<f64>() / n as f64;
let y_centered: Vec<f64> = y.iter().map(|&yi| yi - mu_y).collect();
let gamma_signal: Vec<f64> = lag_grid
.iter()
.enumerate()
.map(|(l, _)| {
let x_col: Vec<f64> = (0..n).map(|i| x_lag[i][l]).collect();
let x_mean = x_col.iter().sum::<f64>() / n as f64;
let xx: f64 = x_col.iter().map(|&v| (v - x_mean).powi(2)).sum();
let xy: f64 = x_col
.iter()
.zip(&y_centered)
.map(|(&x, &yc)| (x - x_mean) * yc)
.sum();
if xx > 1e-15 {
xy / xx
} else {
0.0
}
})
.collect();
let h_gamma = if config.bandwidth > 0.0 {
config.bandwidth
} else {
let bw_result = optim_bandwidth(
&lag_grid,
&gamma_signal,
None,
CvCriterion::Gcv,
&config.kernel,
20,
);
bw_result.h_opt.max(delta_u) };
let gamma_raw = nadaraya_watson(&lag_grid, &gamma_signal, &lag_grid, h_gamma, &config.kernel)?;
let norm_sq: f64 = gamma_raw.iter().map(|&g| g * g).sum::<f64>() * delta_u;
let norm = norm_sq.sqrt();
let gamma: Vec<f64> = if norm > 1e-15 {
gamma_raw.iter().map(|&g| g / norm).collect()
} else {
vec![1.0 / (n_lags as f64).sqrt(); n_lags]
};
let history_scores: Vec<f64> = (0..n)
.map(|i| {
gamma
.iter()
.enumerate()
.map(|(l, &g)| g * x_lag[i][l] * delta_u)
.sum()
})
.collect();
let score_mean = history_scores.iter().sum::<f64>() / n as f64;
let sxx: f64 = history_scores
.iter()
.map(|&s| (s - score_mean).powi(2))
.sum();
let sxy: f64 = history_scores
.iter()
.zip(y.iter())
.map(|(&s, &yi)| (s - score_mean) * yi)
.sum();
let slope = if sxx > 1e-15 { sxy / sxx } else { 0.0 };
let intercept = mu_y - slope * score_mean;
let fitted_values: Vec<f64> = history_scores
.iter()
.map(|&s| intercept + slope * s)
.collect();
let residuals: Vec<f64> = y
.iter()
.zip(&fitted_values)
.map(|(&yi, &yh)| yi - yh)
.collect();
let (r_squared, _) = super::compute_r_squared(y, &residuals, 2);
Ok(HistoryIndexResult {
fitted_values,
residuals,
intercept,
slope,
gamma,
lag_grid,
history_scores,
r_squared,
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::uniform_grid;
fn make_sine_data(n: usize, m: usize, freq_scale: f64) -> FdMatrix {
let data: Vec<f64> = (0..n)
.flat_map(|i| {
(0..m).map(move |j| {
let t = j as f64 / (m - 1) as f64;
(freq_scale * (i as f64 + 1.0) * t).sin()
})
})
.collect();
let mut cm = vec![0.0_f64; n * m];
for i in 0..n {
for j in 0..m {
cm[j * n + i] = data[i * m + j];
}
}
FdMatrix::from_column_major(cm, n, m).unwrap()
}
#[test]
fn fam_synthetic_recovery() {
let n = 50;
let m = 20;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let fpca = fdata_to_pc_1d(&data, 2, &argvals).unwrap();
let y: Vec<f64> = (0..n)
.map(|i| {
let xi1 = fpca.scores[(i, 0)];
let xi2 = fpca.scores[(i, 1)];
let noise = (i as f64 * 0.31).sin() * 0.05;
xi1.sin() + xi2 * xi2 + noise
})
.collect();
let config = FamConfig {
ncomp: 2,
bandwidth: 0.0,
..Default::default()
};
let result = fam(&data, &y, &argvals, None, &config).unwrap();
assert!(
result.r_squared > 0.75,
"expected R² > 0.75, got {}",
result.r_squared
);
let y_mean = y.iter().sum::<f64>() / n as f64;
let ss_y: f64 = y.iter().map(|&yi| (yi - y_mean).powi(2)).sum::<f64>();
let ss_res: f64 = result.residuals.iter().map(|r| r * r).sum();
let rel_err = (ss_res / ss_y).sqrt();
assert!(
rel_err < 0.30,
"expected relative fitted error < 0.30, got {rel_err:.4}"
);
}
#[test]
fn fam_decomposition_identity() {
let n = 30;
let m = 15;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.5);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.2).cos()).collect();
let config = FamConfig {
ncomp: 2,
..Default::default()
};
let result = fam(&data, &y, &argvals, None, &config).unwrap();
for i in 0..n {
let reconstructed = result.fitted_values[i] + result.residuals[i];
assert!(
(reconstructed - y[i]).abs() < 1e-9,
"decomposition failed at i={i}: fitted={} residual={} sum={} y={}",
result.fitted_values[i],
result.residuals[i],
reconstructed,
y[i]
);
}
}
#[test]
fn fam_output_shapes() {
let n = 25;
let m = 12;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = FamConfig {
ncomp: 3,
..Default::default()
};
let result = fam(&data, &y, &argvals, None, &config).unwrap();
assert_eq!(result.ncomp, 3, "ncomp field should be 3");
assert_eq!(
result.component_fits.len(),
3,
"component_fits.len() should equal ncomp"
);
for (k, cf) in result.component_fits.iter().enumerate() {
assert_eq!(cf.len(), n, "component_fits[{k}] should have length n={n}");
}
assert_eq!(
result.bandwidths.len(),
3,
"bandwidths.len() should equal ncomp"
);
assert_eq!(result.fitted_values.len(), n);
assert_eq!(result.residuals.len(), n);
}
#[test]
fn fam_invalid_dimension() {
let n = 20;
let m = 10;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y_ok: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = FamConfig {
ncomp: 2,
..Default::default()
};
let empty_data = FdMatrix::zeros(0, m);
let err = fam(&empty_data, &y_ok, &argvals, None, &config);
assert!(err.is_err(), "empty data should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { parameter, .. } => {
assert_eq!(parameter, "data");
}
e => panic!("expected InvalidDimension, got {e:?}"),
}
let y_wrong: Vec<f64> = vec![1.0; n + 5];
let err = fam(&data, &y_wrong, &argvals, None, &config);
assert!(err.is_err(), "mismatched y length should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { parameter, .. } => {
assert_eq!(parameter, "y");
}
e => panic!("expected InvalidDimension, got {e:?}"),
}
let argvals_wrong: Vec<f64> = uniform_grid(m + 3);
let err = fam(&data, &y_ok, &argvals_wrong, None, &config);
assert!(err.is_err(), "mismatched argvals should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { parameter, .. } => {
assert_eq!(parameter, "argvals");
}
e => panic!("expected InvalidDimension, got {e:?}"),
}
}
#[test]
fn gkam_r2_synthetic() {
let n = 40;
let m = 15;
let argvals = uniform_grid(m);
let mut cm = vec![0.0_f64; n * m];
for i in 0..n {
let amp = (i as f64 + 1.0) / n as f64; for j in 0..m {
let t = j as f64 / (m - 1) as f64;
cm[j * n + i] = amp * (std::f64::consts::PI * 2.0 * t).sin();
}
}
let data = FdMatrix::from_column_major(cm, n, m).unwrap();
let y: Vec<f64> = (0..n)
.map(|i| {
let amp = (i as f64 + 1.0) / n as f64;
let noise = (i as f64 * 0.23).sin() * 0.002;
amp * amp + noise
})
.collect();
let config = GkamConfig {
max_iter: 20,
epsilon: 1e-4,
..Default::default()
};
let result = fregre_gkam(&[&data], &y, &[&argvals], None, &config).unwrap();
assert!(
result.r_squared > 0.70,
"expected R² > 0.70, got {}",
result.r_squared
);
}
#[test]
fn gkam_convergence() {
let n = 25;
let m = 10;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
let config = GkamConfig {
max_iter: 50,
epsilon: 1e-4,
..Default::default()
};
let result = fregre_gkam(&[&data], &y, &[&argvals], None, &config).unwrap();
assert!(
result.converged,
"expected convergence, got iterations={}",
result.iterations
);
assert!(
result.iterations <= config.max_iter,
"iterations {} > max_iter {}",
result.iterations,
config.max_iter
);
}
#[test]
fn gkam_invalid_inputs() {
let n = 20;
let m = 10;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y_ok: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = GkamConfig::default();
let err = fregre_gkam(&[], &y_ok, &[], None, &config);
assert!(err.is_err(), "empty predictors should return Err");
let data_wrong = make_sine_data(n + 5, m, 1.0);
let err = fregre_gkam(&[&data_wrong], &y_ok, &[&argvals], None, &config);
assert!(err.is_err(), "mismatched n should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { .. } => {}
e => panic!("expected InvalidDimension, got {e:?}"),
}
let err = fregre_gkam(&[&data], &y_ok, &[], None, &config);
assert!(
err.is_err(),
"argvals_list length mismatch should return Err"
);
}
#[test]
fn gsam_matches_fam_identity() {
let n = 40;
let m = 16;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let fpca_ref = fdata_to_pc_1d(&data, 2, &argvals).unwrap();
let y: Vec<f64> = (0..n)
.map(|i| {
let xi1 = fpca_ref.scores[(i, 0)];
let xi2 = fpca_ref.scores[(i, 1)];
xi1 + xi2 * xi2 + (i as f64 * 0.23).sin() * 0.02
})
.collect();
let fam_config = FamConfig {
ncomp: 2,
bandwidth: 0.5, kernel: "gaussian".to_string(),
n_grid_bandwidth: 20,
};
let gsam_config = GsamConfig {
ncomp: 2,
bandwidth: 0.5,
kernel: "gaussian".to_string(),
n_grid_bandwidth: 20,
};
let fam_res = fam(&data, &y, &argvals, None, &fam_config).unwrap();
let gsam_res = fregre_gsam(&data, &y, &argvals, None, &gsam_config).unwrap();
for i in 0..n {
let diff = (fam_res.fitted_values[i] - gsam_res.fitted_values[i]).abs();
assert!(
diff < 1e-6,
"fam vs gsam mismatch at i={i}: fam={} gsam={} diff={diff:.2e}",
fam_res.fitted_values[i],
gsam_res.fitted_values[i]
);
}
}
#[test]
fn gsam_ncomp_too_large() {
let n = 15;
let m = 8;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = GsamConfig {
ncomp: 100,
..Default::default()
};
let err = fregre_gsam(&data, &y, &argvals, None, &config);
assert!(err.is_err(), "ncomp > min(n,m) should return Err");
match err.unwrap_err() {
FdarError::InvalidParameter { parameter, .. } => {
assert_eq!(parameter, "config.ncomp");
}
e => panic!("expected InvalidParameter, got {e:?}"),
}
}
#[test]
fn gsam_output_shapes() {
let n = 30;
let m = 10;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
let config = GsamConfig {
ncomp: 3,
..Default::default()
};
let result = fregre_gsam(&data, &y, &argvals, None, &config).unwrap();
assert_eq!(result.ncomp, 3);
assert_eq!(
result.component_fits.len(),
3,
"component_fits.len() should equal ncomp"
);
assert_eq!(result.fitted_values.len(), n);
}
#[test]
fn varselect_active_subset_recovery() {
let n = 100;
let m = 30;
let argvals = uniform_grid(m);
let make_orth = |p_idx: usize| -> FdMatrix {
let mut cm = vec![0.0_f64; n * m];
for i in 0..n {
let amp = (std::f64::consts::PI * (p_idx + 1) as f64 * i as f64 / n as f64).sin();
for j in 0..m {
let t = j as f64 / (m - 1) as f64;
cm[j * n + i] = amp * (std::f64::consts::PI * 2.0 * t).cos();
}
}
FdMatrix::from_column_major(cm, n, m).unwrap()
};
let preds: Vec<FdMatrix> = (0..5).map(make_orth).collect();
let pred_refs: Vec<&FdMatrix> = preds.iter().collect();
let argvals_list: Vec<&[f64]> = (0..5).map(|_| argvals.as_slice()).collect();
let y: Vec<f64> = (0..n)
.map(|i| {
let a0 = (std::f64::consts::PI * i as f64 / n as f64).sin();
let a2 = (std::f64::consts::PI * 3.0 * i as f64 / n as f64).sin();
5.0 * a0 + 3.0 * a2 + (i as f64 * 0.31).sin() * 0.01
})
.collect();
let config = VarSelectConfig {
ncomp: 1,
lambda_n_grid: 20,
..Default::default()
};
let result = variable_selection(&pred_refs, &y, &argvals_list, None, &config).unwrap();
assert_eq!(
result.active_predictors.len(),
5,
"should have 5 active_predictors entries"
);
assert!(
result.active_predictors[0],
"predictor 0 should be active, got {:?}",
result.active_predictors
);
assert!(
result.active_predictors[2],
"predictor 2 should be active, got {:?}",
result.active_predictors
);
assert!(
!result.active_predictors[1],
"predictor 1 should be inactive, got {:?}",
result.active_predictors
);
assert!(
!result.active_predictors[3],
"predictor 3 should be inactive, got {:?}",
result.active_predictors
);
assert!(
!result.active_predictors[4],
"predictor 4 should be inactive, got {:?}",
result.active_predictors
);
assert!(
result.r_squared > 0.5,
"expected R² > 0.5, got {}",
result.r_squared
);
}
#[test]
fn varselect_lambda_max_zeros() {
let n = 30;
let m = 10;
let argvals = uniform_grid(m);
let preds: Vec<FdMatrix> = (0..3usize).map(|_| make_sine_data(n, m, 1.0)).collect();
let pred_refs: Vec<&FdMatrix> = preds.iter().collect();
let argvals_list: Vec<&[f64]> = (0..3).map(|_| argvals.as_slice()).collect();
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.1).sin()).collect();
let config = VarSelectConfig {
ncomp: 2,
lambda: 1e6, ..Default::default()
};
let result = variable_selection(&pred_refs, &y, &argvals_list, None, &config).unwrap();
assert!(
result.active_predictors.iter().all(|&a| !a),
"expected all inactive at lambda=1e6, got {:?}",
result.active_predictors
);
}
#[test]
fn varselect_invalid_inputs() {
let n = 20;
let m = 10;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y_ok: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = VarSelectConfig {
ncomp: 2,
..Default::default()
};
let err = variable_selection(&[], &y_ok, &[], None, &config);
assert!(err.is_err(), "empty predictors should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { .. } => {}
e => panic!("expected InvalidDimension, got {e:?}"),
}
let data_wrong = make_sine_data(n + 5, m, 1.0);
let err = variable_selection(&[&data_wrong], &y_ok, &[&argvals], None, &config);
assert!(err.is_err(), "mismatched n should return Err");
match err.unwrap_err() {
FdarError::InvalidDimension { .. } => {}
e => panic!("expected InvalidDimension, got {e:?}"),
}
let err = variable_selection(&[&data], &y_ok, &[], None, &config);
assert!(err.is_err(), "argvals_list mismatch should return Err");
let config_mcp = VarSelectConfig {
penalty: VarSelectPenalty::GroupMcp,
..config.clone()
};
let err = variable_selection(&[&data], &y_ok, &[&argvals], None, &config_mcp);
assert!(err.is_err(), "GroupMcp should return Err");
match err.unwrap_err() {
FdarError::InvalidParameter { parameter, .. } => {
assert_eq!(parameter, "config.penalty");
}
e => panic!("expected InvalidParameter, got {e:?}"),
}
}
#[test]
fn perm_seeded_reproducibility() {
let n = 30;
let m = 12;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let fpca = fdata_to_pc_1d(&data, 1, &argvals).unwrap();
let y: Vec<f64> = (0..n)
.map(|i| fpca.scores[(i, 0)] * 2.0 + (i as f64 * 0.31).sin() * 0.05)
.collect();
let fam_cfg = FamConfig {
ncomp: 1,
..Default::default()
};
let perm_cfg = PermTestConfig {
n_perm: 19,
seed: 42,
statistic: PermTestStatistic::R2,
};
let r1 = permutation_test_fam(&data, &y, &argvals, None, &fam_cfg, &perm_cfg).unwrap();
let r2 = permutation_test_fam(&data, &y, &argvals, None, &fam_cfg, &perm_cfg).unwrap();
assert_eq!(
r1.p_value, r2.p_value,
"same seed should give same p_value: {} vs {}",
r1.p_value, r2.p_value
);
assert_eq!(
r1.null_statistics, r2.null_statistics,
"same seed should give same null distribution"
);
}
#[test]
fn perm_pvalue_range() {
let n = 20;
let m = 8;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| i as f64).collect();
let fam_cfg = FamConfig {
ncomp: 1,
..Default::default()
};
let perm_cfg = PermTestConfig {
n_perm: 9,
seed: 0,
statistic: PermTestStatistic::FittedNorm,
};
let result = permutation_test_fam(&data, &y, &argvals, None, &fam_cfg, &perm_cfg).unwrap();
assert!(
(0.0..=1.0).contains(&result.p_value),
"p_value out of [0,1]: {}",
result.p_value
);
}
#[test]
fn perm_detects_true_effect() {
let n = 40;
let m = 15;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let fpca = fdata_to_pc_1d(&data, 1, &argvals).unwrap();
let y_signal: Vec<f64> = (0..n)
.map(|i| {
let xi1 = fpca.scores[(i, 0)];
2.0 * xi1 + (i as f64 * 0.17).sin() * 0.02
})
.collect();
let y_null: Vec<f64> = (0..n).map(|i| (i as f64 * 0.37).sin() * 0.3).collect();
let fam_cfg = FamConfig {
ncomp: 1,
..Default::default()
};
let perm_cfg = PermTestConfig {
n_perm: 99,
seed: 42,
statistic: PermTestStatistic::R2,
};
let r_signal =
permutation_test_fam(&data, &y_signal, &argvals, None, &fam_cfg, &perm_cfg).unwrap();
let r_null =
permutation_test_fam(&data, &y_null, &argvals, None, &fam_cfg, &perm_cfg).unwrap();
assert!(
r_signal.p_value < 0.1,
"expected p < 0.1 under true effect, got p={}",
r_signal.p_value
);
assert!(
r_null.p_value > 0.1,
"expected p > 0.1 under the null, got p={}",
r_null.p_value
);
}
#[test]
fn history_index_synthetic_recovery() {
let n = 50;
let m = 30;
let argvals: Vec<f64> = (0..m).map(|j| j as f64 / (m - 1) as f64).collect();
let mut cm = vec![0.0_f64; n * m];
for i in 0..n {
let amp = (i as f64 + 1.0) / n as f64;
for j in 0..m {
let t = j as f64 / (m - 1) as f64;
cm[j * n + i] = amp * (std::f64::consts::PI * 2.0 * t).sin();
}
}
let data = FdMatrix::from_column_major(cm, n, m).unwrap();
let window = 0.5_f64;
let n_lags = 10;
let delta_u = window / n_lags as f64;
let big_t = argvals.last().copied().unwrap();
let y: Vec<f64> = (0..n)
.map(|i| {
(0..n_lags)
.map(|l| {
let u_l = l as f64 * delta_u;
let t_target = big_t - u_l;
let j = argvals
.partition_point(|&v| v < t_target)
.saturating_sub(1)
.min(m - 1);
data[(i, j)] * delta_u
})
.sum::<f64>()
})
.collect();
let config = HistoryIndexConfig {
window,
n_lags,
bandwidth: 0.0,
kernel: "gaussian".to_string(),
};
let result = history_index(&data, &y, &argvals, &config).unwrap();
assert!(
result.r_squared > 0.70,
"expected R² > 0.70, got {}",
result.r_squared
);
let g_mean = result.gamma.iter().sum::<f64>() / n_lags as f64;
let g_std = (result
.gamma
.iter()
.map(|&g| (g - g_mean).powi(2))
.sum::<f64>()
/ n_lags as f64)
.sqrt();
let cv = if g_mean.abs() > 1e-10 {
g_std / g_mean.abs()
} else {
0.0
};
assert!(
cv < 2.0,
"gamma should be approximately uniform (CV < 2.0), got CV={}",
cv
);
}
#[test]
fn history_index_window_too_large() {
let n = 20;
let m = 10;
let argvals = uniform_grid(m); let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| i as f64).collect();
let config = HistoryIndexConfig {
window: 2.0,
n_lags: 10,
..Default::default()
};
let err = history_index(&data, &y, &argvals, &config);
assert!(err.is_err(), "window > argvals range should return Err");
match err.unwrap_err() {
FdarError::InvalidParameter { parameter, .. } => {
assert_eq!(parameter, "config.window");
}
e => panic!("expected InvalidParameter, got {e:?}"),
}
}
#[test]
fn history_index_output_shapes() {
let n = 25;
let m = 15;
let argvals = uniform_grid(m); let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.2).cos()).collect();
let n_lags = 12;
let config = HistoryIndexConfig {
window: 0.5,
n_lags,
..Default::default()
};
let result = history_index(&data, &y, &argvals, &config).unwrap();
assert_eq!(
result.gamma.len(),
n_lags,
"gamma.len() should equal n_lags"
);
assert_eq!(
result.lag_grid.len(),
n_lags,
"lag_grid.len() should equal n_lags"
);
assert_eq!(
result.fitted_values.len(),
n,
"fitted_values.len() should equal n"
);
assert_eq!(
result.history_scores.len(),
n,
"history_scores.len() should equal n"
);
}
#[test]
fn gkam_empty_y_returns_err() {
let m = 10;
let argvals = uniform_grid(m);
let empty_data = FdMatrix::zeros(0, m);
let y_empty: Vec<f64> = vec![];
let config = GkamConfig::default();
let result = fregre_gkam(&[&empty_data], &y_empty, &[&argvals], None, &config);
assert!(
result.is_err(),
"fregre_gkam with empty y should return Err, got Ok"
);
match result.unwrap_err() {
FdarError::InvalidDimension { parameter, .. } => {
assert_eq!(parameter, "y", "error should report parameter='y'");
}
e => panic!("expected InvalidDimension(y), got {e:?}"),
}
}
#[test]
fn fam_scalar_covariates_component_fits_len() {
let n = 30;
let m = 12;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.2).cos()).collect();
let p_scalar = 2_usize;
let sc_vals: Vec<f64> = (0..n * p_scalar).map(|k| (k as f64 * 0.1).sin()).collect();
let mut sc_cm = vec![0.0_f64; n * p_scalar];
for row in 0..n {
for col in 0..p_scalar {
sc_cm[col * n + row] = sc_vals[row * p_scalar + col];
}
}
let sc = FdMatrix::from_column_major(sc_cm, n, p_scalar).unwrap();
let ncomp = 2;
let config = FamConfig {
ncomp,
..Default::default()
};
let result = fam(&data, &y, &argvals, Some(&sc), &config).unwrap();
let expected_len = ncomp + p_scalar;
assert_eq!(
result.component_fits.len(),
expected_len,
"component_fits.len() should be ncomp + p_scalar = {expected_len}, got {}",
result.component_fits.len()
);
assert_eq!(
result.bandwidths.len(),
expected_len,
"bandwidths.len() should be ncomp + p_scalar = {expected_len}, got {}",
result.bandwidths.len()
);
}
#[test]
fn gsam_scalar_covariates_component_fits_len() {
let n = 30;
let m = 12;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| (i as f64 * 0.3).sin()).collect();
let p_scalar = 2_usize;
let mut sc_cm = vec![0.0_f64; n * p_scalar];
for row in 0..n {
for col in 0..p_scalar {
sc_cm[col * n + row] = ((row * p_scalar + col) as f64 * 0.15).cos();
}
}
let sc = FdMatrix::from_column_major(sc_cm, n, p_scalar).unwrap();
let ncomp = 2;
let config = GsamConfig {
ncomp,
..Default::default()
};
let result = fregre_gsam(&data, &y, &argvals, Some(&sc), &config).unwrap();
let expected_len = ncomp + p_scalar;
assert_eq!(
result.component_fits.len(),
expected_len,
"component_fits.len() should be ncomp + p_scalar = {expected_len}, got {}",
result.component_fits.len()
);
assert_eq!(
result.bandwidths.len(),
expected_len,
"bandwidths.len() should be ncomp + p_scalar = {expected_len}, got {}",
result.bandwidths.len()
);
}
#[test]
fn perm_zero_nperm_returns_err() {
let n = 20;
let m = 8;
let argvals = uniform_grid(m);
let data = make_sine_data(n, m, 1.0);
let y: Vec<f64> = (0..n).map(|i| i as f64).collect();
let fam_cfg = FamConfig {
ncomp: 1,
..Default::default()
};
let perm_cfg = PermTestConfig {
n_perm: 0,
seed: 42,
statistic: PermTestStatistic::R2,
};
let result = permutation_test_fam(&data, &y, &argvals, None, &fam_cfg, &perm_cfg);
assert!(result.is_err(), "n_perm=0 should return Err");
match result.unwrap_err() {
FdarError::InvalidParameter { parameter, .. } => {
assert_eq!(parameter, "perm_config.n_perm");
}
e => panic!("expected InvalidParameter, got {e:?}"),
}
}
}