facett_plotters/

violin.rs

//! **Violin plot** — a mirrored kernel-density silhouette per group. The KDE
//! ([`kde`]) is a pure function (Gaussian kernel, Silverman bandwidth) sampled
//! on an evenly-spaced grid; [`ViolinView`] paints the mirrored density as a
//! filled polygon and overlays the median line.

use egui::{Align2, FontId, Pos2, Sense, Stroke, Ui, vec2};
use facett_core::{Facet, FacetCaps, theme};

use crate::boxplot::{Quartiles, paint_tooltip};

/// One density curve: `grid[i]` (the value axis) carries density `dens[i]`.
#[derive(Clone, Debug, PartialEq)]
pub struct Density {
    pub grid: Vec<f64>,
    pub dens: Vec<f64>,
    /// Peak density (for normalising the silhouette width).
    pub peak: f64,
}

/// Silverman's rule-of-thumb bandwidth: `1.06 · σ · n^(−1/5)`. Falls back to a
/// small positive width on a degenerate (zero-variance) sample so the kernel
/// stays well-defined.
pub fn silverman_bandwidth(values: &[f64]) -> f64 {
    let n = values.len();
    if n < 2 {
        return 1.0;
    }
    let mean = values.iter().sum::<f64>() / n as f64;
    let var = values.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (n as f64 - 1.0);
    let sd = var.sqrt();
    let bw = 1.06 * sd * (n as f64).powf(-0.2);
    if bw > 0.0 { bw } else { 1.0 }
}

/// Gaussian-kernel density estimate of `values`, evaluated on `bins` evenly
/// spaced points spanning `[min, max]` padded by one bandwidth. `None` for an
/// empty sample.
pub fn kde(values: &[f64], bins: usize) -> Option<Density> {
    let v: Vec<f64> = values.iter().copied().filter(|x| x.is_finite()).collect();
    if v.is_empty() {
        return None;
    }
    let bins = bins.max(2);
    let bw = silverman_bandwidth(&v);
    let (mut lo, mut hi) = (v[0], v[0]);
    for &x in &v {
        lo = lo.min(x);
        hi = hi.max(x);
    }
    // Pad so the tails of the kernel are visible.
    lo -= bw;
    hi += bw;
    let span = (hi - lo).max(1e-9);
    let inv = 1.0 / (bw * (std::f64::consts::TAU).sqrt());
    let n = v.len() as f64;
    let mut grid = Vec::with_capacity(bins);
    let mut dens = Vec::with_capacity(bins);
    let mut peak = 0.0_f64;
    for i in 0..bins {
        let x = lo + span * (i as f64) / (bins as f64 - 1.0);
        let mut d = 0.0;
        for &xi in &v {
            let u = (x - xi) / bw;
            d += (-0.5 * u * u).exp();
        }
        d *= inv / n;
        peak = peak.max(d);
        grid.push(x);
        dens.push(d);
    }
    Some(Density { grid, dens, peak })
}

/// A violin plot: named groups of raw samples.
pub struct ViolinView {
    pub title: String,
    pub groups: Vec<(String, Vec<f64>)>,
    /// KDE resolution (grid points per group).
    pub bins: usize,
}

impl ViolinView {
    pub fn new(title: impl Into<String>, groups: Vec<(String, Vec<f64>)>) -> Self {
        Self { title: title.into(), groups, bins: 64 }
    }

    fn densities(&self) -> Vec<(String, Density, f64)> {
        self.groups
            .iter()
            .filter_map(|(name, vals)| {
                let d = kde(vals, self.bins)?;
                let median = Quartiles::compute(vals)?.median;
                Some((name.clone(), d, median))
            })
            .collect()
    }

    fn value_range(&self, dens: &[(String, Density, f64)]) -> Option<(f64, f64)> {
        let mut r: Option<(f64, f64)> = None;
        for (_, d, _) in dens {
            for &g in &d.grid {
                r = Some(match r {
                    None => (g, g),
                    Some((lo, hi)) => (lo.min(g), hi.max(g)),
                });
            }
        }
        r
    }

    fn to_csv(&self) -> String {
        let mut s = String::from("group,n,median,peak_density\n");
        for (name, d, median) in self.densities() {
            let n = self.groups.iter().find(|(g, _)| g == &name).map(|(_, v)| v.len()).unwrap_or(0);
            s.push_str(&format!("{name},{n},{median},{}\n", d.peak));
        }
        s
    }
}

impl Facet for ViolinView {
    fn title(&self) -> &str {
        &self.title
    }
    fn caps(&self) -> FacetCaps {
        FacetCaps::NONE.themeable().copyable().resizable()
    }
    fn copy(&mut self) -> Option<String> {
        Some(self.to_csv())
    }
    fn ui(&mut self, ui: &mut Ui) {
        let th = theme(ui);
        let dens = self.densities();
        let (rect, resp) = ui.allocate_exact_size(ui.available_size(), Sense::hover());
        let painter = ui.painter_at(rect);
        painter.rect_filled(rect, egui::CornerRadius::ZERO, th.bg);
        let Some((vmin, vmax)) = self.value_range(&dens) else {
            painter.text(rect.center(), Align2::CENTER_CENTER, "no data", FontId::proportional(13.0), th.text_dim);
            return;
        };
        let plot = rect.shrink2(vec2(28.0, 24.0));
        let span = (vmax - vmin).max(1e-9);
        let y_of = |v: f64| plot.bottom() - ((v - vmin) / span) as f32 * plot.height();
        let n = dens.len().max(1);
        let slot = plot.width() / n as f32;
        let half_w = (slot * 0.42).min(80.0);
        let hover = resp.hover_pos();
        let mut tip: Option<(Pos2, String)> = None;

        for (i, (name, d, median)) in dens.iter().enumerate() {
            let cx = plot.left() + slot * (i as f32 + 0.5);
            let color = facett_core::hash_color(name);
            let peak = d.peak.max(1e-12);
            // Build the mirrored silhouette: right edge down, left edge up.
            let mut poly: Vec<Pos2> = Vec::with_capacity(d.grid.len() * 2);
            for k in 0..d.grid.len() {
                let w = (d.dens[k] / peak) as f32 * half_w;
                poly.push(Pos2::new(cx + w, y_of(d.grid[k])));
            }
            for k in (0..d.grid.len()).rev() {
                let w = (d.dens[k] / peak) as f32 * half_w;
                poly.push(Pos2::new(cx - w, y_of(d.grid[k])));
            }
            painter.add(egui::Shape::convex_polygon(poly.clone(), color.gamma_multiply(0.40), Stroke::new(1.2, th.accent)));
            // Median line across the body.
            let mw = {
                // density at the median, for the line half-width
                let idx = d.grid.iter().enumerate().min_by(|a, b| {
                    (a.1 - median).abs().partial_cmp(&(b.1 - median).abs()).unwrap()
                }).map(|(i, _)| i).unwrap_or(0);
                (d.dens[idx] / peak) as f32 * half_w
            };
            painter.line_segment([Pos2::new(cx - mw, y_of(*median)), Pos2::new(cx + mw, y_of(*median))], Stroke::new(2.0, th.glow));
            painter.text(Pos2::new(cx, plot.bottom() + 12.0), Align2::CENTER_CENTER, name, FontId::proportional(11.0), th.text);

            if let Some(hp) = hover {
                if hp.x >= cx - slot * 0.5 && hp.x < cx + slot * 0.5 && plot.y_range().contains(hp.y) {
                    // map cursor y back to a value
                    let v = vmin + ((plot.bottom() - hp.y) / plot.height()) as f64 * span;
                    tip = Some((hp, format!("{name}\nvalue {v:.3}\nmedian {median:.3}\npeak {:.4}", d.peak)));
                }
            }
        }
        if let Some((at, text)) = tip {
            paint_tooltip(&painter, &th, at, &text);
        }
    }
    fn state_json(&self) -> serde_json::Value {
        serde_json::json!({
            "kind": "violin",
            "bins": self.bins,
            "groups": self.densities().iter().map(|(n, d, m)| serde_json::json!({
                "name": n, "median": m, "peak": d.peak, "grid": d.grid.len(),
            })).collect::<Vec<_>>(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn kde_grid_and_density_lengths_match_bins() {
        let d = kde(&[1.0, 2.0, 2.0, 3.0, 5.0], 32).unwrap();
        assert_eq!(d.grid.len(), 32);
        assert_eq!(d.dens.len(), 32);
        assert!(d.peak > 0.0);
    }

    #[test]
    fn kde_density_is_nonnegative_and_peaks_near_mode() {
        // bimodal-ish: cluster at 0 and at 10
        let mut data = vec![];
        for _ in 0..20 { data.push(0.0); }
        for _ in 0..20 { data.push(10.0); }
        let d = kde(&data, 100).unwrap();
        assert!(d.dens.iter().all(|&x| x >= 0.0));
        // the densest grid point should sit near one of the two clusters
        let (peak_i, _) = d.dens.iter().enumerate().max_by(|a, b| a.1.partial_cmp(b.1).unwrap()).unwrap();
        let peak_x = d.grid[peak_i];
        assert!(peak_x.abs() < 2.0 || (peak_x - 10.0).abs() < 2.0, "peak at {peak_x} should be near 0 or 10");
    }

    #[test]
    fn kde_integrates_to_about_one() {
        // A proper density integrates to ~1 over a wide grid.
        let data: Vec<f64> = (0..200).map(|i| (i as f64) * 0.05).collect();
        let d = kde(&data, 400).unwrap();
        let dx = d.grid[1] - d.grid[0];
        let area: f64 = d.dens.iter().sum::<f64>() * dx;
        assert!((area - 1.0).abs() < 0.05, "area {area} should be ~1");
    }

    #[test]
    fn kde_empty_is_none() {
        assert!(kde(&[], 16).is_none());
        assert!(kde(&[f64::NAN], 16).is_none());
    }

    #[test]
    fn silverman_positive_even_for_constant_sample() {
        assert!(silverman_bandwidth(&[5.0, 5.0, 5.0]) > 0.0);
        assert!(silverman_bandwidth(&[]) > 0.0);
    }

    #[test]
    fn csv_header_and_rows() {
        let v = ViolinView::new("t", vec![("a".into(), vec![1.0, 2.0, 3.0])]);
        let csv = v.to_csv();
        assert!(csv.starts_with("group,n,median,peak_density\n"));
        assert_eq!(csv.lines().count(), 2);
    }
}