plotters-statistical 0.2.0

//! Gaussian kernel density estimation, used to shape violin plots.

use super::{sorted_finite, StatsError};
use std::f64::consts::PI;

/// An evaluated kernel density estimate: paired `x` positions and density
/// values, plus the bandwidth that produced them.
#[derive(Debug, Clone, PartialEq)]
pub struct KdeCurve {
    /// Evaluation positions along the value axis, ascending.
    pub xs: Vec<f64>,
    /// Estimated density at each corresponding `xs` position. Integrates to ~1.
    pub density: Vec<f64>,
    /// The bandwidth (kernel standard deviation) actually used.
    pub bandwidth: f64,
}

impl KdeCurve {
    /// The largest density value on the curve — handy for scaling a violin's
    /// half-width so its widest point maps to a fixed pixel/coordinate extent.
    pub fn max_density(&self) -> f64 {
        self.density.iter().copied().fold(0.0_f64, f64::max)
    }
}

/// Silverman's rule-of-thumb bandwidth:
/// `h = 0.9 * min(std, IQR/1.34) * n^(-1/5)`.
///
/// Uses the sample standard deviation (n − 1 denominator). Returns
/// [`StatsError::InvalidBandwidth`] when the estimate is not finite and
/// positive — e.g. fewer than two points, or an all-identical sample where both
/// the spread and the IQR are zero.
pub fn silverman_bandwidth(data: &[f64]) -> Result<f64, StatsError> {
    let sorted = sorted_finite(data);
    let n = sorted.len();
    if n < 2 {
        return Err(StatsError::InvalidBandwidth);
    }
    let nf = n as f64;
    let mean = sorted.iter().sum::<f64>() / nf;
    let var = sorted.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (nf - 1.0);
    let std = var.sqrt();

    // IQR via the same type-7 percentile used for box plots.
    let q1 = super::percentile_sorted(&sorted, 0.25);
    let q3 = super::percentile_sorted(&sorted, 0.75);
    let iqr = q3 - q1;

    let spread = if iqr > 0.0 { std.min(iqr / 1.34) } else { std };
    let h = 0.9 * spread * nf.powf(-1.0 / 5.0);
    if h.is_finite() && h > 0.0 {
        Ok(h)
    } else {
        Err(StatsError::InvalidBandwidth)
    }
}

fn gaussian_kernel(u: f64) -> f64 {
    (-0.5 * u * u).exp() / (2.0 * PI).sqrt()
}

/// Evaluate a Gaussian KDE of `data` on a uniform grid.
///
/// * `bandwidth` — `Some(h)` to force a bandwidth, or `None` to derive one via
///   [`silverman_bandwidth`].
/// * `n_points` — number of grid points; higher gives a smoother violin outline
///   (values around 200 render cleanly). Clamped to at least 2.
/// * `cut` — how far past the data range to extend, in bandwidths (a value of
///   `3.0` captures the Gaussian tails without a visible cutoff).
///
/// Returns [`StatsError::EmptyInput`] for no finite data, or
/// [`StatsError::InvalidBandwidth`] for a non-positive bandwidth (explicit or
/// derived).
pub fn kde_curve(
    data: &[f64],
    bandwidth: Option<f64>,
    n_points: usize,
    cut: f64,
) -> Result<KdeCurve, StatsError> {
    let sorted = sorted_finite(data);
    if sorted.is_empty() {
        return Err(StatsError::EmptyInput);
    }
    let h = match bandwidth {
        Some(b) if b.is_finite() && b > 0.0 => b,
        Some(_) => return Err(StatsError::InvalidBandwidth),
        None => silverman_bandwidth(&sorted)?,
    };

    let n_points = n_points.max(2);
    let lo = sorted[0] - cut * h;
    let hi = sorted[sorted.len() - 1] + cut * h;
    let step = (hi - lo) / (n_points as f64 - 1.0);

    let nf = sorted.len() as f64;
    let mut xs = Vec::with_capacity(n_points);
    let mut density = Vec::with_capacity(n_points);
    for i in 0..n_points {
        let x = lo + step * i as f64;
        let d = sorted
            .iter()
            .map(|&xi| gaussian_kernel((x - xi) / h))
            .sum::<f64>()
            / (nf * h);
        xs.push(x);
        density.push(d);
    }

    Ok(KdeCurve {
        xs,
        density,
        bandwidth: h,
    })
}