1use crate::maybe_par_chunks_mut_enumerate;
24use rand::prelude::*;
25use rand_distr::Normal;
26use std::f64::consts::PI;
27
28#[derive(Clone, Copy, Debug, PartialEq)]
30pub enum EFunType {
31 Fourier = 0,
33 Poly = 1,
35 PolyHigh = 2,
37 Wiener = 3,
39}
40
41impl EFunType {
42 pub fn from_i32(value: i32) -> Option<Self> {
44 match value {
45 0 => Some(EFunType::Fourier),
46 1 => Some(EFunType::Poly),
47 2 => Some(EFunType::PolyHigh),
48 3 => Some(EFunType::Wiener),
49 _ => None,
50 }
51 }
52}
53
54#[derive(Clone, Copy, Debug, PartialEq)]
56pub enum EValType {
57 Linear = 0,
59 Exponential = 1,
61 Wiener = 2,
63}
64
65impl EValType {
66 pub fn from_i32(value: i32) -> Option<Self> {
68 match value {
69 0 => Some(EValType::Linear),
70 1 => Some(EValType::Exponential),
71 2 => Some(EValType::Wiener),
72 _ => None,
73 }
74 }
75}
76
77pub fn fourier_eigenfunctions(t: &[f64], m: usize) -> Vec<f64> {
95 let n = t.len();
96 let mut phi = vec![0.0; n * m];
97 let sqrt2 = 2.0_f64.sqrt();
98
99 for (i, &ti) in t.iter().enumerate() {
100 phi[i] = 1.0;
102
103 let mut k = 1; let mut freq = 1; while k < m {
107 if k < m {
109 phi[i + k * n] = sqrt2 * (2.0 * PI * freq as f64 * ti).sin();
110 k += 1;
111 }
112 if k < m {
114 phi[i + k * n] = sqrt2 * (2.0 * PI * freq as f64 * ti).cos();
115 k += 1;
116 }
117 freq += 1;
118 }
119 }
120 phi
121}
122
123pub fn legendre_eigenfunctions(t: &[f64], m: usize, high: bool) -> Vec<f64> {
136 let n = t.len();
137 let mut phi = vec![0.0; n * m];
138 let start_deg = if high { 2 } else { 0 };
139
140 for (i, &ti) in t.iter().enumerate() {
141 let x = 2.0 * ti - 1.0;
143
144 for j in 0..m {
145 let deg = start_deg + j;
146 let p = legendre_p(x, deg);
148 let norm = ((2 * deg + 1) as f64).sqrt();
150 phi[i + j * n] = p * norm;
151 }
152 }
153 phi
154}
155
156fn legendre_p(x: f64, n: usize) -> f64 {
161 if n == 0 {
162 return 1.0;
163 }
164 if n == 1 {
165 return x;
166 }
167
168 let mut p_prev = 1.0;
169 let mut p_curr = x;
170
171 for k in 2..=n {
172 let p_next = ((2 * k - 1) as f64 * x * p_curr - (k - 1) as f64 * p_prev) / k as f64;
173 p_prev = p_curr;
174 p_curr = p_next;
175 }
176 p_curr
177}
178
179pub fn wiener_eigenfunctions(t: &[f64], m: usize) -> Vec<f64> {
193 let n = t.len();
194 let mut phi = vec![0.0; n * m];
195 let sqrt2 = 2.0_f64.sqrt();
196
197 for (i, &ti) in t.iter().enumerate() {
198 for j in 0..m {
199 let k = (j + 1) as f64;
200 phi[i + j * n] = sqrt2 * ((k - 0.5) * PI * ti).sin();
202 }
203 }
204 phi
205}
206
207pub fn eigenfunctions(t: &[f64], m: usize, efun_type: EFunType) -> Vec<f64> {
217 match efun_type {
218 EFunType::Fourier => fourier_eigenfunctions(t, m),
219 EFunType::Poly => legendre_eigenfunctions(t, m, false),
220 EFunType::PolyHigh => legendre_eigenfunctions(t, m, true),
221 EFunType::Wiener => wiener_eigenfunctions(t, m),
222 }
223}
224
225pub fn eigenvalues_linear(m: usize) -> Vec<f64> {
233 (1..=m).map(|k| 1.0 / k as f64).collect()
234}
235
236pub fn eigenvalues_exponential(m: usize) -> Vec<f64> {
240 (1..=m).map(|k| (-(k as f64)).exp()).collect()
241}
242
243pub fn eigenvalues_wiener(m: usize) -> Vec<f64> {
249 (1..=m)
250 .map(|k| {
251 let denom = (k as f64 - 0.5) * PI;
252 1.0 / (denom * denom)
253 })
254 .collect()
255}
256
257pub fn eigenvalues(m: usize, eval_type: EValType) -> Vec<f64> {
266 match eval_type {
267 EValType::Linear => eigenvalues_linear(m),
268 EValType::Exponential => eigenvalues_exponential(m),
269 EValType::Wiener => eigenvalues_wiener(m),
270 }
271}
272
273pub fn sim_kl(
294 n: usize,
295 phi: &[f64],
296 m: usize,
297 big_m: usize,
298 lambda: &[f64],
299 seed: Option<u64>,
300) -> Vec<f64> {
301 let mut rng = match seed {
303 Some(s) => StdRng::seed_from_u64(s),
304 None => StdRng::from_entropy(),
305 };
306
307 let normal = Normal::new(0.0, 1.0).unwrap();
308
309 let mut xi = vec![0.0; n * big_m];
312 for k in 0..big_m {
313 let sd = lambda[k].sqrt();
314 for i in 0..n {
315 xi[i + k * n] = rng.sample::<f64, _>(normal) * sd;
316 }
317 }
318
319 let mut data = vec![0.0; n * m];
322
323 maybe_par_chunks_mut_enumerate!(data, n, |(j, col)| {
325 for i in 0..n {
326 let mut sum = 0.0;
327 for k in 0..big_m {
328 sum += xi[i + k * n] * phi[j + k * m];
331 }
332 col[i] = sum;
333 }
334 });
335
336 data
337}
338
339pub fn sim_fundata(
355 n: usize,
356 t: &[f64],
357 big_m: usize,
358 efun_type: EFunType,
359 eval_type: EValType,
360 seed: Option<u64>,
361) -> Vec<f64> {
362 let m = t.len();
363 let phi = eigenfunctions(t, big_m, efun_type);
364 let lambda = eigenvalues(big_m, eval_type);
365 sim_kl(n, &phi, m, big_m, &lambda, seed)
366}
367
368pub fn add_error_pointwise(
386 data: &[f64],
387 _n: usize,
388 _m: usize,
389 sd: f64,
390 seed: Option<u64>,
391) -> Vec<f64> {
392 let mut rng = match seed {
393 Some(s) => StdRng::seed_from_u64(s),
394 None => StdRng::from_entropy(),
395 };
396
397 let normal = Normal::new(0.0, sd).unwrap();
398
399 data.iter()
400 .map(|&x| x + rng.sample::<f64, _>(normal))
401 .collect()
402}
403
404pub fn add_error_curve(data: &[f64], n: usize, m: usize, sd: f64, seed: Option<u64>) -> Vec<f64> {
419 let mut rng = match seed {
420 Some(s) => StdRng::seed_from_u64(s),
421 None => StdRng::from_entropy(),
422 };
423
424 let normal = Normal::new(0.0, sd).unwrap();
425
426 let curve_noise: Vec<f64> = (0..n).map(|_| rng.sample::<f64, _>(normal)).collect();
428
429 let mut result = data.to_vec();
431 for j in 0..m {
432 for i in 0..n {
433 result[i + j * n] += curve_noise[i];
434 }
435 }
436 result
437}
438
439#[cfg(test)]
440mod tests {
441 use super::*;
442
443 #[test]
444 fn test_fourier_eigenfunctions_dimensions() {
445 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
446 let phi = fourier_eigenfunctions(&t, 5);
447 assert_eq!(phi.len(), 100 * 5);
448 }
449
450 #[test]
451 fn test_fourier_eigenfunctions_first_is_constant() {
452 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
453 let phi = fourier_eigenfunctions(&t, 3);
454
455 for i in 0..100 {
457 assert!((phi[i] - 1.0).abs() < 1e-10);
458 }
459 }
460
461 #[test]
462 fn test_eigenvalues_linear() {
463 let lambda = eigenvalues_linear(5);
464 assert_eq!(lambda.len(), 5);
465 assert!((lambda[0] - 1.0).abs() < 1e-10);
466 assert!((lambda[1] - 0.5).abs() < 1e-10);
467 assert!((lambda[2] - 1.0 / 3.0).abs() < 1e-10);
468 }
469
470 #[test]
471 fn test_eigenvalues_exponential() {
472 let lambda = eigenvalues_exponential(3);
473 assert_eq!(lambda.len(), 3);
474 assert!((lambda[0] - (-1.0_f64).exp()).abs() < 1e-10);
475 assert!((lambda[1] - (-2.0_f64).exp()).abs() < 1e-10);
476 }
477
478 #[test]
479 fn test_sim_kl_dimensions() {
480 let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
481 let phi = fourier_eigenfunctions(&t, 5);
482 let lambda = eigenvalues_linear(5);
483
484 let data = sim_kl(10, &phi, 50, 5, &lambda, Some(42));
485 assert_eq!(data.len(), 10 * 50);
486 }
487
488 #[test]
489 fn test_sim_fundata_dimensions() {
490 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
491 let data = sim_fundata(20, &t, 5, EFunType::Fourier, EValType::Linear, Some(42));
492 assert_eq!(data.len(), 20 * 100);
493 }
494
495 #[test]
496 fn test_add_error_pointwise() {
497 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let noisy = add_error_pointwise(&data, 2, 3, 0.1, Some(42));
499 assert_eq!(noisy.len(), 6);
500 for i in 0..6 {
502 assert!((noisy[i] - data[i]).abs() < 1.0);
503 }
504 }
505
506 #[test]
507 fn test_legendre_orthonormality() {
508 let n = 1000;
510 let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
511 let m = 5;
512 let phi = legendre_eigenfunctions(&t, m, false);
513 let dt = 1.0 / (n - 1) as f64;
514
515 for j1 in 0..m {
517 for j2 in 0..m {
518 let mut integral = 0.0;
519 for i in 0..n {
520 integral += phi[i + j1 * n] * phi[i + j2 * n] * dt;
521 }
522 let expected = if j1 == j2 { 1.0 } else { 0.0 };
523 assert!(
524 (integral - expected).abs() < 0.05,
525 "Orthonormality check failed for ({}, {}): {} vs {}",
526 j1,
527 j2,
528 integral,
529 expected
530 );
531 }
532 }
533 }
534
535 #[test]
540 fn test_wiener_eigenfunctions_dimensions() {
541 let t: Vec<f64> = (0..100).map(|i| i as f64 / 99.0).collect();
542 let phi = wiener_eigenfunctions(&t, 7);
543 assert_eq!(phi.len(), 100 * 7);
544 }
545
546 #[test]
547 fn test_wiener_eigenfunctions_orthonormality() {
548 let n = 1000;
551 let t: Vec<f64> = (0..n).map(|i| i as f64 / (n - 1) as f64).collect();
552 let m = 5;
553 let phi = wiener_eigenfunctions(&t, m);
554 let dt = 1.0 / (n - 1) as f64;
555
556 for j1 in 0..m {
557 for j2 in 0..m {
558 let mut integral = 0.0;
559 for i in 0..n {
560 integral += phi[i + j1 * n] * phi[i + j2 * n] * dt;
561 }
562 let expected = if j1 == j2 { 1.0 } else { 0.0 };
563 assert!(
564 (integral - expected).abs() < 0.05,
565 "Wiener orthonormality failed for ({}, {}): {} vs {}",
566 j1,
567 j2,
568 integral,
569 expected
570 );
571 }
572 }
573 }
574
575 #[test]
576 fn test_wiener_eigenfunctions_analytical_form() {
577 let t = vec![0.0, 0.25, 0.5, 0.75, 1.0];
579 let phi = wiener_eigenfunctions(&t, 2);
580 let sqrt2 = 2.0_f64.sqrt();
581
582 for (i, &ti) in t.iter().enumerate() {
584 let expected = sqrt2 * (0.5 * PI * ti).sin();
585 assert!(
586 (phi[i] - expected).abs() < 1e-10,
587 "k=1 at t={}: got {} expected {}",
588 ti,
589 phi[i],
590 expected
591 );
592 }
593
594 for (i, &ti) in t.iter().enumerate() {
596 let expected = sqrt2 * (1.5 * PI * ti).sin();
597 assert!(
598 (phi[i + t.len()] - expected).abs() < 1e-10,
599 "k=2 at t={}: got {} expected {}",
600 ti,
601 phi[i + t.len()],
602 expected
603 );
604 }
605 }
606
607 #[test]
612 fn test_eigenvalues_wiener_decay_rate() {
613 let lambda = eigenvalues_wiener(5);
615 assert_eq!(lambda.len(), 5);
616
617 for k in 1..=5 {
618 let denom = (k as f64 - 0.5) * PI;
619 let expected = 1.0 / (denom * denom);
620 assert!(
621 (lambda[k - 1] - expected).abs() < 1e-12,
622 "Wiener eigenvalue k={}: got {} expected {}",
623 k,
624 lambda[k - 1],
625 expected
626 );
627 }
628 }
629
630 #[test]
631 fn test_eigenvalues_wiener_decreasing() {
632 let lambda = eigenvalues_wiener(10);
634
635 for i in 1..lambda.len() {
636 assert!(
637 lambda[i] < lambda[i - 1],
638 "Eigenvalues not decreasing at {}: {} >= {}",
639 i,
640 lambda[i],
641 lambda[i - 1]
642 );
643 }
644 }
645
646 #[test]
651 fn test_add_error_curve_properties() {
652 let data = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]; let n = 2;
655 let m = 3;
656 let noisy = add_error_curve(&data, n, m, 0.5, Some(42));
657
658 assert_eq!(noisy.len(), 6);
659
660 let diff0_j0 = noisy[0] - data[0]; let diff0_j1 = noisy[n] - data[n]; let diff0_j2 = noisy[2 * n] - data[2 * n]; assert!(
667 (diff0_j0 - diff0_j1).abs() < 1e-10,
668 "Curve 0 noise differs: {} vs {}",
669 diff0_j0,
670 diff0_j1
671 );
672 assert!(
673 (diff0_j0 - diff0_j2).abs() < 1e-10,
674 "Curve 0 noise differs: {} vs {}",
675 diff0_j0,
676 diff0_j2
677 );
678
679 let diff1_j0 = noisy[1] - data[1];
681 assert!(
684 (diff0_j0 - diff1_j0).abs() > 1e-10,
685 "Different curves got same noise"
686 );
687 }
688
689 #[test]
690 fn test_add_error_curve_reproducibility() {
691 let data = vec![1.0, 2.0, 3.0, 4.0];
692 let noisy1 = add_error_curve(&data, 2, 2, 1.0, Some(123));
693 let noisy2 = add_error_curve(&data, 2, 2, 1.0, Some(123));
694
695 for i in 0..4 {
696 assert!(
697 (noisy1[i] - noisy2[i]).abs() < 1e-10,
698 "Reproducibility failed at {}: {} vs {}",
699 i,
700 noisy1[i],
701 noisy2[i]
702 );
703 }
704 }
705
706 #[test]
711 fn test_efun_type_from_i32() {
712 assert_eq!(EFunType::from_i32(0), Some(EFunType::Fourier));
713 assert_eq!(EFunType::from_i32(1), Some(EFunType::Poly));
714 assert_eq!(EFunType::from_i32(2), Some(EFunType::PolyHigh));
715 assert_eq!(EFunType::from_i32(3), Some(EFunType::Wiener));
716 assert_eq!(EFunType::from_i32(-1), None);
717 assert_eq!(EFunType::from_i32(4), None);
718 assert_eq!(EFunType::from_i32(100), None);
719 }
720
721 #[test]
722 fn test_eval_type_from_i32() {
723 assert_eq!(EValType::from_i32(0), Some(EValType::Linear));
724 assert_eq!(EValType::from_i32(1), Some(EValType::Exponential));
725 assert_eq!(EValType::from_i32(2), Some(EValType::Wiener));
726 assert_eq!(EValType::from_i32(-1), None);
727 assert_eq!(EValType::from_i32(3), None);
728 assert_eq!(EValType::from_i32(99), None);
729 }
730
731 #[test]
732 fn test_eigenfunctions_dispatcher() {
733 let t: Vec<f64> = (0..50).map(|i| i as f64 / 49.0).collect();
734 let m = 4;
735
736 let phi_fourier = eigenfunctions(&t, m, EFunType::Fourier);
738 let phi_fourier_direct = fourier_eigenfunctions(&t, m);
739 assert_eq!(phi_fourier, phi_fourier_direct);
740
741 let phi_poly = eigenfunctions(&t, m, EFunType::Poly);
742 let phi_poly_direct = legendre_eigenfunctions(&t, m, false);
743 assert_eq!(phi_poly, phi_poly_direct);
744
745 let phi_poly_high = eigenfunctions(&t, m, EFunType::PolyHigh);
746 let phi_poly_high_direct = legendre_eigenfunctions(&t, m, true);
747 assert_eq!(phi_poly_high, phi_poly_high_direct);
748
749 let phi_wiener = eigenfunctions(&t, m, EFunType::Wiener);
750 let phi_wiener_direct = wiener_eigenfunctions(&t, m);
751 assert_eq!(phi_wiener, phi_wiener_direct);
752 }
753}