optimization_solvers/
plotter_3d.rs

1use super::*;
2use plotly::common::Mode;
3use plotly::layout::Layout;
4use plotly::{Plot, Surface};
5
6pub struct Plotter3d {
7    mesh_size: usize,
8    mesh_x: Vec<f64>,
9    mesh_y: Vec<f64>,
10    plot: Plot,
11}
12impl Plotter3d {
13    pub fn new(xmin: f64, xmax: f64, ymin: f64, ymax: f64, mesh_size: usize) -> Self {
14        let (mesh_x, mesh_y) = (0..mesh_size)
15            .map(|i| {
16                (
17                    xmin + (xmax - xmin) * (i as f64) / (mesh_size as f64),
18                    ymin + (ymax - ymin) * (i as f64) / (mesh_size as f64),
19                )
20            })
21            .unzip();
22        let mut plot = Plot::new();
23        plot.set_layout(Layout::new().width(1600).height(1000));
24        Plotter3d {
25            mesh_size,
26            mesh_x,
27            mesh_y,
28            plot,
29        }
30    }
31    pub fn with_mesh_x(mut self, mesh_x: Vec<f64>) -> Self {
32        self.mesh_x = mesh_x;
33        self
34    }
35    pub fn with_mesh_y(mut self, mesh_y: Vec<f64>) -> Self {
36        self.mesh_y = mesh_y;
37        self
38    }
39    pub fn append_plot(
40        mut self,
41        oracle: &mut impl FnMut(&DVector<f64>) -> FuncEvalMultivariate,
42        title: &str,
43        opacity: f64,
44    ) -> Self {
45        let (n, (x, y)) = (self.mesh_size, (self.mesh_x.clone(), self.mesh_y.clone()));
46
47        let mut z = vec![vec![0.0; n]; n];
48
49        for (i, x) in x.iter().enumerate() {
50            for (j, y) in y.iter().enumerate() {
51                let input = DVector::from_vec(vec![*x, *y]);
52                z[j][i] = *oracle(&input).f();
53            }
54        }
55
56        let surface = Surface::new(z)
57            .x(x)
58            .y(y)
59            .name(title)
60            .show_scale(true)
61            .opacity(opacity);
62        self.plot.add_trace(surface);
63        self
64    }
65    pub fn append_scatter_points(
66        mut self,
67        // f: & mut impl FnMut(f64, f64) -> f64,
68        oracle: &mut impl FnMut(&DVector<f64>) -> FuncEvalMultivariate,
69        points: &[DVector<f64>],
70        title: &str,
71    ) -> Self {
72        let n = points.len();
73        let mut z = vec![0.0; n];
74        let labels = (0..n)
75            .map(|i| format!("Point {}", i))
76            .collect::<Vec<String>>();
77        for (i, input) in points.iter().enumerate() {
78            // let input = DVector::from_vec(vec![*x, *y]);
79            z[i] = *oracle(input).f();
80        }
81        let (x, y) = points.iter().cloned().map(|v| (v[0], v[1])).unzip();
82        let scatter = plotly::Scatter3D::new(x, y, z)
83            .mode(Mode::Markers)
84            .name(title)
85            .text_array(labels);
86        self.plot.add_trace(scatter);
87        self
88    }
89    pub fn set_title(mut self, title: &str) -> Self {
90        let layout = self.plot.layout().clone().title(title).show_legend(true);
91
92        self.plot.set_layout(layout);
93        self
94    }
95    pub fn set_layout_size(mut self, width: usize, height: usize) -> Self {
96        let layout = self.plot.layout().clone().width(width).height(height);
97        self.plot.set_layout(layout);
98        self
99    }
100    pub fn build(&self, filename: &str) {
101        self.plot.write_html(filename);
102    }
103}