|
1 | | -use std::collections::HashMap; |
2 | | -use itertools::chain; |
| 1 | +use indexmap::IndexMap; |
| 2 | +use num_traits::ToPrimitive; |
3 | 3 |
|
4 | 4 | // Utility for creating simple comparison plots for tests without having to aggregate the data manually. |
5 | 5 | // Create an instance and add points for various plots as needed. When the instance goes out of scope, |
6 | 6 | // the plot images are written to the target directory, named after the currently running test. |
7 | 7 | // Inspired by https://github.com/fabianboesiger/debug-plotter |
8 | 8 |
|
| 9 | +// TODO: Replace plotter module with plotter2 |
| 10 | + |
9 | 11 | pub struct Plotter { |
10 | | - plots: HashMap<PlotInfo, PlotData> |
| 12 | + plots: IndexMap<String, PlotData> // IndexMap preserves order of insertion |
11 | 13 | } |
12 | 14 |
|
13 | 15 | impl Plotter { |
14 | 16 | pub fn new() -> Self { |
15 | 17 | Self { |
16 | | - plots: HashMap::new() |
| 18 | + plots: Default::default() |
17 | 19 | } |
18 | 20 | } |
19 | 21 |
|
20 | | - pub fn add_point(&mut self, point: (f64, f64), point_ref: (f64, f64), name: &str, x_label: &str, y_label: &str) { |
21 | | - self.data_mut(name, x_label, y_label).add_point(point, point_ref); |
| 22 | + pub fn add_points<const N: usize, X, Y>(&mut self, title: &str, xlabel: &str, ylabel: &str, points: [(&str, X, Y); N]) |
| 23 | + where X: ToPrimitive + Copy, |
| 24 | + Y: ToPrimitive + Copy |
| 25 | + { |
| 26 | + for i in 0..N { |
| 27 | + let (series, x, y) = points[i]; |
| 28 | + self.data_mut(title, xlabel, ylabel, series, i).add_point((x, y)); // Color according to index |
| 29 | + } |
22 | 30 | } |
23 | 31 |
|
24 | | - pub fn add_points<I1, I2>(&mut self, points: I1, points_ref: I2, name: &str, x_label: &str, y_label: &str) |
25 | | - where I1: IntoIterator<Item=(f64, f64)>, I2: IntoIterator<Item=(f64, f64)> |
| 32 | + pub fn add_point<X, Y>(&mut self, title: &str, xlabel: &str, ylabel: &str, series: &str, point: (X, Y)) |
| 33 | + where X: ToPrimitive, |
| 34 | + Y: ToPrimitive |
26 | 35 | { |
27 | | - self.data_mut(name, x_label, y_label).add_points(points, points_ref); |
| 36 | + self.data_mut(title, xlabel, ylabel, series, 0).add_point(point); // Color zero |
28 | 37 | } |
29 | 38 |
|
30 | | - fn data_mut(&mut self, name: &str, x_label: &str, y_label: &str) -> &mut PlotData { |
31 | | - let info = PlotInfo { |
32 | | - name: name.into(), |
33 | | - x_label: x_label.into(), |
34 | | - y_label: y_label.into() |
35 | | - }; |
| 39 | + fn data_mut(&mut self, title: &str, xlabel: &str, ylabel: &str, series: &str, color: usize) -> &mut SeriesData { |
| 40 | + let plot = self.plots.entry(title.into()).or_insert(PlotData { |
| 41 | + series: Default::default(), |
| 42 | + xlabel: xlabel.into(), |
| 43 | + ylabel: ylabel.into(), |
| 44 | + }); |
36 | 45 |
|
37 | | - self.plots.entry(info).or_default() |
| 46 | + plot.series.entry(series.into()).or_insert(SeriesData { |
| 47 | + points: Vec::new(), |
| 48 | + color, |
| 49 | + }) |
38 | 50 | } |
39 | 51 | } |
40 | 52 |
|
@@ -66,89 +78,107 @@ impl Drop for Plotter { |
66 | 78 | } |
67 | 79 | } |
68 | 80 |
|
69 | | -#[derive(Hash, Eq, PartialEq)] |
70 | | -struct PlotInfo { |
71 | | - name: String, |
72 | | - x_label: String, |
73 | | - y_label: String |
| 81 | +#[allow(dead_code)] |
| 82 | +struct PlotData { |
| 83 | + series: IndexMap<String, SeriesData>, // IndexMap preserves order of insertion |
| 84 | + xlabel: String, |
| 85 | + ylabel: String, |
74 | 86 | } |
75 | 87 |
|
76 | | -#[derive(Default)] |
77 | | -struct PlotData { |
| 88 | +#[allow(dead_code)] |
| 89 | +struct SeriesData { |
78 | 90 | points: Vec<(f64, f64)>, |
79 | | - points_ref: Vec<(f64, f64)>, |
| 91 | + color: usize |
80 | 92 | } |
81 | 93 |
|
82 | | -#[allow(dead_code)] |
83 | | -impl PlotData { |
84 | | - fn add_point(&mut self, point: (f64, f64), point_ref: (f64, f64)) { |
85 | | - self.points.push(point); |
86 | | - self.points_ref.push(point_ref); |
| 94 | +impl SeriesData { |
| 95 | + fn add_point<X, Y>(&mut self, point: (X, Y)) |
| 96 | + where X: ToPrimitive, |
| 97 | + Y: ToPrimitive |
| 98 | + { |
| 99 | + let x = point.0.to_f64().unwrap(); |
| 100 | + let y = point.1.to_f64().unwrap(); |
| 101 | + self.points.push((x, y)); |
87 | 102 | } |
88 | 103 |
|
89 | | - fn add_points<I1, I2>(&mut self, points: I1, points_ref: I2) |
90 | | - where I1: IntoIterator<Item=(f64, f64)>, I2: IntoIterator<Item=(f64, f64)> |
91 | | - { |
92 | | - self.points.extend(points); |
93 | | - self.points_ref.extend(points_ref); |
| 104 | + fn x_min(&self) -> f64 { |
| 105 | + self.points.iter().map(|p| p.0).fold(f64::NAN, f64::min) |
| 106 | + } |
| 107 | + |
| 108 | + fn x_max(&self) -> f64 { |
| 109 | + self.points.iter().map(|p| p.0).fold(f64::NAN, f64::max) |
94 | 110 | } |
95 | 111 |
|
| 112 | + fn y_min(&self) -> f64 { |
| 113 | + self.points.iter().map(|p| p.1).fold(f64::NAN, f64::min) |
| 114 | + } |
| 115 | + |
| 116 | + fn y_max(&self) -> f64 { |
| 117 | + self.points.iter().map(|p| p.1).fold(f64::NAN, f64::max) |
| 118 | + } |
| 119 | +} |
| 120 | + |
| 121 | +#[allow(dead_code)] |
| 122 | +impl PlotData { |
96 | 123 | fn x_min(&self) -> f64 { |
97 | | - chain!(&self.points, &self.points_ref).map(|pt| pt.0).fold(f64::NAN, f64::min) |
| 124 | + self.series.values().map(|s| s.x_min()).fold(f64::NAN, f64::min) |
98 | 125 | } |
99 | 126 |
|
100 | 127 | fn x_max(&self) -> f64 { |
101 | | - chain!(&self.points, &self.points_ref).map(|pt| pt.0).fold(f64::NAN, f64::max) |
| 128 | + self.series.values().map(|s| s.x_max()).fold(f64::NAN, f64::max) |
102 | 129 | } |
103 | 130 |
|
104 | 131 | fn y_min(&self) -> f64 { |
105 | | - chain!(&self.points, &self.points_ref).map(|pt| pt.1).fold(f64::NAN, f64::min) |
| 132 | + self.series.values().map(|s| s.y_min()).fold(f64::NAN, f64::min) |
106 | 133 | } |
107 | 134 |
|
108 | 135 | fn y_max(&self) -> f64 { |
109 | | - chain!(&self.points, &self.points_ref).map(|pt| pt.1).fold(f64::NAN, f64::max) |
| 136 | + self.series.values().map(|s| s.y_max()).fold(f64::NAN, f64::max) |
110 | 137 | } |
111 | 138 | } |
112 | 139 |
|
113 | 140 | // Actually creates the plot file from the given info and data (only if the optional "plotters" dependency is enabled) |
114 | 141 | // The output directory is determined from the name of the current thread, which is named after the test method |
115 | 142 | #[cfg(feature = "plotters")] |
116 | | -fn create_plot(output_path: &str, info: &PlotInfo, data: &PlotData) { |
| 143 | +fn create_plot(output_path: &str, title: &str, data: &PlotData) { |
117 | 144 | use plotters::backend::BitMapBackend; |
118 | 145 | use plotters::chart::ChartBuilder; |
119 | 146 | use plotters::drawing::IntoDrawingArea; |
120 | | - use plotters::element::PathElement; |
121 | 147 | use plotters::series::LineSeries; |
122 | | - use plotters::style::WHITE; |
123 | | - use plotters::style::BLACK; |
124 | | - use plotters::style::BLUE; |
125 | | - use plotters::style::RED; |
| 148 | + use plotters::element::PathElement; |
| 149 | + use plotters::style::{BLACK, WHITE, BLUE, RED, MAGENTA, CYAN, GREEN, YELLOW}; |
126 | 150 |
|
127 | | - let file_path = format!("{}/{}.png", output_path, info.name.to_lowercase().replace(" ", "_")); |
| 151 | + let file_path = format!("{}/{}.png", output_path, title.to_lowercase().replace(" ", "_")); |
128 | 152 | let root_area = BitMapBackend::new(&file_path, (1200, 800)).into_drawing_area(); |
129 | 153 | root_area.fill(&WHITE).unwrap(); |
130 | 154 |
|
| 155 | + |
| 156 | + let colors = [BLUE, RED, MAGENTA, CYAN, GREEN, YELLOW]; |
| 157 | + let color_from_index = |index: usize| { |
| 158 | + let idx = index % colors.len(); |
| 159 | + colors[idx] |
| 160 | + }; |
| 161 | + |
131 | 162 | let mut ctx = ChartBuilder::on(&root_area) |
132 | 163 | .margin(30) |
133 | 164 | .x_label_area_size(30) |
134 | 165 | .y_label_area_size(60) |
135 | | - .caption(&info.name, 20) |
| 166 | + .caption(&title, 20) |
136 | 167 | .build_cartesian_2d(data.x_min()..data.x_max(), data.y_min()..data.y_max()) |
137 | 168 | .unwrap(); |
138 | 169 |
|
139 | 170 | ctx.configure_mesh() |
140 | | - .x_desc(&info.x_label) |
141 | | - .y_desc(&info.y_label) |
| 171 | + .x_desc(&data.xlabel) |
| 172 | + .y_desc(&data.ylabel) |
142 | 173 | .draw() |
143 | 174 | .unwrap(); |
144 | 175 |
|
145 | | - ctx.draw_series(LineSeries::new(data.points.iter().copied(), BLUE)).unwrap() |
146 | | - .label("Actual") |
147 | | - .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE)); |
148 | | - |
149 | | - ctx.draw_series(LineSeries::new(data.points_ref.iter().copied(), RED)).unwrap() |
150 | | - .label("Reference") |
151 | | - .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED)); |
| 176 | + for (name, series) in &data.series { |
| 177 | + let color = color_from_index(series.color); |
| 178 | + ctx.draw_series(LineSeries::new(series.points.iter().copied(), color)).unwrap() |
| 179 | + .label(name) |
| 180 | + .legend(move |(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], color)); |
| 181 | + } |
152 | 182 |
|
153 | 183 | ctx.configure_series_labels() |
154 | 184 | .border_style(BLACK) |
|
0 commit comments