plotters-statistical 0.1.0

//! Empirical CDF series — a right-continuous step curve, optionally with a DKW
//! confidence band and step markers.

use plotters::element::{Drawable, PointCollection};
use plotters::style::RGBColor;
use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};

use crate::stats::{dkw_epsilon, ecdf, StatsError};
use crate::style::{stroke_style, translucent_fill};

const DEFAULT_COLOR: RGBColor = RGBColor(0, 114, 178); // Okabe–Ito blue

/// An empirical cumulative distribution function as a drawable step series
/// (coordinate space `(f64, f64)` = `(value, cumulative proportion)`).
#[derive(Debug, Clone)]
pub struct Ecdf {
    ex: Vec<f64>,
    ep: Vec<f64>,
    n: usize,
    complementary: bool,
    ci_alpha: Option<f64>,
    // Rendered geometry: [step vertices (n_step)] then, if a band, [band polygon
    // (2 * n_step)].
    points: Vec<(f64, f64)>,
    n_step: usize,
    has_band: bool,
    color: RGBColor,
    stroke_width: u32,
    band_opacity: f64,
    marker_radius: u32,
    show_markers: bool,
}

impl Ecdf {
    /// Build from a raw `data` sample. Errors if no finite values remain.
    pub fn from_data(data: &[f64]) -> Result<Self, StatsError> {
        let e = ecdf(data)?;
        let mut this = Self {
            ex: e.x,
            ep: e.p,
            n: e.n,
            complementary: false,
            ci_alpha: None,
            points: Vec::new(),
            n_step: 0,
            has_band: false,
            color: DEFAULT_COLOR,
            stroke_width: 2,
            band_opacity: 0.15,
            marker_radius: 3,
            show_markers: false,
        };
        this.rebuild();
        Ok(this)
    }

    /// Plot the complementary ECDF (survival function, `1 - F`) instead.
    pub fn complementary(mut self, yes: bool) -> Self {
        self.complementary = yes;
        self.rebuild();
        self
    }

    /// Add a Dvoretzky–Kiefer–Wolfowitz confidence band at level `1 - alpha`
    /// (e.g. `alpha = 0.05` for 95%).
    pub fn confidence_band(mut self, alpha: f64) -> Self {
        self.ci_alpha = Some(alpha);
        self.rebuild();
        self
    }

    /// Draw a marker at each observed step.
    pub fn markers(mut self, show: bool) -> Self {
        self.show_markers = show;
        self
    }

    /// Set the line color.
    pub fn color(mut self, color: RGBColor) -> Self {
        self.color = color;
        self
    }

    /// Set the line stroke width in pixels.
    pub fn stroke_width(mut self, width: u32) -> Self {
        self.stroke_width = width;
        self
    }

    fn py(&self, i: usize) -> f64 {
        if self.complementary {
            1.0 - self.ep[i]
        } else {
            self.ep[i]
        }
    }

    fn rebuild(&mut self) {
        let m = self.ex.len();
        let base = if self.complementary { 1.0 } else { 0.0 };
        let mut verts: Vec<(f64, f64)> = Vec::with_capacity(2 * m);
        verts.push((self.ex[0], base));
        verts.push((self.ex[0], self.py(0)));
        for i in 1..m {
            verts.push((self.ex[i], self.py(i - 1)));
            verts.push((self.ex[i], self.py(i)));
        }
        self.n_step = verts.len();

        self.points = verts.clone();
        self.has_band = false;
        if let Some(alpha) = self.ci_alpha {
            let eps = dkw_epsilon(self.n, alpha);
            let mut band: Vec<(f64, f64)> = Vec::with_capacity(2 * self.n_step);
            for &(x, y) in &verts {
                band.push((x, (y + eps).clamp(0.0, 1.0)));
            }
            for &(x, y) in verts.iter().rev() {
                band.push((x, (y - eps).clamp(0.0, 1.0)));
            }
            self.points.extend(band);
            self.has_band = true;
        }
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a Ecdf {
    type Point = &'a (f64, f64);
    type IntoIter = &'a [(f64, f64)];
    fn point_iter(self) -> &'a [(f64, f64)] {
        &self.points
    }
}

impl<DB: DrawingBackend> Drawable<DB> for Ecdf {
    fn draw<I: Iterator<Item = BackendCoord>>(
        &self,
        points: I,
        backend: &mut DB,
        _parent_dim: (u32, u32),
    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
        let pix: Vec<BackendCoord> = points.collect();
        if pix.len() < self.n_step {
            return Ok(());
        }
        if self.has_band && pix.len() >= self.n_step * 3 {
            let band = &pix[self.n_step..self.n_step * 3];
            backend.fill_polygon(
                band.iter().copied(),
                &translucent_fill(&self.color, self.band_opacity),
            )?;
        }
        let step = &pix[..self.n_step];
        backend.draw_path(
            step.iter().copied(),
            &stroke_style(&self.color, self.stroke_width),
        )?;
        if self.show_markers {
            let fill = translucent_fill(&self.color, 0.9);
            for k in (1..self.n_step).step_by(2) {
                backend.draw_circle(step[k], self.marker_radius, &fill, true)?;
            }
        }
        Ok(())
    }
}