ecoord_core/
transform_tree.rs

1use 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    /// Computes transforms at all dynamic sample timestamps along a transform path.
187    ///
188    /// This method identifies all dynamic transforms in the path from the parent frame to the
189    /// child frame specified by `transform_id`, collects all their sample timestamps, and
190    /// computes the complete transform at each of those timestamps.
191    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    /// Checks if a transform path contains only static transforms.
223    ///
224    /// This method determines whether the entire path from the parent frame to the
225    /// child frame specified by `transform_id` consists exclusively of static transforms,
226    /// with no dynamic (time-varying) transforms.
227    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}