use num_complex::Complex;
use num_traits::Float;
use rustfft::{FftNum, FftPlanner};
use super::Validation;
impl<T: Float + FftNum> Validation<T> {
pub fn cross_correlation(&self, max_lag: usize, z: T) -> CrossCorrelation<T> {
let (start, end) = (self.start.min(self.residual.len()), self.residual.len());
let n = T::from(end - start).unwrap();
let mean = |x: &[T]| x.iter().fold(T::zero(), |acc, &v| acc + v) / n;
let (e_mean, u_mean) = (mean(&self.residual[start..]), mean(&self.input[start..]));
let e: Vec<T> = self.residual.iter().map(|&v| v - e_mean).collect();
let u: Vec<T> = self.input.iter().map(|&v| v - u_mean).collect();
let e_evaluated: Vec<T> = (0..end).map(|k| if k >= start { e[k] } else { T::zero() }).collect();
let r_eu: Vec<T> = correlation(&e_evaluated, &u, max_lag).iter().map(|&c| c / n).collect();
let autocorrelation = |x: &[T]| -> Vec<T> { correlation(x, x, max_lag)[max_lag..].iter().map(|&c| c / n).collect() };
let (r_e, r_u) = (autocorrelation(&e[start..]), autocorrelation(&u[start..]));
let max_lag = max_lag as isize;
let lags: Vec<isize> = (-max_lag..=max_lag).collect();
let scale = (r_e[0] * r_u[0]).sqrt();
let p = (1..=max_lag as usize).fold(r_e[0] * r_u[0], |acc, k| {
let w = T::one() - T::from(k).unwrap() / T::from(max_lag + 1).unwrap();
acc + (T::one() + T::one()) * w * r_e[k] * r_u[k]
});
CrossCorrelation {
lags,
correlation: r_eu.iter().map(|&r| r / scale).collect(),
bound: z * (p.max(T::zero()) / n).sqrt() / scale,
}
}
}
#[derive(Clone, Debug)]
pub struct CrossCorrelation<T> {
pub lags: Vec<isize>,
pub correlation: Vec<T>,
pub bound: T,
}
impl<T: Float> CrossCorrelation<T> {
pub fn outside(&self) -> Vec<isize> {
self.lags.iter().zip(&self.correlation).filter(|(_, r)| r.abs() > self.bound).map(|(&tau, _)| tau).collect()
}
pub fn fraction_outside(&self) -> T {
T::from(self.outside().len()).unwrap() / T::from(self.lags.len()).unwrap()
}
pub fn max_ratio(&self) -> T {
self.lags.iter().zip(&self.correlation).filter(|(tau, _)| **tau >= 0).fold(T::zero(), |acc, (_, r)| acc.max(r.abs())) / self.bound
}
}
fn correlation<T: Float + FftNum>(a: &[T], b: &[T], max_lag: usize) -> Vec<T> {
let len = a.len().max(b.len());
let size = (len + max_lag).next_power_of_two();
let mut planner = FftPlanner::new();
let (forward, inverse) = (planner.plan_fft_forward(size), planner.plan_fft_inverse(size));
let padded = |x: &[T]| {
let mut buffer = vec![Complex::new(T::zero(), T::zero()); size];
for (b, &v) in buffer.iter_mut().zip(x) {
b.re = v;
}
buffer
};
let (mut fa, mut fb) = (padded(a), padded(b));
forward.process(&mut fa);
forward.process(&mut fb);
let mut product: Vec<Complex<T>> = fa.iter().zip(&fb).map(|(x, y)| x * y.conj()).collect();
inverse.process(&mut product);
let scale = T::from(size).unwrap();
(0..=2 * max_lag).map(|i| product[(i + size - max_lag) % size].re / scale).collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn correlation_by_fft_matches_the_direct_sum() {
let a: Vec<f64> = (0..300).map(|k| ((k * 7919) % 97) as f64 / 97.0 - 0.5).collect();
let b: Vec<f64> = (0..300).map(|k| ((k * 104729) % 89) as f64 / 89.0 - 0.5).collect();
let max_lag = 40;
let c = correlation(&a, &b, max_lag);
for (i, tau) in (-(max_lag as isize)..=max_lag as isize).enumerate() {
let direct: f64 = (0..a.len() as isize)
.filter(|&k| k - tau >= 0 && k - tau < b.len() as isize)
.map(|k| a[k as usize] * b[(k - tau) as usize])
.sum();
assert!((c[i] - direct).abs() < 1e-10, "τ = {tau}: {} vs {direct}", c[i]);
}
}
}