use crate::matrix::FdMatrix;
use crate::maybe_par_chunks_mut_enumerate;
use rand::prelude::*;
use rand_distr::Normal;
use std::f64::consts::PI;
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub enum EFunType {
Fourier = 0,
Poly = 1,
PolyHigh = 2,
Wiener = 3,
}
impl EFunType {
pub fn from_i32(value: i32) -> Result<Self, crate::FdarError> {
match value {
0 => Ok(EFunType::Fourier),
1 => Ok(EFunType::Poly),
2 => Ok(EFunType::PolyHigh),
3 => Ok(EFunType::Wiener),
_ => Err(crate::FdarError::InvalidEnumValue {
enum_name: "EFunType",
value,
}),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq)]
#[non_exhaustive]
pub enum EValType {
Linear = 0,
Exponential = 1,
Wiener = 2,
}
impl EValType {
pub fn from_i32(value: i32) -> Result<Self, crate::FdarError> {
match value {
0 => Ok(EValType::Linear),
1 => Ok(EValType::Exponential),
2 => Ok(EValType::Wiener),
_ => Err(crate::FdarError::InvalidEnumValue {
enum_name: "EValType",
value,
}),
}
}
}
pub fn fourier_eigenfunctions(t: &[f64], m: usize) -> FdMatrix {
let n = t.len();
let mut phi = FdMatrix::zeros(n, m);
let sqrt2 = 2.0_f64.sqrt();
for (i, &ti) in t.iter().enumerate() {
phi[(i, 0)] = 1.0;
let mut k = 1; let mut freq = 1;
while k < m {
if k < m {
phi[(i, k)] = sqrt2 * (2.0 * PI * f64::from(freq) * ti).sin();
k += 1;
}
if k < m {
phi[(i, k)] = sqrt2 * (2.0 * PI * f64::from(freq) * ti).cos();
k += 1;
}
freq += 1;
}
}
phi
}
pub fn legendre_eigenfunctions(t: &[f64], m: usize, high: bool) -> FdMatrix {
let n = t.len();
let mut phi = FdMatrix::zeros(n, m);
let start_deg = if high { 2 } else { 0 };
for (i, &ti) in t.iter().enumerate() {
let x = 2.0 * ti - 1.0;
for j in 0..m {
let deg = start_deg + j;
let p = legendre_p(x, deg);
let norm = ((2 * deg + 1) as f64).sqrt();
phi[(i, j)] = p * norm;
}
}
phi
}
fn legendre_p(x: f64, n: usize) -> f64 {
if n == 0 {
return 1.0;
}
if n == 1 {
return x;
}
let mut p_prev = 1.0;
let mut p_curr = x;
for k in 2..=n {
let p_next = ((2 * k - 1) as f64 * x * p_curr - (k - 1) as f64 * p_prev) / k as f64;
p_prev = p_curr;
p_curr = p_next;
}
p_curr
}
pub fn wiener_eigenfunctions(t: &[f64], m: usize) -> FdMatrix {
let n = t.len();
let mut phi = FdMatrix::zeros(n, m);
let sqrt2 = 2.0_f64.sqrt();
for (i, &ti) in t.iter().enumerate() {
for j in 0..m {
let k = (j + 1) as f64;
phi[(i, j)] = sqrt2 * ((k - 0.5) * PI * ti).sin();
}
}
phi
}
pub fn eigenfunctions(t: &[f64], m: usize, efun_type: EFunType) -> FdMatrix {
match efun_type {
EFunType::Fourier => fourier_eigenfunctions(t, m),
EFunType::Poly => legendre_eigenfunctions(t, m, false),
EFunType::PolyHigh => legendre_eigenfunctions(t, m, true),
EFunType::Wiener => wiener_eigenfunctions(t, m),
}
}
pub fn eigenvalues_linear(m: usize) -> Vec<f64> {
(1..=m).map(|k| 1.0 / k as f64).collect()
}
pub fn eigenvalues_exponential(m: usize) -> Vec<f64> {
(1..=m).map(|k| (-(k as f64)).exp()).collect()
}
pub fn eigenvalues_wiener(m: usize) -> Vec<f64> {
(1..=m)
.map(|k| {
let denom = (k as f64 - 0.5) * PI;
1.0 / (denom * denom)
})
.collect()
}
pub fn eigenvalues(m: usize, eval_type: EValType) -> Vec<f64> {
match eval_type {
EValType::Linear => eigenvalues_linear(m),
EValType::Exponential => eigenvalues_exponential(m),
EValType::Wiener => eigenvalues_wiener(m),
}
}
pub fn sim_kl(
n: usize,
phi: &FdMatrix,
big_m: usize,
lambda: &[f64],
seed: Option<u64>,
) -> FdMatrix {
let m = phi.nrows();
let mut rng = match seed {
Some(s) => StdRng::seed_from_u64(s),
None => StdRng::from_entropy(),
};
let normal = Normal::new(0.0, 1.0).expect("valid distribution parameters");
let mut xi = vec![0.0; n * big_m];
for k in 0..big_m {
let sd = lambda[k].sqrt();
for i in 0..n {
xi[i + k * n] = rng.sample::<f64, _>(normal) * sd;
}
}
let mut data = vec![0.0; n * m];
maybe_par_chunks_mut_enumerate!(data, n, |(j, col)| {
for i in 0..n {
let mut sum = 0.0;
for k in 0..big_m {
sum += xi[i + k * n] * phi[(j, k)];
}
col[i] = sum;
}
});
FdMatrix::from_column_major(data, n, m).expect("dimension invariant: data.len() == n * m")
}
pub fn sim_fundata(
n: usize,
t: &[f64],
big_m: usize,
efun_type: EFunType,
eval_type: EValType,
seed: Option<u64>,
) -> FdMatrix {
let phi = eigenfunctions(t, big_m, efun_type);
let lambda = eigenvalues(big_m, eval_type);
sim_kl(n, &phi, big_m, &lambda, seed)
}
pub fn add_error_pointwise(data: &FdMatrix, sd: f64, seed: Option<u64>) -> FdMatrix {
let mut rng = match seed {
Some(s) => StdRng::seed_from_u64(s),
None => StdRng::from_entropy(),
};
let normal = Normal::new(0.0, sd).expect("valid distribution parameters: sd > 0");
let noisy: Vec<f64> = data
.as_slice()
.iter()
.map(|&x| x + rng.sample::<f64, _>(normal))
.collect();
FdMatrix::from_column_major(noisy, data.nrows(), data.ncols())
.expect("dimension invariant: data.len() == n * m")
}
pub fn add_error_curve(data: &FdMatrix, sd: f64, seed: Option<u64>) -> FdMatrix {
let n = data.nrows();
let m = data.ncols();
let mut rng = match seed {
Some(s) => StdRng::seed_from_u64(s),
None => StdRng::from_entropy(),
};
let normal = Normal::new(0.0, sd).expect("valid distribution parameters: sd > 0");
let curve_noise: Vec<f64> = (0..n).map(|_| rng.sample::<f64, _>(normal)).collect();
let mut result = data.as_slice().to_vec();
for j in 0..m {
for i in 0..n {
result[i + j * n] += curve_noise[i];
}
}
FdMatrix::from_column_major(result, n, m).expect("dimension invariant: data.len() == n * m")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_fourier_eigenfunctions_dimensions() {
let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
let phi = fourier_eigenfunctions(&t, 5);
assert_eq!(phi.nrows(), 100);
assert_eq!(phi.ncols(), 5);
assert_eq!(phi.len(), 100 * 5);
}
#[test]
fn test_fourier_eigenfunctions_first_is_constant() {
let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
let phi = fourier_eigenfunctions(&t, 3);
for i in 0..100 {
assert!((phi[(i, 0)] - 1.0).abs() < 1e-10);
}
}
#[test]
fn test_eigenvalues_linear() {
let lambda = eigenvalues_linear(5);
assert_eq!(lambda.len(), 5);
assert!((lambda[0] - 1.0).abs() < 1e-10);
assert!((lambda[1] - 0.5).abs() < 1e-10);
assert!((lambda[2] - 1.0 / 3.0).abs() < 1e-10);
}
#[test]
fn test_eigenvalues_exponential() {
let lambda = eigenvalues_exponential(3);
assert_eq!(lambda.len(), 3);
assert!((lambda[0] - (-1.0_f64).exp()).abs() < 1e-10);
assert!((lambda[1] - (-2.0_f64).exp()).abs() < 1e-10);
}
#[test]
fn test_sim_kl_dimensions() {
let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
let phi = fourier_eigenfunctions(&t, 5);
let lambda = eigenvalues_linear(5);
let data = sim_kl(10, &phi, 5, &lambda, Some(42));
assert_eq!(data.nrows(), 10);
assert_eq!(data.ncols(), 50);
assert_eq!(data.len(), 10 * 50);
}
#[test]
fn test_sim_fundata_dimensions() {
let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
let data = sim_fundata(20, &t, 5, EFunType::Fourier, EValType::Linear, Some(42));
assert_eq!(data.nrows(), 20);
assert_eq!(data.ncols(), 100);
assert_eq!(data.len(), 20 * 100);
}
#[test]
fn test_add_error_pointwise() {
let raw = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let data = FdMatrix::from_column_major(raw.clone(), 2, 3).unwrap();
let noisy = add_error_pointwise(&data, 0.1, Some(42));
assert_eq!(noisy.len(), 6);
let noisy_slice = noisy.as_slice();
for i in 0..6 {
assert!((noisy_slice[i] - raw[i]).abs() < 1.0);
}
}
#[test]
fn test_legendre_orthonormality() {
let n = 1000;
let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
let m = 5;
let phi = legendre_eigenfunctions(&t, m, false);
let dt = 1.0 / (n - 1) as f64;
for j1 in 0..m {
for j2 in 0..m {
let mut integral = 0.0;
for i in 0..n {
integral += phi[(i, j1)] * phi[(i, j2)] * dt;
}
let expected = if j1 == j2 { 1.0 } else { 0.0 };
assert!(
(integral - expected).abs() < 0.05,
"Orthonormality check failed for ({}, {}): {} vs {}",
j1,
j2,
integral,
expected
);
}
}
}
#[test]
fn test_wiener_eigenfunctions_dimensions() {
let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
let phi = wiener_eigenfunctions(&t, 7);
assert_eq!(phi.nrows(), 100);
assert_eq!(phi.ncols(), 7);
assert_eq!(phi.len(), 100 * 7);
}
#[test]
fn test_wiener_eigenfunctions_orthonormality() {
let n = 1000;
let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
let m = 5;
let phi = wiener_eigenfunctions(&t, m);
let dt = 1.0 / (n - 1) as f64;
for j1 in 0..m {
for j2 in 0..m {
let mut integral = 0.0;
for i in 0..n {
integral += phi[(i, j1)] * phi[(i, j2)] * dt;
}
let expected = if j1 == j2 { 1.0 } else { 0.0 };
assert!(
(integral - expected).abs() < 0.05,
"Wiener orthonormality failed for ({}, {}): {} vs {}",
j1,
j2,
integral,
expected
);
}
}
}
#[test]
fn test_wiener_eigenfunctions_analytical_form() {
let t = vec![0.0, 0.25, 0.5, 0.75, 1.0];
let phi = wiener_eigenfunctions(&t, 2);
let sqrt2 = 2.0_f64.sqrt();
for (i, &ti) in t.iter().enumerate() {
let expected = sqrt2 * (0.5 * PI * ti).sin();
assert!(
(phi[(i, 0)] - expected).abs() < 1e-10,
"k=1 at t={}: got {} expected {}",
ti,
phi[(i, 0)],
expected
);
}
for (i, &ti) in t.iter().enumerate() {
let expected = sqrt2 * (1.5 * PI * ti).sin();
assert!(
(phi[(i, 1)] - expected).abs() < 1e-10,
"k=2 at t={}: got {} expected {}",
ti,
phi[(i, 1)],
expected
);
}
}
#[test]
fn test_eigenvalues_wiener_decay_rate() {
let lambda = eigenvalues_wiener(5);
assert_eq!(lambda.len(), 5);
for k in 1..=5 {
let denom = (k as f64 - 0.5) * PI;
let expected = 1.0 / (denom * denom);
assert!(
(lambda[k - 1] - expected).abs() < 1e-12,
"Wiener eigenvalue k={}: got {} expected {}",
k,
lambda[k - 1],
expected
);
}
}
#[test]
fn test_eigenvalues_wiener_decreasing() {
let lambda = eigenvalues_wiener(10);
for i in 1..lambda.len() {
assert!(
lambda[i] < lambda[i - 1],
"Eigenvalues not decreasing at {}: {} >= {}",
i,
lambda[i],
lambda[i - 1]
);
}
}
#[test]
fn test_add_error_curve_properties() {
let raw = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let n = 2;
let data = FdMatrix::from_column_major(raw.clone(), n, 3).unwrap();
let noisy = add_error_curve(&data, 0.5, Some(42));
assert_eq!(noisy.len(), 6);
let diff0_j0 = noisy[(0, 0)] - raw[0]; let diff0_j1 = noisy[(0, 1)] - raw[n]; let diff0_j2 = noisy[(0, 2)] - raw[2 * n];
assert!(
(diff0_j0 - diff0_j1).abs() < 1e-10,
"Curve 0 noise differs: {} vs {}",
diff0_j0,
diff0_j1
);
assert!(
(diff0_j0 - diff0_j2).abs() < 1e-10,
"Curve 0 noise differs: {} vs {}",
diff0_j0,
diff0_j2
);
let diff1_j0 = noisy[(1, 0)] - raw[1];
assert!(
(diff0_j0 - diff1_j0).abs() > 1e-10,
"Different curves got same noise"
);
}
#[test]
fn test_add_error_curve_reproducibility() {
let raw = vec![1.0, 2.0, 3.0, 4.0];
let data = FdMatrix::from_column_major(raw, 2, 2).unwrap();
let noisy1 = add_error_curve(&data, 1.0, Some(123));
let noisy2 = add_error_curve(&data, 1.0, Some(123));
let s1 = noisy1.as_slice();
let s2 = noisy2.as_slice();
for i in 0..4 {
assert!(
(s1[i] - s2[i]).abs() < 1e-10,
"Reproducibility failed at {}: {} vs {}",
i,
s1[i],
s2[i]
);
}
}
#[test]
fn test_efun_type_from_i32() {
assert_eq!(EFunType::from_i32(0), Ok(EFunType::Fourier));
assert_eq!(EFunType::from_i32(1), Ok(EFunType::Poly));
assert_eq!(EFunType::from_i32(2), Ok(EFunType::PolyHigh));
assert_eq!(EFunType::from_i32(3), Ok(EFunType::Wiener));
assert!(EFunType::from_i32(-1).is_err());
assert!(EFunType::from_i32(4).is_err());
assert!(EFunType::from_i32(100).is_err());
}
#[test]
fn test_eval_type_from_i32() {
assert_eq!(EValType::from_i32(0), Ok(EValType::Linear));
assert_eq!(EValType::from_i32(1), Ok(EValType::Exponential));
assert_eq!(EValType::from_i32(2), Ok(EValType::Wiener));
assert!(EValType::from_i32(-1).is_err());
assert!(EValType::from_i32(3).is_err());
assert!(EValType::from_i32(99).is_err());
}
#[test]
fn test_eigenfunctions_dispatcher() {
let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
let m = 4;
let phi_fourier = eigenfunctions(&t, m, EFunType::Fourier);
let phi_fourier_direct = fourier_eigenfunctions(&t, m);
assert_eq!(phi_fourier, phi_fourier_direct);
let phi_poly = eigenfunctions(&t, m, EFunType::Poly);
let phi_poly_direct = legendre_eigenfunctions(&t, m, false);
assert_eq!(phi_poly, phi_poly_direct);
let phi_poly_high = eigenfunctions(&t, m, EFunType::PolyHigh);
let phi_poly_high_direct = legendre_eigenfunctions(&t, m, true);
assert_eq!(phi_poly_high, phi_poly_high_direct);
let phi_wiener = eigenfunctions(&t, m, EFunType::Wiener);
let phi_wiener_direct = wiener_eigenfunctions(&t, m);
assert_eq!(phi_wiener, phi_wiener_direct);
}
#[test]
fn test_sigma_zero_error() {
let t: Vec<f64> = (0..20).map(|i| i as f64 / 19.0).collect();
let data = sim_fundata(5, &t, 3, EFunType::Fourier, EValType::Exponential, Some(42));
let noisy = add_error_pointwise(&data, 0.0, Some(42));
for i in 0..5 {
for j in 0..20 {
assert!(
(noisy[(i, j)] - data[(i, j)]).abs() < 1e-12,
"Zero-sigma error should not change data"
);
}
}
}
#[test]
fn test_ncomp1_eigenfunctions() {
let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
let phi = fourier_eigenfunctions(&t, 1);
assert_eq!(phi.nrows(), t.len());
assert_eq!(phi.ncols(), 1);
let first_val = phi[(0, 0)];
for i in 1..t.len() {
assert!((phi[(i, 0)] - first_val).abs() < 1e-10);
}
}
#[test]
fn test_deterministic_seed() {
let t: Vec<f64> = (0..30).map(|i| i as f64 / 29.0).collect();
let d1 = sim_fundata(10, &t, 3, EFunType::Fourier, EValType::Linear, Some(123));
let d2 = sim_fundata(10, &t, 3, EFunType::Fourier, EValType::Linear, Some(123));
for i in 0..10 {
for j in 0..30 {
assert!(
(d1[(i, j)] - d2[(i, j)]).abs() < 1e-12,
"Same seed should produce identical results"
);
}
}
}
}