spatialrust_registration/
point_to_plane_reduction.rs

1use super::{delta_from_solution, update_magnitude, PointToPlaneIcp};
2use crate::RegistrationResult;
3use spatialrust_core::{
4    HasNormals3, HasPositions3, PointCloud, SpatialError, SpatialResult, TransferDirection,
5    TransferStats,
6};
7use spatialrust_gpu::{GpuPointToPlaneEquations, GpuPointToPlaneInput, WgpuRuntime};
8use spatialrust_math::{solve_linear_system, Isometry3, LeastSquaresResult};
9use spatialrust_search::KdTree;
10
11/// One completed update in the explicit resident point-to-plane ICP precision path.
12#[derive(Clone, Debug, PartialEq)]
13pub struct GpuPointToPlaneIteration {
14    /// One-based completed update number.
15    pub iteration: usize,
16    /// Gated pairs used by the normal-equation solve.
17    pub correspondences: u32,
18    /// Gated pairs used for post-update raw-normal plane fitness.
19    pub evaluated_correspondences: u32,
20    /// Post-update raw-normal plane mean squared distance, f64::MAX when unmatched.
21    pub fitness: f64,
22    /// Norm of the world rotation/translation linearized solution, as in CPU point-to-plane ICP.
23    pub update_magnitude: f64,
24}
25
26/// Explicit point-to-plane ICP resident reduction result, trace and full transfer receipt.
27#[derive(Clone, Debug, PartialEq)]
28pub struct GpuPointToPlaneReductionDiagnostics {
29    /// Estimated pose, raw-normal plane fitness and stopping result.
30    pub result: RegistrationResult,
31    /// Completed update rows from this device precision path.
32    pub history: Vec<GpuPointToPlaneIteration>,
33    /// Position/tree/normal upload and fixed-size equations readback.
34    pub transfers: TransferStats,
35    /// Completed resident evaluations (updates plus one).
36    pub evaluation_batches: usize,
37}
38
39fn invalid() -> SpatialError {
40    SpatialError::InvalidArgument("GPU point-to-plane ICP reduction requires nonempty clouds, at least one iteration, minimum six correspondences, a finite positive distance square, finite nonnegative stopping thresholds and a finite unit initial pose".to_owned())
41}
42fn valid_pose(pose: Isometry3<f32>) -> bool {
43    let q = pose.rotation();
44    let t = pose.translation();
45    let norm = f64::from(q.x).hypot(f64::from(q.y)).hypot(f64::from(q.z)).hypot(f64::from(q.w));
46    t.x.is_finite()
47        && t.y.is_finite()
48        && t.z.is_finite()
49        && norm.is_finite()
50        && (norm - 1.0).abs() <= 1e-4
51}
52fn record(total: &mut TransferStats, part: TransferStats) {
53    total.record(TransferDirection::HostToDevice, part.host_to_device_bytes());
54    total.record(TransferDirection::DeviceToHost, part.device_to_host_bytes());
55}
56fn fitness(equations: &GpuPointToPlaneEquations) -> f64 {
57    if equations.fitness_correspondences == 0 {
58        f64::MAX
59    } else {
60        equations.squared_error_sum / f64::from(equations.fitness_correspondences)
61    }
62}
63
64impl PointToPlaneIcp {
65    /// Retains source/target/tree/raw normals on GPU for transforms, nearest
66    /// matching, plane Hessian/gradient and raw-normal plane fitness reductions.
67    /// CPU solves normalized 6x6 equations and converts the linearized step to
68    /// the existing world update/stopping convention. Each evaluation explicitly
69    /// downloads 256 bytes; post-update equations feed the next solve/final fitness.
70    /// No CPU fallback occurs. Raw normals keep relative length weights and need
71    /// not be unit vectors. f32 arithmetic may change matches, pose and stopping.
72    pub fn align_with_gpu_reduction(
73        &self,
74        runtime: &WgpuRuntime,
75        source: &PointCloud,
76        target: &PointCloud,
77    ) -> SpatialResult<GpuPointToPlaneReductionDiagnostics> {
78        if source.is_empty() || target.is_empty() {
79            return Err(invalid());
80        }
81        let c = self.config;
82        let distance_squared = c.max_correspondence_distance * c.max_correspondence_distance;
83        if c.max_iterations == 0
84            || c.min_correspondences < 6
85            || !c.max_correspondence_distance.is_finite()
86            || c.max_correspondence_distance <= 0.0
87            || !distance_squared.is_finite()
88            || distance_squared == 0.0
89            || !c.transformation_epsilon.is_finite()
90            || c.transformation_epsilon < 0.0
91            || !c.fitness_epsilon.is_finite()
92            || c.fitness_epsilon < 0.0
93            || !valid_pose(c.initial_guess)
94        {
95            return Err(invalid());
96        }
97        let (sx, sy, sz) = source.positions3()?;
98        let (tx, ty, tz) = target.positions3()?;
99        if sx.iter().chain(sy).chain(sz).chain(tx).chain(ty).chain(tz).any(|v| !v.is_finite()) {
100            return Err(SpatialError::InvalidArgument(
101                "GPU point-to-plane ICP requires finite XYZ".to_owned(),
102            ));
103        }
104        let tree = KdTree::from_slices(tx, ty, tz);
105        let mut transfers = TransferStats::default();
106        let (nx, ny, nz) = target.normals3()?;
107        let positions: Vec<_> = (0..source.len()).map(|i| [sx[i], sy[i], sz[i]]).collect();
108        let normals: Vec<_> = (0..target.len()).map(|i| [nx[i], ny[i], nz[i]]).collect();
109        let input = GpuPointToPlaneInput::upload(runtime, &positions, &tree, &normals)?;
110        record(&mut transfers, input.upload_transfers());
111        let mut evaluation_batches = 0;
112        let mut evaluate = |pose| -> SpatialResult<GpuPointToPlaneEquations> {
113            let equations = input
114                .equations_resident(runtime, pose, c.max_correspondence_distance)?
115                .download(runtime)?;
116            record(&mut transfers, equations.transfers);
117            evaluation_batches += 1;
118            Ok(equations)
119        };
120        let mut transform = c.initial_guess;
121        let mut equations = evaluate(transform)?;
122        let mut converged = false;
123        let mut history = Vec::new();
124        for iteration in 1..=c.max_iterations {
125            let count = equations.correspondences;
126            if (count as usize) < c.min_correspondences {
127                return Err(SpatialError::InvalidArgument(format!(
128                    "point-to-plane ICP found only {count} correspondences, minimum is {}",
129                    c.min_correspondences
130                )));
131            }
132            let rows = equations.hessian.iter().map(|row| row.to_vec()).collect();
133            let gradient = equations.gradient.iter().map(|g| -g).collect();
134            let LeastSquaresResult::Solved(solution) = solve_linear_system(rows, gradient) else {
135                return Err(SpatialError::InvalidArgument(
136                    "point-to-plane ICP normal equations were singular".to_owned(),
137                ));
138            };
139            let solution: [f64; 6] = solution.try_into().map_err(|_| invalid())?;
140            let world = equations.world_step(solution)?;
141            transform = delta_from_solution(&world).compose(transform);
142            if !valid_pose(transform) {
143                return Err(invalid());
144            }
145            equations = evaluate(transform)?;
146            let fitness = fitness(&equations);
147            let update_magnitude = update_magnitude(&world);
148            history.push(GpuPointToPlaneIteration {
149                iteration,
150                correspondences: count,
151                evaluated_correspondences: equations.fitness_correspondences,
152                fitness,
153                update_magnitude,
154            });
155            if update_magnitude < c.transformation_epsilon || fitness < c.fitness_epsilon {
156                converged = true;
157                break;
158            }
159        }
160        Ok(GpuPointToPlaneReductionDiagnostics {
161            result: RegistrationResult {
162                transform,
163                fitness: fitness(&equations),
164                iterations: history.len(),
165                converged,
166            },
167            history,
168            transfers,
169            evaluation_batches,
170        })
171    }
172}
spatialrust_registration/point_to_plane_reduction.rs

spatialrust_registration/
point_to_plane_reduction.rs
