use crate::slice_maybe_parallel;
#[cfg(feature = "parallel")]
use rayon::iter::ParallelIterator;
use super::IrregFdata;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum KernelType {
Epanechnikov,
Gaussian,
}
#[inline]
pub(crate) fn kernel_epanechnikov(u: f64) -> f64 {
if u.abs() <= 1.0 {
0.75 * (1.0 - u * u)
} else {
0.0
}
}
#[inline]
pub(crate) fn kernel_gaussian(u: f64) -> f64 {
(-0.5 * u * u).exp() / (2.0 * std::f64::consts::PI).sqrt()
}
impl KernelType {
#[inline]
pub(crate) fn as_fn(self) -> fn(f64) -> f64 {
match self {
KernelType::Epanechnikov => kernel_epanechnikov,
KernelType::Gaussian => kernel_gaussian,
}
}
}
pub fn mean_irreg(
ifd: &IrregFdata,
target_argvals: &[f64],
bandwidth: f64,
kernel_type: KernelType,
) -> Vec<f64> {
let n = ifd.n_obs();
let kernel = kernel_type.as_fn();
slice_maybe_parallel!(target_argvals)
.map(|&t| {
let mut sum_weights = 0.0;
let mut sum_values = 0.0;
for i in 0..n {
let (obs_t, obs_x) = ifd.get_obs(i);
for (&ti, &xi) in obs_t.iter().zip(obs_x.iter()) {
let u = (ti - t) / bandwidth;
let w = kernel(u);
sum_weights += w;
sum_values += w * xi;
}
}
if sum_weights > 0.0 {
sum_values / sum_weights
} else {
f64::NAN
}
})
.collect()
}