plotters-statistical 0.1.0

//! ROC curve series.
//!
//! Construct from raw `(scores, labels)` or from precomputed `(FPR, TPR)`
//! points. Renders as a line with an optional area fill (AUC shading) and an
//! optional random-chance diagonal. The computed AUC is exposed so it can be
//! folded into a legend label — see [`RocCurve::legend_label`].

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

use crate::stats::roc::auc_trapezoid;
use crate::stats::{roc_curve as compute_roc, StatsError};
use crate::style::{stroke_style, translucent_fill};

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

/// A single ROC curve as a drawable series (coordinate space `(f64, f64)` =
/// `(FPR, TPR)`).
#[derive(Debug, Clone)]
pub struct RocCurve {
    // [curve points (n_curve)] then one baseline anchor at (1.0, 0.0) used to
    // close the AUC shading polygon down to the x-axis.
    points: Vec<(f64, f64)>,
    n_curve: usize,
    auc: f64,
    color: RGBColor,
    stroke_width: u32,
    shade: bool,
    shade_alpha: f64,
    baseline: bool,
}

impl RocCurve {
    /// Build from raw predicted `scores` and true binary `labels`.
    pub fn from_scores(scores: &[f64], labels: &[bool]) -> Result<Self, StatsError> {
        let data = compute_roc(scores, labels)?;
        let pts: Vec<(f64, f64)> = data.points.iter().map(|p| (p.fpr, p.tpr)).collect();
        Ok(Self::new(pts, data.auc))
    }

    /// Build directly from precomputed `(FPR, TPR)` points (assumed ascending in
    /// FPR). AUC is computed from them by trapezoidal integration.
    pub fn from_points(points: Vec<(f64, f64)>) -> Self {
        let auc = auc_trapezoid(&points);
        Self::new(points, auc)
    }

    fn new(mut points: Vec<(f64, f64)>, auc: f64) -> Self {
        let n_curve = points.len();
        points.push((1.0, 0.0)); // baseline anchor for shading
        Self {
            points,
            n_curve,
            auc,
            color: DEFAULT_COLOR,
            stroke_width: 2,
            shade: false,
            shade_alpha: 0.15,
            baseline: false,
        }
    }

    /// The area under this ROC curve.
    pub fn auc(&self) -> f64 {
        self.auc
    }

    /// A legend label of the form `"{name} (AUC = 0.87)"`. Build this *before*
    /// moving the curve into `draw_series`, since that consumes it.
    pub fn legend_label(&self, name: &str) -> String {
        format!("{name} (AUC = {:.2})", self.auc)
    }

    /// Set the line color (also used for the shading and legend key).
    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
    }

    /// Fill the area under the curve (AUC shading).
    pub fn shade_area(mut self, shade: bool) -> Self {
        self.shade = shade;
        self
    }

    /// Draw the `y = x` random-chance diagonal. A one-line opt-in, since this
    /// reference line is standard on ROC charts.
    pub fn with_baseline(mut self) -> Self {
        self.baseline = true;
        self
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a RocCurve {
    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 RocCurve {
    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 + 1 || self.n_curve < 2 {
            return Ok(());
        }
        let curve = &pix[..self.n_curve];
        let anchor = pix[self.n_curve]; // pixel of (1.0, 0.0)

        // AUC shading: curve, down to the baseline anchor, back to the origin.
        if self.shade {
            let mut poly = curve.to_vec();
            poly.push(anchor);
            poly.push(curve[0]); // (0,0)
            backend.fill_polygon(poly, &translucent_fill(&self.color, self.shade_alpha))?;
        }

        // Random-chance diagonal from (0,0) to (1,1).
        if self.baseline {
            let dashed: ShapeStyle = stroke_style(&BASELINE_GRAY, 1);
            backend.draw_line(curve[0], curve[self.n_curve - 1], &dashed)?;
        }

        backend.draw_path(
            curve.iter().copied(),
            &stroke_style(&self.color, self.stroke_width),
        )?;
        Ok(())
    }
}