spatialrust_registration/
icp_reduction.rs

1use super::{
2    transform_delta_magnitudes, validate_icp_transform, IcpDiagnostics, IcpIteration,
3    IcpRegistration, IcpStopReason,
4};
5use crate::{kabsch::estimate_rigid_transform_from_moments, RegistrationResult};
6use spatialrust_core::{
7    HasPositions3, PointCloud, SpatialError, SpatialResult, TransferDirection, TransferStats,
8};
9use spatialrust_gpu::{GpuIcpInput, GpuIcpTrimPrecision, GpuIcpTrimmedStatistics, WgpuRuntime};
10use spatialrust_search::KdTree;
11
12/// Explicit resident-transform/reduction ICP result and fixed-size transfer receipt.
13#[derive(Clone, Debug, PartialEq)]
14pub struct GpuIcpReductionDiagnostics {
15    /// Pose, convergence and per-update measurements from this GPU precision path.
16    pub diagnostics: IcpDiagnostics,
17    /// Source/target/tree upload plus fixed-size moments readbacks.
18    pub transfers: TransferStats,
19    /// Completed transformation/search/reduction evaluations (updates plus one).
20    pub evaluation_batches: usize,
21}
22
23impl IcpRegistration {
24    /// Aligns with explicitly uploaded source/target retained across iterations.
25    /// GPU transforms original source points, finds bounded correspondences and
26    /// reduces centered cross-covariance/centroids/fitness. Only 176 moment/status
27    /// bytes are downloaded per evaluation; pose solving and convergence stay CPU.
28    /// Reuses post-update moments for the next update and final fitness. No point
29    /// or correspondence readback, repeated payload upload or CPU fallback occurs.
30    /// Requires trim_fraction=1; `align_with_gpu_trimmed_reduction` supports
31    /// resident trimming. f32 transforms/reductions can change matches, pose, fitness and
32    /// stopping near numerical/geometry degeneracy compared with CPU f64 sums.
33    #[cfg(feature = "register-icp-gpu")]
34    pub fn align_with_gpu_reduction(
35        &self,
36        runtime: &WgpuRuntime,
37        source: &PointCloud,
38        target: &PointCloud,
39    ) -> SpatialResult<GpuIcpReductionDiagnostics> {
40        self.align_reduction_impl(runtime, source, target, false, GpuIcpTrimPrecision::DoubleSingle)
41    }
42
43    /// Aligns using GPU stable distance/source-index trimming and resident
44    /// moments. floor(trim_fraction * gated_count) pairs contribute to each
45    /// pose solve. Post-update fitness/stopping use ALL gated pairs. Only 176
46    /// bytes cross to CPU per evaluation; source/target/search/selection remain
47    /// resident. CPU target-tree construction, Horn solve and convergence remain.
48    /// A fraction of one uses the canonical untrimmed GPU evaluation. f32 device
49    /// matching/moments can differ from CPU near gates, ties or degeneracy.
50    #[cfg(feature = "register-icp-gpu")]
51    pub fn align_with_gpu_trimmed_reduction(
52        &self,
53        runtime: &WgpuRuntime,
54        source: &PointCloud,
55        target: &PointCloud,
56    ) -> SpatialResult<GpuIcpReductionDiagnostics> {
57        self.align_with_gpu_trimmed_reduction_with_precision(
58            runtime,
59            source,
60            target,
61            GpuIcpTrimPrecision::DoubleSingle,
62        )
63    }
64
65    /// Aligns with an explicit retained-moment precision choice. DoubleSingle
66    /// (the default resident trimmed route) carries normalized quotient/product
67    /// residuals in two f32 words. CompensatedF32 retains the earlier lower-cost
68    /// arithmetic. Transform/search precision, all-gated fitness, CPU solving
69    /// and explicit fixed-size transfer receipts are the same for both choices.
70    #[cfg(feature = "register-icp-gpu")]
71    pub fn align_with_gpu_trimmed_reduction_with_precision(
72        &self,
73        runtime: &WgpuRuntime,
74        source: &PointCloud,
75        target: &PointCloud,
76        precision: GpuIcpTrimPrecision,
77    ) -> SpatialResult<GpuIcpReductionDiagnostics> {
78        self.align_reduction_impl(runtime, source, target, true, precision)
79    }
80
81    fn align_reduction_impl(
82        &self,
83        runtime: &WgpuRuntime,
84        source: &PointCloud,
85        target: &PointCloud,
86        trimmed: bool,
87        precision: GpuIcpTrimPrecision,
88    ) -> SpatialResult<GpuIcpReductionDiagnostics> {
89        self.validate_inputs(source, target)?;
90        if !trimmed && self.trim_fraction != 1.0 {
91            return Err(SpatialError::InvalidArgument("GPU ICP reduction requires trim_fraction=1; use the explicit correspondence API for trimming".to_owned()));
92        }
93        let (sx, sy, sz) = source.positions3()?;
94        let (tx, ty, tz) = target.positions3()?;
95        let source_positions: Vec<_> = (0..source.len()).map(|i| [sx[i], sy[i], sz[i]]).collect();
96        let tree = KdTree::from_slices(tx, ty, tz);
97        let input = GpuIcpInput::upload(runtime, &source_positions, &tree)?;
98        let mut transfers = input.upload_transfers();
99        let mut evaluation_batches = 0;
100        let mut evaluate = |pose| -> SpatialResult<GpuIcpTrimmedStatistics> {
101            let statistics = input
102                .trimmed_moments_resident_with_precision(
103                    runtime,
104                    pose,
105                    self.config.max_correspondence_distance,
106                    self.trim_fraction,
107                    precision,
108                )?
109                .download(runtime)?;
110            transfers.record(
111                TransferDirection::DeviceToHost,
112                statistics.moments.transfers.device_to_host_bytes(),
113            );
114            evaluation_batches += 1;
115            Ok(statistics)
116        };
117        let mut transform = self.config.initial_guess;
118        let mut moments = evaluate(transform)?;
119        let mut history = Vec::new();
120        let mut previous_fitness = None;
121        let mut converged = false;
122        let mut stop_reason = IcpStopReason::IterationLimit;
123        for iteration in 1..=self.config.max_iterations {
124            let evaluated = moments.evaluated_correspondences as usize;
125            let count = moments.moments.correspondences as usize;
126            if evaluated < self.config.min_correspondences {
127                return Err(SpatialError::InvalidArgument(format!(
128                    "ICP found only {evaluated} correspondences, minimum is {}",
129                    self.config.min_correspondences
130                )));
131            }
132            if count < self.config.min_correspondences {
133                return Err(SpatialError::InvalidArgument(format!(
134                    "trim_fraction retains only {count} of {evaluated} correspondences, minimum is {}",
135                    self.config.min_correspondences
136                )));
137            }
138            let delta = estimate_rigid_transform_from_moments(
139                count,
140                moments.moments.source_centroid,
141                moments.moments.target_centroid,
142                moments.moments.cross_covariance,
143            )
144            .ok_or_else(|| {
145                SpatialError::InvalidArgument("ICP failed to estimate a rigid transform".to_owned())
146            })?;
147            validate_icp_transform(delta)?;
148            transform = delta.compose(transform);
149            validate_icp_transform(transform)?;
150            moments = evaluate(transform)?;
151            let fitness = fitness(&moments);
152            let (translation_delta, rotation_delta_radians) = transform_delta_magnitudes(delta);
153            history.push(IcpIteration {
154                iteration,
155                correspondences: count,
156                evaluated_correspondences: moments.evaluated_correspondences as usize,
157                fitness,
158                fitness_change: previous_fitness.map(|previous| previous - fitness),
159                translation_delta,
160                rotation_delta_radians,
161            });
162            previous_fitness = Some(fitness);
163            if fitness < self.criteria.fitness_epsilon {
164                converged = true;
165                stop_reason = IcpStopReason::FitnessThreshold;
166                break;
167            }
168            if translation_delta < self.criteria.translation_epsilon
169                && rotation_delta_radians < self.criteria.rotation_epsilon
170            {
171                converged = true;
172                stop_reason = IcpStopReason::TransformThreshold;
173                break;
174            }
175        }
176        Ok(GpuIcpReductionDiagnostics {
177            diagnostics: IcpDiagnostics {
178                result: RegistrationResult {
179                    transform,
180                    fitness: fitness(&moments),
181                    iterations: history.len(),
182                    converged,
183                },
184                history,
185                stop_reason,
186            },
187            transfers,
188            evaluation_batches,
189        })
190    }
191}
192
193fn fitness(statistics: &GpuIcpTrimmedStatistics) -> f64 {
194    if statistics.evaluated_correspondences == 0 {
195        f64::MAX
196    } else {
197        statistics.moments.squared_error_sum / f64::from(statistics.evaluated_correspondences)
198    }
199}
spatialrust_registration/icp_reduction.rs

spatialrust_registration/
icp_reduction.rs
