use numr::algorithm::fft::{FftAlgorithms, FftNormalization};
use numr::dtype::Complex128;
use numr::runtime::cpu::{CpuClient, CpuDevice, CpuRuntime};
use numr::tensor::Tensor;
use crate::error::Result;
pub(super) struct Correlator {
client: CpuClient,
device: CpuDevice,
padded: usize,
signal: Vec<f64>,
product: Vec<Complex128>,
lags: Vec<f32>,
}
impl Correlator {
pub(super) fn new(segment_len: usize) -> Self {
let device = CpuDevice::new();
let client = CpuClient::new(device.clone());
let padded = 2 * segment_len;
Self {
client,
device,
padded,
signal: vec![0.0; padded],
product: Vec::new(),
lags: Vec::new(),
}
}
pub(super) fn spectrum(&mut self, head: &[f32]) -> Result<Vec<Complex128>> {
self.signal.fill(0.0);
for (dst, &src) in self.signal.iter_mut().zip(head) {
*dst = f64::from(src);
}
let input = Tensor::<CpuRuntime>::from_slice(&self.signal, &[self.padded], &self.device)?;
let spectrum = self.client.rfft(&input, FftNormalization::None)?;
Ok(spectrum.try_to_vec::<Complex128>()?)
}
pub(super) fn correlate(
&mut self,
earlier: &[Complex128],
later: &[Complex128],
) -> Result<&[f32]> {
self.product.clear();
self.product
.extend(earlier.iter().zip(later).map(|(a, b)| Complex128 {
re: a.re * b.re + a.im * b.im,
im: a.im * b.re - a.re * b.im,
}));
let input =
Tensor::<CpuRuntime>::from_slice(&self.product, &[self.product.len()], &self.device)?;
let lags = self
.client
.irfft(&input, Some(self.padded), FftNormalization::Backward)?;
self.lags.clear();
self.lags
.extend(lags.try_to_vec::<f64>()?.into_iter().map(|v| v as f32));
Ok(&self.lags)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn matches_direct_sum() -> Result<()> {
let len = 16;
let earlier: Vec<f32> = (0..len).map(|n| ((n * 7) % 5) as f32 - 2.0).collect();
let later: Vec<f32> = (0..len).map(|n| ((n * 3) % 4) as f32 - 1.5).collect();
let mut correlator = Correlator::new(len);
let a = correlator.spectrum(&earlier)?;
let b = correlator.spectrum(&later)?;
let got = correlator.correlate(&a, &b)?;
assert_eq!(got.len(), 2 * len);
for (lag, &value) in got.iter().enumerate().take(len) {
let want: f32 = (0..len - lag).map(|n| earlier[n + lag] * later[n]).sum();
assert!(
(value - want).abs() < 1e-3,
"lag {lag}: got {value}, want {want}"
);
}
Ok(())
}
}