rustmodel/src/plot.rs

108 lines
2.3 KiB
Rust

use std::cmp::Ordering;
use crate::util::*;
use nalgebra::Matrix;
use plotters::prelude::*;
#[derive(Clone, Debug, derive_builder::Builder)]
#[builder(setter(into), default)]
pub struct Plot {
dt: f64,
size: (u32, u32),
title: Option<String>,
x_label: Option<String>,
x_max: Option<f64>,
x_min: Option<f64>,
y_max: Option<f64>,
y_min: Option<f64>,
}
impl Default for Plot {
fn default() -> Self {
Self {
dt: 1.0,
size: (640, 480),
title: None,
x_label: None,
x_max: None,
x_min: None,
y_max: None,
y_min: None,
}
}
}
impl Plot {
pub fn plot<const R: usize>(
&self,
path: impl AsRef<std::path::Path>,
data: &[Vect<f64, R>],
series: &[(&str, RGBColor); R],
) {
let root = BitMapBackend::new(path.as_ref(), self.size).into_drawing_area();
root.fill(&WHITE).unwrap();
let mut chart = ChartBuilder::on(&root);
if let Some(title) = &self.title {
chart.caption(
title,
FontDesc::new(FontFamily::SansSerif, 22.0, FontStyle::Normal),
);
}
let x_range =
self.x_min.unwrap_or(0.0)..self.x_max.unwrap_or_else(|| data.len() as f64 * self.dt);
let y_range = self
.y_min
.unwrap_or_else(|| find_extremum(data.iter().map(Matrix::min), Ordering::Less).unwrap())
..self.y_max.unwrap_or_else(|| {
find_extremum(data.iter().map(Matrix::max), Ordering::Greater).unwrap()
});
let mut chart = chart
.margin_right(12)
.y_label_area_size(30)
.x_label_area_size(30)
.build_cartesian_2d(x_range, y_range)
.unwrap();
let mut x_axis = chart.configure_mesh();
if let Some(x_label) = &self.x_label {
x_axis.x_desc(x_label);
}
x_axis.draw().unwrap();
for (i, (label, color)) in series.into_iter().enumerate() {
chart
.draw_series(LineSeries::new(
data.iter()
.enumerate()
.map(|(x, y)| (x as f64 * self.dt, y[i])),
color,
))
.unwrap()
.label(*label)
.legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], color.clone()));
}
chart
.configure_series_labels()
.background_style(WHITE.mix(0.8))
.border_style(BLACK)
.draw()
.unwrap();
}
}
fn find_extremum<T: PartialOrd, I: IntoIterator<Item = T>>(iter: I, ord: Ordering) -> Option<T> {
let mut iter = iter.into_iter();
let extremum = iter.next()?;
Some(iter.fold(extremum, |extremum, i| {
if i.partial_cmp(&extremum) == Some(ord) {
i
} else {
extremum
}
}))
}