Skip to main content

plotters_statistical/series/
precision_recall_curve.rs

1//! Precision–recall curve series.
2//!
3//! Mirrors [`RocCurve`](crate::RocCurve)'s construction and rendering, but the baseline is the
4//! positive-class prevalence (a *horizontal* line), not a diagonal — getting
5//! this right is the whole reason PR is a separate type rather than a copy of
6//! ROC. The computed average precision is exposed for the legend.
7
8use plotters::element::{Drawable, PointCollection};
9use plotters::style::RGBColor;
10use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};
11
12use crate::stats::{precision_recall_curve as compute_pr, StatsError};
13use crate::style::{stroke_style, translucent_fill};
14
15const DEFAULT_COLOR: RGBColor = RGBColor(213, 94, 0); // Okabe–Ito vermillion
16const BASELINE_GRAY: RGBColor = RGBColor(150, 150, 150);
17
18/// A single precision–recall curve as a drawable series (coordinate space
19/// `(f64, f64)` = `(recall, precision)`).
20#[derive(Debug, Clone)]
21pub struct PrecisionRecallCurve {
22    // [curve points (n_curve)] then anchors:
23    //   [n_curve]     = (max_recall, 0.0)  shading corner
24    //   [n_curve + 1] = (0.0, 0.0)         shading corner
25    //   [n_curve + 2] = (0.0, baseline)    baseline endpoint
26    //   [n_curve + 3] = (1.0, baseline)    baseline endpoint
27    points: Vec<(f64, f64)>,
28    n_curve: usize,
29    average_precision: f64,
30    baseline_value: f64,
31    color: RGBColor,
32    stroke_width: u32,
33    shade: bool,
34    shade_alpha: f64,
35    baseline: bool,
36}
37
38impl PrecisionRecallCurve {
39    /// Build from raw predicted `scores` and true binary `labels`.
40    pub fn from_scores(scores: &[f64], labels: &[bool]) -> Result<Self, StatsError> {
41        let data = compute_pr(scores, labels)?;
42        let mut pts: Vec<(f64, f64)> = data
43            .points
44            .iter()
45            .map(|p| (p.recall, p.precision))
46            .collect();
47        // Start the drawn curve at the left axis for a clean line.
48        if let Some(&(_, p0)) = pts.first() {
49            pts.insert(0, (0.0, p0));
50        }
51        Ok(Self::new(pts, data.average_precision, data.baseline))
52    }
53
54    /// Build directly from precomputed `(recall, precision)` points plus the
55    /// average precision and positive-class prevalence baseline.
56    pub fn from_points(points: Vec<(f64, f64)>, average_precision: f64, baseline: f64) -> Self {
57        Self::new(points, average_precision, baseline)
58    }
59
60    fn new(mut points: Vec<(f64, f64)>, average_precision: f64, baseline_value: f64) -> Self {
61        let n_curve = points.len();
62        let max_recall = points.iter().map(|p| p.0).fold(0.0_f64, f64::max);
63        points.push((max_recall, 0.0));
64        points.push((0.0, 0.0));
65        points.push((0.0, baseline_value));
66        points.push((1.0, baseline_value));
67        Self {
68            points,
69            n_curve,
70            average_precision,
71            baseline_value,
72            color: DEFAULT_COLOR,
73            stroke_width: 2,
74            shade: false,
75            shade_alpha: 0.15,
76            baseline: false,
77        }
78    }
79
80    /// The average precision (area under the PR curve, step definition).
81    pub fn average_precision(&self) -> f64 {
82        self.average_precision
83    }
84
85    /// The positive-class prevalence used as the chance baseline.
86    pub fn baseline_value(&self) -> f64 {
87        self.baseline_value
88    }
89
90    /// A legend label of the form `"{name} (AP = 0.83)"`. Build before moving
91    /// the curve into `draw_series`.
92    pub fn legend_label(&self, name: &str) -> String {
93        format!("{name} (AP = {:.2})", self.average_precision)
94    }
95
96    /// Set the line color.
97    pub fn color(mut self, color: RGBColor) -> Self {
98        self.color = color;
99        self
100    }
101
102    /// Set the line stroke width in pixels.
103    pub fn stroke_width(mut self, width: u32) -> Self {
104        self.stroke_width = width;
105        self
106    }
107
108    /// Fill the area under the curve.
109    pub fn shade_area(mut self, shade: bool) -> Self {
110        self.shade = shade;
111        self
112    }
113
114    /// Draw the horizontal prevalence baseline (the PR chance line).
115    pub fn with_baseline(mut self) -> Self {
116        self.baseline = true;
117        self
118    }
119}
120
121impl<'a> PointCollection<'a, (f64, f64)> for &'a PrecisionRecallCurve {
122    type Point = &'a (f64, f64);
123    type IntoIter = &'a [(f64, f64)];
124    fn point_iter(self) -> &'a [(f64, f64)] {
125        &self.points
126    }
127}
128
129impl<DB: DrawingBackend> Drawable<DB> for PrecisionRecallCurve {
130    fn draw<I: Iterator<Item = BackendCoord>>(
131        &self,
132        points: I,
133        backend: &mut DB,
134        _parent_dim: (u32, u32),
135    ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
136        let pix: Vec<BackendCoord> = points.collect();
137        if pix.len() < self.n_curve + 4 || self.n_curve < 2 {
138            return Ok(());
139        }
140        let curve = &pix[..self.n_curve];
141        let corner_right = pix[self.n_curve];
142        let corner_left = pix[self.n_curve + 1];
143        let base_left = pix[self.n_curve + 2];
144        let base_right = pix[self.n_curve + 3];
145
146        if self.shade {
147            let mut poly = curve.to_vec();
148            poly.push(corner_right);
149            poly.push(corner_left);
150            backend.fill_polygon(poly, &translucent_fill(&self.color, self.shade_alpha))?;
151        }
152
153        if self.baseline {
154            backend.draw_line(base_left, base_right, &stroke_style(&BASELINE_GRAY, 1))?;
155        }
156
157        backend.draw_path(
158            curve.iter().copied(),
159            &stroke_style(&self.color, self.stroke_width),
160        )?;
161        Ok(())
162    }
163}