liftof_lib/
sine_fitter.rs

1use ndarray::Array1;
2//, Array, ArrayBase, Data, Ix1};
3//use ndarray_linalg::solve::Inverse;
4use rustfft::{FftPlanner, num_complex::Complex};
5use std::f64::consts::PI;
6use statrs::distribution::ChiSquared;
7//use ndarray_rand::RandomExt;
8//use ndarray_rand::rand_distr::Uniform;
9use argmin::prelude::*;
10//use argmin::ArgminOp;
11// use argmin::core::Executor;
12// use argmin::core::Error;
13use argmin::solver::neldermead::NelderMead;
14use statrs::distribution::ContinuousCDF;
15
16struct SinFunc<'a> {
17    tt: &'a Array1<f64>,
18    yy: &'a Array1<f64>,
19}
20
21impl ArgminOp for SinFunc<'_> {
22    type Param = Vec<f64>;
23    type Output = f64;
24    type Hessian = ();
25    type Jacobian = ();
26    type Float = f64;
27
28    fn apply(&self, p: &Self::Param) -> Result<Self::Output, Error> {
29        #[allow(non_snake_case)]
30        let (A, w, ph, c) = (p[0], p[1], p[2], p[3]);
31        let sinfunc = |t: f64| A * (w * t + ph).sin() + c;
32        let residuals: f64 = self.tt.iter().zip(self.yy.iter())
33            .map(|(&t, &y)| (y - sinfunc(t)).powi(2))
34            .sum();
35        Ok(residuals)
36    }
37}
38
39pub fn fit_sine(nanoseconds: &Array1<f64>, volts: &Array1<f64>) -> Result<(f64, f64, f64, f64, f64, f64), Box<dyn std::error::Error>> {
40    let tt = nanoseconds;
41    let yy = volts;
42
43    // FFT
44    let mut planner = FftPlanner::new();
45    let fft = planner.plan_fft_forward(tt.len());
46    let mut signal: Vec<Complex<f64>> = yy.iter().map(|&y| Complex::new(y, 0.0)).collect();
47    fft.process(&mut signal);
48
49    let mut magnitudes: Vec<f64> = signal.iter().map(|c| c.norm()).collect();
50    magnitudes[0] = 0.0; // exclude zero frequency peak
51
52    let guess_freq = (magnitudes.iter().cloned().enumerate().max_by(|a, b| a.1.partial_cmp(&b.1).unwrap()).unwrap().0 as f64) / tt.len() as f64;
53    let guess_amp = yy.std(0.0) * 2f64.sqrt();
54    let guess_offset = yy.mean().unwrap();
55    let guess = vec![guess_amp, 2.0 * PI * guess_freq, 0.0, guess_offset];
56
57    // Optimization
58    let problem = SinFunc { tt, yy };
59    let solver = NelderMead::new();
60    let res = Executor::new(problem, solver, guess)
61        .max_iters(1000)
62        .run()?;
63
64    let result = &res.state().best_param;
65    #[allow(non_snake_case)]
66    let (A, w, p, c) = (result[0], result[1], result[2], result[3]);
67    let phase_multiple_pi = p / PI;
68
69    let fitfunc = |t: f64| -> f64 {
70        A * (w * t + p).sin() + c
71    };
72
73    let residuals: Vec<f64> = tt.iter().zip(yy.iter()).map(|(&t, &y)| y - fitfunc(t)).collect();
74    let _ss_res = residuals.iter().map(|&r| r.powi(2)).sum::<f64>();
75    let _ss_tot = yy.iter().map(|&y| (y - yy.mean().unwrap()).powi(2)).sum::<f64>();
76    //let r_squared = 1.0 - (ss_res / ss_tot);
77
78    // Chi-squared calculation
79    let expected_values: Vec<f64> = tt.iter().map(|&t| fitfunc(t)).collect();
80    let observed_values = yy.to_vec();
81    let chi_squared_stat: f64 = observed_values.iter().zip(expected_values.iter())
82        .map(|(&o, &e)| (o - e).powi(2) / e).sum();
83
84    let df = tt.len() as f64 - result.len() as f64;
85    let reduced_chi_squared = chi_squared_stat / df;
86
87    let chi_squared = ChiSquared::new(df).unwrap();
88    let _p_value = 1.0 - chi_squared.cdf(chi_squared_stat);
89
90    let f = w / (2.0 * PI);
91    let _period = 1.0 / f;
92
93
94    Ok((A, f, p, phase_multiple_pi, c, reduced_chi_squared))
95    // println!("amp: {:.2}", A);
96    // println!("omega: {:.2}", w);
97    // println!("phase: {:.2}", p);
98    // println!("phase_formatted: {:.2}π", phase_multiple_pi);
99    // println!("offset: {:.2}", c);
100    // println!("freq: {:.2}", f);
101    // println!("period: {:.2}", period);
102    // println!("r_squared: {:.2}", r_squared);
103    // println!("chi_squared_stat: {:.2}", chi_squared_stat);
104    // println!("p_value: {:.2}", p_value);
105    // println!("reduced_chi_squared: {:.2}", reduced_chi_squared);
106
107}
liftof_lib/sine_fitter.rs

liftof_lib/
sine_fitter.rs
