Skip to main content

spatialrust_sync/
frame_graph.rs

1//! Calibrated frame graph with parent→child isometries.
2
3use std::collections::{HashMap, HashSet, VecDeque};
4
5use spatialrust_core::{
6    FieldSemantic, FrameId, HasNormals3, HasPositions3, PointBuffer, PointBufferSet, PointCloud,
7};
8use spatialrust_math::{Isometry3, TransformPoint, Vec3};
9use spatialrust_records::SpatialRecord;
10
11use crate::{SyncError, SyncResult};
12
13/// One directed edge: `child_T_parent` transform.
14#[derive(Clone, Debug, PartialEq)]
15pub struct FrameEdge {
16    /// Parent frame.
17    pub parent: FrameId,
18    /// Child frame.
19    pub child: FrameId,
20    /// Transform that maps parent coordinates into child coordinates.
21    pub child_t_parent: Isometry3<f32>,
22}
23
24/// Directed calibrated frame graph.
25#[derive(Clone, Debug, Default, PartialEq)]
26pub struct FrameGraph {
27    /// Adjacency: parent → (child, child_T_parent).
28    edges: HashMap<String, Vec<(String, Isometry3<f32>)>>,
29    /// Reverse adjacency for lookups towards parents.
30    reverse: HashMap<String, Vec<(String, Isometry3<f32>)>>,
31}
32
33impl FrameGraph {
34    /// Creates an empty frame graph.
35    #[must_use]
36    pub fn new() -> Self {
37        Self::default()
38    }
39
40    /// Inserts or replaces a parent→child edge.
41    pub fn insert_edge(&mut self, edge: FrameEdge) -> SyncResult<()> {
42        if edge.parent.0 == edge.child.0 {
43            return Err(SyncError::InvalidConfiguration(
44                "frame edge parent and child must differ".into(),
45            ));
46        }
47        self.edges
48            .entry(edge.parent.0.clone())
49            .or_default()
50            .retain(|(child, _)| child != &edge.child.0);
51        self.edges
52            .entry(edge.parent.0.clone())
53            .or_default()
54            .push((edge.child.0.clone(), edge.child_t_parent));
55
56        let parent_t_child = edge.child_t_parent.inverse();
57        self.reverse
58            .entry(edge.child.0.clone())
59            .or_default()
60            .retain(|(parent, _)| parent != &edge.parent.0);
61        self.reverse.entry(edge.child.0.clone()).or_default().push((edge.parent.0, parent_t_child));
62        Ok(())
63    }
64
65    /// Looks up a transform that maps `from` coordinates into `to` coordinates.
66    pub fn lookup(&self, from: &FrameId, to: &FrameId) -> SyncResult<Isometry3<f32>> {
67        if from == to {
68            return Ok(Isometry3::identity());
69        }
70        let mut queue = VecDeque::from([(from.0.clone(), Isometry3::identity())]);
71        let mut visited = HashSet::from([from.0.clone()]);
72        while let Some((node, acc)) = queue.pop_front() {
73            for (next, edge) in self.neighbors(&node) {
74                if !visited.insert(next.clone()) {
75                    continue;
76                }
77                let composed = edge.compose(acc);
78                if next == to.0 {
79                    return Ok(composed);
80                }
81                queue.push_back((next, composed));
82            }
83        }
84        Err(SyncError::NoTransformPath { from: from.0.clone(), to: to.0.clone() })
85    }
86
87    /// Transforms a record from its metadata frame into `target`.
88    ///
89    /// Positions and complete normal triplets are transformed by the rigid
90    /// frame path. Other columns, timestamps, schema, and source provenance
91    /// are preserved. The returned record owns a new CPU cloud, so no hidden
92    /// device transfer or in-place mutation occurs.
93    pub fn transform_record_to(
94        &self,
95        record: &SpatialRecord,
96        target: &FrameId,
97    ) -> SyncResult<SpatialRecord> {
98        let source = &record.metadata().frame_id;
99        let transform = self.lookup(source, target)?;
100        let cloud = record.cloud();
101        let (x, y, z) = cloud.positions3()?;
102        let mut positions = Vec::with_capacity(cloud.len());
103        for index in 0..cloud.len() {
104            positions.push(transform.transform_point(Vec3::new(x[index], y[index], z[index])));
105        }
106
107        let normal_names = normal_field_names(cloud)?;
108        let normals = if normal_names.is_some() {
109            let (x, y, z) = cloud.normals3()?;
110            Some(
111                (0..cloud.len())
112                    .map(|index| {
113                        transform
114                            .transform_vector(Vec3::new(x[index], y[index], z[index]))
115                            .normalize()
116                    })
117                    .collect::<Vec<_>>(),
118            )
119        } else {
120            None
121        };
122
123        let position_names = position_field_names(cloud);
124        let mut buffers = PointBufferSet::new();
125        for field in cloud.schema().fields() {
126            let buffer = if field.name == position_names.0 {
127                PointBuffer::from_f32(positions.iter().map(|point| point.x).collect())
128            } else if field.name == position_names.1 {
129                PointBuffer::from_f32(positions.iter().map(|point| point.y).collect())
130            } else if field.name == position_names.2 {
131                PointBuffer::from_f32(positions.iter().map(|point| point.z).collect())
132            } else if let Some((names, values)) = normal_names.as_ref().zip(normals.as_ref()) {
133                if field.name == names.0 {
134                    PointBuffer::from_f32(values.iter().map(|normal| normal.x).collect())
135                } else if field.name == names.1 {
136                    PointBuffer::from_f32(values.iter().map(|normal| normal.y).collect())
137                } else if field.name == names.2 {
138                    PointBuffer::from_f32(values.iter().map(|normal| normal.z).collect())
139                } else {
140                    clone_buffer(cloud.field(&field.name)?)
141                }
142            } else {
143                clone_buffer(cloud.field(&field.name)?)
144            };
145            buffers.insert(field.name.clone(), buffer);
146        }
147
148        let mut metadata = cloud.metadata().clone();
149        metadata.frame_id = target.clone();
150        metadata.sensor_origin =
151            metadata.sensor_origin.map(|origin| transform.transform_point(origin));
152        let transformed_cloud =
153            PointCloud::try_from_parts(cloud.schema().clone(), buffers, metadata)?;
154        Ok(SpatialRecord::try_new_with_provenance(
155            record.schema().clone(),
156            transformed_cloud,
157            record.provenance().clone(),
158        )?)
159    }
160
161    fn neighbors(&self, node: &str) -> Vec<(String, Isometry3<f32>)> {
162        let mut out = Vec::new();
163        if let Some(forward) = self.edges.get(node) {
164            out.extend(forward.iter().cloned());
165        }
166        if let Some(back) = self.reverse.get(node) {
167            out.extend(back.iter().cloned());
168        }
169        out
170    }
171}
172
173fn position_field_names(cloud: &PointCloud) -> (String, String, String) {
174    (
175        cloud
176            .schema()
177            .find_semantic(FieldSemantic::PositionX)
178            .expect("positions3 validates position fields")
179            .name
180            .clone(),
181        cloud
182            .schema()
183            .find_semantic(FieldSemantic::PositionY)
184            .expect("positions3 validates position fields")
185            .name
186            .clone(),
187        cloud
188            .schema()
189            .find_semantic(FieldSemantic::PositionZ)
190            .expect("positions3 validates position fields")
191            .name
192            .clone(),
193    )
194}
195
196fn normal_field_names(cloud: &PointCloud) -> SyncResult<Option<(String, String, String)>> {
197    let names = [
198        (FieldSemantic::NormalX, "normal x"),
199        (FieldSemantic::NormalY, "normal y"),
200        (FieldSemantic::NormalZ, "normal z"),
201    ]
202    .map(|(semantic, label)| {
203        let fields: Vec<&str> = cloud
204            .schema()
205            .fields()
206            .iter()
207            .filter(|field| field.semantic == semantic)
208            .map(|field| field.name.as_str())
209            .collect();
210        (label, fields)
211    });
212    let present = names.iter().filter(|(_, fields)| !fields.is_empty()).count();
213    if present == 0 {
214        return Ok(None);
215    }
216    if names.iter().any(|(_, fields)| fields.len() != 1) {
217        return Err(SyncError::InvalidConfiguration(
218            "normal fields must contain exactly one x/y/z semantic each".into(),
219        ));
220    }
221    Ok(Some((names[0].1[0].to_owned(), names[1].1[0].to_owned(), names[2].1[0].to_owned())))
222}
223
224fn clone_buffer(buffer: &PointBuffer) -> PointBuffer {
225    match buffer {
226        PointBuffer::F32(values) => PointBuffer::from_f32(values.clone()),
227        PointBuffer::F64(values) => PointBuffer::F64(values.clone()),
228        PointBuffer::U8(values) => PointBuffer::U8(values.clone()),
229        PointBuffer::U16(values) => PointBuffer::U16(values.clone()),
230        PointBuffer::U32(values) => PointBuffer::U32(values.clone()),
231        PointBuffer::I32(values) => PointBuffer::I32(values.clone()),
232    }
233}
234
235#[cfg(test)]
236mod tests {
237    use super::{FrameEdge, FrameGraph};
238    use spatialrust_core::{
239        FrameId, HasNormals3, HasPositions3, PointBuffer, PointBufferSet, PointCloud,
240        SpatialMetadata, StandardSchemas, Timestamp,
241    };
242    use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
243    use spatialrust_records::{RecordProvenance, SchemaVersion, SpatialRecord};
244
245    #[test]
246    fn composes_chain_base_to_lidar() {
247        let mut graph = FrameGraph::new();
248        graph
249            .insert_edge(FrameEdge {
250                parent: FrameId::new("base"),
251                child: FrameId::new("sensor"),
252                child_t_parent: Isometry3::new(
253                    Quat::new(0.0, 0.0, 0.0, 1.0),
254                    Vec3::new(1.0, 0.0, 0.0),
255                ),
256            })
257            .unwrap();
258        graph
259            .insert_edge(FrameEdge {
260                parent: FrameId::new("sensor"),
261                child: FrameId::new("lidar"),
262                child_t_parent: Isometry3::new(
263                    Quat::new(0.0, 0.0, 0.0, 1.0),
264                    Vec3::new(0.0, 2.0, 0.0),
265                ),
266            })
267            .unwrap();
268        let t = graph.lookup(&FrameId::new("base"), &FrameId::new("lidar")).unwrap();
269        let p = t.transform_point(Vec3::new(0.0, 0.0, 0.0));
270        assert!((p.x - 1.0).abs() < 1e-5);
271        assert!((p.y - 2.0).abs() < 1e-5);
272    }
273
274    #[test]
275    fn transforms_record_and_preserves_non_geometry_data() {
276        let mut graph = FrameGraph::new();
277        graph
278            .insert_edge(FrameEdge {
279                parent: FrameId::new("base"),
280                child: FrameId::new("sensor"),
281                child_t_parent: Isometry3::new(
282                    Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), std::f32::consts::FRAC_PI_2),
283                    Vec3::new(1.0, 2.0, 0.0),
284                ),
285            })
286            .unwrap();
287
288        let mut buffers = PointBufferSet::new();
289        buffers.insert("x", PointBuffer::from_f32(vec![1.0]));
290        buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
291        buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
292        buffers.insert("intensity", PointBuffer::from_f32(vec![7.0]));
293        buffers.insert("normal_x", PointBuffer::from_f32(vec![1.0]));
294        buffers.insert("normal_y", PointBuffer::from_f32(vec![0.0]));
295        buffers.insert("normal_z", PointBuffer::from_f32(vec![0.0]));
296        let mut metadata = SpatialMetadata::new("base", Timestamp::from_nanos(42));
297        metadata.sensor_origin = Some(Vec3::new(0.0, 0.0, 0.0));
298        let cloud =
299            PointCloud::try_from_parts(StandardSchemas::point_xyzinormal(), buffers, metadata)
300                .unwrap();
301        let provenance = RecordProvenance::try_new("bag")
302            .unwrap()
303            .with_stream_id("/lidar")
304            .with_sequence(Some(9));
305        let record = SpatialRecord::try_from_cloud_with_provenance(
306            "point",
307            SchemaVersion::new(1, 0),
308            cloud,
309            provenance.clone(),
310        )
311        .unwrap();
312
313        let transformed = graph.transform_record_to(&record, &FrameId::new("sensor")).unwrap();
314        let (x, y, z) = transformed.cloud().positions3().unwrap();
315        assert!((x[0] - 1.0).abs() < 1e-5);
316        assert!((y[0] - 3.0).abs() < 1e-5);
317        assert!((z[0]).abs() < 1e-5);
318        let (nx, ny, nz) = transformed.cloud().normals3().unwrap();
319        assert!(nx[0].abs() < 1e-5);
320        assert!((ny[0] - 1.0).abs() < 1e-5);
321        assert!(nz[0].abs() < 1e-5);
322        assert_eq!(transformed.metadata().frame_id, FrameId::new("sensor"));
323        assert_eq!(transformed.metadata().sensor_origin, Some(Vec3::new(1.0, 2.0, 0.0)));
324        assert_eq!(transformed.cloud().field("intensity").unwrap().as_f32().unwrap(), &[7.0]);
325        assert_eq!(transformed.provenance(), &provenance);
326    }
327
328    #[test]
329    fn transform_record_rejects_missing_path() {
330        let mut buffers = PointBufferSet::new();
331        buffers.insert("x", PointBuffer::from_f32(vec![0.0]));
332        buffers.insert("y", PointBuffer::from_f32(vec![0.0]));
333        buffers.insert("z", PointBuffer::from_f32(vec![0.0]));
334        let cloud = PointCloud::try_from_parts(
335            StandardSchemas::point_xyz(),
336            buffers,
337            SpatialMetadata::new("source", Timestamp::from_nanos(0)),
338        )
339        .unwrap();
340        let record =
341            SpatialRecord::try_from_cloud("point", SchemaVersion::new(1, 0), cloud).unwrap();
342        let error =
343            FrameGraph::new().transform_record_to(&record, &FrameId::new("target")).unwrap_err();
344        assert!(matches!(error, crate::SyncError::NoTransformPath { .. }));
345    }
346}