ecoord_core/
transform_tree.rs1use crate::error::Error;
2use crate::frames::{FrameId, FrameInfo};
3use crate::transform::TransformId;
4
5use crate::Error::ContainsDynamicTransform;
6use crate::frame_graph::FrameGraph;
7use crate::transform_edge::TransformEdge;
8use crate::{DynamicTransform, StaticTransform, TimedTransform, Transform};
9use chrono::{DateTime, Utc};
10use nalgebra::Isometry3;
11use rayon::iter::ParallelIterator;
12use std::collections::{HashMap, HashSet};
13
14#[derive(Debug, Clone, Default)]
15pub struct TransformTree {
16 pub edges: HashMap<TransformId, TransformEdge>,
17 pub frames: HashMap<FrameId, FrameInfo>,
18 frame_graph: FrameGraph,
19}
20
21impl TransformTree {
22 pub fn new(edges: Vec<TransformEdge>, frames: Vec<FrameInfo>) -> Result<Self, Error> {
23 let edges: HashMap<TransformId, TransformEdge> = edges
24 .into_iter()
25 .map(|x| (x.transform_id(), x))
26 .collect::<HashMap<_, _>>();
27 let mut frames: HashMap<FrameId, FrameInfo> = frames
28 .into_iter()
29 .map(|x| (x.id.clone(), x))
30 .collect::<HashMap<_, _>>();
31
32 for transform_id in edges.keys() {
33 if !frames.contains_key(&transform_id.parent_frame_id) {
34 frames.insert(
35 transform_id.parent_frame_id.clone(),
36 FrameInfo::new(
37 transform_id.parent_frame_id.clone(),
38 Default::default(),
39 Default::default(),
40 ),
41 );
42 }
43 if !frames.contains_key(&transform_id.child_frame_id) {
44 frames.insert(
45 transform_id.child_frame_id.clone(),
46 FrameInfo::new(
47 transform_id.child_frame_id.clone(),
48 Default::default(),
49 Default::default(),
50 ),
51 );
52 }
53 }
54
55 let frame_graph = FrameGraph::new(edges.keys().cloned().collect())?;
56
57 Ok(Self {
58 edges,
59 frames,
60 frame_graph,
61 })
62 }
63
64 pub fn is_empty(&self) -> bool {
65 self.edges.is_empty()
66 }
67
68 pub fn edges(&self) -> &HashMap<TransformId, TransformEdge> {
69 &self.edges
70 }
71
72 pub fn frames(&self) -> &HashMap<FrameId, FrameInfo> {
73 &self.frames
74 }
75
76 pub fn insert_edge(&mut self, edge: TransformEdge) {
77 if !self.frames.contains_key(edge.parent_frame_id()) {
78 let frame = FrameInfo::new(
79 edge.parent_frame_id().clone(),
80 Default::default(),
81 Default::default(),
82 );
83 self.frames.insert(edge.parent_frame_id().clone(), frame);
84 }
85 if !self.frames.contains_key(edge.child_frame_id()) {
86 let frame = FrameInfo::new(
87 edge.child_frame_id().clone(),
88 Default::default(),
89 Default::default(),
90 );
91 self.frames.insert(edge.child_frame_id().clone(), frame);
92 }
93 self.edges.insert(edge.transform_id().clone(), edge);
94 }
95
96 pub fn root_frames(&self) -> HashSet<FrameId> {
97 self.frame_graph.root_frames()
98 }
99
100 pub fn child_frames(&self) -> HashSet<FrameId> {
101 self.frame_graph.child_frames()
102 }
103}
104
105impl TransformTree {
106 pub fn get_frame_ids(&self) -> HashSet<FrameId> {
107 self.frames.keys().cloned().collect()
108 }
109
110 pub fn contains_frame(&self, frame_id: &FrameId) -> bool {
111 self.frames.contains_key(frame_id)
112 }
113
114 pub fn contains_transform(&self, transform_id: &TransformId) -> bool {
115 self.edges.contains_key(transform_id)
116 }
117
118 pub fn remove_transform(&mut self, transform_id: &TransformId) {
119 self.edges.remove(transform_id);
120 }
121}
122
123impl TransformTree {
124 pub fn static_snapshot_at(&self, timestamp: DateTime<Utc>) -> Result<TransformTree, Error> {
125 let transform_edges: Vec<TransformEdge> = self
126 .edges
127 .values()
128 .map(|x| {
129 let transform = x.at_time(timestamp);
130
131 let static_transform = StaticTransform::new(
132 x.parent_frame_id().clone(),
133 x.child_frame_id().clone(),
134 transform,
135 );
136 TransformEdge::Static(static_transform)
137 })
138 .collect();
139
140 TransformTree::new(transform_edges, self.frames.values().cloned().collect())
141 }
142
143 pub fn get_transform_at_time(
144 &self,
145 transform_id: &TransformId,
146 timestamp: DateTime<Utc>,
147 ) -> Result<Transform, Error> {
148 let transform_id_path = self.frame_graph.get_transform_id_path(transform_id)?;
149 let transforms: Vec<Transform> = transform_id_path
150 .into_iter()
151 .map(|x| self.edges.get(&x).expect("must exist").at_time(timestamp))
152 .collect();
153
154 let isometry: Isometry3<f64> =
155 transforms
156 .into_iter()
157 .fold(Isometry3::identity(), |acc, t| {
158 let iso: Isometry3<f64> = t.isometry();
159 acc * iso
160 });
161
162 Ok(Transform::from(isometry))
163 }
164
165 pub fn get_static_transform(&self, transform_id: &TransformId) -> Result<Transform, Error> {
166 let transform_id_path = self.frame_graph.get_transform_id_path(transform_id)?;
167 let transforms: Vec<Transform> = transform_id_path
168 .into_iter()
169 .map(|x| match self.edges.get(&x).expect("must exist") {
170 TransformEdge::Static(x) => Ok(x.transform),
171 TransformEdge::Dynamic(_x) => Err(ContainsDynamicTransform()),
172 })
173 .collect::<Result<Vec<Transform>, Error>>()?;
174
175 let isometry: Isometry3<f64> =
176 transforms
177 .into_iter()
178 .fold(Isometry3::identity(), |acc, t| {
179 let iso: Isometry3<f64> = t.isometry();
180 acc * iso
181 });
182
183 Ok(Transform::from(isometry))
184 }
185
186 pub fn compute_timed_transforms_for_all_samples(
192 &self,
193 transform_id: &TransformId,
194 ) -> Result<Vec<TimedTransform>, Error> {
195 let transform_id_path = self.frame_graph.get_transform_id_path(transform_id)?;
196 let dynamic_transforms: Vec<&DynamicTransform> = transform_id_path
197 .into_iter()
198 .filter_map(|x| match self.edges.get(&x).expect("must exist") {
199 TransformEdge::Dynamic(dynamic) => Some(dynamic),
200 TransformEdge::Static(_) => None,
201 })
202 .collect();
203
204 let all_sample_timestamps: Vec<DateTime<Utc>> = dynamic_transforms
205 .into_iter()
206 .flat_map(|x| x.sample_timestamps())
207 .collect();
208
209 let timed_transforms = all_sample_timestamps
210 .into_iter()
211 .map(|x| {
212 let timed_transform =
213 TimedTransform::new(x, self.get_transform_at_time(transform_id, x)?);
214
215 Ok(timed_transform)
216 })
217 .collect::<Result<Vec<TimedTransform>, Error>>()?;
218
219 Ok(timed_transforms)
220 }
221
222 pub fn is_transform_path_static(&self, transform_id: &TransformId) -> Result<bool, Error> {
228 let transform_id_path = self.frame_graph.get_transform_id_path(transform_id)?;
229
230 let is_static =
231 transform_id_path
232 .into_iter()
233 .all(|x| match self.edges.get(&x).expect("must exist") {
234 TransformEdge::Static(_) => true,
235 TransformEdge::Dynamic(_) => false,
236 });
237
238 Ok(is_static)
239 }
240}