plotters-statistical 0.1.0

//! Cumulative gain and lift charts for ranking/classification evaluation.

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

use crate::stats::{gain_curve, GainPoint, StatsError};
use crate::style::stroke_style;

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

/// Whether a [`GainChart`] shows the cumulative gain curve or the lift curve.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GainMode {
    /// Fraction targeted vs fraction of positives captured; baseline is `y = x`.
    Gain,
    /// Fraction targeted vs lift (`gain / fraction`); baseline is `y = 1`.
    Lift,
}

/// A cumulative gain or lift chart as a drawable series (coordinate space
/// `(f64, f64)`).
#[derive(Debug, Clone)]
pub struct GainChart {
    gain: Vec<GainPoint>,
    mode: GainMode,
    // [curve points (n_curve)] then [2 baseline endpoints].
    points: Vec<(f64, f64)>,
    n_curve: usize,
    color: RGBColor,
    stroke_width: u32,
    show_baseline: bool,
}

impl GainChart {
    /// Build from predicted `scores` and true binary `labels`, defaulting to the
    /// cumulative-gain view.
    pub fn from_scores(scores: &[f64], labels: &[bool]) -> Result<Self, StatsError> {
        let gain = gain_curve(scores, labels)?;
        let mut this = Self {
            gain,
            mode: GainMode::Gain,
            points: Vec::new(),
            n_curve: 0,
            color: DEFAULT_COLOR,
            stroke_width: 2,
            show_baseline: true,
        };
        this.rebuild();
        Ok(this)
    }

    /// Switch between the gain and lift views.
    pub fn mode(mut self, mode: GainMode) -> Self {
        self.mode = mode;
        self.rebuild();
        self
    }

    /// Show/hide the chance baseline.
    pub fn baseline(mut self, show: bool) -> Self {
        self.show_baseline = 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 rebuild(&mut self) {
        let mut curve: Vec<(f64, f64)> = match self.mode {
            GainMode::Gain => self.gain.iter().map(|g| (g.fraction, g.gain)).collect(),
            // Lift is undefined at fraction 0; skip that leading point.
            GainMode::Lift => self
                .gain
                .iter()
                .filter(|g| g.fraction > 0.0 && g.lift.is_finite())
                .map(|g| (g.fraction, g.lift))
                .collect(),
        };
        self.n_curve = curve.len();
        // Baseline endpoints.
        match self.mode {
            GainMode::Gain => {
                curve.push((0.0, 0.0));
                curve.push((1.0, 1.0));
            }
            GainMode::Lift => {
                curve.push((0.0, 1.0));
                curve.push((1.0, 1.0));
            }
        }
        self.points = curve;
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a GainChart {
    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 GainChart {
    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_curve + 2 || self.n_curve < 2 {
            return Ok(());
        }
        if self.show_baseline {
            backend.draw_line(
                pix[self.n_curve],
                pix[self.n_curve + 1],
                &stroke_style(&BASELINE_GRAY, 1),
            )?;
        }
        backend.draw_path(
            pix[..self.n_curve].iter().copied(),
            &stroke_style(&self.color, self.stroke_width),
        )?;
        Ok(())
    }
}