1use crate::matrix::FdMatrix;
24use crate::maybe_par_chunks_mut_enumerate;
25use rand::prelude::*;
26use rand_distr::Normal;
27use std::f64::consts::PI;
28
29#[derive(Clone, Copy, Debug, PartialEq)]
31pub enum EFunType {
32 Fourier = 0,
34 Poly = 1,
36 PolyHigh = 2,
38 Wiener = 3,
40}
41
42impl EFunType {
43 pub fn from_i32(value: i32) -> Result<Self, crate::FdarError> {
45 match value {
46 0 => Ok(EFunType::Fourier),
47 1 => Ok(EFunType::Poly),
48 2 => Ok(EFunType::PolyHigh),
49 3 => Ok(EFunType::Wiener),
50 _ => Err(crate::FdarError::InvalidEnumValue {
51 enum_name: "EFunType",
52 value,
53 }),
54 }
55 }
56}
57
58#[derive(Clone, Copy, Debug, PartialEq)]
60pub enum EValType {
61 Linear = 0,
63 Exponential = 1,
65 Wiener = 2,
67}
68
69impl EValType {
70 pub fn from_i32(value: i32) -> Result<Self, crate::FdarError> {
72 match value {
73 0 => Ok(EValType::Linear),
74 1 => Ok(EValType::Exponential),
75 2 => Ok(EValType::Wiener),
76 _ => Err(crate::FdarError::InvalidEnumValue {
77 enum_name: "EValType",
78 value,
79 }),
80 }
81 }
82}
83
84pub fn fourier_eigenfunctions(t: &[f64], m: usize) -> FdMatrix {
102 let n = t.len();
103 let mut phi = FdMatrix::zeros(n, m);
104 let sqrt2 = 2.0_f64.sqrt();
105
106 for (i, &ti) in t.iter().enumerate() {
107 phi[(i, 0)] = 1.0;
109
110 let mut k = 1; let mut freq = 1; while k < m {
114 if k < m {
116 phi[(i, k)] = sqrt2 * (2.0 * PI * f64::from(freq) * ti).sin();
117 k += 1;
118 }
119 if k < m {
121 phi[(i, k)] = sqrt2 * (2.0 * PI * f64::from(freq) * ti).cos();
122 k += 1;
123 }
124 freq += 1;
125 }
126 }
127 phi
128}
129
130pub fn legendre_eigenfunctions(t: &[f64], m: usize, high: bool) -> FdMatrix {
143 let n = t.len();
144 let mut phi = FdMatrix::zeros(n, m);
145 let start_deg = if high { 2 } else { 0 };
146
147 for (i, &ti) in t.iter().enumerate() {
148 let x = 2.0 * ti - 1.0;
150
151 for j in 0..m {
152 let deg = start_deg + j;
153 let p = legendre_p(x, deg);
155 let norm = ((2 * deg + 1) as f64).sqrt();
157 phi[(i, j)] = p * norm;
158 }
159 }
160 phi
161}
162
163fn legendre_p(x: f64, n: usize) -> f64 {
168 if n == 0 {
169 return 1.0;
170 }
171 if n == 1 {
172 return x;
173 }
174
175 let mut p_prev = 1.0;
176 let mut p_curr = x;
177
178 for k in 2..=n {
179 let p_next = ((2 * k - 1) as f64 * x * p_curr - (k - 1) as f64 * p_prev) / k as f64;
180 p_prev = p_curr;
181 p_curr = p_next;
182 }
183 p_curr
184}
185
186pub fn wiener_eigenfunctions(t: &[f64], m: usize) -> FdMatrix {
200 let n = t.len();
201 let mut phi = FdMatrix::zeros(n, m);
202 let sqrt2 = 2.0_f64.sqrt();
203
204 for (i, &ti) in t.iter().enumerate() {
205 for j in 0..m {
206 let k = (j + 1) as f64;
207 phi[(i, j)] = sqrt2 * ((k - 0.5) * PI * ti).sin();
209 }
210 }
211 phi
212}
213
214pub fn eigenfunctions(t: &[f64], m: usize, efun_type: EFunType) -> FdMatrix {
224 match efun_type {
225 EFunType::Fourier => fourier_eigenfunctions(t, m),
226 EFunType::Poly => legendre_eigenfunctions(t, m, false),
227 EFunType::PolyHigh => legendre_eigenfunctions(t, m, true),
228 EFunType::Wiener => wiener_eigenfunctions(t, m),
229 }
230}
231
232pub fn eigenvalues_linear(m: usize) -> Vec<f64> {
240 (1..=m).map(|k| 1.0 / k as f64).collect()
241}
242
243pub fn eigenvalues_exponential(m: usize) -> Vec<f64> {
247 (1..=m).map(|k| (-(k as f64)).exp()).collect()
248}
249
250pub fn eigenvalues_wiener(m: usize) -> Vec<f64> {
256 (1..=m)
257 .map(|k| {
258 let denom = (k as f64 - 0.5) * PI;
259 1.0 / (denom * denom)
260 })
261 .collect()
262}
263
264pub fn eigenvalues(m: usize, eval_type: EValType) -> Vec<f64> {
273 match eval_type {
274 EValType::Linear => eigenvalues_linear(m),
275 EValType::Exponential => eigenvalues_exponential(m),
276 EValType::Wiener => eigenvalues_wiener(m),
277 }
278}
279
280pub fn sim_kl(
300 n: usize,
301 phi: &FdMatrix,
302 big_m: usize,
303 lambda: &[f64],
304 seed: Option<u64>,
305) -> FdMatrix {
306 let m = phi.nrows();
307
308 let mut rng = match seed {
310 Some(s) => StdRng::seed_from_u64(s),
311 None => StdRng::from_entropy(),
312 };
313
314 let normal = Normal::new(0.0, 1.0).expect("valid distribution parameters");
315
316 let mut xi = vec![0.0; n * big_m];
319 for k in 0..big_m {
320 let sd = lambda[k].sqrt();
321 for i in 0..n {
322 xi[i + k * n] = rng.sample::<f64, _>(normal) * sd;
323 }
324 }
325
326 let mut data = vec![0.0; n * m];
329
330 maybe_par_chunks_mut_enumerate!(data, n, |(j, col)| {
332 for i in 0..n {
333 let mut sum = 0.0;
334 for k in 0..big_m {
335 sum += xi[i + k * n] * phi[(j, k)];
338 }
339 col[i] = sum;
340 }
341 });
342
343 FdMatrix::from_column_major(data, n, m).expect("dimension invariant: data.len() == n * m")
344}
345
346pub fn sim_fundata(
362 n: usize,
363 t: &[f64],
364 big_m: usize,
365 efun_type: EFunType,
366 eval_type: EValType,
367 seed: Option<u64>,
368) -> FdMatrix {
369 let phi = eigenfunctions(t, big_m, efun_type);
370 let lambda = eigenvalues(big_m, eval_type);
371 sim_kl(n, &phi, big_m, &lambda, seed)
372}
373
374pub fn add_error_pointwise(data: &FdMatrix, sd: f64, seed: Option<u64>) -> FdMatrix {
390 let mut rng = match seed {
391 Some(s) => StdRng::seed_from_u64(s),
392 None => StdRng::from_entropy(),
393 };
394
395 let normal = Normal::new(0.0, sd).expect("valid distribution parameters: sd > 0");
396
397 let noisy: Vec<f64> = data
398 .as_slice()
399 .iter()
400 .map(|&x| x + rng.sample::<f64, _>(normal))
401 .collect();
402
403 FdMatrix::from_column_major(noisy, data.nrows(), data.ncols())
404 .expect("dimension invariant: data.len() == n * m")
405}
406
407pub fn add_error_curve(data: &FdMatrix, sd: f64, seed: Option<u64>) -> FdMatrix {
420 let n = data.nrows();
421 let m = data.ncols();
422
423 let mut rng = match seed {
424 Some(s) => StdRng::seed_from_u64(s),
425 None => StdRng::from_entropy(),
426 };
427
428 let normal = Normal::new(0.0, sd).expect("valid distribution parameters: sd > 0");
429
430 let curve_noise: Vec<f64> = (0..n).map(|_| rng.sample::<f64, _>(normal)).collect();
432
433 let mut result = data.as_slice().to_vec();
435 for j in 0..m {
436 for i in 0..n {
437 result[i + j * n] += curve_noise[i];
438 }
439 }
440 FdMatrix::from_column_major(result, n, m).expect("dimension invariant: data.len() == n * m")
441}
442
443#[cfg(test)]
444mod tests {
445 use super::*;
446
447 #[test]
448 fn test_fourier_eigenfunctions_dimensions() {
449 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
450 let phi = fourier_eigenfunctions(&t, 5);
451 assert_eq!(phi.nrows(), 100);
452 assert_eq!(phi.ncols(), 5);
453 assert_eq!(phi.len(), 100 * 5);
454 }
455
456 #[test]
457 fn test_fourier_eigenfunctions_first_is_constant() {
458 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
459 let phi = fourier_eigenfunctions(&t, 3);
460
461 for i in 0..100 {
463 assert!((phi[(i, 0)] - 1.0).abs() < 1e-10);
464 }
465 }
466
467 #[test]
468 fn test_eigenvalues_linear() {
469 let lambda = eigenvalues_linear(5);
470 assert_eq!(lambda.len(), 5);
471 assert!((lambda[0] - 1.0).abs() < 1e-10);
472 assert!((lambda[1] - 0.5).abs() < 1e-10);
473 assert!((lambda[2] - 1.0 / 3.0).abs() < 1e-10);
474 }
475
476 #[test]
477 fn test_eigenvalues_exponential() {
478 let lambda = eigenvalues_exponential(3);
479 assert_eq!(lambda.len(), 3);
480 assert!((lambda[0] - (-1.0_f64).exp()).abs() < 1e-10);
481 assert!((lambda[1] - (-2.0_f64).exp()).abs() < 1e-10);
482 }
483
484 #[test]
485 fn test_sim_kl_dimensions() {
486 let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
487 let phi = fourier_eigenfunctions(&t, 5);
488 let lambda = eigenvalues_linear(5);
489
490 let data = sim_kl(10, &phi, 5, &lambda, Some(42));
491 assert_eq!(data.nrows(), 10);
492 assert_eq!(data.ncols(), 50);
493 assert_eq!(data.len(), 10 * 50);
494 }
495
496 #[test]
497 fn test_sim_fundata_dimensions() {
498 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
499 let data = sim_fundata(20, &t, 5, EFunType::Fourier, EValType::Linear, Some(42));
500 assert_eq!(data.nrows(), 20);
501 assert_eq!(data.ncols(), 100);
502 assert_eq!(data.len(), 20 * 100);
503 }
504
505 #[test]
506 fn test_add_error_pointwise() {
507 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();
509 let noisy = add_error_pointwise(&data, 0.1, Some(42));
510 assert_eq!(noisy.len(), 6);
511 let noisy_slice = noisy.as_slice();
513 for i in 0..6 {
514 assert!((noisy_slice[i] - raw[i]).abs() < 1.0);
515 }
516 }
517
518 #[test]
519 fn test_legendre_orthonormality() {
520 let n = 1000;
522 let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
523 let m = 5;
524 let phi = legendre_eigenfunctions(&t, m, false);
525 let dt = 1.0 / (n - 1) as f64;
526
527 for j1 in 0..m {
529 for j2 in 0..m {
530 let mut integral = 0.0;
531 for i in 0..n {
532 integral += phi[(i, j1)] * phi[(i, j2)] * dt;
533 }
534 let expected = if j1 == j2 { 1.0 } else { 0.0 };
535 assert!(
536 (integral - expected).abs() < 0.05,
537 "Orthonormality check failed for ({}, {}): {} vs {}",
538 j1,
539 j2,
540 integral,
541 expected
542 );
543 }
544 }
545 }
546
547 #[test]
552 fn test_wiener_eigenfunctions_dimensions() {
553 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
554 let phi = wiener_eigenfunctions(&t, 7);
555 assert_eq!(phi.nrows(), 100);
556 assert_eq!(phi.ncols(), 7);
557 assert_eq!(phi.len(), 100 * 7);
558 }
559
560 #[test]
561 fn test_wiener_eigenfunctions_orthonormality() {
562 let n = 1000;
565 let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
566 let m = 5;
567 let phi = wiener_eigenfunctions(&t, m);
568 let dt = 1.0 / (n - 1) as f64;
569
570 for j1 in 0..m {
571 for j2 in 0..m {
572 let mut integral = 0.0;
573 for i in 0..n {
574 integral += phi[(i, j1)] * phi[(i, j2)] * dt;
575 }
576 let expected = if j1 == j2 { 1.0 } else { 0.0 };
577 assert!(
578 (integral - expected).abs() < 0.05,
579 "Wiener orthonormality failed for ({}, {}): {} vs {}",
580 j1,
581 j2,
582 integral,
583 expected
584 );
585 }
586 }
587 }
588
589 #[test]
590 fn test_wiener_eigenfunctions_analytical_form() {
591 let t = vec![0.0, 0.25, 0.5, 0.75, 1.0];
593 let phi = wiener_eigenfunctions(&t, 2);
594 let sqrt2 = 2.0_f64.sqrt();
595
596 for (i, &ti) in t.iter().enumerate() {
598 let expected = sqrt2 * (0.5 * PI * ti).sin();
599 assert!(
600 (phi[(i, 0)] - expected).abs() < 1e-10,
601 "k=1 at t={}: got {} expected {}",
602 ti,
603 phi[(i, 0)],
604 expected
605 );
606 }
607
608 for (i, &ti) in t.iter().enumerate() {
610 let expected = sqrt2 * (1.5 * PI * ti).sin();
611 assert!(
612 (phi[(i, 1)] - expected).abs() < 1e-10,
613 "k=2 at t={}: got {} expected {}",
614 ti,
615 phi[(i, 1)],
616 expected
617 );
618 }
619 }
620
621 #[test]
626 fn test_eigenvalues_wiener_decay_rate() {
627 let lambda = eigenvalues_wiener(5);
629 assert_eq!(lambda.len(), 5);
630
631 for k in 1..=5 {
632 let denom = (k as f64 - 0.5) * PI;
633 let expected = 1.0 / (denom * denom);
634 assert!(
635 (lambda[k - 1] - expected).abs() < 1e-12,
636 "Wiener eigenvalue k={}: got {} expected {}",
637 k,
638 lambda[k - 1],
639 expected
640 );
641 }
642 }
643
644 #[test]
645 fn test_eigenvalues_wiener_decreasing() {
646 let lambda = eigenvalues_wiener(10);
648
649 for i in 1..lambda.len() {
650 assert!(
651 lambda[i] < lambda[i - 1],
652 "Eigenvalues not decreasing at {}: {} >= {}",
653 i,
654 lambda[i],
655 lambda[i - 1]
656 );
657 }
658 }
659
660 #[test]
665 fn test_add_error_curve_properties() {
666 let raw = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let n = 2;
669 let data = FdMatrix::from_column_major(raw.clone(), n, 3).unwrap();
670 let noisy = add_error_curve(&data, 0.5, Some(42));
671
672 assert_eq!(noisy.len(), 6);
673
674 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!(
681 (diff0_j0 - diff0_j1).abs() < 1e-10,
682 "Curve 0 noise differs: {} vs {}",
683 diff0_j0,
684 diff0_j1
685 );
686 assert!(
687 (diff0_j0 - diff0_j2).abs() < 1e-10,
688 "Curve 0 noise differs: {} vs {}",
689 diff0_j0,
690 diff0_j2
691 );
692
693 let diff1_j0 = noisy[(1, 0)] - raw[1];
695 assert!(
698 (diff0_j0 - diff1_j0).abs() > 1e-10,
699 "Different curves got same noise"
700 );
701 }
702
703 #[test]
704 fn test_add_error_curve_reproducibility() {
705 let raw = vec![1.0, 2.0, 3.0, 4.0];
706 let data = FdMatrix::from_column_major(raw, 2, 2).unwrap();
707 let noisy1 = add_error_curve(&data, 1.0, Some(123));
708 let noisy2 = add_error_curve(&data, 1.0, Some(123));
709
710 let s1 = noisy1.as_slice();
711 let s2 = noisy2.as_slice();
712 for i in 0..4 {
713 assert!(
714 (s1[i] - s2[i]).abs() < 1e-10,
715 "Reproducibility failed at {}: {} vs {}",
716 i,
717 s1[i],
718 s2[i]
719 );
720 }
721 }
722
723 #[test]
728 fn test_efun_type_from_i32() {
729 assert_eq!(EFunType::from_i32(0), Ok(EFunType::Fourier));
730 assert_eq!(EFunType::from_i32(1), Ok(EFunType::Poly));
731 assert_eq!(EFunType::from_i32(2), Ok(EFunType::PolyHigh));
732 assert_eq!(EFunType::from_i32(3), Ok(EFunType::Wiener));
733 assert!(EFunType::from_i32(-1).is_err());
734 assert!(EFunType::from_i32(4).is_err());
735 assert!(EFunType::from_i32(100).is_err());
736 }
737
738 #[test]
739 fn test_eval_type_from_i32() {
740 assert_eq!(EValType::from_i32(0), Ok(EValType::Linear));
741 assert_eq!(EValType::from_i32(1), Ok(EValType::Exponential));
742 assert_eq!(EValType::from_i32(2), Ok(EValType::Wiener));
743 assert!(EValType::from_i32(-1).is_err());
744 assert!(EValType::from_i32(3).is_err());
745 assert!(EValType::from_i32(99).is_err());
746 }
747
748 #[test]
749 fn test_eigenfunctions_dispatcher() {
750 let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
751 let m = 4;
752
753 let phi_fourier = eigenfunctions(&t, m, EFunType::Fourier);
755 let phi_fourier_direct = fourier_eigenfunctions(&t, m);
756 assert_eq!(phi_fourier, phi_fourier_direct);
757
758 let phi_poly = eigenfunctions(&t, m, EFunType::Poly);
759 let phi_poly_direct = legendre_eigenfunctions(&t, m, false);
760 assert_eq!(phi_poly, phi_poly_direct);
761
762 let phi_poly_high = eigenfunctions(&t, m, EFunType::PolyHigh);
763 let phi_poly_high_direct = legendre_eigenfunctions(&t, m, true);
764 assert_eq!(phi_poly_high, phi_poly_high_direct);
765
766 let phi_wiener = eigenfunctions(&t, m, EFunType::Wiener);
767 let phi_wiener_direct = wiener_eigenfunctions(&t, m);
768 assert_eq!(phi_wiener, phi_wiener_direct);
769 }
770
771 #[test]
772 fn test_sigma_zero_error() {
773 let t: Vec<f64> = (0..20).map(|i| i as f64 / 19.0).collect();
774 let data = sim_fundata(5, &t, 3, EFunType::Fourier, EValType::Exponential, Some(42));
775 let noisy = add_error_pointwise(&data, 0.0, Some(42));
776 for i in 0..5 {
778 for j in 0..20 {
779 assert!(
780 (noisy[(i, j)] - data[(i, j)]).abs() < 1e-12,
781 "Zero-sigma error should not change data"
782 );
783 }
784 }
785 }
786
787 #[test]
788 fn test_ncomp1_eigenfunctions() {
789 let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
790 let phi = fourier_eigenfunctions(&t, 1);
791 assert_eq!(phi.nrows(), t.len());
792 assert_eq!(phi.ncols(), 1);
793 let first_val = phi[(0, 0)];
795 for i in 1..t.len() {
796 assert!((phi[(i, 0)] - first_val).abs() < 1e-10);
797 }
798 }
799
800 #[test]
801 fn test_deterministic_seed() {
802 let t: Vec<f64> = (0..30).map(|i| i as f64 / 29.0).collect();
803 let d1 = sim_fundata(10, &t, 3, EFunType::Fourier, EValType::Linear, Some(123));
804 let d2 = sim_fundata(10, &t, 3, EFunType::Fourier, EValType::Linear, Some(123));
805 for i in 0..10 {
806 for j in 0..30 {
807 assert!(
808 (d1[(i, j)] - d2[(i, j)]).abs() < 1e-12,
809 "Same seed should produce identical results"
810 );
811 }
812 }
813 }
814}