Mesh: Add differentiable shape-to-shape fitting

CHANGELOG.md

@@ -52,6 +52,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   ``Mesh.free_form_deform`` / ``DomainMesh.free_form_deform`` with Bernstein
   (classic FFD) and locally supported uniform cubic B-spline bases, plus
   node-interpolating `linear`, `cubic_hermite`, and `quintic_hermite` modes.
+- Adds differentiable template-to-surface fitting with tensor-level
+  `arap_energy`, Torch/NVIDIA Warp `point_to_mesh_distance`, and
+  `fit_template_points`, plus the connectivity-preserving `Mesh.fit_template`
+  wrapper.
 - Adds `uniform_grid_divergence`, `uniform_grid_curl`, and
   `uniform_grid_laplacian` to `physicsnemo.nn.functional`, with Torch and fused
   Warp implementations for periodic Cartesian grids.

benchmarks/physicsnemo/nn/functional/registry.py

@@ -41,12 +41,15 @@
     ViewAsComplex,
 )
 from physicsnemo.nn.functional.geometry import (
+    ARAPEnergy,
     DisplacePoints,
     FarthestPointSampling,
+    FitTemplatePoints,
     FreeFormDeformPoints,
     MeshPoissonDiskSample,
     MeshToVoxelFraction,
     MorphPoints,
+    PointToMeshDistance,
     RayMeshIntersect,
     SignedDistanceField,
 )
@@ -83,9 +86,12 @@
     UniformGridCurl,
     UniformGridLaplacian,
     # Geometry.
+    ARAPEnergy,
     DisplacePoints,
+    FitTemplatePoints,
     MorphPoints,
     FreeFormDeformPoints,
+    PointToMeshDistance,
     FarthestPointSampling,
     MeshPoissonDiskSample,
     MeshToVoxelFraction,

docs/api/mesh/transformations.rst

@@ -41,14 +41,16 @@ Deformations
 
 .. currentmodule:: physicsnemo.mesh.transformations.deform
 
-Dense displacement, sparse control-point morphing, and lattice free-form
-deformation are available from the ``deform`` namespace and as methods on
-:class:`~physicsnemo.mesh.mesh.Mesh`.
+Dense displacement, sparse control-point morphing, lattice free-form
+deformation, and template shape fitting are available from the ``deform``
+namespace and as methods on :class:`~physicsnemo.mesh.mesh.Mesh`.
 
 The mesh methods wrap the tensor-level
 :func:`~physicsnemo.nn.functional.displace_points`,
 :func:`~physicsnemo.nn.functional.morph_points`, and
 :func:`~physicsnemo.nn.functional.free_form_deform_points` operations.
+Template fitting wraps
+:func:`~physicsnemo.nn.functional.fit_template_points`.
 
 Dense displacement accepts a tensor or a point-data key (including a nested
 tuple key). The operation returns a new mesh without changing ``mesh.points``.
@@ -218,6 +220,41 @@ produce two local bulges and a local indentation.
    :alt: Sphere with control lattice, Bernstein taper, and local B-spline sculpt
    :width: 100%
 
+Template Shape Fitting
+^^^^^^^^^^^^^^^^^^^^^^
+
+Template fitting moves the vertices of one prealigned 3D triangle mesh toward
+another triangle surface. The target may use different connectivity; the
+returned mesh keeps the template's connectivity, point ordering, and attached
+data.
+
+.. code:: python
+
+    fitted = template.fit_template(
+        target,
+        fit_weight=1.0,
+        arap_weight=0.1,
+        steps=10,
+    )
+
+The fitting term is one-way from template vertices to closest points on the
+target. ``fit_weight`` controls target adherence, while ``arap_weight`` resists
+local stretching and shearing. A larger ARAP weight favors shape preservation
+over an exact surface match. Rigid registration is a separate preprocessing
+step; the two meshes should already be approximately aligned.
+
+The solver performs a fixed number of local/global steps. First-order gradients
+propagate from ``fitted.points`` to both input point tensors for the selected
+nearest-surface and rotation branches. Those branches are discrete, so
+gradients are valid almost everywhere and higher-order differentiation is not
+supported.
+
+Template fitting is intentionally defined for one
+:class:`~physicsnemo.mesh.mesh.Mesh` and one target surface. It is not exposed
+on :class:`~physicsnemo.mesh.domain_mesh.DomainMesh`, whose interior and
+boundary components may have different manifold dimensions and do not define a
+single unambiguous fitting surface.
+
 Domain Meshes
 ^^^^^^^^^^^^^
 
@@ -291,6 +328,8 @@ caches and recompute them lazily. They retain topology caches.
 
 .. autofunction:: free_form_deform
 
+.. autofunction:: fit_template
+
 .. autofunction:: morph
 
 Projections

docs/api/nn/functionals/geometry.rst

@@ -155,6 +155,100 @@ For connectivity-preserving object APIs, use
 :meth:`~physicsnemo.mesh.mesh.Mesh.free_form_deform` or
 :meth:`~physicsnemo.mesh.domain_mesh.DomainMesh.free_form_deform`.
 
+As-Rigid-As-Possible Energy
+---------------------------
+
+.. autofunction:: physicsnemo.nn.functional.arap_energy
+
+``arap_energy(reference_points, deformed_points, edges)`` measures local shape
+distortion while allowing each point neighborhood to rotate. It is therefore
+zero, up to numerical precision, for a global rigid transformation. The
+functional accepts coordinate and connectivity tensors rather than
+:class:`~physicsnemo.mesh.mesh.Mesh` objects, so callers choose the undirected
+edge graph and optional edge weights explicitly.
+
+The fitted local rotations use the first-order envelope derivative of the
+minimized energy. Rotation selection is not unique for a degenerate
+neighborhood, and higher derivatives through that selection are outside the
+functional's contract.
+
+.. rubric:: Visualization
+
+The visualization compares a reference edge graph with rigid and non-rigid
+deformations measured by the ARAP energy.
+
+.. figure:: /img/nn/functional/geometry/arap_energy/arap_energy.png
+   :alt: Reference mesh and deformations with different ARAP energies
+   :width: 85%
+
+Point-to-Mesh Distance
+----------------------
+
+.. autofunction:: physicsnemo.nn.functional.point_to_mesh_distance
+
+The target mesh tensors come first and query points come last:
+
+.. code:: python
+
+    distance, closest = point_to_mesh_distance(
+        target_points,
+        target_triangles,
+        query_points,
+        squared=True,
+    )
+
+Nearest-face and triangle-feature choices are discrete. Once those choices are
+fixed, both outputs retain gradients to the query coordinates and the selected
+target vertices. Gradients are therefore branchwise and valid almost
+everywhere, but not at correspondence ties or triangle-region boundaries.
+The Warp backend rebuilds its target acceleration structure on each call so
+the public tensor API remains stateless. CUDA Graph capture is not part of the
+current contract.
+
+.. rubric:: Visualization
+
+The visualization shows query points, their closest points on the target
+surface, and the corresponding distances.
+
+.. figure:: /img/nn/functional/geometry/point_to_mesh_distance/point_to_mesh_distance.png
+   :alt: Query points and their closest points on a triangle mesh
+   :width: 85%
+
+Template Shape Fitting
+----------------------
+
+.. autofunction:: physicsnemo.nn.functional.fit_template_points
+
+``fit_template_points`` deforms a prealigned 3D triangle template toward a
+target triangle surface. The fitting term is one-way: every template vertex is
+matched to its closest point on the target. The returned tensor retains the
+template's point ordering, and the target may use different connectivity.
+
+``fit_weight`` controls target adherence, while ``arap_weight`` resists local
+stretching and shearing. Increasing ``arap_weight`` generally preserves the
+template shape more strongly at the cost of a less exact surface fit. The
+operation runs a fixed number of local/global steps and provides first-order
+gradients through that fixed-step result. Correspondence and rotation branches
+are discrete, and higher-order differentiation is not supported.
+
+The initial contract is unbatched, 3D, and prealigned. Warp accelerates only
+nearest-face selection; local rotations, ARAP assembly, and the matrix-free
+linear solves remain in Torch. The target acceleration structure is rebuilt
+for every fitting step, and CUDA Graph capture is not part of the current
+contract.
+
+For a connectivity-preserving object API, use
+:meth:`~physicsnemo.mesh.mesh.Mesh.fit_template`.
+
+.. rubric:: Visualization
+
+The visualization compares a prealigned template, its target surface, and the
+fitted template while retaining the template connectivity.
+
+.. figure:: /img/nn/functional/geometry/fit_template_points/fit_template_points.png
+   :alt: Template mesh fitted to a target surface with ARAP regularization
+   :width: 90%
+
 Mesh Poisson Disk Sample
 ------------------------
 

physicsnemo/mesh/mesh.py

@@ -39,7 +39,12 @@
 
 from physicsnemo.mesh.geometry._cell_areas import compute_cell_areas
 from physicsnemo.mesh.geometry._cell_normals import compute_cell_normals
-from physicsnemo.mesh.transformations.deform import displace, free_form_deform, morph
+from physicsnemo.mesh.transformations.deform import (
+    displace,
+    fit_template,
+    free_form_deform,
+    morph,
+)
 from physicsnemo.mesh.transformations.deform.ffd import _FFDBasis
 from physicsnemo.mesh.transformations.geometric import (
     rotate,
@@ -2759,6 +2764,39 @@ def free_form_deform(
             implementation=implementation,
         )
 
+    def fit_template(
+        self,
+        target: "Mesh",
+        *,
+        fit_weight: builtins.float = 1.0,
+        arap_weight: builtins.float = 0.1,
+        steps: int = 10,
+        cg_tolerance: builtins.float = 1.0e-6,
+        cg_max_iterations: int = 256,
+        implementation: Literal["torch", "warp"] | None = None,
+    ) -> "Mesh":
+        """Fit this prealigned triangle template to a target surface.
+
+        Convenience wrapper for
+        :func:`physicsnemo.mesh.transformations.deform.fit_template`, which
+        documents all parameters and numerical behavior.
+
+        Returns
+        -------
+        Mesh
+            New fitted mesh with this template's connectivity and fields.
+        """
+        return fit_template(
+            self,
+            target,
+            fit_weight=fit_weight,
+            arap_weight=arap_weight,
+            steps=steps,
+            cg_tolerance=cg_tolerance,
+            cg_max_iterations=cg_max_iterations,
+            implementation=implementation,
+        )
+
     def rotate(
         self,
         angle: float,

physicsnemo/mesh/transformations/deform/__init__.py

@@ -18,6 +18,7 @@
 
 from physicsnemo.mesh.transformations.deform.displace import displace
 from physicsnemo.mesh.transformations.deform.ffd import free_form_deform
+from physicsnemo.mesh.transformations.deform.fit_template import fit_template
 from physicsnemo.mesh.transformations.deform.morph import morph
 
-__all__ = ["displace", "free_form_deform", "morph"]
+__all__ = ["displace", "fit_template", "free_form_deform", "morph"]

physicsnemo/mesh/transformations/deform/fit_template.py

@@ -0,0 +1,123 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shape-to-shape fitting for triangle surface meshes."""
+
+from typing import TYPE_CHECKING, Literal
+
+from physicsnemo.mesh.transformations.deform._utils import (
+    _mesh_with_deformed_points,
+)
+
+if TYPE_CHECKING:
+    from physicsnemo.mesh.mesh import Mesh
+
+
+def _validate_triangle_surface(mesh: object, argument_name: str) -> "Mesh":
+    """Require a triangle surface mesh embedded in three dimensions."""
+
+    from physicsnemo.mesh.mesh import Mesh
+
+    if not isinstance(mesh, Mesh):
+        raise TypeError(f"{argument_name} must be a Mesh, got {type(mesh).__name__}")
+    if mesh.n_manifold_dims != 2 or mesh.n_spatial_dims != 3:
+        raise ValueError(
+            f"{argument_name} must be a triangle surface mesh embedded in 3D, got "
+            f"n_manifold_dims={mesh.n_manifold_dims} and "
+            f"n_spatial_dims={mesh.n_spatial_dims}"
+        )
+    return mesh
+
+
+def fit_template(
+    template: "Mesh",
+    target: "Mesh",
+    *,
+    fit_weight: float = 1.0,
+    arap_weight: float = 0.1,
+    steps: int = 10,
+    cg_tolerance: float = 1.0e-6,
+    cg_max_iterations: int = 256,
+    implementation: Literal["torch", "warp"] | None = None,
+) -> "Mesh":
+    """Fit a prealigned triangle template to a target surface.
+
+    The template's vertices move toward closest points on the target while an
+    as-rigid-as-possible term discourages local distortion. Only template point
+    coordinates change; its connectivity and attached fields are retained.
+
+    Parameters
+    ----------
+    template : Mesh
+        Prealigned 3D triangle surface mesh to deform. The source mesh is not
+        modified.
+    target : Mesh
+        Prealigned 3D triangle surface mesh defining the fitting surface. Its
+        topology may differ from the template and it is not modified.
+    fit_weight : float, optional
+        Positive closest-surface fitting weight. Default is ``1.0``.
+    arap_weight : float, optional
+        Nonnegative as-rigid-as-possible regularization weight. Default is
+        ``0.1``.
+    steps : int, optional
+        Number of fixed local/global fitting iterations. Zero returns a new
+        mesh with cloned template points. Default is ``10``.
+    cg_tolerance : float, optional
+        Positive scaled residual tolerance for the linear solves. Default is
+        ``1e-6``.
+    cg_max_iterations : int, optional
+        Positive maximum iteration count for each linear solve. Default is
+        ``256``.
+    implementation : {"warp", "torch"} or None, optional
+        Correspondence-search backend. ``None`` selects Warp for CUDA float32
+        coordinates when available and Torch otherwise.
+
+    Returns
+    -------
+    Mesh
+        New mesh with fitted points and the template's connectivity and
+        attached fields.
+
+    Notes
+    -----
+    Inputs are assumed to be prealigned. Attached fields are treated as
+    Lagrangian data and are not pushed forward. Geometry-dependent caches are
+    invalidated and template topology caches are retained. The fit does not
+    guarantee freedom from inverted, degenerate, or self-intersecting cells;
+    call :meth:`~physicsnemo.mesh.mesh.Mesh.validate` explicitly when needed.
+    """
+
+    template = _validate_triangle_surface(template, "template")
+    target = _validate_triangle_surface(target, "target")
+
+    from physicsnemo.nn.functional import fit_template_points
+
+    points = fit_template_points(
+        template.points,
+        template.cells,
+        target.points,
+        target.cells,
+        fit_weight=fit_weight,
+        arap_weight=arap_weight,
+        steps=steps,
+        cg_tolerance=cg_tolerance,
+        cg_max_iterations=cg_max_iterations,
+        implementation=implementation,
+    )
+    return _mesh_with_deformed_points(template, points)
+
+
+__all__ = ["fit_template"]