Skip to main content

spatialrust_mapping/
pose_graph.rs

1//! Pose graph nodes, relative edges, and loop-closure candidates.
2
3use std::collections::HashMap;
4
5use spatialrust_math::Isometry3;
6
7use crate::{MappingError, MappingResult, StampedPose};
8
9/// Stable pose-graph node identifier.
10#[derive(Clone, Debug, PartialEq, Eq, Hash)]
11pub struct PoseNodeId(pub String);
12
13impl PoseNodeId {
14    /// Creates a node id.
15    #[must_use]
16    pub fn new(value: impl Into<String>) -> Self {
17        Self(value.into())
18    }
19}
20
21impl From<&str> for PoseNodeId {
22    fn from(value: &str) -> Self {
23        Self(value.to_owned())
24    }
25}
26
27/// Relative pose constraint between two nodes: `to_T_from`.
28#[derive(Clone, Debug, PartialEq)]
29pub struct PoseGraphEdge {
30    /// Source node.
31    pub from: PoseNodeId,
32    /// Target node.
33    pub to: PoseNodeId,
34    /// Transform mapping `from` coordinates into `to` coordinates.
35    pub to_t_from: Isometry3<f32>,
36    /// Whether this edge is a loop closure.
37    pub loop_closure: bool,
38}
39
40/// Pose graph with absolute node estimates and relative constraints.
41#[derive(Clone, Debug, Default, PartialEq)]
42pub struct PoseGraph {
43    nodes: HashMap<String, StampedPose>,
44    edges: Vec<PoseGraphEdge>,
45}
46
47impl PoseGraph {
48    /// Creates an empty pose graph.
49    #[must_use]
50    pub fn new() -> Self {
51        Self::default()
52    }
53
54    /// Inserts or replaces a node estimate.
55    pub fn upsert_node(&mut self, id: impl Into<PoseNodeId>, pose: StampedPose) {
56        self.nodes.insert(id.into().0, pose);
57    }
58
59    /// Adds a relative constraint.
60    pub fn add_edge(&mut self, edge: PoseGraphEdge) -> MappingResult<()> {
61        if !self.nodes.contains_key(&edge.from.0) || !self.nodes.contains_key(&edge.to.0) {
62            return Err(MappingError::Missing("pose graph endpoint".into()));
63        }
64        self.edges.push(edge);
65        Ok(())
66    }
67
68    /// Returns node estimates.
69    #[must_use]
70    pub fn nodes(&self) -> &HashMap<String, StampedPose> {
71        &self.nodes
72    }
73
74    /// Returns relative edges.
75    #[must_use]
76    pub fn edges(&self) -> &[PoseGraphEdge] {
77        &self.edges
78    }
79
80    /// Propagates absolute poses along non-loop edges from a fixed root using composition.
81    pub fn localize_from_root(&mut self, root: &PoseNodeId) -> MappingResult<()> {
82        let root_pose = self
83            .nodes
84            .get(&root.0)
85            .cloned()
86            .ok_or_else(|| MappingError::Missing(format!("root node `{}`", root.0)))?;
87        let mut changed = true;
88        let mut guard = 0;
89        while changed && guard < self.edges.len().saturating_mul(2).saturating_add(1) {
90            changed = false;
91            guard += 1;
92            for edge in &self.edges {
93                if edge.loop_closure {
94                    continue;
95                }
96                let Some(from_pose) = self.nodes.get(&edge.from.0).cloned() else {
97                    continue;
98                };
99                let predicted =
100                    spatialrust_math::Pose3::new(edge.to_t_from.compose(from_pose.pose.isometry));
101                let Some(existing) = self.nodes.get_mut(&edge.to.0) else {
102                    continue;
103                };
104                let delta = (existing.pose.isometry.translation()
105                    - predicted.isometry.translation())
106                .length();
107                if delta > 1e-4 {
108                    existing.pose = predicted;
109                    // Keep target stamp; overwrite pose only.
110                    changed = true;
111                }
112            }
113        }
114        // Ensure root remains fixed.
115        self.nodes.insert(root.0.clone(), root_pose);
116        Ok(())
117    }
118
119    /// Suggests loop closures for nodes within `max_distance` translation of each other.
120    pub fn loop_closure_candidates(&self, max_distance: f32) -> Vec<(PoseNodeId, PoseNodeId)> {
121        let mut ids: Vec<_> = self.nodes.keys().cloned().collect();
122        ids.sort();
123        let mut out = Vec::new();
124        for i in 0..ids.len() {
125            for j in (i + 1)..ids.len() {
126                let a = &self.nodes[&ids[i]];
127                let b = &self.nodes[&ids[j]];
128                let delta =
129                    (a.pose.isometry.translation() - b.pose.isometry.translation()).length();
130                if delta <= max_distance {
131                    out.push((PoseNodeId(ids[i].clone()), PoseNodeId(ids[j].clone())));
132                }
133            }
134        }
135        out
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use super::{PoseGraph, PoseGraphEdge, PoseNodeId};
142    use crate::StampedPose;
143    use spatialrust_core::Timestamp;
144    use spatialrust_math::{Isometry3, Pose3, Quat, Vec3};
145    use spatialrust_sync::{ClockDomain, StampedTime};
146
147    fn stamped(x: f32, nanos: u64) -> StampedPose {
148        StampedPose::new(
149            StampedTime::exact("host", ClockDomain::HostSteady, Timestamp::from_nanos(nanos)),
150            Pose3::new(Isometry3::new(Quat::new(0.0, 0.0, 0.0, 1.0), Vec3::new(x, 0.0, 0.0))),
151        )
152    }
153
154    #[test]
155    fn localizes_child_from_relative_edge() {
156        let mut graph = PoseGraph::new();
157        graph.upsert_node("a", stamped(0.0, 0));
158        graph.upsert_node("b", stamped(0.0, 1));
159        graph
160            .add_edge(PoseGraphEdge {
161                from: PoseNodeId::new("a"),
162                to: PoseNodeId::new("b"),
163                to_t_from: Isometry3::new(Quat::new(0.0, 0.0, 0.0, 1.0), Vec3::new(2.0, 0.0, 0.0)),
164                loop_closure: false,
165            })
166            .unwrap();
167        graph.localize_from_root(&PoseNodeId::new("a")).unwrap();
168        assert!((graph.nodes()["b"].pose.isometry.translation().x - 2.0).abs() < 1e-4);
169        let loops = graph.loop_closure_candidates(0.1);
170        assert!(loops.is_empty());
171    }
172}