plotters-statistical 0.1.0

//! Precision–recall curve series.
//!
//! Mirrors [`RocCurve`](crate::RocCurve)'s construction and rendering, but the baseline is the
//! positive-class prevalence (a *horizontal* line), not a diagonal — getting
//! this right is the whole reason PR is a separate type rather than a copy of
//! ROC. The computed average precision is exposed for the legend.

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

use crate::stats::{precision_recall_curve as compute_pr, StatsError};
use crate::style::{stroke_style, translucent_fill};

const DEFAULT_COLOR: RGBColor = RGBColor(213, 94, 0); // Okabe–Ito vermillion
const BASELINE_GRAY: RGBColor = RGBColor(150, 150, 150);

/// A single precision–recall curve as a drawable series (coordinate space
/// `(f64, f64)` = `(recall, precision)`).
#[derive(Debug, Clone)]
pub struct PrecisionRecallCurve {
    // [curve points (n_curve)] then anchors:
    //   [n_curve]     = (max_recall, 0.0)  shading corner
    //   [n_curve + 1] = (0.0, 0.0)         shading corner
    //   [n_curve + 2] = (0.0, baseline)    baseline endpoint
    //   [n_curve + 3] = (1.0, baseline)    baseline endpoint
    points: Vec<(f64, f64)>,
    n_curve: usize,
    average_precision: f64,
    baseline_value: f64,
    color: RGBColor,
    stroke_width: u32,
    shade: bool,
    shade_alpha: f64,
    baseline: bool,
}

impl PrecisionRecallCurve {
    /// Build from raw predicted `scores` and true binary `labels`.
    pub fn from_scores(scores: &[f64], labels: &[bool]) -> Result<Self, StatsError> {
        let data = compute_pr(scores, labels)?;
        let mut pts: Vec<(f64, f64)> = data
            .points
            .iter()
            .map(|p| (p.recall, p.precision))
            .collect();
        // Start the drawn curve at the left axis for a clean line.
        if let Some(&(_, p0)) = pts.first() {
            pts.insert(0, (0.0, p0));
        }
        Ok(Self::new(pts, data.average_precision, data.baseline))
    }

    /// Build directly from precomputed `(recall, precision)` points plus the
    /// average precision and positive-class prevalence baseline.
    pub fn from_points(points: Vec<(f64, f64)>, average_precision: f64, baseline: f64) -> Self {
        Self::new(points, average_precision, baseline)
    }

    fn new(mut points: Vec<(f64, f64)>, average_precision: f64, baseline_value: f64) -> Self {
        let n_curve = points.len();
        let max_recall = points.iter().map(|p| p.0).fold(0.0_f64, f64::max);
        points.push((max_recall, 0.0));
        points.push((0.0, 0.0));
        points.push((0.0, baseline_value));
        points.push((1.0, baseline_value));
        Self {
            points,
            n_curve,
            average_precision,
            baseline_value,
            color: DEFAULT_COLOR,
            stroke_width: 2,
            shade: false,
            shade_alpha: 0.15,
            baseline: false,
        }
    }

    /// The average precision (area under the PR curve, step definition).
    pub fn average_precision(&self) -> f64 {
        self.average_precision
    }

    /// The positive-class prevalence used as the chance baseline.
    pub fn baseline_value(&self) -> f64 {
        self.baseline_value
    }

    /// A legend label of the form `"{name} (AP = 0.83)"`. Build before moving
    /// the curve into `draw_series`.
    pub fn legend_label(&self, name: &str) -> String {
        format!("{name} (AP = {:.2})", self.average_precision)
    }

    /// 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
    }

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

    /// Draw the horizontal prevalence baseline (the PR chance line).
    pub fn with_baseline(mut self) -> Self {
        self.baseline = true;
        self
    }
}

impl<'a> PointCollection<'a, (f64, f64)> for &'a PrecisionRecallCurve {
    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 PrecisionRecallCurve {
    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 + 4 || self.n_curve < 2 {
            return Ok(());
        }
        let curve = &pix[..self.n_curve];
        let corner_right = pix[self.n_curve];
        let corner_left = pix[self.n_curve + 1];
        let base_left = pix[self.n_curve + 2];
        let base_right = pix[self.n_curve + 3];

        if self.shade {
            let mut poly = curve.to_vec();
            poly.push(corner_right);
            poly.push(corner_left);
            backend.fill_polygon(poly, &translucent_fill(&self.color, self.shade_alpha))?;
        }

        if self.baseline {
            backend.draw_line(base_left, base_right, &stroke_style(&BASELINE_GRAY, 1))?;
        }

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