spatialrust_mapping/
pose_graph.rs1use std::collections::HashMap;
4
5use spatialrust_math::Isometry3;
6
7use crate::{MappingError, MappingResult, StampedPose};
8
9#[derive(Clone, Debug, PartialEq, Eq, Hash)]
11pub struct PoseNodeId(pub String);
12
13impl PoseNodeId {
14 #[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#[derive(Clone, Debug, PartialEq)]
29pub struct PoseGraphEdge {
30 pub from: PoseNodeId,
32 pub to: PoseNodeId,
34 pub to_t_from: Isometry3<f32>,
36 pub loop_closure: bool,
38}
39
40#[derive(Clone, Debug, Default, PartialEq)]
42pub struct PoseGraph {
43 nodes: HashMap<String, StampedPose>,
44 edges: Vec<PoseGraphEdge>,
45}
46
47impl PoseGraph {
48 #[must_use]
50 pub fn new() -> Self {
51 Self::default()
52 }
53
54 pub fn upsert_node(&mut self, id: impl Into<PoseNodeId>, pose: StampedPose) {
56 self.nodes.insert(id.into().0, pose);
57 }
58
59 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 #[must_use]
70 pub fn nodes(&self) -> &HashMap<String, StampedPose> {
71 &self.nodes
72 }
73
74 #[must_use]
76 pub fn edges(&self) -> &[PoseGraphEdge] {
77 &self.edges
78 }
79
80 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 changed = true;
111 }
112 }
113 }
114 self.nodes.insert(root.0.clone(), root_pose);
116 Ok(())
117 }
118
119 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}