Skip to main content

rust_robotics_viz/
visualizer.rs

1//! Visualization utilities for rust_robotics
2//!
3//! Provides a unified interface for plotting using gnuplot.
4
5use gnuplot::{AutoOption, AxesCommon, Caption, Color, Figure, LineWidth, PointSize, PointSymbol};
6use rust_robotics_core::{Obstacles, Path2D, Point2D, Pose2D};
7
8/// Color palette for consistent styling
9pub mod colors {
10    pub const BLACK: &str = "#000000";
11    pub const RED: &str = "#FF0000";
12    pub const GREEN: &str = "#00FF00";
13    pub const BLUE: &str = "#0000FF";
14    pub const YELLOW: &str = "#FFFF00";
15    pub const CYAN: &str = "#00FFFF";
16    pub const MAGENTA: &str = "#FF00FF";
17    pub const ORANGE: &str = "#FFA500";
18    pub const PURPLE: &str = "#800080";
19    pub const GRAY: &str = "#808080";
20
21    // Semantic colors
22    pub const OBSTACLE: &str = BLACK;
23    pub const START: &str = GREEN;
24    pub const GOAL: &str = BLUE;
25    pub const PATH: &str = RED;
26    pub const ROBOT: &str = CYAN;
27    pub const ESTIMATED: &str = "#35C788";
28    pub const GROUND_TRUTH: &str = BLUE;
29    pub const MEASUREMENT: &str = "#DD3355";
30    pub const DEAD_RECKONING: &str = YELLOW;
31}
32
33/// Style for path rendering
34#[derive(Debug, Clone)]
35pub struct PathStyle {
36    pub color: String,
37    pub line_width: f64,
38    pub caption: String,
39}
40
41impl PathStyle {
42    pub fn new(color: &str, caption: &str) -> Self {
43        Self {
44            color: color.to_string(),
45            line_width: 2.0,
46            caption: caption.to_string(),
47        }
48    }
49
50    pub fn with_line_width(mut self, width: f64) -> Self {
51        self.line_width = width;
52        self
53    }
54}
55
56impl Default for PathStyle {
57    fn default() -> Self {
58        Self {
59            color: colors::PATH.to_string(),
60            line_width: 2.0,
61            caption: "Path".to_string(),
62        }
63    }
64}
65
66/// Style for point rendering
67#[derive(Debug, Clone)]
68pub struct PointStyle {
69    pub color: String,
70    pub size: f64,
71    pub symbol: char,
72    pub caption: String,
73}
74
75impl PointStyle {
76    pub fn new(color: &str, caption: &str) -> Self {
77        Self {
78            color: color.to_string(),
79            size: 1.0,
80            symbol: 'O',
81            caption: caption.to_string(),
82        }
83    }
84
85    pub fn with_size(mut self, size: f64) -> Self {
86        self.size = size;
87        self
88    }
89
90    pub fn with_symbol(mut self, symbol: char) -> Self {
91        self.symbol = symbol;
92        self
93    }
94}
95
96/// Plot command to be rendered
97enum PlotCommand {
98    Path {
99        x: Vec<f64>,
100        y: Vec<f64>,
101        style: PathStyle,
102    },
103    Points {
104        x: Vec<f64>,
105        y: Vec<f64>,
106        style: PointStyle,
107    },
108}
109
110/// Main visualizer struct
111pub struct Visualizer {
112    figure: Figure,
113    title: String,
114    x_label: String,
115    y_label: String,
116    x_range: Option<(f64, f64)>,
117    y_range: Option<(f64, f64)>,
118    aspect_ratio: Option<f64>,
119    commands: Vec<PlotCommand>,
120}
121
122impl Visualizer {
123    /// Create a new visualizer
124    pub fn new() -> Self {
125        Self {
126            figure: Figure::new(),
127            title: String::new(),
128            x_label: "X [m]".to_string(),
129            y_label: "Y [m]".to_string(),
130            x_range: None,
131            y_range: None,
132            aspect_ratio: Some(1.0),
133            commands: Vec::new(),
134        }
135    }
136
137    /// Set the plot title
138    pub fn set_title(&mut self, title: &str) -> &mut Self {
139        self.title = title.to_string();
140        self
141    }
142
143    /// Set X axis label
144    pub fn set_x_label(&mut self, label: &str) -> &mut Self {
145        self.x_label = label.to_string();
146        self
147    }
148
149    /// Set Y axis label
150    pub fn set_y_label(&mut self, label: &str) -> &mut Self {
151        self.y_label = label.to_string();
152        self
153    }
154
155    /// Set X axis range
156    pub fn set_x_range(&mut self, min: f64, max: f64) -> &mut Self {
157        self.x_range = Some((min, max));
158        self
159    }
160
161    /// Set Y axis range
162    pub fn set_y_range(&mut self, min: f64, max: f64) -> &mut Self {
163        self.y_range = Some((min, max));
164        self
165    }
166
167    /// Set aspect ratio (None for auto)
168    pub fn set_aspect_ratio(&mut self, ratio: Option<f64>) -> &mut Self {
169        self.aspect_ratio = ratio;
170        self
171    }
172
173    /// Get mutable reference to the internal figure for advanced usage
174    pub fn figure_mut(&mut self) -> &mut Figure {
175        &mut self.figure
176    }
177
178    /// Plot a path
179    pub fn plot_path(&mut self, path: &Path2D, style: &PathStyle) -> &mut Self {
180        let x: Vec<f64> = path.points.iter().map(|p| p.x).collect();
181        let y: Vec<f64> = path.points.iter().map(|p| p.y).collect();
182        self.commands.push(PlotCommand::Path {
183            x,
184            y,
185            style: style.clone(),
186        });
187        self
188    }
189
190    /// Plot a path from x,y vectors
191    pub fn plot_path_xy(&mut self, x: &[f64], y: &[f64], style: &PathStyle) -> &mut Self {
192        self.commands.push(PlotCommand::Path {
193            x: x.to_vec(),
194            y: y.to_vec(),
195            style: style.clone(),
196        });
197        self
198    }
199
200    /// Plot obstacles
201    pub fn plot_obstacles(&mut self, obstacles: &Obstacles) -> &mut Self {
202        let x: Vec<f64> = obstacles.points.iter().map(|p| p.x).collect();
203        let y: Vec<f64> = obstacles.points.iter().map(|p| p.y).collect();
204        self.commands.push(PlotCommand::Points {
205            x,
206            y,
207            style: PointStyle {
208                color: colors::OBSTACLE.to_string(),
209                size: 0.5,
210                symbol: 'S',
211                caption: "Obstacles".to_string(),
212            },
213        });
214        self
215    }
216
217    /// Plot obstacles from x,y vectors
218    pub fn plot_obstacles_xy(&mut self, ox: &[f64], oy: &[f64]) -> &mut Self {
219        self.commands.push(PlotCommand::Points {
220            x: ox.to_vec(),
221            y: oy.to_vec(),
222            style: PointStyle {
223                color: colors::OBSTACLE.to_string(),
224                size: 0.5,
225                symbol: 'S',
226                caption: "Obstacles".to_string(),
227            },
228        });
229        self
230    }
231
232    /// Plot a single point (start, goal, etc.)
233    pub fn plot_point(&mut self, point: Point2D, style: &PointStyle) -> &mut Self {
234        self.commands.push(PlotCommand::Points {
235            x: vec![point.x],
236            y: vec![point.y],
237            style: style.clone(),
238        });
239        self
240    }
241
242    /// Plot multiple points
243    pub fn plot_points(&mut self, points: &[Point2D], style: &PointStyle) -> &mut Self {
244        let x: Vec<f64> = points.iter().map(|p| p.x).collect();
245        let y: Vec<f64> = points.iter().map(|p| p.y).collect();
246        self.commands.push(PlotCommand::Points {
247            x,
248            y,
249            style: style.clone(),
250        });
251        self
252    }
253
254    /// Plot points from x,y vectors
255    pub fn plot_points_xy(&mut self, x: &[f64], y: &[f64], style: &PointStyle) -> &mut Self {
256        self.commands.push(PlotCommand::Points {
257            x: x.to_vec(),
258            y: y.to_vec(),
259            style: style.clone(),
260        });
261        self
262    }
263
264    /// Plot robot pose with direction indicator
265    pub fn plot_robot(&mut self, pose: &Pose2D, size: f64) -> &mut Self {
266        // Plot robot position
267        self.commands.push(PlotCommand::Points {
268            x: vec![pose.x],
269            y: vec![pose.y],
270            style: PointStyle {
271                color: colors::ROBOT.to_string(),
272                size,
273                symbol: 'O',
274                caption: "Robot".to_string(),
275            },
276        });
277
278        // Plot direction line (arrow substitute)
279        let arrow_len = size * 0.5;
280        let end_x = pose.x + arrow_len * pose.yaw.cos();
281        let end_y = pose.y + arrow_len * pose.yaw.sin();
282        self.commands.push(PlotCommand::Path {
283            x: vec![pose.x, end_x],
284            y: vec![pose.y, end_y],
285            style: PathStyle {
286                color: colors::ROBOT.to_string(),
287                line_width: 2.0,
288                caption: String::new(),
289            },
290        });
291        self
292    }
293
294    /// Plot start position
295    pub fn plot_start(&mut self, point: Point2D) -> &mut Self {
296        self.plot_point(
297            point,
298            &PointStyle::new(colors::START, "Start").with_size(1.5),
299        )
300    }
301
302    /// Plot goal position
303    pub fn plot_goal(&mut self, point: Point2D) -> &mut Self {
304        self.plot_point(point, &PointStyle::new(colors::GOAL, "Goal").with_size(1.5))
305    }
306
307    /// Finalize and show the plot
308    pub fn show(&mut self) -> Result<(), String> {
309        self.apply_settings();
310        self.figure.show().map_err(|e| e.to_string()).map(|_| ())
311    }
312
313    /// Save plot to PNG file
314    pub fn save_png(&mut self, path: &str, width: u32, height: u32) -> Result<(), String> {
315        self.apply_settings();
316        self.figure
317            .save_to_png(path, width, height)
318            .map_err(|e| e.to_string())
319    }
320
321    /// Save plot to SVG file
322    pub fn save_svg(&mut self, path: &str, width: u32, height: u32) -> Result<(), String> {
323        self.apply_settings();
324        self.figure
325            .save_to_svg(path, width, height)
326            .map_err(|e| e.to_string())
327    }
328
329    fn apply_settings(&mut self) {
330        let axes = self.figure.axes2d();
331
332        // Apply all plot commands to the same axes
333        for cmd in &self.commands {
334            match cmd {
335                PlotCommand::Path { x, y, style } => {
336                    axes.lines(
337                        x,
338                        y,
339                        &[
340                            Caption(&style.caption),
341                            Color(&style.color),
342                            LineWidth(style.line_width),
343                        ],
344                    );
345                }
346                PlotCommand::Points { x, y, style } => {
347                    axes.points(
348                        x,
349                        y,
350                        &[
351                            Caption(&style.caption),
352                            Color(&style.color),
353                            PointSymbol(style.symbol),
354                            PointSize(style.size),
355                        ],
356                    );
357                }
358            }
359        }
360
361        if !self.title.is_empty() {
362            axes.set_title(&self.title, &[]);
363        }
364        axes.set_x_label(&self.x_label, &[]);
365        axes.set_y_label(&self.y_label, &[]);
366
367        if let Some((min, max)) = self.x_range {
368            axes.set_x_range(AutoOption::Fix(min), AutoOption::Fix(max));
369        }
370        if let Some((min, max)) = self.y_range {
371            axes.set_y_range(AutoOption::Fix(min), AutoOption::Fix(max));
372        }
373        if let Some(ratio) = self.aspect_ratio {
374            axes.set_aspect_ratio(AutoOption::Fix(ratio));
375        }
376    }
377}
378
379impl Default for Visualizer {
380    fn default() -> Self {
381        Self::new()
382    }
383}
384
385/// Quick plot function for simple path visualization
386pub fn quick_plot_path(
387    path: &Path2D,
388    obstacles: Option<&Obstacles>,
389    start: Option<Point2D>,
390    goal: Option<Point2D>,
391    title: &str,
392) -> Visualizer {
393    let mut vis = Visualizer::new();
394    vis.set_title(title);
395
396    if let Some(obs) = obstacles {
397        vis.plot_obstacles(obs);
398    }
399    if let Some(s) = start {
400        vis.plot_start(s);
401    }
402    if let Some(g) = goal {
403        vis.plot_goal(g);
404    }
405    vis.plot_path(path, &PathStyle::default());
406
407    vis
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    #[test]
415    fn test_visualizer_creation() {
416        let vis = Visualizer::new();
417        assert!(vis.aspect_ratio.is_some());
418    }
419
420    #[test]
421    fn test_path_style() {
422        let style = PathStyle::new(colors::RED, "Test Path").with_line_width(3.0);
423        assert_eq!(style.line_width, 3.0);
424        assert_eq!(style.color, colors::RED);
425    }
426}