use crate::float_trait::Float;
use crate::time_series::TimeSeries;
use conv::ConvAsUtil;
mod fft;
mod fft_thread_local;
pub use fft_thread_local::FloatSupportedByFft;
mod freq;
pub use freq::FreqGrid;
pub use freq::{AverageNyquistFreq, MedianNyquistFreq, NyquistFreq, QuantileNyquistFreq};
mod power;
pub use power::PeriodogramPower;
mod power_fft;
pub use power_fft::PeriodogramPowerFft;
mod power_direct;
pub use power_direct::PeriodogramPowerDirect;
pub mod recurrent_sin_cos;
pub struct Periodogram<T> {
freq_grid: FreqGrid<T>,
periodogram_power: Box<dyn PeriodogramPower<T>>,
}
impl<T> Periodogram<T>
where
T: Float,
{
pub fn new(periodogram_power: Box<dyn PeriodogramPower<T>>, freq_grid: FreqGrid<T>) -> Self {
assert!(freq_grid.step.is_sign_positive() && freq_grid.step.is_finite());
assert!(freq_grid.size > 0);
Self {
freq_grid,
periodogram_power,
}
}
pub fn set_periodogram_power(
&mut self,
periodogram_power: Box<dyn PeriodogramPower<T>>,
) -> &mut Self {
self.periodogram_power = periodogram_power;
self
}
pub fn init_thread_local_fft_plans(n: &[usize]) {
for size in n {
T::init_fft_plan(size.next_power_of_two());
}
}
#[allow(clippy::borrowed_box)]
pub fn from_t(
periodogram_power: Box<dyn PeriodogramPower<T>>,
t: &[T],
resolution: f32,
max_freq_factor: f32,
nyquist: &Box<dyn NyquistFreq<T>>,
) -> Self {
Self::new(
periodogram_power,
FreqGrid::from_t(t, resolution, max_freq_factor, nyquist),
)
}
pub fn freq(&self, i: usize) -> T {
self.freq_grid.step * (i + 1).approx().unwrap()
}
pub fn power(&self, ts: &mut TimeSeries<T>) -> Vec<T> {
self.periodogram_power.power(&self.freq_grid, ts)
}
}
#[cfg(test)]
#[allow(clippy::unreadable_literal)]
#[allow(clippy::excessive_precision)]
mod tests {
use super::*;
use crate::statistics::Statistics;
use light_curve_common::{all_close, linspace};
use rand::prelude::*;
#[test]
fn compr_direct_with_scipy() {
const OMEGA_SIN: f64 = 0.07;
const N: usize = 100;
let t = linspace(0.0, 99.0, N);
let m: Vec<_> = t.iter().map(|&x| f64::sin(OMEGA_SIN * x)).collect();
let mut ts = TimeSeries::new(&t[..], &m[..], None);
let mut periodogram = Periodogram::new(
Box::new(PeriodogramPowerDirect),
FreqGrid {
step: OMEGA_SIN,
size: 1,
},
);
periodogram.set_periodogram_power(Box::new(PeriodogramPowerDirect));
all_close(
&[periodogram.power(&mut ts)[0] * 2.0 / (N as f64 - 1.0)],
&[1.0],
1.0 / (N as f64),
);
let freq_grid = FreqGrid {
step: 0.01,
size: 5,
};
let periodogram = Periodogram::new(Box::new(PeriodogramPowerDirect), freq_grid.clone());
all_close(
&linspace(
freq_grid.step,
freq_grid.step * freq_grid.size as f64,
freq_grid.size,
),
&(0..freq_grid.size)
.map(|i| periodogram.freq(i))
.collect::<Vec<_>>(),
1e-12,
);
let desired = [
16.99018018,
18.57722516,
21.96049738,
28.15056806,
36.66519435,
];
let actual = periodogram.power(&mut ts);
all_close(&actual[..], &desired[..], 1e-6);
}
#[test]
fn direct_vs_fft_one_to_one() {
const OMEGA: f64 = 0.472;
const N: usize = 64;
const RESOLUTION: f32 = 1.0;
const MAX_FREQ_FACTOR: f32 = 1.0;
let t = linspace(0.0, (N - 1) as f64, N);
let m: Vec<_> = t.iter().map(|&x| f64::sin(OMEGA * x)).collect();
let mut ts = TimeSeries::new(&t[..], &m[..], None);
let nyquist: Box<dyn NyquistFreq<f64>> = Box::new(AverageNyquistFreq);
let direct = Periodogram::from_t(
Box::new(PeriodogramPowerDirect),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
let fft = Periodogram::from_t(
Box::new(PeriodogramPowerFft),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
all_close(&fft[..direct.len() - 1], &direct[..direct.len() - 1], 1e-8);
}
#[test]
fn direct_vs_fft_uniform_sin_cos() {
const OMEGA1: f64 = 0.472;
const OMEGA2: f64 = 1.222;
const AMPLITUDE2: f64 = 2.0;
const N: usize = 100;
const RESOLUTION: f32 = 4.0;
const MAX_FREQ_FACTOR: f32 = 1.0;
let t = linspace(0.0, (N - 1) as f64, N);
let m: Vec<_> = t
.iter()
.map(|&x| f64::sin(OMEGA1 * x) + AMPLITUDE2 * f64::cos(OMEGA2 * x))
.collect();
let mut ts = TimeSeries::new(&t[..], &m[..], None);
let nyquist: Box<dyn NyquistFreq<f64>> = Box::new(AverageNyquistFreq);
let direct = Periodogram::from_t(
Box::new(PeriodogramPowerDirect),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
let fft = Periodogram::from_t(
Box::new(PeriodogramPowerFft),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
assert_eq!(
&fft.peak_indices_reverse_sorted()[..2],
&direct.peak_indices_reverse_sorted()[..2]
);
}
#[test]
fn direct_vs_fft_unevenly_sin_cos() {
const OMEGA1: f64 = 0.472;
const OMEGA2: f64 = 1.222;
const AMPLITUDE2: f64 = 2.0;
const NOISE_AMPLITUDE: f64 = 1.0;
const N: usize = 100;
const RESOLUTION: f32 = 6.0;
const MAX_FREQ_FACTOR: f32 = 1.0;
let mut rng = StdRng::seed_from_u64(0);
let t = (0..N)
.map(|_| rng.gen::<f64>() * (N - 1) as f64)
.collect::<Vec<_>>()
.sorted();
let m: Vec<_> = t
.iter()
.map(|&x| {
f64::sin(OMEGA1 * x)
+ AMPLITUDE2 * f64::cos(OMEGA2 * x)
+ NOISE_AMPLITUDE * rng.gen::<f64>()
})
.collect();
let mut ts = TimeSeries::new(&t[..], &m[..], None);
let nyquist: Box<dyn NyquistFreq<f64>> = Box::new(MedianNyquistFreq);
let direct = Periodogram::from_t(
Box::new(PeriodogramPowerDirect),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
let fft = Periodogram::from_t(
Box::new(PeriodogramPowerFft),
&t,
RESOLUTION,
MAX_FREQ_FACTOR,
&nyquist,
)
.power(&mut ts);
assert_eq!(
&fft.peak_indices_reverse_sorted()[..2],
&direct.peak_indices_reverse_sorted()[..2]
);
}
}