spatialrust_registration/
kabsch.rs

1use spatialrust_math::{Isometry3, Quat, Vec3};
2
3/// Estimates the least-squares proper rigid transform mapping paired source to target.
4///
5/// Uses f64 centroids and Horn's unit-quaternion formulation with a normalized
6/// symmetric eigensystem. Returns None for short, mismatched or nonfinite inputs,
7/// or a transform that cannot be represented in f32. Collinear/coincident inputs
8/// can admit multiple equally valid rotations; the returned pose is not a
9/// certificate of observability or of correct correspondences.
10#[must_use]
11pub fn estimate_rigid_transform(
12    source: &[Vec3<f32>],
13    target: &[Vec3<f32>],
14) -> Option<Isometry3<f32>> {
15    if source.len() != target.len() || source.len() < 3 {
16        return None;
17    }
18    let mut mean_source = [0.0_f64; 3];
19    let mut mean_target = [0.0_f64; 3];
20    for (src, dst) in source.iter().zip(target) {
21        for (i, (s, t)) in [src.x, src.y, src.z].into_iter().zip([dst.x, dst.y, dst.z]).enumerate()
22        {
23            if !s.is_finite() || !t.is_finite() {
24                return None;
25            }
26            mean_source[i] += f64::from(s);
27            mean_target[i] += f64::from(t);
28        }
29    }
30    for i in 0..3 {
31        mean_source[i] /= source.len() as f64;
32        mean_target[i] /= source.len() as f64;
33    }
34    let mut h = [[0.0_f64; 3]; 3];
35    for (src, dst) in source.iter().zip(target) {
36        let s = [f64::from(src.x), f64::from(src.y), f64::from(src.z)];
37        let t = [f64::from(dst.x), f64::from(dst.y), f64::from(dst.z)];
38        for i in 0..3 {
39            for j in 0..3 {
40                h[i][j] += (s[i] - mean_source[i]) * (t[j] - mean_target[j]);
41            }
42        }
43    }
44    estimate_rigid_transform_from_moments(source.len(), mean_source, mean_target, h)
45}
46
47pub(crate) fn estimate_rigid_transform_from_moments(
48    count: usize,
49    mean_source: [f64; 3],
50    mean_target: [f64; 3],
51    mut h: [[f64; 3]; 3],
52) -> Option<Isometry3<f32>> {
53    if count < 3
54        || mean_source.iter().chain(&mean_target).chain(h.iter().flatten()).any(|v| !v.is_finite())
55    {
56        return None;
57    }
58    let scale = h.iter().flatten().fold(0.0_f64, |largest, v| largest.max(v.abs()));
59    let rotation = if scale == 0.0 {
60        Quat::<f32>::identity()
61    } else {
62        for row in &mut h {
63            for v in row {
64                *v /= scale;
65            }
66        }
67        let [[xx, xy, xz], [yx, yy, yz], [zx, zy, zz]] = h;
68        // Quaternion coordinates are [w, x, y, z]. The largest eigenvector
69        // maximizes trace(R H) and always represents a proper rotation.
70        let n = [
71            [xx + yy + zz, yz - zy, zx - xz, xy - yx],
72            [yz - zy, xx - yy - zz, xy + yx, zx + xz],
73            [zx - xz, xy + yx, -xx + yy - zz, yz + zy],
74            [xy - yx, zx + xz, yz + zy, -xx - yy + zz],
75        ];
76        let q = largest_eigenvector(n)?;
77        Quat::new(q[1] as f32, q[2] as f32, q[3] as f32, q[0] as f32).normalize()
78    };
79    // Use the actual returned f32 rotation for the centroid correction.
80    let r = rotation.to_mat3();
81    let mut t = [0.0_f32; 3];
82    for i in 0..3 {
83        t[i] = (mean_target[i] - (0..3).map(|j| f64::from(r.m[i][j]) * mean_source[j]).sum::<f64>())
84            as f32;
85        if !t[i].is_finite() {
86            return None;
87        }
88    }
89    Some(Isometry3::new(rotation, Vec3::new(t[0], t[1], t[2])))
90}
91
92fn largest_eigenvector(mut a: [[f64; 4]; 4]) -> Option<[f64; 4]> {
93    let mut vectors = [[0.0_f64; 4]; 4];
94    for (i, row) in vectors.iter_mut().enumerate() {
95        row[i] = 1.0;
96    }
97    for _ in 0..64 {
98        let (mut p, mut q, mut largest) = (0, 1, 0.0_f64);
99        for (i, row) in a.iter().enumerate() {
100            for (j, value) in row.iter().enumerate().skip(i + 1) {
101                if value.abs() > largest {
102                    p = i;
103                    q = j;
104                    largest = value.abs();
105                }
106            }
107        }
108        let diagonal_scale = (0..4).fold(0.0_f64, |s, i| s.max(a[i][i].abs()));
109        if largest <= 16.0 * f64::EPSILON * diagonal_scale.max(1.0) {
110            let mut index = 0;
111            for i in 1..4 {
112                if a[i][i] > a[index][index] {
113                    index = i;
114                }
115            }
116            return Some(std::array::from_fn(|i| vectors[i][index]));
117        }
118        let tau = (a[q][q] - a[p][p]) / (2.0 * a[p][q]);
119        let tangent = 1.0_f64.copysign(tau) / (tau.abs() + tau.hypot(1.0));
120        let cosine = 1.0 / tangent.hypot(1.0);
121        let sine = tangent * cosine;
122        a[p][p] -= tangent * a[p][q];
123        a[q][q] += tangent * a[p][q];
124        a[p][q] = 0.0;
125        a[q][p] = 0.0;
126        for k in [0, 1, 2, 3] {
127            if k != p && k != q {
128                let (kp, kq) = (a[k][p], a[k][q]);
129                a[k][p] = cosine * kp - sine * kq;
130                a[p][k] = a[k][p];
131                a[k][q] = sine * kp + cosine * kq;
132                a[q][k] = a[k][q];
133            }
134        }
135        for row in &mut vectors {
136            let (vp, vq) = (row[p], row[q]);
137            row[p] = cosine * vp - sine * vq;
138            row[q] = sine * vp + cosine * vq;
139        }
140    }
141    None // Do not silently use an eigensystem that failed to converge.
142}
143
144#[cfg(test)]
145mod tests {
146    use super::estimate_rigid_transform;
147    use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
148
149    #[test]
150    fn recovers_pure_translation() {
151        let source: Vec<Vec3<f32>> = (0..20).map(|i| Vec3::new(i as f32 * 0.1, 0.0, 0.0)).collect();
152        let offset = Vec3::new(0.5, -0.2, 0.1);
153        let target: Vec<Vec3<f32>> = source.iter().map(|point| *point + offset).collect();
154
155        let transform = estimate_rigid_transform(&source, &target).unwrap();
156        assert!((transform.translation().x - offset.x).abs() < 1e-4);
157        assert!((transform.translation().y - offset.y).abs() < 1e-4);
158        assert!((transform.translation().z - offset.z).abs() < 1e-4);
159    }
160
161    #[test]
162    fn recovers_known_rigid_transform() {
163        let target: Vec<Vec3<f32>> = (0..4)
164            .flat_map(|x| {
165                (0..4).flat_map(move |y| {
166                    (0..3).map(move |z| Vec3::new(x as f32, y as f32, z as f32 * 0.2))
167                })
168            })
169            .collect();
170        let misalignment = Isometry3::new(
171            Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.2),
172            Vec3::new(0.3, -0.1, 0.05),
173        );
174        let source: Vec<Vec3<f32>> =
175            target.iter().map(|point| misalignment.transform_point(*point)).collect();
176
177        let estimated = estimate_rigid_transform(&source, &target).unwrap();
178        let composed = estimated.compose(misalignment);
179        let probe = Vec3::new(1.0, 2.0, 0.0);
180        let restored = composed.transform_point(probe);
181        assert!((restored.x - probe.x).abs() < 1e-3);
182        assert!((restored.y - probe.y).abs() < 1e-3);
183        assert!((restored.z - probe.z).abs() < 1e-3);
184    }
185
186    #[test]
187    fn rigid_estimation_is_scale_invariant_for_volume_and_plane() {
188        for scale in [1e-4_f32, 1.0, 1e4] {
189            for planar in [false, true] {
190                let source: Vec<_> = (0..4)
191                    .flat_map(|x| {
192                        (0..5).map(move |y| {
193                            Vec3::new(
194                                (x as f32 - 1.5) * scale,
195                                (y as f32 - 2.0) * scale,
196                                if planar { 0.0 } else { ((x * y % 3) as f32 - 1.0) * scale },
197                            )
198                        })
199                    })
200                    .collect();
201                let truth = Isometry3::new(
202                    Quat::from_axis_angle(Vec3::new(1.0, 2.0, 3.0), 0.45),
203                    Vec3::new(0.1 * scale, -0.2 * scale, 0.3 * scale),
204                );
205                let target: Vec<_> = source.iter().map(|p| truth.transform_point(*p)).collect();
206                let estimate = estimate_rigid_transform(&source, &target).unwrap();
207                for (p, expected) in source.iter().zip(&target) {
208                    let restored = estimate.transform_point(*p);
209                    let error = restored - *expected;
210                    assert!(
211                        error.x.hypot(error.y).hypot(error.z) / scale < 1e-5,
212                        "scale={scale}, planar={planar}, estimated={estimate:?}"
213                    );
214                }
215            }
216        }
217    }
218
219    #[test]
220    fn many_points_at_large_origin_retain_translation_precision() {
221        let source: Vec<_> = (0..20000)
222            .map(|i| {
223                Vec3::new(
224                    10000.0 + ((i % 11) as f32 - 5.0) * 0.25,
225                    -20000.0 + ((i % 7) as f32 - 3.0) * 0.5,
226                    30000.0 + ((i % 5) as f32 - 2.0) * 0.25,
227                )
228            })
229            .collect();
230        let shift = Vec3::new(0.125, -0.25, 0.5);
231        let target: Vec<_> = source.iter().map(|p| *p + shift).collect();
232        let estimated = estimate_rigid_transform(&source, &target).unwrap();
233        let t = estimated.translation();
234        assert!((t.x - shift.x).abs() < 1e-5);
235        assert!((t.y - shift.y).abs() < 1e-5);
236        assert!((t.z - shift.z).abs() < 1e-5);
237        for (source, target) in source.iter().zip(&target) {
238            assert_eq!(estimated.transform_point(*source), *target);
239        }
240    }
241
242    #[test]
243    fn rejects_nonfinite_and_unrepresentable_results() {
244        let source = [Vec3::new(0.0, 0.0, 0.0), Vec3::new(1.0, 0.0, 0.0), Vec3::new(0.0, 1.0, 0.0)];
245        for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
246            let mut invalid = source;
247            invalid[1].z = bad;
248            assert!(estimate_rigid_transform(&invalid, &source).is_none());
249            assert!(estimate_rigid_transform(&source, &invalid).is_none());
250        }
251        assert!(estimate_rigid_transform(&source[..2], &source[..2]).is_none());
252        assert!(estimate_rigid_transform(&source[..2], &source).is_none());
253        let negative = [Vec3::new(-f32::MAX, 0.0, 0.0); 3];
254        let positive = [Vec3::new(f32::MAX, 0.0, 0.0); 3];
255        assert!(estimate_rigid_transform(&negative, &positive).is_none());
256    }
257
258    #[test]
259    fn returns_proper_rotation_for_reflection_and_half_turn() {
260        let source = [
261            Vec3::new(1.0, 0.0, 0.0),
262            Vec3::new(0.0, 2.0, 0.0),
263            Vec3::new(0.0, 0.0, 3.0),
264            Vec3::new(-1.0, -2.0, -3.0),
265        ];
266        let reflected: Vec<_> = source.iter().map(|p| Vec3::new(-p.x, p.y, p.z)).collect();
267        let estimated = estimate_rigid_transform(&source, &reflected).unwrap();
268        let r = estimated.rotation().to_mat3().m;
269        let determinant = r[0][0] * (r[1][1] * r[2][2] - r[1][2] * r[2][1])
270            - r[0][1] * (r[1][0] * r[2][2] - r[1][2] * r[2][0])
271            + r[0][2] * (r[1][0] * r[2][1] - r[1][1] * r[2][0]);
272        assert!((determinant - 1.0).abs() < 1e-5);
273        let truth = Isometry3::new(
274            Quat::from_axis_angle(Vec3::new(1.0, 2.0, 3.0), std::f32::consts::PI),
275            Vec3::new(0.0, 0.0, 0.0),
276        );
277        let target: Vec<_> = source.iter().map(|p| truth.transform_point(*p)).collect();
278        let estimated = estimate_rigid_transform(&source, &target).unwrap();
279        for (p, expected) in source.iter().zip(target) {
280            let error = estimated.transform_point(*p) - expected;
281            assert!(error.x.hypot(error.y).hypot(error.z) < 1e-5);
282        }
283    }
284}
spatialrust_registration/kabsch.rs

spatialrust_registration/
kabsch.rs
