plotters_statistical/series/
gain_chart.rs1use plotters::element::{Drawable, PointCollection};
4use plotters::style::RGBColor;
5use plotters_backend::{BackendCoord, DrawingBackend, DrawingErrorKind};
6
7use crate::stats::{gain_curve, GainPoint, StatsError};
8use crate::style::stroke_style;
9
10const DEFAULT_COLOR: RGBColor = RGBColor(0, 114, 178); const BASELINE_GRAY: RGBColor = RGBColor(150, 150, 150);
12
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
15pub enum GainMode {
16 Gain,
18 Lift,
20}
21
22#[derive(Debug, Clone)]
25pub struct GainChart {
26 gain: Vec<GainPoint>,
27 mode: GainMode,
28 points: Vec<(f64, f64)>,
30 n_curve: usize,
31 color: RGBColor,
32 stroke_width: u32,
33 show_baseline: bool,
34}
35
36impl GainChart {
37 pub fn from_scores(scores: &[f64], labels: &[bool]) -> Result<Self, StatsError> {
40 let gain = gain_curve(scores, labels)?;
41 let mut this = Self {
42 gain,
43 mode: GainMode::Gain,
44 points: Vec::new(),
45 n_curve: 0,
46 color: DEFAULT_COLOR,
47 stroke_width: 2,
48 show_baseline: true,
49 };
50 this.rebuild();
51 Ok(this)
52 }
53
54 pub fn mode(mut self, mode: GainMode) -> Self {
56 self.mode = mode;
57 self.rebuild();
58 self
59 }
60
61 pub fn baseline(mut self, show: bool) -> Self {
63 self.show_baseline = show;
64 self
65 }
66
67 pub fn color(mut self, color: RGBColor) -> Self {
69 self.color = color;
70 self
71 }
72
73 pub fn stroke_width(mut self, width: u32) -> Self {
75 self.stroke_width = width;
76 self
77 }
78
79 fn rebuild(&mut self) {
80 let mut curve: Vec<(f64, f64)> = match self.mode {
81 GainMode::Gain => self.gain.iter().map(|g| (g.fraction, g.gain)).collect(),
82 GainMode::Lift => self
84 .gain
85 .iter()
86 .filter(|g| g.fraction > 0.0 && g.lift.is_finite())
87 .map(|g| (g.fraction, g.lift))
88 .collect(),
89 };
90 self.n_curve = curve.len();
91 match self.mode {
93 GainMode::Gain => {
94 curve.push((0.0, 0.0));
95 curve.push((1.0, 1.0));
96 }
97 GainMode::Lift => {
98 curve.push((0.0, 1.0));
99 curve.push((1.0, 1.0));
100 }
101 }
102 self.points = curve;
103 }
104}
105
106impl<'a> PointCollection<'a, (f64, f64)> for &'a GainChart {
107 type Point = &'a (f64, f64);
108 type IntoIter = &'a [(f64, f64)];
109 fn point_iter(self) -> &'a [(f64, f64)] {
110 &self.points
111 }
112}
113
114impl<DB: DrawingBackend> Drawable<DB> for GainChart {
115 fn draw<I: Iterator<Item = BackendCoord>>(
116 &self,
117 points: I,
118 backend: &mut DB,
119 _parent_dim: (u32, u32),
120 ) -> Result<(), DrawingErrorKind<DB::ErrorType>> {
121 let pix: Vec<BackendCoord> = points.collect();
122 if pix.len() < self.n_curve + 2 || self.n_curve < 2 {
123 return Ok(());
124 }
125 if self.show_baseline {
126 backend.draw_line(
127 pix[self.n_curve],
128 pix[self.n_curve + 1],
129 &stroke_style(&BASELINE_GRAY, 1),
130 )?;
131 }
132 backend.draw_path(
133 pix[..self.n_curve].iter().copied(),
134 &stroke_style(&self.color, self.stroke_width),
135 )?;
136 Ok(())
137 }
138}