cartesian_tree/
frame.rs

1use crate::CartesianTreeError;
2use crate::Pose;
3use crate::orientation::IntoOrientation;
4use crate::tree::Walking;
5use crate::tree::{HasChildren, HasParent, NodeEquality};
6
7use nalgebra::UnitQuaternion;
8use nalgebra::{Isometry3, Translation3, Vector3};
9use std::cell::RefCell;
10use std::rc::{Rc, Weak};
11
12use serde::{Deserialize, Serialize};
13use serde_json;
14
15/// Represents a coordinate frame in a Cartesian tree structure.
16///
17/// Each frame can have one parent and multiple children. The frame stores its
18/// transformation (position and orientation) relative to its parent.
19///
20/// Root frames (created via `Frame::new_origin`) have no parent and use the identity transform.
21#[derive(Clone, Debug)]
22pub struct Frame {
23    pub(crate) data: Rc<RefCell<FrameData>>,
24}
25
26#[derive(Debug)]
27pub(crate) struct FrameData {
28    /// The name of the frame (must be unique among siblings).
29    pub(crate) name: String,
30    /// Reference to the parent frame.
31    parent: Option<Weak<RefCell<FrameData>>>,
32    /// Transformation from this frame to its parent frame.
33    transform_to_parent: Isometry3<f64>,
34    /// Child frames directly connected to this frame.
35    children: Vec<Frame>,
36}
37
38#[derive(Serialize, Deserialize, Debug, Clone)]
39struct SerialFrame {
40    name: String,
41    position: Vector3<f64>,
42    orientation: UnitQuaternion<f64>,
43    children: Vec<SerialFrame>,
44}
45
46impl Frame {
47    /// Creates a new root frame (origin) with the given name.
48    ///
49    /// The origin has no parent and uses the identity transform.
50    /// # Arguments
51    /// - `name`: The name of the root frame.
52    ///
53    /// # Example
54    /// ```
55    /// use cartesian_tree::Frame;
56    ///
57    /// let origin = Frame::new_origin("world");
58    /// ```
59    pub fn new_origin(name: impl Into<String>) -> Self {
60        Frame {
61            data: Rc::new(RefCell::new(FrameData {
62                name: name.into(),
63                parent: None,
64                children: Vec::new(),
65                transform_to_parent: Isometry3::identity(),
66            })),
67        }
68    }
69
70    pub(crate) fn borrow(&self) -> std::cell::Ref<'_, FrameData> {
71        self.data.borrow()
72    }
73
74    fn borrow_mut(&self) -> std::cell::RefMut<'_, FrameData> {
75        self.data.borrow_mut()
76    }
77
78    pub(crate) fn downgrade(&self) -> Weak<RefCell<FrameData>> {
79        Rc::downgrade(&self.data)
80    }
81
82    pub(crate) fn walk_up_and_transform(
83        &self,
84        target: &Frame,
85    ) -> Result<Isometry3<f64>, CartesianTreeError> {
86        let mut transform = Isometry3::identity();
87        let mut current = self.clone();
88
89        while !current.is_same(target) {
90            let transform_to_its_parent = {
91                // Scope borrow
92                let current_data = current.borrow();
93
94                // If current frame is root and not target, then target is not an ancestor.
95                if current_data.parent.is_none() {
96                    return Err(CartesianTreeError::IsNoAncestor(target.name(), self.name()));
97                }
98                current_data.transform_to_parent
99            };
100
101            transform = transform_to_its_parent * transform;
102
103            let parent_frame_opt = current.parent();
104            current = parent_frame_opt
105                .ok_or_else(|| CartesianTreeError::IsNoAncestor(target.name(), self.name()))?;
106        }
107
108        Ok(transform)
109    }
110
111    /// Returns the name of the frame.
112    pub fn name(&self) -> String {
113        self.borrow().name.clone()
114    }
115
116    /// Returns the transformation from this frame to its parent frame.
117    ///
118    /// # Returns
119    /// - `Ok(Isometry3<f64>)` if the frame has a parent.
120    /// - `Err(String)` if the frame has no parent.
121    pub fn transform_to_parent(&self) -> Result<Isometry3<f64>, CartesianTreeError> {
122        if self.parent().is_none() {
123            return Err(CartesianTreeError::RootHasNoParent(self.name()));
124        }
125        Ok(self.borrow().transform_to_parent)
126    }
127
128    /// Updates the frame's transformation relative to its parent.
129    ///
130    /// This method modifies the frame's position and orientation relative to its parent frame.
131    /// It fails if the frame is a root frame (i.e., has no parent).
132    ///
133    /// # Arguments
134    /// - `position`: A 3D vector representing the new translational offset from the parent.
135    /// - `orientation`: An orientation convertible into a unit quaternion for new orientational offset from the parent.
136    ///
137    /// # Returns
138    /// - `Ok(())` if the transformation was updated successfully.
139    /// - `Err(String)` if the frame has no parent.
140    ///
141    /// # Example
142    /// ```
143    /// use cartesian_tree::Frame;
144    /// use nalgebra::{Vector3, UnitQuaternion};
145    ///
146    /// let root = Frame::new_origin("root");
147    /// let child = root
148    ///     .add_child("camera", Vector3::new(0.0, 0.0, 1.0), UnitQuaternion::identity())
149    ///     .unwrap();
150    /// child.update_transform(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
151    ///     .unwrap();
152    /// ```
153    pub fn update_transform<O>(
154        &self,
155        position: Vector3<f64>,
156        orientation: O,
157    ) -> Result<(), CartesianTreeError>
158    where
159        O: IntoOrientation,
160    {
161        if self.parent().is_none() {
162            return Err(CartesianTreeError::CannotUpdateRootTransform(self.name()));
163        }
164        self.borrow_mut().transform_to_parent =
165            Isometry3::from_parts(Translation3::from(position), orientation.into_orientation());
166        Ok(())
167    }
168
169    /// Adds a new child frame to the current frame.
170    ///
171    /// The child is positioned and oriented relative to this frame.
172    ///
173    /// Returns an error if a child with the same name already exists.
174    ///
175    /// # Arguments
176    /// - `name`: The name of the new child frame.
177    /// - `position`: A 3D vector representing the translational offset from the parent.
178    /// - `orientation`: An orientation convertible into a unit quaternion.
179    ///
180    /// # Returns
181    /// - `Ok(Rc<Frame>)` the newly added child frame.
182    /// - `Err(String)` if a child with the same name already exists.
183    ///
184    /// # Example
185    /// ```
186    /// use cartesian_tree::Frame;
187    /// use nalgebra::{Vector3, UnitQuaternion};
188    ///
189    /// let root = Frame::new_origin("base");
190    /// let child = root
191    ///     .add_child("camera", Vector3::new(0.0, 0.0, 1.0), UnitQuaternion::identity())
192    ///     .unwrap();
193    /// ```
194    pub fn add_child<O>(
195        &self,
196        name: impl Into<String>,
197        position: Vector3<f64>,
198        orientation: O,
199    ) -> Result<Frame, CartesianTreeError>
200    where
201        O: IntoOrientation,
202    {
203        let child_name = name.into();
204        {
205            let frame = self.borrow();
206            if frame
207                .children
208                .iter()
209                .any(|child| child.borrow().name == child_name)
210            {
211                return Err(CartesianTreeError::ChildNameConflict(
212                    child_name,
213                    self.name(),
214                ));
215            }
216        }
217        let quat = orientation.into_orientation();
218        let transform = Isometry3::from_parts(Translation3::from(position), quat);
219
220        let child = Frame {
221            data: Rc::new(RefCell::new(FrameData {
222                name: child_name,
223                parent: Some(Rc::downgrade(&self.data)),
224                children: Vec::new(),
225                transform_to_parent: transform,
226            })),
227        };
228
229        self.borrow_mut().children.push(child.clone());
230        Ok(child)
231    }
232
233    /// Adds a new child frame calibrated such that a reference pose, when expressed in the new frame,
234    /// matches the desired position and orientation.
235    ///
236    /// # Arguments
237    /// - `name`: The name of the new child frame.
238    /// - `desired_position`: The desired position of the reference pose in the new frame.
239    /// - `desired_orientation`: The desired orientation of the reference pose in the new frame.
240    /// - `reference_pose`: The existing pose (in some frame A) used as the calibration reference.
241    ///
242    /// # Returns
243    /// - `Ok(Frame)` the new child frame if successful.
244    /// - `Err(CartesianTreeError)` if the reference frame is invalid, no common ancestor exists,
245    ///   or a child name conflict occurs.
246    ///
247    /// # Example
248    /// ```
249    /// use cartesian_tree::Frame;
250    /// use nalgebra::{Vector3, UnitQuaternion};
251    ///
252    /// let root = Frame::new_origin("root");
253    /// let reference_pose = root.add_pose(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity());
254    /// let calibrated_child = root.calibrate_child(
255    ///     "calibrated",
256    ///     Vector3::zeros(),
257    ///     UnitQuaternion::identity(),
258    ///     &reference_pose,
259    /// ).unwrap();
260    /// ```
261    pub fn calibrate_child<O>(
262        &self,
263        name: impl Into<String>,
264        desired_position: Vector3<f64>,
265        desired_orientation: O,
266        reference_pose: &Pose,
267    ) -> Result<Frame, CartesianTreeError>
268    where
269        O: IntoOrientation,
270    {
271        let reference_frame = reference_pose.frame().ok_or_else(|| {
272            CartesianTreeError::FrameDropped("Reference pose frame has been dropped".to_string())
273        })?;
274
275        let ancestor = self.lca_with(&reference_frame).ok_or_else(|| {
276            CartesianTreeError::NoCommonAncestor(self.name(), reference_frame.name())
277        })?;
278
279        let t_reference_to_ancestor = reference_frame.walk_up_and_transform(&ancestor)?;
280        let t_pose_to_reference = reference_pose.transformation();
281        let t_pose_to_ancestor = t_reference_to_ancestor * t_pose_to_reference;
282
283        let t_parent_to_ancestor = self.walk_up_and_transform(&ancestor)?;
284        let t_ancestor_to_parent = t_parent_to_ancestor.inverse();
285
286        let desired_pose = Isometry3::from_parts(
287            Translation3::from(desired_position),
288            desired_orientation.into_orientation(),
289        );
290
291        let t_calibrated_to_parent =
292            t_pose_to_ancestor * desired_pose.inverse() * t_ancestor_to_parent;
293
294        self.add_child(
295            name,
296            t_calibrated_to_parent.translation.vector,
297            t_calibrated_to_parent.rotation,
298        )
299    }
300
301    /// Adds a pose to the current frame.
302    ///
303    /// # Arguments
304    /// - `position`: The translational part of the pose.
305    /// - `orientation`: The orientational part of the pose.
306    ///
307    /// # Returns
308    /// - The newly added pose.
309    ///
310    /// # Example
311    /// ```
312    /// use cartesian_tree::Frame;
313    /// use nalgebra::{Vector3, UnitQuaternion};
314    ///
315    /// let frame = Frame::new_origin("base");
316    /// let pose = frame.add_pose(Vector3::new(0.5, 0.0, 0.0), UnitQuaternion::identity());
317    /// ```
318    pub fn add_pose<O>(&self, position: Vector3<f64>, orientation: O) -> Pose
319    where
320        O: IntoOrientation,
321    {
322        Pose::new(self.downgrade(), position, orientation)
323    }
324
325    /// Serializes the frame tree to a JSON string.
326    ///
327    /// This recursively serializes the hierarchy starting from this frame (ideally the root).
328    /// Transforms for root frames are set to identity.
329    ///
330    /// # Returns
331    /// - `Ok(String)` the pretty-printed JSON.
332    /// - `Err(CartesianTreeError)` on serialization failure.
333    pub fn to_json(&self) -> Result<String, CartesianTreeError> {
334        let serial = self.to_serial();
335        Ok(serde_json::to_string_pretty(&serial)?)
336    }
337
338    /// Helper function to convert the frame and its children recursively into a serializable structure.
339    ///
340    /// This is used internally for JSON serialization.
341    fn to_serial(&self) -> SerialFrame {
342        let (position, orientation) = if self.parent().is_some() {
343            let iso = self.transform_to_parent().unwrap_or(Isometry3::identity());
344            (iso.translation.vector, iso.rotation)
345        } else {
346            (Vector3::zeros(), UnitQuaternion::identity())
347        };
348
349        SerialFrame {
350            name: self.name(),
351            position,
352            orientation,
353            children: self.children().into_iter().map(|c| c.to_serial()).collect(),
354        }
355    }
356
357    /// Applies a JSON config to this frame tree by updating matching transforms.
358    ///
359    /// Deserializes the JSON to a temporary structure, then recursively updates transforms
360    /// where names match (partial apply; ignores unmatched nodes in config).
361    /// Skips updating root frames (identity assumed) - assumes this frame is the root.
362    ///
363    /// # Arguments
364    /// - `json`: The JSON string to apply.
365    ///
366    /// # Returns
367    /// - `Ok(())` if applied successfully (even if partial).
368    /// - `Err(CartesianTreeError)` on deserialization or mismatch (e.g., root names differ).
369    pub fn apply_config(&self, json: &str) -> Result<(), CartesianTreeError> {
370        let serial: SerialFrame = serde_json::from_str(json)?;
371        self.apply_serial(&serial)
372    }
373
374    fn apply_serial(&self, serial: &SerialFrame) -> Result<(), CartesianTreeError> {
375        if self.name() != serial.name {
376            return Err(CartesianTreeError::Mismatch(format!(
377                "Frame names do not match: {} vs {}",
378                self.name(),
379                serial.name
380            )));
381        }
382
383        // only update if frame has parent
384        if self.parent().is_some() {
385            self.update_transform(serial.position, serial.orientation)?;
386        }
387
388        for potential_child in &serial.children {
389            if let Some(child) = self
390                .children()
391                .into_iter()
392                .find(|c| c.name() == potential_child.name)
393            {
394                child.apply_serial(potential_child)?;
395            }
396        }
397
398        Ok(())
399    }
400}
401
402impl HasParent for Frame {
403    type Node = Frame;
404
405    fn parent(&self) -> Option<Self::Node> {
406        self.borrow()
407            .parent
408            .clone()
409            .and_then(|data_weak| data_weak.upgrade().map(|data_rc| Frame { data: data_rc }))
410    }
411}
412
413impl NodeEquality for Frame {
414    fn is_same(&self, other: &Self) -> bool {
415        Rc::ptr_eq(&self.data, &other.data)
416    }
417}
418
419impl HasChildren for Frame {
420    type Node = Frame;
421    fn children(&self) -> Vec<Frame> {
422        self.borrow().children.clone()
423    }
424}
425
426#[cfg(test)]
427mod tests {
428    use super::*;
429    use nalgebra::{UnitQuaternion, Vector3};
430
431    #[test]
432    fn create_origin_frame() {
433        let root = Frame::new_origin("world");
434        let root_borrow = root.borrow();
435        assert_eq!(root_borrow.name, "world");
436        assert!(root_borrow.parent.is_none());
437        assert_eq!(root_borrow.children.len(), 0);
438    }
439
440    #[test]
441    fn add_child_frame_with_quaternion() {
442        let root = Frame::new_origin("world");
443        let child = root
444            .add_child(
445                "dummy",
446                Vector3::new(1.0, 0.0, 0.0),
447                UnitQuaternion::identity(),
448            )
449            .unwrap();
450
451        let root_borrow = root.borrow();
452        assert_eq!(root_borrow.children.len(), 1);
453
454        let child_borrow = child.borrow();
455        assert_eq!(child_borrow.name, "dummy");
456        assert!(child_borrow.parent.is_some());
457
458        let parent_name = child_borrow
459            .parent
460            .as_ref()
461            .unwrap()
462            .upgrade()
463            .unwrap()
464            .borrow()
465            .name
466            .clone();
467        assert_eq!(parent_name, "world");
468    }
469
470    #[test]
471    fn add_child_frame_with_rpy() {
472        let root = Frame::new_origin("world");
473        let child = root
474            .add_child(
475                "dummy",
476                Vector3::new(0.0, 1.0, 0.0),
477                (0.0, 0.0, std::f64::consts::FRAC_PI_2),
478            )
479            .unwrap();
480
481        let child_borrow = child.borrow();
482        assert_eq!(child_borrow.name, "dummy");
483
484        let rotation = child_borrow.transform_to_parent.rotation;
485        let expected = UnitQuaternion::from_euler_angles(0.0, 0.0, std::f64::consts::FRAC_PI_2);
486        assert!((rotation.angle() - expected.angle()).abs() < 1e-10);
487    }
488
489    #[test]
490    fn multiple_child_frames() {
491        let root = Frame::new_origin("world");
492
493        let a = root
494            .add_child("a", Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
495            .unwrap();
496        let b = root
497            .add_child("b", Vector3::new(0.0, 1.0, 0.0), UnitQuaternion::identity())
498            .unwrap();
499
500        let root_borrow = root.borrow();
501        assert_eq!(root_borrow.children.len(), 2);
502
503        let a_borrow = a.borrow();
504        let b_borrow = b.borrow();
505
506        assert_eq!(
507            a_borrow
508                .parent
509                .as_ref()
510                .unwrap()
511                .upgrade()
512                .unwrap()
513                .borrow()
514                .name,
515            "world"
516        );
517        assert_eq!(
518            b_borrow
519                .parent
520                .as_ref()
521                .unwrap()
522                .upgrade()
523                .unwrap()
524                .borrow()
525                .name,
526            "world"
527        );
528    }
529
530    #[test]
531    fn reject_duplicate_child_name() {
532        let root = Frame::new_origin("world");
533
534        let _ = root
535            .add_child(
536                "duplicate",
537                Vector3::new(1.0, 0.0, 0.0),
538                UnitQuaternion::identity(),
539            )
540            .unwrap();
541
542        let result = root.add_child(
543            "duplicate",
544            Vector3::new(2.0, 0.0, 0.0),
545            UnitQuaternion::identity(),
546        );
547        assert!(result.is_err());
548    }
549
550    #[test]
551    #[should_panic(expected = "already borrowed")]
552    fn test_borrow_conflict() {
553        let frame = Frame::new_origin("root");
554        let _borrow = frame.borrow(); // Immutable borrow
555        frame.borrow_mut(); // Should panic
556    }
557
558    #[test]
559    fn test_add_pose_to_frame() {
560        let frame = Frame::new_origin("dummy");
561        let pose = frame.add_pose(Vector3::new(1.0, 2.0, 3.0), UnitQuaternion::identity());
562
563        assert_eq!(pose.frame().unwrap().name(), "dummy");
564    }
565
566    #[test]
567    fn test_update_transform() {
568        let root = Frame::new_origin("root");
569        let child = root
570            .add_child(
571                "dummy",
572                Vector3::new(0.0, 0.0, 1.0),
573                UnitQuaternion::identity(),
574            )
575            .unwrap();
576        child
577            .update_transform(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
578            .unwrap();
579        assert_eq!(
580            child.transform_to_parent().unwrap().translation.vector,
581            Vector3::new(1.0, 0.0, 0.0)
582        );
583
584        // Test root frame error
585        assert!(
586            root.update_transform(Vector3::new(1.0, 0.0, 0.0), UnitQuaternion::identity())
587                .is_err()
588        );
589    }
590
591    #[test]
592    fn test_pose_transformation_between_frames() {
593        let root = Frame::new_origin("root");
594
595        let f1 = root
596            .add_child(
597                "f1",
598                Vector3::new(1.0, 0.0, 0.0),
599                UnitQuaternion::identity(),
600            )
601            .unwrap();
602
603        let f2 = f1
604            .add_child(
605                "f2",
606                Vector3::new(0.0, 2.0, 0.0),
607                UnitQuaternion::identity(),
608            )
609            .unwrap();
610
611        let pose_in_f2 = f2.add_pose(Vector3::new(1.0, 1.0, 0.0), UnitQuaternion::identity());
612
613        let pose_in_root = pose_in_f2.in_frame(&root).unwrap();
614        let pos = pose_in_root.transformation().translation.vector;
615
616        // Total offset should be: f2 (0,2,0) + pose (1,1,0) + f1 (1,0,0)
617        assert!((pos - Vector3::new(2.0, 3.0, 0.0)).norm() < 1e-6);
618    }
619
620    #[test]
621    fn test_calibrate_child() {
622        let root = Frame::new_origin("root");
623
624        let reference_pose = root.add_pose(
625            Vector3::new(1.0, 2.0, 3.0),
626            UnitQuaternion::from_euler_angles(0.0, 0.0, std::f64::consts::FRAC_PI_2),
627        );
628
629        // Calibrate a child where the reference pose should appear at (0,0,0) with identity orientation.
630        let calibrated_frame = root
631            .calibrate_child(
632                "calibrated",
633                Vector3::zeros(),
634                UnitQuaternion::identity(),
635                &reference_pose,
636            )
637            .unwrap();
638
639        let pose_in_calibrated = reference_pose.in_frame(&calibrated_frame).unwrap();
640        let transformation = pose_in_calibrated.transformation();
641
642        assert!((transformation.translation.vector - Vector3::zeros()).norm() < 1e-6);
643        assert!((transformation.rotation.angle() - 0.0).abs() < 1e-6);
644
645        // Verify the child's transform matches the reference pose's original transform.
646        let calibrated_transformation = calibrated_frame.transform_to_parent().unwrap();
647        assert!(
648            (calibrated_transformation.translation.vector - Vector3::new(1.0, 2.0, 3.0)).norm()
649                < 1e-6
650        );
651        assert!(
652            (calibrated_transformation.rotation.angle() - std::f64::consts::FRAC_PI_2).abs() < 1e-6
653        );
654    }
655
656    #[test]
657    fn test_to_json_and_apply_config() {
658        let root = Frame::new_origin("root");
659        let _ = root
660            .add_child(
661                "child",
662                Vector3::new(1.0, 2.0, 3.0),
663                UnitQuaternion::from_euler_angles(0.1, 0.2, 0.3),
664            )
665            .unwrap();
666
667        let json = root.to_json().unwrap();
668        // roughly verify JSON structure
669        assert!(json.contains(r#""name": "root""#));
670        assert!(json.contains(r#""name": "child""#));
671
672        // Create a default tree with different transforms
673        let default_root = Frame::new_origin("root");
674        default_root
675            .add_child(
676                "child",
677                Vector3::new(0.0, 0.0, 0.0),
678                UnitQuaternion::identity(),
679            )
680            .unwrap();
681
682        // Apply config
683        default_root.apply_config(&json).unwrap();
684
685        // Verify child transform updated
686        let updated_child = default_root
687            .children()
688            .into_iter()
689            .find(|c| c.name() == "child")
690            .unwrap();
691        let iso = updated_child.transform_to_parent().unwrap();
692        assert_eq!(iso.translation.vector, Vector3::new(1.0, 2.0, 3.0));
693        let (r, p, y) = iso.rotation.euler_angles();
694        assert!((r - 0.1).abs() < 1e-6);
695        assert!((p - 0.2).abs() < 1e-6);
696        assert!((y - 0.3).abs() < 1e-6);
697
698        // Test partial: If config has extra, ignore it
699        let partial_json = r#"
700        {
701            "name": "root",
702            "position": [0.0, 0.0, 0.0],
703            "orientation": [0.0, 0.0, 0.0, 1.0],
704            "children": [
705                {
706                    "name": "child",
707                    "position": [4.0, 5.0, 6.0],
708                    "orientation": [0.0, 0.0, 0.0, 1.0],
709                    "children": []
710                },
711                {
712                    "name": "extra",
713                    "position": [0.0, 0.0, 0.0],
714                    "orientation": [0.0, 0.0, 0.0, 1.0],
715                    "children": []
716                }
717            ]
718        }
719        "#;
720        default_root.apply_config(partial_json).unwrap();
721        let updated_child = default_root
722            .children()
723            .into_iter()
724            .find(|c| c.name() == "child")
725            .unwrap();
726        assert_eq!(
727            updated_child
728                .transform_to_parent()
729                .unwrap()
730                .translation
731                .vector,
732            Vector3::new(4.0, 5.0, 6.0)
733        );
734
735        // Test mismatch
736        let mismatch_json = r#"
737        {
738            "name": "wrong_root",
739            "position": [0.0, 0.0, 0.0],
740            "orientation": [0.0, 0.0, 0.0, 1.0],
741            "children": []
742        }
743        "#;
744        assert!(default_root.apply_config(mismatch_json).is_err());
745    }
746}