Skip to main content

asap_rs/utils/
acf.rs

1use crate::statistics::Metrics;
2use crate::fft::{transform, inverse_transform};
3
4pub struct ACF {
5    mean: f64,
6    values: Vec<f64>,
7    pub correlations: Vec<f64>,
8    corr_thresh: f64,
9    pub max_acf: f64,
10}
11
12impl ACF {
13    pub fn new(values: Vec<f64>, max_lag: usize) -> Self {
14        // Handle empty values case
15        if values.is_empty() {
16            return ACF {
17                mean: 0.0,
18                values,
19                correlations: vec![0.0], // Only lag 0
20                corr_thresh: 0.2,
21                max_acf: 0.0,
22            };
23        }
24        
25        // For non-empty arrays, ensure max_lag is not larger than the data length - 1
26        let adjusted_max_lag = max_lag.min(values.len() - 1);
27        
28        let mut acf = ACF {
29            mean: Metrics::mean(&values),
30            values,
31            correlations: vec![0.0; adjusted_max_lag + 1], // +1 to include lag 0
32            corr_thresh: 0.2,
33            max_acf: 0.0,
34        };
35        
36        // Skip calculation if max_lag is 0
37        if adjusted_max_lag > 0 {
38            acf.calculate();
39        }
40        
41        acf
42    }
43
44    fn calculate(&mut self) {
45        // Skip calculation if values is empty
46        if self.values.is_empty() {
47            return;
48        }
49        
50        // Padding to the closest power of 2
51        let len = 2_usize.pow((self.values.len() as f64).log2().ceil() as u32);
52        let mut fft_real = vec![0.0; len];
53        let mut fft_imag = vec![0.0; len];
54
55        for (i, &value) in self.values.iter().enumerate() {
56            fft_real[i] = value - self.mean;
57        }
58
59        // F_R(f) = FFT(X)
60        if let Ok(()) = transform(&mut fft_real, &mut fft_imag) {
61            // S(f) = F_R(f)F_R*(f)
62            for i in 0..fft_real.len() {
63                fft_real[i] = fft_real[i].powi(2) + fft_imag[i].powi(2);
64                fft_imag[i] = 0.0;
65            }
66
67            // R(t) = IFFT(S(f))
68            if let Ok(()) = inverse_transform(&mut fft_real, &mut fft_imag) {
69                // Ensure correlations[0] is not zero to avoid division by zero
70                if fft_real[0].abs() > 1e-10 {
71                    // Fill correlations array (only up to correlations.len())
72                    for i in 1..self.correlations.len() {
73                        self.correlations[i] = fft_real[i] / fft_real[0];
74                    }
75                }
76            }
77        }
78    }
79
80    pub fn find_peaks(&mut self) -> Vec<usize> {
81        let mut peak_indices = Vec::new();
82
83        if self.correlations.len() > 2 {  // Need at least 3 elements (lag 0, 1, 2)
84            // Start at correlations[1] and check against correlations[0]
85            let mut positive = self.correlations.get(1).map_or(false, |&v| v > self.correlations[0]);
86            let mut max = 1;
87
88            for i in 2..self.correlations.len() {
89                if !positive && self.correlations[i] > self.correlations[i - 1] {
90                    max = i;
91                    positive = !positive;
92                } else if positive && self.correlations[i] > self.correlations[max] {
93                    max = i;
94                } else if positive && self.correlations[i] < self.correlations[i - 1] {
95                    if max > 1 && self.correlations[max] > self.corr_thresh {
96                        peak_indices.push(max);
97                        if self.correlations[max] > self.max_acf {
98                            self.max_acf = self.correlations[max];
99                        }
100                    }
101                    positive = !positive;
102                }
103            }
104        }
105
106        // If there is no autocorrelation peak within the MAX_WINDOW boundary,
107        // try windows from the largest to the smallest
108        if peak_indices.len() <= 1 && self.correlations.len() > 2 {
109            peak_indices.extend(2..self.correlations.len());
110        }
111
112        peak_indices
113    }
114}