registration

triwarp.registration

Point cloud registration: Procrustes analysis and ICP.

procrustes

procrustes(
    a: array[vec3],
    b: array[vec3],
    weights: array[float32] | None = None,
    reflection: bool = True,
    translation: bool = True,
    scale: bool = True,
    return_cost: Literal[True] = True,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]

procrustes(
    a: array[vec3],
    b: array[vec3],
    weights: array[float32] | None = None,
    reflection: bool = True,
    translation: bool = True,
    scale: bool = True,
    return_cost: Literal[False] = False,
) -> wp.array[wp.mat44]

procrustes(
    a: array[vec3],
    b: array[vec3],
    weights: array[float32] | None = None,
    reflection: bool = True,
    translation: bool = True,
    scale: bool = True,
    return_cost: bool = True,
) -> (
    tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]
    | wp.array[wp.mat44]
)

procrustes(
    a: array[vec3],
    b: array[vec3],
    weights: array[float32] | None = None,
    reflection: bool = True,
    translation: bool = True,
    scale: bool = True,
    return_cost: bool = True,
) -> (
    tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]
    | wp.array[wp.mat44]
)

Find the optimal transform mapping point cloud a onto b.

Warp port of trimesh.registration.procrustes. Accepts and returns GPU-resident arrays; no NumPy in the hot path.

Parameters:

Name	Type	Description	Default
a	array[vec3]	Corresponding point clouds, shape (n,), dtype wp.vec3.	required
b	array[vec3]	Corresponding point clouds, shape (n,), dtype wp.vec3.	required
weights	array[float32] | None	Per-point weights, shape (n,), dtype wp.float32. Uniform weights are assumed when None.	None
reflection	bool	Allow reflections in the rotation component.	True
translation	bool	Allow translation.	True
scale	bool	Allow uniform scaling.	True
return_cost	bool	When True (default) also return the transformed points and cost.	True

Returns:

Name	Type	Description
matrix	tuple[array[mat44], array[vec3], float] | array[mat44]	wp.array[wp.mat44] of shape (1,) encoding the transform.
transformed	tuple[array[mat44], array[vec3], float] | array[mat44]	wp.array[wp.vec3] of shape (n,) — image of a under the transform. Only returned when return_cost=True.
cost	tuple[array[mat44], array[vec3], float] | array[mat44]	Weighted sum of squared distances between transformed and b. Only returned when return_cost=True.

Raises:

Type	Description
RuntimeError	If a, b and weights are not all on one device.
ValueError	If b (or a non-empty weights) does not have the same length as a, or a non-empty weights sums to zero.

Notes

Latency-bound at every size that matters: the fit is two kernel launches (four with return_cost), one allocation for the packed moment accumulator and one readback, so the cost is nearly flat in n. Callers in a loop should use the workspace form — see icp, which allocates once outside its iteration.

The returned matrix is column-vector, which is the transpose of pytorch3d's. pytorch3d.ops.corresponding_points_alignment solves the row-vector form s X R + T = Y, so its R is this matrix's linear block divided by the scale and transposed.

Source code in triwarp/registration.py

def procrustes(
    a: wp.array[wp.vec3],
    b: wp.array[wp.vec3],
    weights: wp.array[wp.float32] | None = None,
    reflection: bool = True,
    translation: bool = True,
    scale: bool = True,
    return_cost: bool = True,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float] | wp.array[wp.mat44]:
    """
    Find the optimal transform mapping point cloud *a* onto *b*.

    Warp port of [`trimesh.registration.procrustes`][]. Accepts and returns
    GPU-resident arrays; no NumPy in the hot path.

    Parameters
    ----------
    a, b:
        Corresponding point clouds, shape ``(n,)``, dtype ``wp.vec3``.
    weights:
        Per-point weights, shape ``(n,)``, dtype ``wp.float32``.
        Uniform weights are assumed when ``None``.
    reflection:
        Allow reflections in the rotation component.
    translation:
        Allow translation.
    scale:
        Allow uniform scaling.
    return_cost:
        When ``True`` (default) also return the transformed points and cost.

    Returns
    -------
    matrix:
        ``wp.array[wp.mat44]`` of shape ``(1,)`` encoding the transform.
    transformed:
        ``wp.array[wp.vec3]`` of shape ``(n,)`` — image of *a* under the transform.
        Only returned when ``return_cost=True``.
    cost:
        Weighted sum of squared distances between *transformed* and *b*.
        Only returned when ``return_cost=True``.

    Raises
    ------
    RuntimeError
        If ``a``, ``b`` and ``weights`` are not all on one device.
    ValueError
        If ``b`` (or a non-empty ``weights``) does not have the same length as ``a``, or a
        non-empty ``weights`` sums to zero.

    Notes
    -----
    Latency-bound at every size that matters: the fit is two kernel launches (four with
    ``return_cost``), one allocation for the packed moment accumulator and one readback, so the
    cost is nearly flat in ``n``. Callers in a loop should use the workspace form — see
    [`icp`][triwarp.registration.icp], which allocates once outside its iteration.

    **The returned matrix is column-vector, which is the transpose of pytorch3d's.**
    ``pytorch3d.ops.corresponding_points_alignment`` solves the row-vector form ``s X R + T = Y``,
    so its ``R`` is this matrix's linear block divided by the scale and *transposed*.
    """
    require_same_device(a=a, b=b, weights=weights)
    n = int(a.shape[0])
    device = a.device
    if int(b.shape[0]) != n:
        raise ValueError(f"a and b must have the same length, got {n} and {b.shape[0]}")
    # A zero-length weights array is the kernels' own "uniform weights" sentinel (see
    # ``_uniform_weights``), not a caller mistake -- only a non-empty mismatch is a real error.
    if weights is not None and int(weights.shape[0]) not in (0, n):
        raise ValueError(f"weights must have the same length as a, got {weights.shape[0]} and {n}")

    if n == 0:
        matrix = _identity_mat44(device)
        if not return_cost:
            return matrix
        return matrix, wp.clone(a), 0.0

    # An all-zero non-empty ``weights`` divides by zero inside the kernel (the accumulated weight
    # sum is the denominator of every centroid and, with ``scale=True``, of the singular values
    # too), silently returning a matrix of NaN with no exception. ``icp`` avoids this by checking
    # the same sum itself and breaking its loop before ever reaching the shared machinery below --
    # this entry point has no loop to break, so it raises instead.
    if weights is not None and int(weights.shape[0]) == n and float(tw.reduce.sum(weights)) == 0.0:
        raise ValueError("weights sum to zero: no point carries any weight")

    workspace = _procrustes_workspace(n, device, return_cost=return_cost)
    return _procrustes_into(a, b, weights, reflection, translation, scale, return_cost, workspace)

icp

icp(
    a: array[vec3],
    target_vertices: array[vec3],
    target_faces: array[int32] | None = None,
    *,
    initial: array[mat44] | mat44 | None = None,
    max_iterations: int = 20,
    threshold: float = 1e-05,
    max_distance: float | None = None,
    reflection: bool = False,
    translation: bool = True,
    scale: bool = False,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]

Point-to-point iterative closest point registration.

Aligns the source cloud a onto a target mesh or point cloud by repeatedly matching each source point to its nearest target counterpart and solving the resulting correspondence with procrustes. Following the pytorch3d formulation, every iteration re-aligns the original a to the current correspondences, so the returned transform is a single map from a to the target (no incremental-composition drift). Only a rough initial alignment is required for a good result.

Warp port combining trimesh.registration.icp and pytorch3d.ops.iterative_closest_point. Correspondence search runs on the GPU via closest_point_on_mesh (mesh target) or query_nearest (point-cloud target); the alignment is the GPU tiled SVD of procrustes.

Its R is the transpose of the linear block returned here, because it solves the row-vector form s X R + T = Y; the converged transform is what is comparable, not the iteration count, since relative_rmse_thr is its own stopping rule.

Parameters:

Name	Type	Description	Default
a	array[vec3]	Source point cloud, shape (n,), dtype wp.vec3.	required
target_vertices	array[vec3]	Target vertex positions, shape (m,), dtype wp.vec3. Used as the target point cloud when target_faces is None.	required
target_faces	array[int32] | None	Flat (f * 3,) triangle index buffer. When given, the internally built wp.Mesh aliases target_vertices and target_faces rather than copying them; do not mutate them for the duration of the call. When given (and non-empty), correspondences are the closest points on the triangle surface; otherwise the target is the point cloud target_vertices.	None
initial	array[mat44] | mat44 | None	Initial transform seeding the first correspondence search, as a (1,) wp.mat44 array or a scalar wp.mat44. Identity when None.	None
max_iterations	int	Maximum number of ICP iterations.	20
threshold	float	Stop when the cost decreases by less than this between iterations.	1e-05
max_distance	float | None	Reject correspondences farther than this (weight 0, excluded from the fit). No rejection when None.	None
reflection	bool	Forwarded to procrustes. Defaults are rigid (reflection=False, scale=False); set scale=True for a similarity transform.	False
translation	bool	Forwarded to procrustes. Defaults are rigid (reflection=False, scale=False); set scale=True for a similarity transform.	False
scale	bool	Forwarded to procrustes. Defaults are rigid (reflection=False, scale=False); set scale=True for a similarity transform.	False

Returns:

Type	Description
matrix	(1,) wp.mat44 transform mapping a onto the target.
transformed	(n,) wp.vec3 image of a under matrix.
cost	Weighted mean squared correspondence distance at the final iteration.

Raises:

Type	Description
RuntimeError	If a, target_vertices, target_faces and initial are not all on one device.

See Also

procrustes icp_point_to_plane trimesh.registration.icp

Source code in triwarp/registration.py

def icp(
    a: wp.array[wp.vec3],
    target_vertices: wp.array[wp.vec3],
    target_faces: wp.array[wp.int32] | None = None,
    *,
    initial: wp.array[wp.mat44] | wp.mat44 | None = None,
    max_iterations: int = 20,
    threshold: float = 1e-5,
    max_distance: float | None = None,
    reflection: bool = False,
    translation: bool = True,
    scale: bool = False,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]:
    """
    Point-to-point iterative closest point registration.

    Aligns the source cloud *a* onto a target mesh or point cloud by repeatedly
    matching each source point to its nearest target counterpart and solving the
    resulting correspondence with [`procrustes`][triwarp.registration.procrustes].
    Following the ``pytorch3d`` formulation, every iteration re-aligns the
    *original* *a* to the current correspondences, so the returned transform is a
    single map from *a* to the target (no incremental-composition drift). Only a
    rough initial alignment is required for a good result.

    Warp port combining [`trimesh.registration.icp`][] and
    ``pytorch3d.ops.iterative_closest_point``. Correspondence search runs on the
    GPU via [`closest_point_on_mesh`][triwarp.proximity.closest_point_on_mesh]
    (mesh target) or [`query_nearest`][triwarp.neighbors.query_nearest]
    (point-cloud target); the alignment is the GPU tiled SVD of
    [`procrustes`][triwarp.registration.procrustes].

    Its ``R`` is the **transpose** of the linear block returned here, because it solves the
    row-vector form ``s X R + T = Y``; the *converged transform* is what is comparable, not the
    iteration count, since ``relative_rmse_thr`` is its own stopping rule.

    Parameters
    ----------
    a
        Source point cloud, shape ``(n,)``, dtype ``wp.vec3``.
    target_vertices
        Target vertex positions, shape ``(m,)``, dtype ``wp.vec3``. Used as the
        target point cloud when ``target_faces`` is ``None``.
    target_faces
        Flat ``(f * 3,)`` triangle index buffer. When given, the internally built ``wp.Mesh``
        aliases ``target_vertices`` and ``target_faces`` rather than copying them; do not mutate
        them for the duration of the call. When given (and non-empty),
        correspondences are the closest points on the triangle surface; otherwise
        the target is the point cloud ``target_vertices``.
    initial
        Initial transform seeding the first correspondence search, as a ``(1,)``
        ``wp.mat44`` array or a scalar ``wp.mat44``. Identity when ``None``.
    max_iterations
        Maximum number of ICP iterations.
    threshold
        Stop when the cost decreases by less than this between iterations.
    max_distance
        Reject correspondences farther than this (weight ``0``, excluded from the
        fit). No rejection when ``None``.
    reflection, translation, scale
        Forwarded to [`procrustes`][triwarp.registration.procrustes]. Defaults are
        rigid (`reflection=False, scale=False`); set ``scale=True`` for a
        similarity transform.

    Returns
    -------
    matrix
        ``(1,)`` ``wp.mat44`` transform mapping *a* onto the target.
    transformed
        ``(n,)`` ``wp.vec3`` image of *a* under *matrix*.
    cost
        Weighted mean squared correspondence distance at the final iteration.

    Raises
    ------
    RuntimeError
        If ``a``, ``target_vertices``, ``target_faces`` and ``initial`` are not all on one device.

    See Also
    --------
    [`procrustes`][triwarp.registration.procrustes]
    [`icp_point_to_plane`][triwarp.registration.icp_point_to_plane]
    [`trimesh.registration.icp`][]
    """
    require_same_device(
        a=a, target_vertices=target_vertices, target_faces=target_faces, initial=initial
    )
    device = a.device
    n = int(a.shape[0])

    if n == 0 or int(target_vertices.shape[0]) == 0:
        return _identity_mat44(device), wp.clone(a), math.inf

    initial_matrix, current = _seed_transform(a, initial, device)
    total = initial_matrix
    transformed = wp.clone(current)
    cost = math.inf

    mesh, query_max, target_index = _resolve_icp_target(
        target_vertices, target_faces, current, max_distance, "icp"
    )

    # Correspondence and weight buffers are allocated once and refilled every iteration, and so is
    # the Procrustes workspace — the fit is latency-bound, so its ~10 per-call allocations would
    # otherwise dominate an iteration that is already down to a handful of launches.
    closest = wp.empty(n, dtype=wp.vec3, device=device)
    distance_mesh = wp.empty(n, dtype=wp.float32, device=device)
    triangle_id_mesh = wp.empty(n, dtype=wp.int32, device=device)
    weights: wp.array[wp.float32] | None = (
        wp.empty(n, dtype=wp.float32, device=device) if max_distance is not None else None
    )
    workspace = _procrustes_workspace(n, device, return_cost=True)

    old_cost = math.inf
    for _ in range(max_iterations):
        distance, triangle_id = _correspondences(
            mesh,
            target_vertices,
            target_index,
            current,
            query_max,
            closest,
            distance_mesh,
            triangle_id_mesh,
        )

        if max_distance is not None and weights is not None:
            wp.map(
                kernel_registration.distance_threshold_weight,
                distance,
                triangle_id,
                wp.float32(max_distance),
                out=weights,
            )
            # Readback: the loop's own early-exit test, no cheaper than the sum it reads.
            if float(tw.reduce.sum(weights)) == 0.0:
                break

        total, transformed, cost = cast(
            tuple[wp.array[wp.mat44], wp.array[wp.vec3], float],
            _procrustes_into(a, closest, weights, reflection, translation, scale, True, workspace),
        )
        current = transformed
        if old_cost - cost < threshold:
            break
        old_cost = cost

    return total, transformed, cost

icp_point_to_plane

icp_point_to_plane(
    a: array[vec3],
    target_vertices: array[vec3],
    target_faces: array[int32] | None = None,
    *,
    target_normals: array[vec3] | None = None,
    initial: array[mat44] | mat44 | None = None,
    max_iterations: int = 20,
    threshold: float = 1e-06,
    max_distance: float | None = None,
    robust_kernel: Literal[
        "none", "huber", "tukey"
    ] = "none",
    robust_scale: float | None = None,
    damping: float = 1e-06,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]

Point-to-plane iterative closest point registration with optional robust loss.

Minimizes the squared distance from each transformed source point to the tangent plane of its target correspondence, which converges faster than the point-to-point metric on smooth surfaces. Each iteration linearizes the residual, assembles the 6x6 Gauss-Newton system on the GPU, solves it, and composes the incremental rigid step into the running transform.

Warp port of the libigl iterative_closest_point (rigid, point-to-plane) and Open3D's robust TransformationEstimationPointToPlane. An optional Huber or Tukey M-estimator down-weights outlier correspondences.

Parameters:

Name	Type	Description	Default
a	array[vec3]	Source point cloud, shape (n,), dtype wp.vec3.	required
target_vertices	array[vec3]	Target vertex positions, shape (m,), dtype wp.vec3.	required
target_faces	array[int32] | None	Flat (f * 3,) triangle index buffer. When given, target normals are the closest triangles' face normals; otherwise the target is the point cloud target_vertices and target_normals is required. When given, the internally built wp.Mesh aliases target_vertices and target_faces rather than copying them; do not mutate them for the duration of the call.	None
target_normals	array[vec3] | None	Per-vertex unit normals for a point-cloud target, shape (m,). Required (and only used) when target_faces is None. Estimate them with estimate_normals if absent.	None
initial	array[mat44] | mat44 | None	Initial transform, as a (1,) wp.mat44 array or scalar wp.mat44. Identity when None.	None
max_iterations	int	Maximum number of ICP iterations.	20
threshold	float	Stop when the cost decreases by less than this between iterations.	1e-06
max_distance	float | None	Reject correspondences farther than this. No rejection when None.	None
robust_kernel	Literal['none', 'huber', 'tukey']	M-estimator applied to residuals: "none", "huber", or "tukey".	'none'
robust_scale	float | None	Tuning constant of the robust kernel (Huber k / Tukey c). When None and a robust kernel is active, derived from the median absolute deviation of the first iteration's residuals.	None
damping	float	Relative Levenberg damping added to the normal-equations diagonal (scaled by its mean magnitude) for conditioning on planar/degenerate targets.	1e-06

Returns:

Type	Description
matrix	(1,) wp.mat44 rigid transform mapping a onto the target.
transformed	(n,) wp.vec3 image of a under matrix.
cost	Robust-weighted sum of squared point-to-plane residuals at the final iteration.

Raises:

Type	Description
ValueError	If robust_kernel is not one of the three names, or if the target is a point cloud (target_faces is None) and target_normals is not provided, or does not have the same length as target_vertices.
RuntimeError	If a, target_vertices, target_faces, target_normals and initial are not all on one device.

See Also

icp estimate_normals normals_at_closest_faces

Source code in triwarp/registration.py

def icp_point_to_plane(
    a: wp.array[wp.vec3],
    target_vertices: wp.array[wp.vec3],
    target_faces: wp.array[wp.int32] | None = None,
    *,
    target_normals: wp.array[wp.vec3] | None = None,
    initial: wp.array[wp.mat44] | wp.mat44 | None = None,
    max_iterations: int = 20,
    threshold: float = 1e-6,
    max_distance: float | None = None,
    robust_kernel: Literal["none", "huber", "tukey"] = "none",
    robust_scale: float | None = None,
    damping: float = 1e-6,
) -> tuple[wp.array[wp.mat44], wp.array[wp.vec3], float]:
    """
    Point-to-plane iterative closest point registration with optional robust loss.

    Minimizes the squared distance from each transformed source point to the
    *tangent plane* of its target correspondence, which converges faster than the
    point-to-point metric on smooth surfaces. Each iteration linearizes the
    residual, assembles the ``6x6`` Gauss-Newton system on the GPU, solves it, and
    composes the incremental rigid step into the running transform.

    Warp port of the ``libigl`` ``iterative_closest_point`` (rigid, point-to-plane)
    and Open3D's robust ``TransformationEstimationPointToPlane``. An optional
    Huber or Tukey M-estimator down-weights outlier correspondences.

    Parameters
    ----------
    a
        Source point cloud, shape ``(n,)``, dtype ``wp.vec3``.
    target_vertices
        Target vertex positions, shape ``(m,)``, dtype ``wp.vec3``.
    target_faces
        Flat ``(f * 3,)`` triangle index buffer. When given, target normals are
        the closest triangles' face normals; otherwise the target is the point
        cloud ``target_vertices`` and ``target_normals`` is required. When given, the internally
        built ``wp.Mesh`` aliases ``target_vertices`` and ``target_faces`` rather than copying
        them; do not mutate them for the duration of the call.
    target_normals
        Per-vertex unit normals for a point-cloud target, shape ``(m,)``. Required
        (and only used) when ``target_faces`` is ``None``. Estimate them with
        [`estimate_normals`][triwarp.points.estimate_normals] if absent.
    initial
        Initial transform, as a ``(1,)`` ``wp.mat44`` array or scalar ``wp.mat44``.
        Identity when ``None``.
    max_iterations
        Maximum number of ICP iterations.
    threshold
        Stop when the cost decreases by less than this between iterations.
    max_distance
        Reject correspondences farther than this. No rejection when ``None``.
    robust_kernel
        M-estimator applied to residuals: ``"none"``, ``"huber"``, or ``"tukey"``.
    robust_scale
        Tuning constant of the robust kernel (Huber ``k`` / Tukey ``c``). When
        ``None`` and a robust kernel is active, derived from the median absolute
        deviation of the first iteration's residuals.
    damping
        Relative Levenberg damping added to the normal-equations diagonal (scaled
        by its mean magnitude) for conditioning on planar/degenerate targets.

    Returns
    -------
    matrix
        ``(1,)`` ``wp.mat44`` rigid transform mapping *a* onto the target.
    transformed
        ``(n,)`` ``wp.vec3`` image of *a* under *matrix*.
    cost
        Robust-weighted sum of squared point-to-plane residuals at the final
        iteration.

    Raises
    ------
    ValueError
        If ``robust_kernel`` is not one of the three names, or if the target is a point cloud
        (``target_faces is None``) and ``target_normals`` is not provided, or does not have the
        same length as ``target_vertices``.
    RuntimeError
        If ``a``, ``target_vertices``, ``target_faces``, ``target_normals`` and ``initial`` are not
        all on one device.

    See Also
    --------
    [`icp`][triwarp.registration.icp]
    [`estimate_normals`][triwarp.points.estimate_normals]
    [`normals_at_closest_faces`][triwarp.proximity.normals_at_closest_faces]
    """
    require_same_device(
        a=a,
        target_vertices=target_vertices,
        target_faces=target_faces,
        target_normals=target_normals,
        initial=initial,
    )
    device = a.device
    n = int(a.shape[0])
    is_mesh = _is_mesh_target(target_faces)
    if robust_kernel not in _ROBUST_KINDS:
        raise ValueError(
            f"robust_kernel must be one of {list(_ROBUST_KINDS)}, got {robust_kernel!r}"
        )
    kind = _ROBUST_KINDS[robust_kernel]

    if not is_mesh and target_normals is None:
        raise ValueError(
            "icp_point_to_plane requires target_normals for a point-cloud target "
            "(target_faces is None)."
        )
    if (
        not is_mesh
        and target_normals is not None
        and int(target_normals.shape[0]) != int(target_vertices.shape[0])
    ):
        # ``target_normals`` feeds ``gather_vec_skip_negative`` below, indexed by a nearest-vertex
        # id that ranges over ``target_vertices``; a shorter buffer is an out-of-bounds read on
        # both devices (CLAUDE.md §12.1), not merely a wrong answer.
        raise ValueError(
            "target_normals must have the same length as target_vertices, got "
            f"{target_normals.shape[0]} and {target_vertices.shape[0]}."
        )

    if n == 0 or int(target_vertices.shape[0]) == 0:
        return _identity_mat44(device), wp.clone(a), math.inf

    initial_matrix, current = _seed_transform(a, initial, device)
    transformed = wp.clone(current)
    cost = math.inf

    mesh, query_max, target_index = _resolve_icp_target(
        target_vertices, target_faces, current, max_distance, "icp_point_to_plane"
    )
    face_normals: wp.array[wp.vec3] | None = None
    if is_mesh:
        assert target_faces is not None
        face_normals, _ = tw.triangles.face_normals_and_areas(target_vertices, target_faces)
    # ``target_index`` is already the point-cloud one from ``_resolve_icp_target`` above when
    # ``not is_mesh`` -- rebuilding it here duplicated a ``wp.Bvh`` build and an ``aabb`` reduction
    # for no behavioral difference.

    max_d = max_distance if max_distance is not None else math.inf
    scale_value = robust_scale
    old_cost = math.inf

    # All per-iteration buffers are allocated once; the accumulators are zeroed in place and
    # ``current``/``updated`` ping-pong. ``total`` is cloned so composing in place never mutates
    # a caller-provided initial transform. The per-iteration cost read stays: it is the
    # stopping criterion (a 4-byte transfer).
    closest = wp.empty(n, dtype=wp.vec3, device=device)
    distance_mesh = wp.empty(n, dtype=wp.float32, device=device)
    triangle_id_mesh = wp.empty(n, dtype=wp.int32, device=device)
    normals = wp.empty(n, dtype=wp.vec3, device=device)
    # Only allocated when correspondences can actually be rejected -- see the all-rejected guard
    # below, which needs a buffer to reduce over.
    valid: wp.array[wp.bool] | None = (
        wp.empty(n, dtype=wp.bool, device=device) if max_distance is not None else None
    )
    jtj = wp.zeros(1, dtype=wp.spatial_matrix, device=device)
    jtr = wp.zeros(1, dtype=wp.spatial_vector, device=device)
    cost_acc = wp.zeros(1, dtype=wp.float32, device=device)
    weight_sum_acc = wp.zeros(1, dtype=wp.float32, device=device)
    step = wp.empty(1, dtype=wp.mat44, device=device)
    updated = wp.empty(n, dtype=wp.vec3, device=device)
    total = wp.clone(initial_matrix)

    for iteration in range(max_iterations):
        # --- correspondence + target normals ---
        distance, triangle_id = _correspondences(
            mesh,
            target_vertices,
            target_index,
            current,
            query_max,
            closest,
            distance_mesh,
            triangle_id_mesh,
        )
        # One gather either way: for a mesh target ``triangle_id`` indexes the face normals, for a
        # cloud it indexes the target's own per-vertex normals.
        normal_source = face_normals if mesh is not None else target_normals
        assert normal_source is not None
        wp.launch(
            kernel_array.gather_vec_skip_negative,
            dim=n,
            inputs=[normal_source, triangle_id, normals],
            device=device,
        )

        # --- bail out once no correspondence survives the distance gate ---
        # Mirrors ``icp``'s ``sum(weights) == 0`` check: without it, every accumulator below stays
        # at its zeroed initial value, the damped 6x6 solve returns a zero step, and the loop
        # reports ``cost=0.0`` -- indistinguishable from a perfect fit -- instead of stopping with
        # the last real cost (or ``math.inf`` if nothing ever matched).
        if valid is not None:
            wp.map(
                kernel_registration.residual_valid,
                triangle_id,
                distance,
                wp.float32(max_d),
                out=valid,
            )
            if not tw.reduce.any(valid):
                break

        # --- resolve robust scale on the first iteration ---
        if kind != 0 and scale_value is None:
            scale_value = _robust_scale_from_residuals(
                current, closest, normals, distance, triangle_id, max_d, kind
            )

        # --- assemble and solve the linearized point-to-plane system ---
        jtj.zero_()
        jtr.zero_()
        cost_acc.zero_()
        weight_sum_acc.zero_()
        wp.launch_tiled(
            kernel_registration.accumulate_point_to_plane,
            dim=kernel_reduce.blocks_1d(n),
            inputs=[
                current,
                closest,
                normals,
                distance,
                triangle_id,
                wp.float32(max_d),
                wp.int32(kind),
                wp.float32(scale_value if scale_value is not None else 0.0),
                jtj,
                jtr,
                cost_acc,
                weight_sum_acc,
            ],
            block_dim=TILE_1D,
            device=device,
        )

        # --- bail out once every in-range correspondence's own robust weight collapsed to zero ---
        # The ``valid`` guard above only catches a correspondence rejected by *distance*; a Tukey
        # kernel can drive every remaining correspondence's weight to exactly zero on its own
        # (``|residual| >= scale``, reachable through an explicit ``robust_scale`` too tight for the
        # residual distribution, or an earlier iteration's step overshooting far past what a
        # first-iteration-derived scale anticipated) while ``residual_valid`` -- and so ``valid`` --
        # still accepts every one of them. Without this, ``jtj``/``jtr`` stay exactly zero, the
        # damped solve returns a near-identity step, and ``cost`` reads ``0.0`` -- indistinguishable
        # from a perfect fit -- for the rest of the run. Reproduced with
        # ``robust_kernel="tukey", robust_scale=1e-9``: every residual exceeds so tight a scale, and
        # the loop returned the identity transform with ``cost=0.0`` on a cloud still offset by
        # (1.0, 0.5, -0.3) from its target.
        if float(read_scalar(weight_sum_acc, 0)) <= 0.0:
            break

        wp.launch(
            kernel_registration.solve_point_to_plane,
            dim=1,
            inputs=[jtj, jtr, wp.float32(damping), step],
            device=device,
        )

        # --- apply the incremental step and compose into the running transform ---
        wp.launch(
            kernel_transform.apply_transform_mat44,
            dim=n,
            inputs=[current, step, updated],
            device=device,
        )
        current, updated = updated, current
        transformed = current
        wp.map(wp.mul, step, total, out=total)

        cost = float(read_scalar(cost_acc, 0))
        if iteration > 0 and old_cost - cost < threshold:
            break
        old_cost = cost

    return total, transformed, cost