apex_solver/io/
mod.rs

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