apex_solver/io/
mod.rs

1use collections::HashMap;
2use nalgebra::{Matrix3, Matrix6, Quaternion, UnitQuaternion, Vector3};
3
4#[cfg(feature = "visualization")]
5use rerun::external::glam::{Quat, Vec3};
6
7use std::{
8    collections, fmt,
9    fmt::{Display, Formatter},
10    io,
11    path::Path,
12};
13use thiserror::Error;
14use tracing::error;
15
16// Import manifold types
17use crate::{
18    core::problem::VariableEnum,
19    manifold::{se2::SE2, se3::SE3},
20};
21
22// Module declarations
23pub mod bal;
24pub mod g2o;
25pub mod toro;
26
27// Re-exports
28pub use bal::{BalCamera, BalDataset, BalLoader, BalObservation, BalPoint};
29pub use g2o::G2oLoader;
30pub use toro::ToroLoader;
31
32/// Errors that can occur during graph file parsing
33#[derive(Error, Debug)]
34pub enum IoError {
35    #[error("IO error: {0}")]
36    Io(#[from] io::Error),
37
38    #[error("Parse error at line {line}: {message}")]
39    Parse { line: usize, message: String },
40
41    #[error("Unsupported vertex type: {0}")]
42    UnsupportedVertexType(String),
43
44    #[error("Unsupported edge type: {0}")]
45    UnsupportedEdgeType(String),
46
47    #[error("Invalid number format at line {line}: {value}")]
48    InvalidNumber { line: usize, value: String },
49
50    #[error("Missing required fields at line {line}")]
51    MissingFields { line: usize },
52
53    #[error("Duplicate vertex ID: {id}")]
54    DuplicateVertex { id: usize },
55
56    #[error("Invalid quaternion at line {line}: norm = {norm:.6}, expected ~1.0")]
57    InvalidQuaternion { line: usize, norm: f64 },
58
59    #[error("Unsupported file format: {0}")]
60    UnsupportedFormat(String),
61
62    #[error("Failed to create file '{path}': {reason}")]
63    FileCreationFailed { path: String, reason: String },
64}
65
66impl IoError {
67    /// Log the error using tracing::error and return self for chaining
68    #[must_use]
69    pub fn log(self) -> Self {
70        error!("{}", self);
71        self
72    }
73
74    /// Log the error with source error information using tracing::error and return self for chaining
75    #[must_use]
76    pub fn log_with_source<E: std::fmt::Debug>(self, source_error: E) -> Self {
77        error!("{} | Source: {:?}", self, source_error);
78        self
79    }
80}
81
82#[derive(Clone, PartialEq)]
83pub struct VertexSE2 {
84    pub id: usize,
85    pub pose: SE2,
86}
87impl VertexSE2 {
88    pub fn new(id: usize, x: f64, y: f64, theta: f64) -> Self {
89        Self {
90            id,
91            pose: SE2::from_xy_angle(x, y, theta),
92        }
93    }
94
95    pub fn from_vector(id: usize, vector: Vector3<f64>) -> Self {
96        Self {
97            id,
98            pose: SE2::from_xy_angle(vector[0], vector[1], vector[2]),
99        }
100    }
101
102    pub fn id(&self) -> usize {
103        self.id
104    }
105
106    pub fn x(&self) -> f64 {
107        self.pose.x()
108    }
109
110    pub fn y(&self) -> f64 {
111        self.pose.y()
112    }
113
114    pub fn theta(&self) -> f64 {
115        self.pose.angle()
116    }
117}
118
119impl Display for VertexSE2 {
120    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
121        write!(f, "VertexSE2 [ id: {}, pose: {} ]", self.id, self.pose)
122    }
123}
124
125impl VertexSE2 {
126    /// Convert to Rerun 2D position with scaling
127    ///
128    /// **Note:** Requires the `visualization` feature to be enabled.
129    ///
130    /// # Arguments
131    /// * `scale` - Scale factor to apply to position
132    ///
133    /// # Returns
134    /// 2D position array [x, y] compatible with Rerun Points2D
135    pub fn to_rerun_position_2d(&self, scale: f32) -> [f32; 2] {
136        [(self.x() as f32) * scale, (self.y() as f32) * scale]
137    }
138
139    /// Convert to Rerun 3D position with scaling and specified height
140    ///
141    /// **Note:** Requires the `visualization` feature to be enabled.
142    ///
143    /// # Arguments
144    /// * `scale` - Scale factor to apply to X and Y
145    /// * `height` - Z coordinate for the 2D point in 3D space
146    ///
147    /// # Returns
148    /// 3D position compatible with Rerun Transform3D or Points3D
149    #[cfg(feature = "visualization")]
150    pub fn to_rerun_position_3d(&self, scale: f32, height: f32) -> Vec3 {
151        Vec3::new((self.x() as f32) * scale, (self.y() as f32) * scale, height)
152    }
153}
154
155/// SE3 vertex with ID (x, y, z, qx, qy, qz, qw)
156#[derive(Clone, PartialEq)]
157pub struct VertexSE3 {
158    pub id: usize,
159    pub pose: SE3,
160}
161
162impl VertexSE3 {
163    pub fn new(id: usize, translation: Vector3<f64>, rotation: UnitQuaternion<f64>) -> Self {
164        Self {
165            id,
166            pose: SE3::new(translation, rotation),
167        }
168    }
169
170    pub fn from_vector(id: usize, vector: [f64; 7]) -> Self {
171        let translation = Vector3::from([vector[0], vector[1], vector[2]]);
172        let rotation = UnitQuaternion::from_quaternion(Quaternion::from([
173            vector[3], vector[4], vector[5], vector[6],
174        ]));
175        Self::new(id, translation, rotation)
176    }
177
178    pub fn from_translation_quaternion(
179        id: usize,
180        translation: Vector3<f64>,
181        quaternion: Quaternion<f64>,
182    ) -> Self {
183        Self {
184            id,
185            pose: SE3::from_translation_quaternion(translation, quaternion),
186        }
187    }
188
189    pub fn id(&self) -> usize {
190        self.id
191    }
192
193    pub fn translation(&self) -> Vector3<f64> {
194        self.pose.translation()
195    }
196
197    pub fn rotation(&self) -> UnitQuaternion<f64> {
198        self.pose.rotation_quaternion()
199    }
200
201    pub fn x(&self) -> f64 {
202        self.pose.x()
203    }
204
205    pub fn y(&self) -> f64 {
206        self.pose.y()
207    }
208
209    pub fn z(&self) -> f64 {
210        self.pose.z()
211    }
212}
213
214impl Display for VertexSE3 {
215    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
216        write!(f, "VertexSE3 [ id: {}, pose: {} ]", self.id, self.pose)
217    }
218}
219
220impl VertexSE3 {
221    /// Convert to Rerun 3D transform components (position and rotation) with scaling
222    ///
223    /// **Note:** Requires the `visualization` feature to be enabled.
224    ///
225    /// # Arguments
226    /// * `scale` - Scale factor to apply to position
227    ///
228    /// # Returns
229    /// Tuple of (position, rotation) compatible with Rerun Transform3D
230    #[cfg(feature = "visualization")]
231    pub fn to_rerun_transform(&self, scale: f32) -> (Vec3, Quat) {
232        // Extract translation and convert to glam Vec3
233        let trans = self.translation();
234        let position = Vec3::new(trans.x as f32, trans.y as f32, trans.z as f32) * scale;
235
236        // Extract rotation quaternion and convert to glam Quat
237        let rot = self.rotation();
238        let nq = rot.as_ref();
239        let rotation = Quat::from_xyzw(nq.i as f32, nq.j as f32, nq.k as f32, nq.w as f32);
240
241        (position, rotation)
242    }
243}
244
245/// 2D edge constraint between two SE2 vertices
246#[derive(Clone, PartialEq)]
247pub struct EdgeSE2 {
248    pub from: usize,
249    pub to: usize,
250    pub measurement: SE2,          // Relative transformation
251    pub information: Matrix3<f64>, // 3x3 information matrix
252}
253
254impl EdgeSE2 {
255    pub fn new(
256        from: usize,
257        to: usize,
258        dx: f64,
259        dy: f64,
260        dtheta: f64,
261        information: Matrix3<f64>,
262    ) -> Self {
263        Self {
264            from,
265            to,
266            measurement: SE2::from_xy_angle(dx, dy, dtheta),
267            information,
268        }
269    }
270}
271
272impl Display for EdgeSE2 {
273    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
274        write!(
275            f,
276            "EdgeSE2 [ from: {}, to: {}, measurement: {}, information: {} ]",
277            self.from, self.to, self.measurement, self.information
278        )
279    }
280}
281
282/// 3D edge constraint between two SE3 vertices
283#[derive(Clone, PartialEq)]
284pub struct EdgeSE3 {
285    pub from: usize,
286    pub to: usize,
287    pub measurement: SE3,          // Relative transformation
288    pub information: Matrix6<f64>, // 6x6 information matrix
289}
290
291impl EdgeSE3 {
292    pub fn new(
293        from: usize,
294        to: usize,
295        translation: Vector3<f64>,
296        rotation: UnitQuaternion<f64>,
297        information: Matrix6<f64>,
298    ) -> Self {
299        Self {
300            from,
301            to,
302            measurement: SE3::new(translation, rotation),
303            information,
304        }
305    }
306}
307
308impl Display for EdgeSE3 {
309    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
310        write!(
311            f,
312            "EdgeSE3 [ from: {}, to: {}, measurement: {}, information: {} ]",
313            self.from, self.to, self.measurement, self.information
314        )
315    }
316}
317
318/// Main graph structure containing vertices and edges
319#[derive(Clone)]
320pub struct Graph {
321    pub vertices_se2: collections::HashMap<usize, VertexSE2>,
322    pub vertices_se3: collections::HashMap<usize, VertexSE3>,
323    pub edges_se2: Vec<EdgeSE2>,
324    pub edges_se3: Vec<EdgeSE3>,
325}
326
327impl Graph {
328    pub fn new() -> Self {
329        Self {
330            vertices_se2: collections::HashMap::new(),
331            vertices_se3: collections::HashMap::new(),
332            edges_se2: Vec::new(),
333            edges_se3: Vec::new(),
334        }
335    }
336
337    pub fn vertex_count(&self) -> usize {
338        self.vertices_se2.len() + self.vertices_se3.len()
339    }
340
341    pub fn edge_count(&self) -> usize {
342        self.edges_se2.len() + self.edges_se3.len()
343    }
344
345    /// Create a new graph from optimized variables, keeping the original edges
346    ///
347    /// This is useful for saving optimization results: vertices are updated with
348    /// optimized poses, while edges (constraints) remain the same.
349    ///
350    /// # Arguments
351    /// * `variables` - collections::HashMap of optimized variable values from solver
352    /// * `original_edges` - Reference to original graph to copy edges from
353    ///
354    /// # Returns
355    /// A new Graph with optimized vertices and original edges
356    pub fn from_optimized_variables(
357        variables: &HashMap<String, VariableEnum>,
358        original_edges: &Self,
359    ) -> Self {
360        use VariableEnum;
361
362        let mut graph = Graph::new();
363
364        // Copy edges from original (they don't change during optimization)
365        graph.edges_se2 = original_edges.edges_se2.clone();
366        graph.edges_se3 = original_edges.edges_se3.clone();
367
368        // Convert optimized variables back to vertices
369        for (var_name, var) in variables {
370            // Extract vertex ID from variable name (format: "x{id}")
371            if let Some(id_str) = var_name.strip_prefix('x')
372                && let Ok(id) = id_str.parse::<usize>()
373            {
374                match var {
375                    VariableEnum::SE2(v) => {
376                        let vertex = VertexSE2 {
377                            id,
378                            pose: v.value.clone(),
379                        };
380                        graph.vertices_se2.insert(id, vertex);
381                    }
382                    VariableEnum::SE3(v) => {
383                        let vertex = VertexSE3 {
384                            id,
385                            pose: v.value.clone(),
386                        };
387                        graph.vertices_se3.insert(id, vertex);
388                    }
389                    _ => {
390                        // Skip other manifold types (SO2, SO3, Rn)
391                        // These are not commonly used in SLAM graphs
392                    }
393                }
394            }
395        }
396
397        graph
398    }
399}
400
401impl Default for Graph {
402    fn default() -> Self {
403        Self::new()
404    }
405}
406
407impl Display for Graph {
408    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
409        write!(
410            f,
411            "Graph [[ vertices_se2: {} (count: {}), vertices_se3: {} (count: {}), edges_se2: {} (count: {}), edges_se3: {} (count: {}) ]]",
412            self.vertices_se2
413                .values()
414                .map(|v| format!("{}", v))
415                .collect::<Vec<_>>()
416                .join(", "),
417            self.vertices_se2.len(),
418            self.vertices_se3
419                .values()
420                .map(|v| format!("{}", v))
421                .collect::<Vec<_>>()
422                .join(", "),
423            self.vertices_se3.len(),
424            self.edges_se2
425                .iter()
426                .map(|e| format!("{}", e))
427                .collect::<Vec<_>>()
428                .join(", "),
429            self.edges_se2.len(),
430            self.edges_se3
431                .iter()
432                .map(|e| format!("{}", e))
433                .collect::<Vec<_>>()
434                .join(", "),
435            self.edges_se3.len()
436        )
437    }
438}
439
440/// Trait for graph file loaders and writers
441pub trait GraphLoader {
442    /// Load a graph from a file
443    fn load<P: AsRef<Path>>(path: P) -> Result<Graph, IoError>;
444
445    /// Write a graph to a file
446    fn write<P: AsRef<Path>>(graph: &Graph, path: P) -> Result<(), IoError>;
447}
448
449/// Convenience function to load any supported format based on file extension
450pub fn load_graph<P: AsRef<Path>>(path: P) -> Result<Graph, IoError> {
451    let path_ref = path.as_ref();
452    let extension = path_ref
453        .extension()
454        .and_then(|ext| ext.to_str())
455        .ok_or_else(|| {
456            IoError::UnsupportedFormat("No file extension".to_string())
457                .log_with_source(format!("File path: {:?}", path_ref))
458        })?;
459
460    match extension.to_lowercase().as_str() {
461        "g2o" => G2oLoader::load(path),
462        "graph" => ToroLoader::load(path),
463        _ => Err(
464            IoError::UnsupportedFormat(format!("Unsupported extension: {extension}"))
465                .log_with_source(format!("File path: {:?}", path_ref)),
466        ),
467    }
468}
469
470#[cfg(test)]
471mod tests {
472    use super::*;
473    use std::{error, io::Write};
474    use tempfile::NamedTempFile;
475
476    #[test]
477    fn test_load_simple_graph() -> Result<(), IoError> {
478        let mut temp_file = NamedTempFile::new().map_err(|e| {
479            IoError::FileCreationFailed {
480                path: "temp_file".to_string(),
481                reason: e.to_string(),
482            }
483            .log()
484        })?;
485        writeln!(temp_file, "VERTEX_SE2 0 0.0 0.0 0.0")?;
486        writeln!(temp_file, "VERTEX_SE2 1 1.0 0.0 0.0")?;
487        writeln!(temp_file, "# This is a comment")?;
488        writeln!(temp_file)?; // Empty line
489        writeln!(temp_file, "VERTEX_SE3:QUAT 2 0.0 0.0 0.0 0.0 0.0 0.0 1.0")?;
490
491        let graph = G2oLoader::load(temp_file.path())?;
492
493        assert_eq!(graph.vertices_se2.len(), 2);
494        assert_eq!(graph.vertices_se3.len(), 1);
495        assert!(graph.vertices_se2.contains_key(&0));
496        assert!(graph.vertices_se2.contains_key(&1));
497        assert!(graph.vertices_se3.contains_key(&2));
498
499        Ok(())
500    }
501
502    #[test]
503    fn test_load_m3500() -> Result<(), Box<dyn error::Error>> {
504        let graph = G2oLoader::load("data/odometry/M3500.g2o")?;
505        assert!(!graph.vertices_se2.is_empty());
506        Ok(())
507    }
508
509    #[test]
510    fn test_load_sphere2500() -> Result<(), Box<dyn error::Error>> {
511        let graph = G2oLoader::load("data/odometry/sphere2500.g2o")?;
512        assert!(!graph.vertices_se3.is_empty());
513        Ok(())
514    }
515
516    #[test]
517    fn test_duplicate_vertex_error() -> Result<(), io::Error> {
518        let mut temp_file = NamedTempFile::new()?;
519        writeln!(temp_file, "VERTEX_SE2 0 0.0 0.0 0.0")?;
520        writeln!(temp_file, "VERTEX_SE2 0 1.0 0.0 0.0")?; // Duplicate ID
521
522        let result = G2oLoader::load(temp_file.path());
523        assert!(matches!(result, Err(IoError::DuplicateVertex { id: 0 })));
524
525        Ok(())
526    }
527
528    #[test]
529    fn test_toro_loader() -> Result<(), IoError> {
530        let mut temp_file = NamedTempFile::new().map_err(|e| {
531            IoError::FileCreationFailed {
532                path: "temp_file".to_string(),
533                reason: e.to_string(),
534            }
535            .log()
536        })?;
537        writeln!(temp_file, "VERTEX2 0 0.0 0.0 0.0")?;
538        writeln!(temp_file, "VERTEX2 1 1.0 0.0 0.0")?;
539
540        let graph = ToroLoader::load(temp_file.path()).map_err(|e| e.log())?;
541        assert_eq!(graph.vertices_se2.len(), 2);
542
543        Ok(())
544    }
545
546    #[test]
547    #[cfg(feature = "visualization")]
548    fn test_se3_to_rerun() {
549        let vertex = VertexSE3::new(0, Vector3::new(1.0, 2.0, 3.0), UnitQuaternion::identity());
550
551        let (pos, rot) = vertex.to_rerun_transform(0.1);
552
553        assert!((pos.x - 0.1).abs() < 1e-6);
554        assert!((pos.y - 0.2).abs() < 1e-6);
555        assert!((pos.z - 0.3).abs() < 1e-6);
556        assert!((rot.w - 1.0).abs() < 1e-6);
557    }
558
559    #[test]
560    fn test_se2_to_rerun_2d() {
561        let vertex = VertexSE2::new(0, 10.0, 20.0, 0.5);
562
563        let pos = vertex.to_rerun_position_2d(0.1);
564
565        assert!((pos[0] - 1.0).abs() < 1e-6);
566        assert!((pos[1] - 2.0).abs() < 1e-6);
567    }
568
569    #[test]
570    #[cfg(feature = "visualization")]
571    fn test_se2_to_rerun_3d() {
572        let vertex = VertexSE2::new(0, 10.0, 20.0, 0.5);
573
574        let pos = vertex.to_rerun_position_3d(0.1, 5.0);
575
576        assert!((pos.x - 1.0).abs() < 1e-6);
577        assert!((pos.y - 2.0).abs() < 1e-6);
578        assert!((pos.z - 5.0).abs() < 1e-6);
579    }
580}