spatialrust_registration/
gicp.rs

1// Explicit row/column indexing reads more clearly than iterators for fixed 3x3
2// and 3x6 linear-algebra kernels.
3#![allow(clippy::needless_range_loop)]
4
5use spatialrust_core::{HasPositions3, PointCloud, SpatialError, SpatialResult};
6use spatialrust_math::{
7    solve_linear_system, symmetric_eigen3, Isometry3, LeastSquaresResult, Mat3, Quat,
8    TransformPoint, Vec3,
9};
10use spatialrust_search::{KdTree, NearestNeighborIndex};
11
12use crate::registration::{PointCloudRegistration, RegistrationResult};
13
14type M3 = [[f64; 3]; 3];
15
16/// Configuration for Generalized ICP (plane-to-plane).
17///
18/// GICP models each point's local surface as an anisotropic Gaussian and
19/// minimizes the Mahalanobis distance between correspondences, combining the
20/// robustness of point-to-plane ICP for both clouds. Local covariances are
21/// estimated from k-nearest neighbors and regularized toward planar disks.
22#[derive(Clone, Copy, Debug, PartialEq)]
23pub struct GicpConfig {
24    /// Maximum number of GICP iterations.
25    pub max_iterations: usize,
26    /// Maximum correspondence distance.
27    pub max_correspondence_distance: f32,
28    /// Number of neighbors used to estimate per-point covariance.
29    pub k_neighbors: usize,
30    /// Planar regularization: eigenvalue assigned along the surface normal.
31    pub epsilon: f64,
32    /// Stop when the transform update is smaller than this threshold.
33    pub transformation_epsilon: f64,
34    /// Stop when the fitness is smaller than this threshold.
35    pub fitness_epsilon: f64,
36    /// Minimum number of correspondences required per iteration.
37    pub min_correspondences: usize,
38    /// Initial transform guess mapping source into target frame.
39    pub initial_guess: Isometry3<f32>,
40    /// When set (and the `register-gicp-gpu` feature is enabled), estimate
41    /// per-point covariances on the GPU using a uniform grid neighbor search of
42    /// this radius instead of the CPU KD-tree. Falls back to CPU on GPU errors.
43    pub covariance_radius: Option<f32>,
44}
45
46impl Default for GicpConfig {
47    fn default() -> Self {
48        Self {
49            max_iterations: 50,
50            max_correspondence_distance: 1.0,
51            k_neighbors: 20,
52            epsilon: 1e-3,
53            transformation_epsilon: 1e-8,
54            fitness_epsilon: 1e-6,
55            min_correspondences: 6,
56            initial_guess: Isometry3::identity(),
57            covariance_radius: None,
58        }
59    }
60}
61
62impl GicpConfig {
63    /// Creates a config with the given correspondence distance.
64    #[must_use]
65    pub fn with_correspondence_distance(max_correspondence_distance: f32) -> Self {
66        Self { max_correspondence_distance, ..Self::default() }
67    }
68}
69
70/// Generalized ICP registration.
71#[derive(Clone, Copy, Debug, PartialEq)]
72pub struct GicpRegistration {
73    config: GicpConfig,
74}
75
76impl GicpRegistration {
77    /// Creates a GICP algorithm from config.
78    #[must_use]
79    pub const fn new(config: GicpConfig) -> Self {
80        Self { config }
81    }
82
83    /// Returns the config.
84    #[must_use]
85    pub const fn config(&self) -> GicpConfig {
86        self.config
87    }
88
89    /// Per-point plane-regularized covariances, on the GPU when a covariance
90    /// radius is configured and the `register-gicp-gpu` feature is enabled.
91    fn point_covariances(&self, x: &[f32], y: &[f32], z: &[f32]) -> Vec<M3> {
92        #[cfg(feature = "register-gicp-gpu")]
93        if let Some(radius) = self.config.covariance_radius {
94            if let Some(cov) = gpu_covariances(x, y, z, radius, self.config.epsilon) {
95                return cov;
96            }
97        }
98        let tree = KdTree::from_slices(x, y, z);
99        covariances(x, y, z, &tree, self.config.k_neighbors, self.config.epsilon)
100    }
101
102    /// Aligns `source` to `target` using Generalized ICP.
103    pub fn align_with_diagnostics(
104        &self,
105        source: &PointCloud,
106        target: &PointCloud,
107    ) -> SpatialResult<RegistrationResult> {
108        if source.is_empty() || target.is_empty() {
109            return Err(SpatialError::InvalidArgument(
110                "GICP requires non-empty source and target point clouds".to_owned(),
111            ));
112        }
113        if self.config.k_neighbors < 3 {
114            return Err(SpatialError::InvalidArgument(
115                "GICP needs k_neighbors >= 3 for covariance estimation".to_owned(),
116            ));
117        }
118
119        let (sx, sy, sz) = source.positions3()?;
120        let (tx, ty, tz) = target.positions3()?;
121        let target_tree = KdTree::from_slices(tx, ty, tz);
122        let source_cov = self.point_covariances(sx, sy, sz);
123        let target_cov = self.point_covariances(tx, ty, tz);
124
125        let max_distance_squared =
126            self.config.max_correspondence_distance * self.config.max_correspondence_distance;
127
128        let mut transform = self.config.initial_guess;
129        let mut transformed = vec![Vec3::new(0.0, 0.0, 0.0); source.len()];
130        apply_transform(&mut transformed, sx, sy, sz, transform);
131
132        let mut iterations = 0usize;
133        let mut converged = false;
134
135        for _ in 0..self.config.max_iterations {
136            iterations += 1;
137            let rot = mat3_to_f64(transform.rotation().to_mat3());
138            let rot_t = transpose(rot);
139
140            let mut hessian = [[0.0_f64; 6]; 6];
141            let mut gradient = [0.0_f64; 6];
142            let mut count = 0usize;
143
144            for (i, point) in transformed.iter().enumerate() {
145                let Some(neighbor) = target_tree.nearest_one(point.x, point.y, point.z) else {
146                    continue;
147                };
148                if neighbor.distance_squared > max_distance_squared {
149                    continue;
150                }
151                let j = neighbor.index;
152
153                // C = C_target + R C_source R^T, then M = C^-1.
154                let rotated_source = mat_mul(mat_mul(rot, source_cov[i]), rot_t);
155                let combined = mat_add(target_cov[j], rotated_source);
156                let Some(m) = inverse3(combined) else {
157                    continue;
158                };
159
160                let e = [
161                    f64::from(point.x) - f64::from(tx[j]),
162                    f64::from(point.y) - f64::from(ty[j]),
163                    f64::from(point.z) - f64::from(tz[j]),
164                ];
165                // Jacobian rows (3x6): [-skew(point) | I].
166                let jac = jacobian([f64::from(point.x), f64::from(point.y), f64::from(point.z)]);
167                let mj = mat3x6_premul(m, &jac);
168                let me = mat_vec(m, e);
169
170                for a in 0..6 {
171                    for k in 0..3 {
172                        gradient[a] += jac[k][a] * me[k];
173                        for b in 0..6 {
174                            hessian[a][b] += jac[k][a] * mj[k][b];
175                        }
176                    }
177                }
178                count += 1;
179            }
180
181            if count < self.config.min_correspondences {
182                return Err(SpatialError::InvalidArgument(format!(
183                    "GICP found only {} correspondences, minimum is {}",
184                    count, self.config.min_correspondences
185                )));
186            }
187
188            let a_rows: Vec<Vec<f64>> = hessian.iter().map(|row| row.to_vec()).collect();
189            let neg_g: Vec<f64> = gradient.iter().map(|value| -value).collect();
190            let LeastSquaresResult::Solved(solution) = solve_linear_system(a_rows, neg_g) else {
191                return Err(SpatialError::InvalidArgument(
192                    "GICP normal equations were singular".to_owned(),
193                ));
194            };
195
196            transform = delta_from_solution(&solution).compose(transform);
197            apply_transform(&mut transformed, sx, sy, sz, transform);
198
199            if update_magnitude(&solution) < self.config.transformation_epsilon {
200                converged = true;
201                break;
202            }
203            if euclidean_fitness(&transformed, tx, ty, tz, &target_tree, max_distance_squared)
204                < self.config.fitness_epsilon
205            {
206                converged = true;
207                break;
208            }
209        }
210
211        Ok(RegistrationResult {
212            transform,
213            fitness: euclidean_fitness(
214                &transformed,
215                tx,
216                ty,
217                tz,
218                &target_tree,
219                max_distance_squared,
220            ),
221            iterations,
222            converged,
223        })
224    }
225}
226
227impl PointCloudRegistration for GicpRegistration {
228    fn name(&self) -> &'static str {
229        "GicpRegistration"
230    }
231
232    fn align(&self, source: &PointCloud, target: &PointCloud) -> SpatialResult<RegistrationResult> {
233        self.align_with_diagnostics(source, target)
234    }
235}
236
237/// GPU path: per-point plane covariances via a uniform-grid neighbor search.
238/// Returns `None` on any GPU error so the caller falls back to the CPU path.
239#[cfg(feature = "register-gicp-gpu")]
240fn gpu_covariances(x: &[f32], y: &[f32], z: &[f32], radius: f32, epsilon: f64) -> Option<Vec<M3>> {
241    use spatialrust_gpu::{estimate_plane_covariances_grid_gpu, WgpuRuntime};
242    let runtime = WgpuRuntime::shared().ok()?;
243    let raw =
244        estimate_plane_covariances_grid_gpu(&runtime, x, y, z, radius, epsilon as f32).ok()?;
245    Some(
246        raw.into_iter()
247            .map(|[c00, c11, c22, c01, c02, c12]| {
248                [
249                    [c00 as f64, c01 as f64, c02 as f64],
250                    [c01 as f64, c11 as f64, c12 as f64],
251                    [c02 as f64, c12 as f64, c22 as f64],
252                ]
253            })
254            .collect(),
255    )
256}
257
258/// Computes per-point plane-regularized covariances from k-nearest neighbors.
259fn covariances(x: &[f32], y: &[f32], z: &[f32], tree: &KdTree, k: usize, epsilon: f64) -> Vec<M3> {
260    let len = x.len();
261    let mut out = Vec::with_capacity(len);
262    for index in 0..len {
263        let neighbors = tree.nearest_k(x[index], y[index], z[index], k);
264        let mut acc = spatialrust_math::CovarianceAccumulator3::new();
265        for neighbor in &neighbors {
266            acc.push(Vec3::new(x[neighbor.index], y[neighbor.index], z[neighbor.index]));
267        }
268        let cov = acc
269            .covariance()
270            .unwrap_or_else(|| Mat3::from_rows([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]));
271        out.push(regularize(cov, epsilon));
272    }
273    out
274}
275
276/// Replaces a covariance's eigenvalues with (epsilon, 1, 1) to model a planar disk.
277fn regularize(cov: Mat3<f64>, epsilon: f64) -> M3 {
278    let eigen = symmetric_eigen3(cov);
279    // Eigenvectors are columns; eigenvalues are ascending, so column 0 is the normal.
280    let v = eigen.eigenvectors.m;
281    let mut result = [[0.0_f64; 3]; 3];
282    for (col, scale) in [(0usize, epsilon), (1, 1.0), (2, 1.0)] {
283        let axis = [v[0][col], v[1][col], v[2][col]];
284        for r in 0..3 {
285            for c in 0..3 {
286                result[r][c] += scale * axis[r] * axis[c];
287            }
288        }
289    }
290    result
291}
292
293fn jacobian(p: [f64; 3]) -> [[f64; 6]; 3] {
294    // [-skew(p) | I3]
295    [
296        [0.0, p[2], -p[1], 1.0, 0.0, 0.0],
297        [-p[2], 0.0, p[0], 0.0, 1.0, 0.0],
298        [p[1], -p[0], 0.0, 0.0, 0.0, 1.0],
299    ]
300}
301
302fn apply_transform(
303    transformed: &mut [Vec3<f32>],
304    x: &[f32],
305    y: &[f32],
306    z: &[f32],
307    transform: Isometry3<f32>,
308) {
309    for (index, point) in transformed.iter_mut().enumerate() {
310        *point = transform.transform_point(Vec3::new(x[index], y[index], z[index]));
311    }
312}
313
314fn delta_from_solution(solution: &[f64]) -> Isometry3<f32> {
315    let (rx, ry, rz) = (solution[0], solution[1], solution[2]);
316    let angle = (rx * rx + ry * ry + rz * rz).sqrt();
317    let rotation = if angle > 1e-12 {
318        let axis = Vec3::new((rx / angle) as f32, (ry / angle) as f32, (rz / angle) as f32);
319        Quat::from_axis_angle(axis, angle as f32)
320    } else {
321        Quat::<f32>::identity()
322    };
323    Isometry3::new(rotation, Vec3::new(solution[3] as f32, solution[4] as f32, solution[5] as f32))
324}
325
326fn update_magnitude(solution: &[f64]) -> f64 {
327    solution.iter().map(|value| value * value).sum::<f64>().sqrt()
328}
329
330fn euclidean_fitness(
331    transformed: &[Vec3<f32>],
332    tx: &[f32],
333    ty: &[f32],
334    tz: &[f32],
335    tree: &KdTree,
336    max_distance_squared: f32,
337) -> f64 {
338    let mut sum = 0.0_f64;
339    let mut count = 0usize;
340    for point in transformed {
341        let Some(neighbor) = tree.nearest_one(point.x, point.y, point.z) else {
342            continue;
343        };
344        if neighbor.distance_squared > max_distance_squared {
345            continue;
346        }
347        let dx = f64::from(point.x - tx[neighbor.index]);
348        let dy = f64::from(point.y - ty[neighbor.index]);
349        let dz = f64::from(point.z - tz[neighbor.index]);
350        sum += dx * dx + dy * dy + dz * dz;
351        count += 1;
352    }
353    if count == 0 {
354        return f64::MAX;
355    }
356    sum / count as f64
357}
358
359fn mat3_to_f64(matrix: Mat3<f32>) -> M3 {
360    let mut out = [[0.0_f64; 3]; 3];
361    for r in 0..3 {
362        for c in 0..3 {
363            out[r][c] = f64::from(matrix.m[r][c]);
364        }
365    }
366    out
367}
368
369fn transpose(a: M3) -> M3 {
370    let mut out = [[0.0_f64; 3]; 3];
371    for r in 0..3 {
372        for c in 0..3 {
373            out[r][c] = a[c][r];
374        }
375    }
376    out
377}
378
379fn mat_mul(a: M3, b: M3) -> M3 {
380    let mut out = [[0.0_f64; 3]; 3];
381    for r in 0..3 {
382        for c in 0..3 {
383            for k in 0..3 {
384                out[r][c] += a[r][k] * b[k][c];
385            }
386        }
387    }
388    out
389}
390
391fn mat_add(a: M3, b: M3) -> M3 {
392    let mut out = [[0.0_f64; 3]; 3];
393    for r in 0..3 {
394        for c in 0..3 {
395            out[r][c] = a[r][c] + b[r][c];
396        }
397    }
398    out
399}
400
401fn mat_vec(a: M3, v: [f64; 3]) -> [f64; 3] {
402    [
403        a[0][0] * v[0] + a[0][1] * v[1] + a[0][2] * v[2],
404        a[1][0] * v[0] + a[1][1] * v[1] + a[1][2] * v[2],
405        a[2][0] * v[0] + a[2][1] * v[1] + a[2][2] * v[2],
406    ]
407}
408
409/// Computes `m * j` where `m` is 3x3 and `j` is 3x6.
410fn mat3x6_premul(m: M3, j: &[[f64; 6]; 3]) -> [[f64; 6]; 3] {
411    let mut out = [[0.0_f64; 6]; 3];
412    for r in 0..3 {
413        for c in 0..6 {
414            for k in 0..3 {
415                out[r][c] += m[r][k] * j[k][c];
416            }
417        }
418    }
419    out
420}
421
422fn inverse3(a: M3) -> Option<M3> {
423    let det = a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1])
424        - a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0])
425        + a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0]);
426    if det.abs() < 1e-18 {
427        return None;
428    }
429    let inv_det = 1.0 / det;
430    let mut out = [[0.0_f64; 3]; 3];
431    out[0][0] = (a[1][1] * a[2][2] - a[1][2] * a[2][1]) * inv_det;
432    out[0][1] = (a[0][2] * a[2][1] - a[0][1] * a[2][2]) * inv_det;
433    out[0][2] = (a[0][1] * a[1][2] - a[0][2] * a[1][1]) * inv_det;
434    out[1][0] = (a[1][2] * a[2][0] - a[1][0] * a[2][2]) * inv_det;
435    out[1][1] = (a[0][0] * a[2][2] - a[0][2] * a[2][0]) * inv_det;
436    out[1][2] = (a[0][2] * a[1][0] - a[0][0] * a[1][2]) * inv_det;
437    out[2][0] = (a[1][0] * a[2][1] - a[1][1] * a[2][0]) * inv_det;
438    out[2][1] = (a[0][1] * a[2][0] - a[0][0] * a[2][1]) * inv_det;
439    out[2][2] = (a[0][0] * a[1][1] - a[0][1] * a[1][0]) * inv_det;
440    Some(out)
441}
442
443#[cfg(test)]
444mod tests {
445    use super::{GicpConfig, GicpRegistration};
446    use crate::registration::PointCloudRegistration;
447    use crate::transform::transform_point_cloud;
448    use spatialrust_core::{PointCloudBuilder, StandardSchemas};
449    use spatialrust_math::{Isometry3, Quat, TransformPoint, Vec3};
450
451    /// Three perpendicular faces (a box corner) giving full 6-DoF constraint.
452    fn box_corner() -> spatialrust_core::PointCloud {
453        let mut builder = PointCloudBuilder::new(StandardSchemas::point_xyz());
454        for i in 0..10 {
455            for j in 0..10 {
456                let (a, b) = (i as f32 * 0.08, j as f32 * 0.08);
457                builder.push_point([a, b, 0.0]).unwrap();
458                builder.push_point([a, 0.0, b + 0.04]).unwrap();
459                builder.push_point([0.0, a + 0.04, b + 0.04]).unwrap();
460            }
461        }
462        builder.build().unwrap()
463    }
464
465    #[test]
466    fn aligns_rotated_and_translated_source() {
467        let target = box_corner();
468        let misalignment = Isometry3::new(
469            Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.06),
470            Vec3::new(0.015, -0.01, 0.012),
471        );
472        let source = transform_point_cloud(&target, misalignment).unwrap();
473
474        let gicp = GicpRegistration::new(GicpConfig {
475            max_correspondence_distance: 0.2,
476            max_iterations: 40,
477            k_neighbors: 12,
478            ..GicpConfig::default()
479        });
480        let result = gicp.align(&source, &target).unwrap();
481        assert!(result.fitness < 1e-4, "fitness too high: {}", result.fitness);
482
483        let composed = result.transform.compose(misalignment);
484        let probe = Vec3::new(0.3, 0.4, 0.2);
485        let restored = composed.transform_point(probe);
486        assert!((restored.x - probe.x).abs() < 1e-2);
487        assert!((restored.y - probe.y).abs() < 1e-2);
488        assert!((restored.z - probe.z).abs() < 1e-2);
489    }
490
491    #[test]
492    fn rejects_tiny_neighbor_count() {
493        let cloud = box_corner();
494        let gicp = GicpRegistration::new(GicpConfig { k_neighbors: 2, ..GicpConfig::default() });
495        assert!(gicp.align(&cloud, &cloud).is_err());
496    }
497
498    #[cfg(feature = "register-gicp-gpu")]
499    #[test]
500    fn aligns_with_gpu_covariances() {
501        // Skip gracefully when no GPU/software adapter is available.
502        if spatialrust_gpu::WgpuRuntime::shared().is_err() {
503            return;
504        }
505        let target = box_corner();
506        let misalignment = Isometry3::new(
507            Quat::from_axis_angle(Vec3::new(0.0, 0.0, 1.0), 0.05),
508            Vec3::new(0.015, -0.01, 0.012),
509        );
510        let source = transform_point_cloud(&target, misalignment).unwrap();
511
512        let gicp = GicpRegistration::new(GicpConfig {
513            max_correspondence_distance: 0.2,
514            max_iterations: 40,
515            covariance_radius: Some(0.16),
516            ..GicpConfig::default()
517        });
518        let result = gicp.align(&source, &target).unwrap();
519        let composed = result.transform.compose(misalignment);
520        let probe = Vec3::new(0.3, 0.4, 0.2);
521        let restored = composed.transform_point(probe);
522        assert!((restored.x - probe.x).abs() < 2e-2, "x off: {}", restored.x);
523        assert!((restored.y - probe.y).abs() < 2e-2, "y off: {}", restored.y);
524        assert!((restored.z - probe.z).abs() < 2e-2, "z off: {}", restored.z);
525    }
526}
spatialrust_registration/gicp.rs

spatialrust_registration/
gicp.rs
