Geometry Functionals#

Point Displacement#

physicsnemo.nn.functional.displace_points(
points: Tensor,
displacement: Tensor,
*,
point_weights: Tensor | None = None,
implementation: Literal['torch'] | None = None,
) → Tensor#

Apply an aligned dense displacement field to points.

The operation is

\[x'_i = x_i + w_i\,d_i,\]

where point_weights contains optional per-point values \(w_i\) and displacement contains vectors \(d_i\). Signed or greater-than-one point weights are permitted.

Inputs may be unbatched (N, D) or batched (B, N, D). Batched inputs are aligned rather than broadcast: points and displacement must have identical shapes. Float32 and float64 are supported.

Parameters:

  • points (torch.Tensor) – Point coordinates with shape (N, D) or (B, N, D).

  • displacement (torch.Tensor) – Dense displacement vectors with exactly the same shape, dtype, and device as points.

  • point_weights (torch.Tensor or None, optional) – Optional bool or floating per-point weights. Accepted shapes are (N,) for unbatched inputs and (B, N) for batched inputs. Floating point weights must match the point dtype and device; bool values act as hard masks. Values are used as supplied without clamping. Default is None.

  • implementation ({"torch"} or None, optional) – Explicit backend. None selects Torch.

Returns:

  • torch.Tensor – Displaced points with the same shape, dtype, and device as points.

  • The operation is implemented with native Torch tensor operations and

  • participates in autograd and torch.compile().

import torch
from physicsnemo.nn.functional import displace_points

points = torch.tensor(
    [[0.0, 0.0], [1.0, 0.0], [2.0, 0.0]], requires_grad=True
)
displacement = torch.tensor(
    [[0.0, 1.0], [0.0, 1.0], [0.0, 1.0]], requires_grad=True
)
point_weights = torch.tensor([0.0, 0.5, 1.0])

moved = displace_points(
    points,
    displacement,
    point_weights=point_weights,
)
moved.square().sum().backward()

Sobolev Point Deformation#

physicsnemo.nn.functional.sobolev_deform_points(
points: Float[Tensor, '*batch num_points num_dims'],
cells: Int[Tensor, 'num_cells num_cell_points'],
displacement: Float[Tensor, '*batch num_points num_dims'],
*,
length_scale: float,
fixed_points: Bool[Tensor, '*batch num_points'] | None = None,
max_iterations: int = 128,
tolerance: float | None = None,
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '*batch num_points num_dims']#

Deform simplicial-mesh points with a smooth Sobolev displacement.

The operation filters a prescribed per-vertex displacement \(d\) by solving the uniform-mass discrete Helmholtz system

\[(M + \ell^2 K)u = M d, \qquad x' = x + u.\]

Here \(M=\bar m I\) is a uniform vertex mass matrix. Its scalar \(\bar m\) is the mean positive lumped P1 vertex mass. \(K\) is the P1 stiffness matrix, and \(\ell\) is length_scale in the same physical units as points. This uniform mass makes the forward filter self-adjoint in standard Euclidean vertex coordinates. The reverse pass therefore applies the same smoothing operator to the displacement adjoint. The solve uses a matrix-free Jacobi-preconditioned conjugate gradient method. Each ambient component is filtered independently.

points and displacement may be unbatched (N, D) or aligned batched (B, N, D) tensors. cells has shape (C, V) and defines one topology shared by every batch entry. It must contain simplices with V - 1 <= D. All floating tensors must use float32 or float64.

Parameters:

  • points (torch.Tensor) – Vertex coordinates with shape (N, D) or (B, N, D).

  • cells (torch.Tensor) – Shared simplex connectivity with shape (C, V) and int32 or int64 dtype.

  • displacement (torch.Tensor) – Raw per-vertex displacement with the same shape, dtype, and device as points.

  • length_scale (float) – Nonnegative physical smoothing length. Zero applies the raw displacement exactly at unfixed points.

  • fixed_points (torch.Tensor or None, optional) – Optional bool mask with shape (N,) or (B, N). True entries use a zero Dirichlet displacement. Other mesh boundaries use the natural homogeneous Neumann condition. Default is None.

  • max_iterations (int, optional) – Maximum PCG iterations. Default is 128.

  • tolerance (float or None, optional) – Positive relative residual tolerance. None selects 1e-6 for float32 and 1e-10 for float64. Default is None.

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU. On CUDA, it selects Warp for segments, triangles, and tetrahedra when available. It otherwise selects Torch, with a one-time RuntimeWarning when Warp is unavailable. The Warp backend requires CUDA tensors.

Returns:

Deformed points with the same shape, dtype, and device as points.

Return type:

torch.Tensor

Raises:

  • TypeError – If argument types or tensor dtypes are unsupported.

  • ValueError – If shapes, devices, indices, scalar options, or simplex geometry are invalid.

  • KeyError – If implementation does not name a registered backend.

  • ImportError – If an explicitly requested backend is unavailable.

  • RuntimeError – If CUDA Graph capture is active or the forward or adjoint PCG solve does not reach tolerance within max_iterations.

Notes

Constant displacements are retained to solver precision when no points are fixed. Isolated points receive their raw displacement. The uniform mass scale is computed over all nonisolated vertices in the supplied topology.

Both backends provide first-order gradients with respect to points and displacement. Their implicit reverse-mode derivatives solve the adjoint of the forward Helmholtz system. This is not an identity-valued surrogate gradient. The Warp backend evaluates the geometry vector-Jacobian product analytically because Warp’s supplied linear solvers do not generate automatic backward kernels. Higher-order gradients are not supported for a positive length scale.

The Warp backend supports segments, triangles, and tetrahedra. The Torch backend also supports higher-dimensional simplices. Default dispatch keeps higher-dimensional CUDA simplices on Torch. Warp CUDA assembly and geometry pullback use atomic accumulation, so results and point gradients may vary at roundoff between runs.

Forward and backward each run at most max_iterations matrix-free PCG steps. A nonconverged solve raises an error instead of returning an inconsistent implicit gradient. At positive length scales, cells must be finite and nondegenerate. The operation does not check for inverted or self-intersecting output cells.

import torch
from physicsnemo.nn.functional import sobolev_deform_points

points = torch.tensor(
    [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]],
    requires_grad=True,
)
cells = torch.tensor([[0, 1, 2], [0, 2, 3]])
displacement = torch.tensor(
    [[0.0, 0.0], [0.0, 0.2], [0.0, -0.2], [0.0, 0.0]],
    requires_grad=True,
)

smooth = sobolev_deform_points(
    points,
    cells,
    displacement,
    length_scale=0.25,
)
smooth.square().mean().backward()

The operation assembles a P1 stiffness matrix and a uniform vertex mass scaled by the mean positive lumped P1 mass. It solves \((M + \ell^2 K)u = Md\) and returns \(x + u\). length_scale is \(\ell\) in the same physical units as points. Zero recovers ordinary dense displacement. Larger values attenuate short-wavelength changes in the raw displacement and its reverse-mode sensitivity. The uniform mass makes this filter self-adjoint in standard Euclidean vertex coordinates.

fixed_points is an optional boolean vertex mask. True entries impose zero Dirichlet displacement. Boundaries without a fixed mask use the natural homogeneous Neumann condition. The Torch and Warp matrix-free solves support one mesh or a batch of point tensors that share cells. Both are differentiable with respect to points and displacement. The Warp backend runs on CUDA and uses an explicit implicit-adjoint backward with an analytic geometry vector-Jacobian product. CUDA segments, triangles, and tetrahedra select Warp by default when it is available. CPU inputs and higher-dimensional simplices select Torch. Positive length scales support first-order reverse-mode differentiation. Higher-order derivatives are not supported. The operation raises an error when a forward or adjoint solve does not converge within max_iterations. CUDA Graph capture is not supported because P1 operator assembly and solver diagnostics are not capture-safe. Warp CUDA results and point gradients may vary at roundoff between runs.

Sparse Control-Point Morphing#

physicsnemo.nn.functional.morph_points(
points: Tensor,
control_points: Tensor,
control_displacements: Tensor,
*,
radius: float | Tensor,
point_weights: Tensor | None = None,
kernel: Literal['wendland_c2'] = 'wendland_c2',
implementation: Literal['warp', 'torch'] | None = None,
) → Tensor#

Morph points from sparse control displacements using compact Shepard blending.

For query point \(x\), control \(c_j\), radius \(r_j\), and normalized distance \(q_j=\lVert x-c_j\rVert/r_j\), define

\[\phi(q) = (1-q)^4(4q+1), \qquad a_j = \frac{\phi(q_j)}{q_j^2}, \quad 0 < q_j < 1.\]

A stationary background of weight one blends the active controls toward zero displacement:

\[u(x) = \frac{\sum_j a_j d_j}{1 + \sum_j a_j}, \qquad x' = x + \mathtt{point\_weights}\,u(x).\]

Each control’s influence vanishes smoothly at its own support boundary. The total field is zero wherever no control support is active; it need not be zero at one control’s boundary if another support overlaps that location. At an exact control coincidence, the unscaled field returns the mean displacement of all controls at that coordinate. point_weights are then applied to that field. The implementation uses a stably scaled equivalent of the quotient near controls.

Inputs may be unbatched (N, D)/(C, D) or aligned batched (B, N, D)/(B, C, D). Point and control batches are not implicitly broadcast. All coordinate and displacement tensors must use float32 or float64 and have the same dtype and device.

Parameters:

  • points (torch.Tensor) – Query points with shape (N, D) or (B, N, D).

  • control_points (torch.Tensor) – World-coordinate control locations with shape (C, D) or (B, C, D). Controls need not be query points.

  • control_displacements (torch.Tensor) – Control displacement vectors, not destination coordinates. The shape, dtype, and device must exactly match control_points.

  • radius (float or torch.Tensor) – Support radius. Accepts a scalar, per-control (C,) tensor, or aligned batched (B, C) tensor. Every value must be positive and finite. Tensor radii must match the control dtype and device; their numerical values are not validated at runtime.

  • point_weights (torch.Tensor or None, optional) – Optional bool or floating per-point weights. Shapes follow DisplacePoints; point weights are not per-control values. Signed and amplifying values are permitted and are used without clamping.

  • kernel ({"wendland_c2"}, optional) – Compact radial kernel used by Shepard blending. The explicit name reserves the algorithm-selection extension point. Default is "wendland_c2".

  • implementation ({"warp", "torch"} or None, optional) – Explicit backend. None selects Torch on CPU and Warp on CUDA when Warp is available, otherwise Torch with a one-time RuntimeWarning.

Returns:

Morphed points with the same shape, dtype, and device as points.

Return type:

torch.Tensor

Notes

Both backends propagate first-order gradients through points, controls, control displacements, tensor-valued radii, and floating-point weights. Only first-order gradients are part of the Warp backend’s public contract.

A learned radius should be parameterized to remain positive, for example as torch.nn.functional.softplus(raw_radius) + eps with a positive eps appropriate to the coordinate scale. With zero controls, the operation is the identity and the numerical value of a scalar radius is unused.

import torch
from physicsnemo.nn.functional import morph_points

x = torch.linspace(0.0, 1.0, 9)
points = torch.stack((x, torch.zeros_like(x)), dim=-1).requires_grad_()
control_points = points.detach()[[0, -1]].clone().requires_grad_()
control_displacements = points.new_tensor(
    [[0.0, 0.25], [0.0, -0.15]], requires_grad=True
)
radii = points.new_tensor([0.8, 0.8])

morphed = morph_points(
    points,
    control_points,
    control_displacements,
    radius=radii,
    kernel="wendland_c2",
)
morphed.square().mean().backward()

This allows an optimizer—or a model producing the control displacements—to learn a deformation from a differentiable objective on morphed.

Global Radial-Basis Deformation#

physicsnemo.nn.functional.radial_basis_function_deform_points(
points: Float[Tensor, '*batch num_points num_dims'],
control_points: Float[Tensor, '*batch num_controls num_dims'],
control_displacements: Float[Tensor, '*batch num_controls num_dims'],
*,
kernel: Literal['thin_plate_spline'] = 'thin_plate_spline',
polynomial: bool = True,
smoothing: float = 0.0,
point_weights: Bool[Tensor, '*batch num_points'] | Float[Tensor, '*batch num_points'] | None = None,
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '*batch num_points num_dims']#

Deform points with a global thin-plate-spline RBF field.

Given controls \(c_j\) with prescribed displacements \(d_j\), the displacement field is

\[u(x) = \sum_j \phi(\lVert x-c_j\rVert)w_j + a_0 + A x, \qquad \phi(r)=r^2\log(r), \quad \phi(0)=0.\]

By default, the radial coefficients \(w_j\) and affine coefficients \((a_0,A)\) are fitted from the standard augmented interpolation system. The affine side constraints make the fit unique for distinct controls in nondegenerate position and reproduce affine displacement fields exactly. With smoothing=0, the field interpolates every control displacement up to solver precision. Positive smoothing relaxes exact interpolation. This formulation follows the thin-plate-spline interpolant described by Bookstein [1].

Inputs may be unbatched (N, D)/(C, D) or aligned batched (B, N, D)/(B, C, D). Point and control batches are not implicitly broadcast. All coordinate and displacement tensors must use float32 or float64 and have the same dtype and device. The formulation is dimensionally generic.

Parameters:

  • points (torch.Tensor) – Query points with shape (N, D) or (B, N, D).

  • control_points (torch.Tensor) – World-coordinate control locations with shape (C, D) or (B, C, D).

  • control_displacements (torch.Tensor) – Prescribed displacement vectors, not destination coordinates. Shape, dtype, and device must exactly match control_points.

  • kernel ({"thin_plate_spline"}, optional) – Radial kernel. Default is "thin_plate_spline".

  • polynomial (bool, optional) – Include the affine polynomial tail and its side constraints. This requires at least D + 1 controls when controls are present. Disabling the tail can make the radial system singular. Default is True.

  • smoothing (float, optional) – Nonnegative diagonal regularization added to the control-kernel block. Zero gives exact interpolation for a nonsingular control layout up to solver precision. Positive values relax interpolation accuracy. Default is 0.0.

  • point_weights (torch.Tensor or None, optional) – Optional bool or floating per-point multipliers with shape (N,) or (B, N). They scale the fitted field after interpolation. Signed and amplifying values are used without clamping. Bool weights must be on the same device as points. Floating weights must have the same dtype and device as points.

  • implementation ({"torch", "warp"} or None, optional) – Field-evaluation backend. Both implementations fit coefficients with a checked differentiable torch.linalg.solve_ex(). None selects Torch on CPU and Warp on CUDA when Warp is available, otherwise Torch with a one-time RuntimeWarning.

Returns:

Deformed points with the same shape, dtype, and device as points.

Return type:

torch.Tensor

Raises:

  • TypeError – If tensor dtypes or Python argument types are unsupported.

  • ValueError – If tensor shapes, devices, control layout, point weights, or RBF options are invalid.

  • KeyError – If implementation does not name a registered backend.

  • ImportError – If an explicitly requested backend is unavailable.

  • RuntimeError – If runtime validation or coefficient fitting fails, including for a singular system or during CUDA Graph capture.

Notes

The dense coefficient solve costs \(O(C^3)\) and stores an \(O(C^2)\) system. Field evaluation costs \(O(NC)\). Reuse of a fit across multiple independent query sets is not currently exposed by this convenience API.

Duplicate controls or controls that do not span the affine basis can make the augmented system singular. In that case the checked linear solve raises an error. With zero controls, the operation is the identity.

Coefficient fitting is not supported inside CUDA Graph capture because the singular-system check requires host interaction.

Both backends propagate first-order gradients through points, control locations, control displacements, and floating point weights. Only first-order gradients are part of the Warp evaluator’s public contract.

References

[1] F. L. Bookstein, “Principal warps: thin-plate splines and the

decomposition of deformations,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 11, no. 6, pp. 567-585, 1989. https://doi.org/10.1109/34.24792

import torch
from physicsnemo.nn.functional import radial_basis_function_deform_points

points = torch.tensor(
    [[0.25, 0.25], [0.75, 0.25], [0.75, 0.75], [0.25, 0.75]],
    requires_grad=True,
)
controls = torch.tensor(
    [[0.0, 0.0], [1.0, 0.0], [1.0, 1.0], [0.0, 1.0]],
    requires_grad=True,
)
control_displacements = torch.tensor(
    [[0.0, 0.0], [0.0, 0.0], [0.15, 0.25], [0.0, 0.0]],
    requires_grad=True,
)

exact = radial_basis_function_deform_points(
    points,
    controls,
    control_displacements,
    kernel="thin_plate_spline",
    polynomial=True,
    smoothing=0.0,
)
exact.square().mean().backward()

With zero smoothing and a nonsingular control layout, the fitted field interpolates every control displacement up to solver precision. The affine polynomial tail also reproduces affine displacement fields. A positive smoothing value adds diagonal regularization and relaxes interpolation. Thin-plate-spline fields have global support, unlike the compact Shepard field used by morph_points(). This formulation follows the thin-plate-spline interpolant described by Bookstein [1].

[1] F. L. Bookstein, “Principal warps: thin-plate splines and the decomposition of deformations,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 11, no. 6, pp. 567-585, 1989. https://doi.org/10.1109/34.24792

Performance and Compilation#

Dense point displacement uses Torch on every device. Compact morphing and radial-basis deformation use Torch by default on CPU and Warp by default on CUDA. If Warp is unavailable, automatic CUDA dispatch falls back to Torch, while explicitly requesting implementation="warp" raises an ImportError. For a repeatedly evaluated, fixed-shape CUDA deformation wrapped in torch.compile(), benchmark implementation="torch" as well. Compiler fusion can make that path faster after its one-time compilation cost. Keep the backend explicit when comparing compiled and eager runs.

Compact morphing and radial-basis deformation evaluate every query/control pair. Their computational cost is proportional to batch_size * n_points * n_controls * n_spatial_dims. Pass all simultaneous controls in one call. For a DomainMesh, the object API combines its interior and boundary queries into one field evaluation before rebuilding the individual component meshes.

Radial-basis deformation additionally solves a dense control system with cubic cost in n_controls. Both backends use the same differentiable PyTorch solve. implementation="warp" selects a fused Warp evaluator for the point/control evaluation phase. The checked coefficient solve is not supported inside CUDA Graph capture. Use torch.compile() when compiled execution is needed.

For connectivity-preserving object APIs, use displace(), morph(), or radial_basis_function_deform(). A DomainMesh provides morph() and radial_basis_function_deform() for shared sparse fields across all components.

Lattice Free-Form Deformation#

physicsnemo.nn.functional.free_form_deform_points(
points: Float[Tensor, '*batch num_points num_dims'],
control_displacements: Float[Tensor, '*lattice_resolution num_dims'],
*,
origin: Float[Tensor, '*box_batch num_dims'] | Sequence[float],
extent: Float[Tensor, '*box_batch num_dims'] | Sequence[float],
basis: Literal['bernstein', 'bspline', 'linear', 'cubic_hermite', 'quintic_hermite'] = 'bernstein',
point_weights: Bool[Tensor, '*batch num_points'] | Float[Tensor, '*batch num_points'] | None = None,
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '*batch num_points num_dims']#

Deform points with a control lattice by free-form deformation.

An \(n_1 \times \dots \times n_D\) array of control displacements defines a field over the axis-aligned box with corner origin \(o\) and edge lengths extent \(e\). For each point inside the box, define the local coordinates as

\[u_d = \frac{x_d - o_d}{e_d} \in [0, 1].\]

The tensor-product interpolation of the control displacements \(\Delta P\) gives the displacement:

\[d(x) = \sum_{i_1=0}^{n_1-1} \dots \sum_{i_D=0}^{n_D-1} \left[\prod_{d=1}^{D} b_{i_d}(u_d)\right] \Delta P_{i_1 \dots i_D}, \qquad x' = x + \mathtt{point\_weights}\,d(x).\]

Points outside the box remain unchanged. All supported bases form a partition of unity, so a lattice of zero displacements is exactly the identity and a constant lattice translates every point inside the box.

basis="bernstein" provides the classic free-form deformation basis from Sederberg and Parry (1986) [1]. Its per-axis functions are the Bernstein polynomials of degree \(p_d = n_d - 1\):

\[b_i(u) = \binom{p_d}{i} u^i (1-u)^{p_d - i},\]

Every lattice node influences every point in the box. basis="bspline" uses uniform cubic B-splines. The axis is divided into \(n_d - 3\) knot spans, and only the four coefficients around the containing span influence a point, independent of the lattice resolution. Coefficient index \(i\) is associated with the Greville coordinate \((i - 1) / (n_d - 3)\). The first and last coefficient planes lie one knot spacing outside the evaluation box.

basis="linear", "cubic_hermite", and "quintic_hermite" are local, node-interpolating alternatives. Each axis is divided into \(n_d - 1\) cells, and only the two nodes bracketing a point contribute along that axis. If \(t \in [0, 1]\) is the coordinate within a cell, the upper-node weights are \(t\), \(3t^2 - 2t^3\), and \(t^3(6t^2 - 15t + 10)\), respectively. The lower-node weight is one minus the upper-node weight. The resulting fields are C0, C1, and C2 across cell boundaries, respectively. Perlin introduced the quintic polynomial in “Improving Noise” [2] to remove the cubic polynomial’s second-derivative discontinuities.

Inputs may be unbatched (points of shape (N, D)) or batched ((B, N, D)). The corresponding control lattice has shape (n_1, ..., n_D, D) or (B, n_1, ..., n_D, D). Points, control lattices, and point weights must be either unbatched or batch-aligned. They are not broadcast. For batched inputs, (D,) box vectors are shared across the batch, while (B, D) box vectors are aligned. Float32 and float64 are supported.

Parameters:

  • points (torch.Tensor) – Query points with shape (N, D) or (B, N, D).

  • control_displacements (torch.Tensor) – Displacement vectors, not destination coordinates, for every lattice node, with shape (n_1, ..., n_D, D) or (B, n_1, ..., n_D, D) and the same dtype and device as points. Each axis needs at least two nodes for "bernstein" and the node-interpolating bases, and four for "bspline". For "bernstein" and the interpolating bases, node (i_1, ..., i_D) sits at coordinate origin_d + extent_d * i_d / (n_d - 1) along each axis. Uniform cubic B-spline coefficient i_d is associated with origin_d + extent_d * (i_d - 1) / (n_d - 3) and is generally not interpolated by the field.

  • origin (torch.Tensor or sequence of float) – Minimum corner of the lattice box with shape (D,) or aligned batched shape (B, D). Tensor values must match the point dtype and device. The operation uses tensor values without runtime validation. It converts and validates sequences on the host. For repeated GPU calls, create device tensors once and reuse them.

  • extent (torch.Tensor or sequence of float) – Edge lengths of the lattice box with the same accepted shapes as origin. Every value must be strictly positive and finite. The operation does not validate tensor values at runtime.

  • basis ({"bernstein", "bspline", "linear", "cubic_hermite", "quintic_hermite"}, optional) – Per-axis basis family. "bernstein" provides classic global-support FFD for coarse lattices. "bspline" uses local four-node-per-axis support and scales to fine lattices. "linear", "cubic_hermite", and "quintic_hermite" are local two-node-per-axis bases that interpolate every lattice node, with progressively smoother transitions. Default is "bernstein".

  • point_weights (torch.Tensor or None, optional) – Optional bool or floating per-point weights. Accepted shapes are (N,) for unbatched inputs and (B, N) for batched inputs. All weights must match the point device. Floating weights must also match the point dtype. Bool values act as hard masks. The operation uses values as supplied without clamping. Default is None.

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU. On CUDA, it selects Warp when available and otherwise Torch with a one-time RuntimeWarning.

Returns:

Deformed points with the same shape, dtype, and device as points.

Return type:

torch.Tensor

Raises:

  • TypeError – If tensor dtypes or Python argument types are unsupported.

  • ValueError – If tensor shapes, devices, lattice parameters, point weights, or basis are invalid.

  • KeyError – If implementation does not name a registered backend.

  • ImportError – If an explicitly requested backend is unavailable.

Notes

Both backends propagate first-order gradients through points, control displacements, and floating point weights. origin and extent are non-differentiable lattice parameters. They must not require gradients. Only first-order gradients are part of the Warp backend’s public contract.

The deformation is generally not continuous across the box boundary. An outside point stays fixed while a neighboring inside point moves. To keep the exterior fixed, set the outermost coefficient plane on every Bernstein or node-interpolating face to zero. For cubic B-splines, set the first and last three coefficient planes on every axis to zero because three planes have nonzero weight at each box face.

Bernstein degree, global support, and evaluation cost grow with the lattice resolution. Use "bspline" for fine lattices or local control. The lattice resolution is a static parameter under torch.compile(). Each distinct resolution compiles its own graph. Eager Torch evaluation chunks query points to keep estimated live FFD temporaries within 256 MiB. Compiled Torch evaluation instead uses one vectorized block because symbolic chunk loops cannot be unrolled, so it does not enforce that budget. Very large Bernstein workloads may therefore require substantially more peak memory when compiled.

References

[1] Sederberg, T. W., and Parry, S. R. (1986). “Free-Form Deformation of Solid Geometric Models.” ACM SIGGRAPH Computer Graphics, 20(4), 151-160. https://doi.org/10.1145/15886.15903

[2] Perlin, K. (2002). “Improving Noise.” ACM Transactions on Graphics, 21(3), 681-682. https://doi.org/10.1145/566654.566636

import torch
from physicsnemo.nn.functional import free_form_deform_points

points = torch.rand(1024, 3)
control_displacements = torch.zeros(4, 4, 4, 3, requires_grad=True)
origin = points.new_zeros(3)
extent = points.new_ones(3)

deformed = free_form_deform_points(
    points,
    control_displacements,
    origin=origin,
    extent=extent,
    basis="bernstein",
)
deformed.square().mean().backward()

With zero control displacements, the operation is exactly the identity, so a lattice initialized at zero is a well-behaved starting point for shape optimization. An optimizer, or a model that produces the lattice displacements, learns the deformation from a differentiable objective on deformed.

For repeated GPU calls, create origin and extent once as device tensors, as shown in the example. Python sequences are convenient for one-off calls. Each invocation with sequence inputs creates and transfers new tensors.

Choosing a basis:

  • "bernstein" provides classic free-form deformation. Every lattice node influences every point in the box, which suits coarse design lattices. The polynomial degree, global support, and evaluation cost grow with the resolution.

  • "bspline" uses uniform cubic B-splines with local four-node-per-axis support. The per-point cost is independent of the lattice resolution, so it scales to fine lattices for local sculpting and registration-style deformation. Along an axis with n coefficients, index i is associated with the Greville coordinate (i - 1) / (n - 3). The first and last coefficient planes therefore lie one knot spacing outside the evaluation box.

  • "linear", "cubic_hermite", and "quintic_hermite" use the two neighboring lattice nodes per axis and exactly reproduce every control-node displacement. "linear" is piecewise multilinear and C0 across cell boundaries. The cubic and quintic Hermite variants are C1 and C2, respectively. These modes suit design parameters whose values must be attained at the lattice nodes.

For "bernstein", the evaluation cost is proportional to batch_size * n_points * prod(resolution) * n_spatial_dims. For "bspline", it is proportional to batch_size * n_points * 4**n_spatial_dims * n_spatial_dims. The node-interpolating modes use 2**n_spatial_dims controls per point. Points outside the lattice box remain unchanged. A sufficient condition for continuity with a fixed exterior is to zero the outermost coefficient plane on every Bernstein or node-interpolating face. For cubic B-splines, zero the first and last three coefficient planes on every axis because three planes have nonzero weight at each box face.

Eager Torch evaluation chunks query points to keep estimated live FFD temporaries within 256 MiB. Under torch.compile(), the Torch backend uses one vectorized block because symbolic chunk loops cannot be unrolled. The eager memory budget is therefore not enforced. Very large Bernstein workloads may require substantially more peak memory when compiled.

For connectivity-preserving object APIs, use free_form_deform() or free_form_deform().

Deformation Energies#

The deformation-energy functionals compare current point coordinates with a fixed reference configuration and topology. They return differentiable penalty objectives for use in an optimization loop. They do not solve a constrained deformation problem or enforce hard constraints.

For a Mesh, the wrappers in physicsnemo.mesh.deformation validate and cache connectivity and build triangle hinges automatically.

physicsnemo.nn.functional.simplex_strain_energy(
points: Float[Tensor, '*batch num_points num_dims'],
reference_points: Float[Tensor, '*batch num_points num_dims'],
cells: Int[Tensor, 'num_cells simplex_vertices'],
*,
lame_lambda: float = 1.0,
shear_modulus: float = 1.0,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Evaluate reference-integrated St. Venant–Kirchhoff simplex energy.

For each edge, triangle, or tetrahedron, an orthonormalized reference edge frame is used to express the right Cauchy–Green tensor \(C\) intrinsically. With Green–Lagrange strain \(\varepsilon=(C-I)/2\), the per-simplex energy is

\[V_0\left[\mu\lVert\varepsilon\rVert_F^2 + \frac{\lambda}{2}\operatorname{tr}(\varepsilon)^2\right].\]

The formulation supports one-, two-, and three-dimensional simplices embedded in any coordinate dimension supported by the selected backend. It is invariant to rigid transformations and reflections. Add simplex_inversion_energy() when element orientation must be preserved.

Parameters:

  • points (torch.Tensor) – Current coordinates with shape (N, D) or (B, N, D).

  • reference_points (torch.Tensor) – Rest coordinates with exactly the same shape, dtype, and device as points. Batches are aligned and are not implicitly broadcast.

  • cells (torch.Tensor) – Shared simplex connectivity with shape (M, m + 1), where m is 1, 2, or 3 and D >= m. Values must be distinct, valid in-range point indices. Int32 and int64 are accepted.

  • lame_lambda (float, optional) – First Lamé parameter. It may be negative for a stable auxetic material, but lame_lambda + 2 * shear_modulus / m must be nonnegative. Default is 1.0.

  • shear_modulus (float, optional) – Nonnegative shear modulus. Default is 1.0.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns one term per simplex. "sum" returns the physically integrated total. "mean" averages all batch/simplex terms. Empty inputs have zero sum and mean. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and Warp on CUDA when Warp is available and D <= 3. Higher coordinate dimensions use Torch.

Returns:

Per-simplex terms or a scalar according to reduction.

Return type:

torch.Tensor

Notes

This function returns an energy term for use in an objective. It does not run an optimizer or impose a hard constraint. Degenerate reference cells produce NaN so invalid rest geometry is not silently regularized. Both coordinate tensors are differentiable. Connectivity is discrete and is not differentiable. Torch supports higher-order derivatives. The Warp backend’s contract is first-order differentiation.

physicsnemo.nn.functional.simplex_measure_energy(
points: Float[Tensor, '*batch num_points num_dims'],
reference_points: Float[Tensor, '*batch num_points num_dims'],
cells: Int[Tensor, 'num_cells simplex_vertices'],
*,
target_ratio: float = 1.0,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Penalize local simplex measure changes relative to rest geometry.

Each term is \(\tfrac12 V_0(q-q_*)^2\), where target_ratio is \(q_*\). For a full-dimensional simplex, \(q=\det(E)/\det(E_0)\) is signed, so a reflected cell does not falsely satisfy the target. For an embedded simplex, \(q=\sqrt{\det(G)/\det(G_0)}\) is the unsigned intrinsic length, area, or volume ratio.

Use total_measure_energy() when local redistribution is allowed and only the aggregate measure should remain near its target.

Parameters:

  • points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • reference_points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • cells (torch.Tensor) – Shared (M, m + 1) connectivity for edges, triangles, or tetrahedra. Values must be distinct, valid in-range indices. Int32 and int64 are accepted.

  • target_ratio (float, optional) – Strictly positive target measure relative to the reference. Default is 1.0.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns (M,) or (B, M). "sum" and "mean" reduce every term to a scalar. Empty inputs have zero sum and mean. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and, when available, Warp on CUDA for coordinate dimensions up to three. Higher-dimensional CUDA inputs use Torch. Explicitly selecting Warp for them is rejected.

Returns:

Per-simplex terms or a scalar according to reduction.

Return type:

torch.Tensor

Notes

This is a penalty energy, not an exact constraint. Reference-degenerate cells produce NaN. Both current and reference coordinates are differentiable. Connectivity is not. At exact collapse, unsigned embedded measure uses a zero current-coordinate subgradient, so this term cannot by itself reopen an already collapsed embedded cell. Warp provides first-order gradients.

physicsnemo.nn.functional.total_measure_energy(
points: Float[Tensor, '*batch num_points num_dims'],
reference_points: Float[Tensor, '*batch num_points num_dims'],
cells: Int[Tensor, 'num_cells simplex_vertices'],
*,
target_ratio: float = 1.0,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Penalize change in the total simplex measure.

The energy for each batch item is \(\tfrac12 V_{0,\mathrm{tot}}(V/V_{0,\mathrm{tot}}-q_*)^2\). Unlike simplex_measure_energy(), local cells may expand and contract while their aggregate measure remains at the target. Full-dimensional cell contributions retain their orientation relative to the corresponding rest cell. Embedded-simplex contributions are unsigned.

Parameters:

  • points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • reference_points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • cells (torch.Tensor) – Nonempty shared (M, m + 1) simplex connectivity. Values must be distinct, valid in-range indices. Int32 and int64 are accepted.

  • target_ratio (float, optional) – Strictly positive target total-measure ratio. Default is 1.0.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns a scalar for unbatched input or (B,) for batched input. "sum" and "mean" reduce across the batch. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and, when available, Warp on CUDA for coordinate dimensions up to three. Higher-dimensional CUDA inputs use Torch. Explicitly selecting Warp for them is rejected.

Returns:

One term per batch item or a scalar according to reduction.

Return type:

torch.Tensor

Notes

This is a soft aggregate constraint. It does not prevent individual cells from collapsing or inverting. Combine it with strain and inversion terms when local validity matters. Empty connectivity is rejected, and a zero or invalid total reference measure produces NaN. At exact collapse, unsigned embedded measure uses a zero current-coordinate subgradient.

physicsnemo.nn.functional.simplex_inversion_energy(
points: Float[Tensor, '*batch num_points num_dims'],
reference_points: Float[Tensor, '*batch num_points num_dims'],
cells: Int[Tensor, 'num_cells simplex_vertices'],
*,
minimum_jacobian: float = 0.1,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Penalize full-dimensional cells below a signed Jacobian threshold.

Each term is \(\tfrac12 V_0\max(0,J_{\min}-J)^2\), where \(J=\det(E)/\det(E_0)\). It is zero above minimum_jacobian and increases quadratically through collapse and inversion.

Parameters:

  • points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • reference_points (torch.Tensor) – Current and rest coordinates with identical (N, D) or (B, N, D) shape, dtype, and device.

  • cells (torch.Tensor) – Shared full-dimensional simplex connectivity: edges in 1D, triangles in 2D, or tetrahedra in 3D. Values must be distinct, valid in-range indices. Int32 and int64 are accepted.

  • minimum_jacobian (float, optional) – Nonnegative signed-Jacobian threshold. Default is 0.1.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns one term per cell. "sum" and "mean" reduce all terms to a scalar. Empty inputs have zero sum and mean. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and Warp on CUDA.

Returns:

Per-cell penalties or a scalar according to reduction.

Return type:

torch.Tensor

Notes

The squared hinge is continuously differentiable, but its second derivative is discontinuous at the threshold. This energy discourages inversion but is not an infinite barrier and cannot guarantee validity under a finite penalty. Embedded simplices are outside this determinant-based term’s domain and are rejected.

physicsnemo.nn.functional.surface_bending_energy(
points: Float[Tensor, '*batch num_points 3'],
reference_points: Float[Tensor, '*batch num_points 3'],
hinges: Int[Tensor, 'num_hinges 4'],
*,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Evaluate reference-relative discrete bending over triangle hinges.

A hinge row (i, j, k, l) denotes the consistently oriented adjacent faces (i, j, k) and (j, i, l). If \(\theta\) and \(\theta_0\) are their signed current and reference dihedral angles, each term is

\[\frac12\frac{\ell_0^2}{A_{0,L}+A_{0,R}} \operatorname{wrap}(\theta-\theta_0)^2.\]
Parameters:

  • points (torch.Tensor) – Current and rest coordinates with identical (N, 3) or (B, N, 3) shape, dtype, and device.

  • reference_points (torch.Tensor) – Current and rest coordinates with identical (N, 3) or (B, N, 3) shape, dtype, and device.

  • hinges (torch.Tensor) – Shared oriented hinge connectivity with shape (H, 4). Values must be distinct, valid in-range indices. Int32 and int64 are accepted. Construct hinges from unique triangles with at most two incident triangles per edge, and omit boundary edges. Build them once from fixed topology and reuse them during optimization. The physicsnemo.mesh.deformation.surface_bending_energy() wrapper validates these conditions, constructs the hinges, and caches them.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns (H,) or (B, H). "sum" and "mean" return a scalar. A surface without interior hinges has zero sum and mean. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and Warp on CUDA.

Returns:

Per-hinge terms or a scalar according to reduction.

Return type:

torch.Tensor

Notes

This is a geometric thin-shell regularizer, not a calibrated physical shell model: material thickness and constitutive scaling are not included. Degenerate current or reference hinge triangles produce NaN. Dihedral wrapping is continuous except at the equivalent-angle branch cut. Both coordinate tensors are differentiable. The discrete hinge construction is not. Warp provides first-order gradients.

physicsnemo.nn.functional.closed_surface_volume_energy(
points: Float[Tensor, '*batch num_points 3'],
reference_points: Float[Tensor, '*batch num_points 3'],
triangles: Int[Tensor, 'num_triangles 3'],
*,
target_ratio: float = 1.0,
reduction: Literal['none', 'sum', 'mean'] = 'sum',
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '...']#

Penalize enclosed-volume change of an oriented closed triangle surface.

With \(y_i=x_i-o\) for one common per-batch origin, the signed volume is assembled from face contributions \(v_f=y_i\cdot(y_j\times y_k)/6\). The closed-surface sum is independent of \(o\). For each batch item, the energy is

\[\frac12 |V_0|\left(V/V_0-q_*\right)^2.\]
Parameters:

  • points (torch.Tensor) – Current and rest coordinates with identical (N, 3) or (B, N, 3) shape, dtype, and device.

  • reference_points (torch.Tensor) – Current and rest coordinates with identical (N, 3) or (B, N, 3) shape, dtype, and device.

  • triangles (torch.Tensor) – Nonempty shared (F, 3) triangle connectivity for one edge-connected, edge-closed, consistently oriented component. Every edge must have two oppositely directed incident triangles. Values must be distinct, valid in-range indices. Int32 and int64 are accepted.

  • target_ratio (float, optional) – Strictly positive target enclosed-volume ratio. Default is 1.0.

  • reduction ({"none", "sum", "mean"}, optional) – "none" returns a scalar for unbatched input or (B,) for batched input. "sum" and "mean" reduce across the batch. Default is "sum".

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend. None selects Torch on CPU and Warp on CUDA.

Returns:

One volume term per batch item or a scalar according to reduction.

Return type:

torch.Tensor

Notes

This function supplies a penalty term, not an exact volume constraint. It assumes the stated connectivity, edge-closure, and orientation contract. Checking those discrete properties belongs outside an iterative optimization loop. Evaluate disconnected components separately. An empty triangle set is rejected. Zero or nonfinite reference volume produces NaN. Vertex-manifoldness and self-intersection are not checked, and local face inversion can cancel in the signed total. Combine this term with local validity energies when needed.

from physicsnemo.nn.functional import (
    simplex_inversion_energy,
    simplex_strain_energy,
    total_measure_energy,
)

loss = (
    simplex_strain_energy(points, reference_points, cells)
    + 40.0 * total_measure_energy(points, reference_points, cells)
    + 10.0 * simplex_inversion_energy(points, reference_points, cells)
)
loss.backward()

The simplex functionals accept unbatched (n_points, n_spatial_dims) or batched (batch_size, n_points, n_spatial_dims) coordinates with shared integer topology. Coordinates must use matching torch.float32 or torch.float64 dtypes. reduction="none" returns one value per simplex or hinge. For total_measure_energy and closed_surface_volume_energy, it returns one global value per batch item. "sum" and "mean" reduce all values to one scalar.

simplex_measure_energy constrains each element separately, whereas total_measure_energy permits local redistribution and constrains only the sum. Full-dimensional measure ratios retain the orientation relative to each reference simplex. Embedded-simplex ratios are unsigned. The St. Venant–Kirchhoff strain formulation used by simplex_strain_energy is reflection-blind. Add simplex_inversion_energy when full-dimensional simplex orientation matters. Closed-surface volume requires one edge-connected, edge-closed, consistently oriented 3D triangle surface. The low-level functional assumes that contract without checking it. Surface bending is a geometric hinge regularizer rather than a material shell model.

The tensor functionals validate topology shape, dtype, and device without scanning index values, which would synchronize a CUDA device on every call. Indices must therefore be distinct and in range as documented. Invalid index values are outside the tensor API contract. Use the mesh wrappers when cached value validation is needed. Direct tensor calls accept int32 or int64 connectivity, but normalize int32 connectivity on every call. Use int64 for repeated direct calls. Mesh wrappers cache this normalization with the topology.

Torch provides higher-order derivatives. Warp provides a first-order CUDA path. Because its backward uses atomic accumulation at shared vertices, Warp gradient results can have small run-to-run floating-point differences.

Nearest-Surface Shrinkwrap#

Shrinkwrap projects each source point to the nearest location on a triangle target surface. Optional point weights scale the displacement between the source and projected positions. A signed offset moves the projection along the selected target face normal.

physicsnemo.nn.functional.shrinkwrap_points(
points: Float[Tensor, '*batch num_points 3'],
target_points: Float[Tensor, 'num_target_points 3'],
target_faces: Int[Tensor, 'num_target_faces 3'],
*,
offset: float | Float[Tensor, ''] = 0.0,
max_distance: float | None = None,
point_weights: Bool[Tensor, '*batch num_points'] | Float[Tensor, '*batch num_points'] | None = None,
implementation: Literal['warp', 'torch'] | None = None,
) → Float[Tensor, '*batch num_points 3']#

Project points onto the nearest locations of a triangle surface.

For every source point \(x_i\), the operation selects the closest location \(p_i\) on the target triangle surface and applies

\[x'_i = x_i + w_i\left(p_i + \delta n_i - x_i\right),\]

where \(w_i\) is an optional point weight, \(\delta\) is offset, and \(n_i\) is the oriented unit normal of the selected target face. Positive offsets follow the target winding. Points with no target triangle within max_distance remain unchanged.

Source points may have shape (N, 3) or (B, N, 3). A single target surface is shared by every source batch. Float32 and float64 are supported.

Parameters:

  • points (torch.Tensor) – Source coordinates with shape (N, 3) or (B, N, 3).

  • target_points (torch.Tensor) – Target-surface vertices with shape (M, 3) and the same dtype and device as points.

  • target_faces (torch.Tensor) – Target triangle connectivity with shape (F, 3), dtype torch.int32 or torch.int64, and the same device as points. Out-of-range indices raise an error. Triangles with nonfinite coordinates or degenerate geometry are ignored, but at least one finite nondegenerate triangle is required.

  • offset (float or torch.Tensor, optional) – Signed scalar offset along the selected target face normal. A scalar tensor must match the source dtype and device and may require gradients. Default is 0.0.

  • max_distance (float or None, optional) – Positive finite search radius measured to the un-offset target. A point exactly at the cutoff is left unchanged. None performs an unbounded nearest-surface search. Values that round to zero in the source dtype are rejected. Default is None.

  • point_weights (torch.Tensor or None, optional) – Optional bool or floating source-point weights with shape (N,) or (B, N). Zero leaves a point unchanged and one applies the full projection. Floating values are not clamped. Default is None.

  • implementation ({"torch", "warp"} or None, optional) – Explicit backend for nearest-face search. None selects Torch on CPU and Warp on CUDA when available.

Returns:

Shrinkwrapped points with the source shape, dtype, and device.

Return type:

torch.Tensor

Notes

Nearest-face selection, closest-feature changes, and max_distance gating are discrete. With those choices fixed, both backends propagate first-order gradients through source points, selected target vertices, floating point weights, and tensor-valued offset. Ties and transitions between target faces, edges, and vertices are nonsmooth.

The Warp backend accelerates only the discrete nearest-face search. Both backends replay the selected point-to-triangle projection with Torch in the original dtype, which supplies identical piecewise geometry derivatives. Float64 targets use the Torch search because Warp searches in float32. Safe float32 coordinates are searched unchanged. Warp falls back to Torch for unsafe coordinate magnitudes or face geometry. Target topology is non-differentiable. Nonzero offsets require consistently oriented target faces. At shared features, the selected face determines the normal. The operation does not prevent source-cell inversion or self-intersection. Shrinkwrap is not supported inside CUDA Graph capture with either backend.

import torch
from physicsnemo.nn.functional import shrinkwrap_points

target_points = torch.tensor(
    [
        [0.0, 0.0, 0.0],
        [1.0, 0.0, 0.0],
        [1.0, 1.0, 0.0],
        [0.0, 1.0, 0.0],
    ],
    requires_grad=True,
)
target_faces = torch.tensor([[0, 1, 2], [0, 2, 3]])
points = torch.tensor(
    [[0.25, 0.25, 0.30], [0.75, 0.75, 0.45]],
    requires_grad=True,
)
weights = torch.tensor([0.5, 1.0], requires_grad=True)
offset = torch.tensor(0.02, requires_grad=True)

wrapped = shrinkwrap_points(
    points,
    target_points,
    target_faces,
    point_weights=weights,
    offset=offset,
)
wrapped.square().mean().backward()

Source points use one of these shapes:

  • (N, 3)

  • (B, N, 3)

One triangle target is shared across the source batches. Keep the following projection rules in mind:

  • A positive offset follows target face winding.

  • A finite max_distance leaves points at or beyond the cutoff unchanged.

Nearest-face selection, closest-feature changes, and distance gating are discrete. Between those transitions, gradients propagate through the following values:

  • Source points

  • Selected target vertices

  • Floating-point weights

  • A tensor-valued offset

Target connectivity is not differentiable.

Torch provides the reference nearest-face search. Warp accelerates that search on CPU and CUDA. Both backends evaluate the selected projection with PyTorch in the input dtype. Exact ties at target edges or vertices can select different adjacent faces. Their closest point is the same, but a nonzero face-normal offset can differ.

Search backend rules differ by dtype and geometry:

  • float64 targets use the Torch search because Warp searches in float32.

  • Warp searches safe float32 coordinates unchanged.

  • Warp falls back to Torch for unsafe coordinate magnitudes or face geometry.

Shrinkwrap performs data-dependent validation and nearest-face search setup. Do not run CUDA executions with either backend inside CUDA Graph capture.

For a connectivity-preserving object API and a triangle-panel example, use shrinkwrap().

Mesh Poisson Disk Sample#

physicsnemo.nn.functional.mesh_poisson_disk_sample(
mesh_vertices: Tensor,
mesh_indices: Tensor,
min_distance: float = 0.02,
per_vertex_radius: Tensor | None = None,
batch_size: int = 131072,
max_points: int = 2000000,
max_iterations: int = 64,
random_seed: int = 42,
hash_grid_resolution: int | Sequence[int] | Tensor = 128,
mode: str = 'dart_throwing',
target_num_points: int | None = None,
*,
implementation: Literal['warp'] | None = None,
) → Tensor#

Generate Poisson-disk samples on a triangle mesh surface with Warp.

This functional supports two sampling modes on triangle meshes:

  1. dart_throwing: iterative parallel dart throwing where each iteration draws area-weighted candidates, rejects points near accepted samples, resolves candidate-candidate conflicts with random-priority MIS, and commits survivors.

  2. weighted_sample_elimination: builds an oversampled Poisson-quality pool, then downsamples to target_num_points using a radius-aware elimination pass.

Both modes produce blue-noise-like sample sets. dart_throwing emphasizes throughput and minimum-distance control; weighted_sample_elimination emphasizes distribution quality at a fixed output count.

Parameters:

  • mesh_vertices (torch.Tensor) – Mesh vertex positions with shape (n_vertices, 3).

  • mesh_indices (torch.Tensor) – Triangle connectivity in shape (n_faces, 3) or flattened shape (3 * n_faces,).

  • min_distance (float, optional) – Minimum Poisson distance for constant-radius mode. Default is 0.02. In weighted_sample_elimination mode this is treated as a lower-bound hint while the algorithm primarily targets target_num_points quality.

  • per_vertex_radius (torch.Tensor | None, optional) – Optional adaptive radius with shape (n_vertices,). If provided, candidate radius is barycentrically interpolated.

  • mode (str, optional) – Sampling mode. "dart_throwing" uses iterative parallel dart throwing. "weighted_sample_elimination" builds an oversampled Poisson pool and then downsamples to target_num_points with radius-aware elimination.

  • batch_size (int, optional) – Number of generated candidates per iteration. Default is 131072.

  • max_points (int, optional) – Maximum number of accepted samples. Default is 2_000_000. For mode="weighted_sample_elimination", this is also the default target_num_points when that argument is omitted.

  • target_num_points (int | None, optional) – Number of output points for mode="weighted_sample_elimination". If None, the mode uses max_points.

  • max_iterations (int, optional) – Iteration cap for the sampler. Default is 64.

  • random_seed (int, optional) – Base random seed for deterministic candidate generation.

  • hash_grid_resolution (int | Sequence[int], optional) – Hash-grid resolution, either scalar or (nx, ny, nz). Default is 128.

  • implementation (str | None, optional) – Explicit implementation name. Defaults to dispatch behavior.

Returns:

Accepted sample positions with shape (n_samples, 3) and dtype torch.float32.

Return type:

torch.Tensor

Notes

  • mode="weighted_sample_elimination" uses Warp kernels and follows Open3D’s Yuksel-style weighting equations.

  • per_vertex_radius is ignored in weighted elimination mode.

  • The output order is implementation-specific and not semantically meaningful.

Visualization

This visualization compares Poisson samples generated by dart_throwing and weighted_sample_elimination on the same Stanford Bunny surface mesh.

Rotating Mesh Poisson disk sampling mode comparison

Mesh To Voxel Fraction#

physicsnemo.nn.functional.mesh_to_voxel_fraction(
mesh_vertices: Tensor,
mesh_indices: Tensor,
origin: Tensor | Sequence[float],
voxel_size: float,
grid_dims: Sequence[int] | Tensor,
n_samples: int = 64,
seed: int = 42,
open_mesh: bool = False,
winding_number_threshold: float = 0.5,
winding_number_accuracy: float = 2.0,
*,
implementation: Literal['warp'] | None = None,
) → Tensor#

Compute mesh-voxel volume fractions on a regular 3D grid.

This functional estimates the fraction of each voxel that lies inside a triangle mesh using Warp kernels and Monte Carlo sampling.

For each voxel, it first performs an AABB-overlap query with mesh triangles. If no triangles overlap the voxel, it classifies only the voxel center as inside or outside. If triangles overlap, it uniformly samples points inside the voxel and estimates the occupancy fraction:

\[f_{ijk} \approx \frac{1}{N_s}\sum_{s=1}^{N_s}\mathbb{1}\left(x_s \in \Omega\right),\]

where \(N_s\) is n_samples and \(\Omega\) is the mesh interior.

Parameters:

  • mesh_vertices (torch.Tensor) – Vertex positions with shape (n_vertices, 3).

  • mesh_indices (torch.Tensor) – Triangle connectivity as shape (n_faces, 3) or flattened shape (3 * n_faces,).

  • origin (torch.Tensor | Sequence[float]) – Lower corner of the voxel grid as a length-3 vector.

  • voxel_size (float) – Edge length of each cubic voxel.

  • grid_dims (Sequence[int]) – Grid resolution (nx, ny, nz).

  • n_samples (int, optional) – Number of Monte Carlo samples per overlapping voxel. Default is 64.

  • seed (int, optional) – Random seed offset used per voxel. Default is 42.

  • open_mesh (bool, optional) – If True, uses winding-number sign queries for open meshes. Default is False.

  • winding_number_threshold (float, optional) – Winding-number threshold used when open_mesh=True.

  • winding_number_accuracy (float, optional) – Winding-number query accuracy used when open_mesh=True.

  • implementation (str | None, optional) – Explicit backend selection. Defaults to dispatch behavior.

Returns:

Volume fractions in [0, 1] with shape (nz, ny, nx) and dtype torch.float32.

Return type:

torch.Tensor

Notes

  • This functional provides a Warp implementation.

  • The operation is stochastic over overlapping voxels; use seed for reproducible runs.

Visualization

This visualization shows a side-by-side rotating view of the Stanford Bunny mesh and the occupied voxels inferred by mesh_to_voxel_fraction.

Mesh to voxel fraction mesh and occupied-voxel rotation animation

Surface Remeshing#

physicsnemo.nn.functional.remeshing(
mesh_vertices: Float[Tensor, 'n_vertices 3'],
mesh_indices: Integer[Tensor, 'n_faces 3'],
n_clusters: int,
*,
max_iterations: int = 4,
vertex_density: Tensor | None = None,
search_radius_scale: float = 1.6,
voxel_width_scale: float = 1.15,
hash_grid_resolution: int = 128,
farthest_point_threshold: int = 256,
farthest_point_oversampling: int = 4,
implementation: Literal['warp'] | None = None,
) → tuple[Float[Tensor, 'n_output_vertices 3'], Int[Tensor, 'n_output_faces 3']]#

Remesh a triangle surface represented by tensors.

This low-level functional performs integration-mass-weighted centroidal clustering, projects cluster centers onto the source surface, and reconstructs compact triangle connectivity. The operation is intentionally non-differentiable. Most users should call physicsnemo.mesh.remeshing.remesh(), which accepts and returns physicsnemo.mesh.Mesh objects.

Parameters:

  • mesh_vertices (torch.Tensor) – Floating-point vertex coordinates with shape (n_vertices, 3).

  • mesh_indices (torch.Tensor) – Integer triangle connectivity with shape (n_faces, 3) on the same device.

  • n_clusters (int) – Target output vertex count between 3 and n_vertices, inclusive.

  • max_iterations (int, optional) – Maximum centroid-relaxation iterations. Default is 4.

  • vertex_density (torch.Tensor or None, optional) – Positive relative integration density with shape (n_vertices,). Larger values allocate more output vertices near the corresponding input vertices. Multiplying the full tensor by a positive constant leaves the objective unchanged. It must use a real floating-point dtype. Default is None for uniform remeshing.

  • search_radius_scale (float, optional) – Base hash-grid query radius relative to sqrt(surface_area / n_clusters). Nonuniform density automatically enlarges the effective radius to cover wider low-density spacing. Default is 1.6.

  • voxel_width_scale (float, optional) – Spatial-stratification voxel width relative to sqrt(surface_area / n_clusters) for the large uniform initializer. It does not affect density-aware initialization. Default is 1.15.

  • hash_grid_resolution (int, optional) – Resolution of each axis of the centroid hash grid. Must be at most 256, which bounds its two dense cell-offset arrays to 128 MiB. Default is 128.

  • farthest_point_threshold (int, optional) – Use farthest-point initialization when n_clusters is at most this value. Set to 0 to disable farthest-point initialization. Default is 256.

  • farthest_point_oversampling (int, optional) – Integration-mass-weighted farthest-point candidate-pool size as a multiple of n_clusters. Default is 4.

  • implementation ({"warp"} | None, optional) – Explicit backend selection. Only "warp" is currently available.

Returns:

Remeshed vertices and triangle indices. Vertex dtype and device match mesh_vertices. Indices use torch.int64 on the same device.

Return type:

tuple[torch.Tensor, torch.Tensor]

Raises:

  • TypeError – If tensor or scalar inputs have invalid types.

  • ValueError – If tensor shapes, devices, counts, geometry, or tuning values are invalid.

  • KeyError – If implementation does not name a registered backend.

  • ImportError – If Warp is unavailable.

  • RuntimeError – If topology reconstruction cannot produce a nonempty manifold surface.

Notes

Remeshing is intentionally non-differentiable. Warp computes in centered and scaled coordinates in float32, then restores the input vertex dtype and coordinate frame. Centroid sampling uses a fixed random seed, although floating-point atomics can still introduce small run-to-run differences. Spatial clustering can weld sheets or thin features separated by less than the mean cluster spacing. Projection can also map distinct centroids to the same surface position. Output vertices are not welded by position.

Ray Mesh Intersect#

physicsnemo.nn.functional.ray_mesh_intersect(
mesh_vertices: Tensor,
mesh_indices: Tensor,
ray_origins: Tensor,
ray_directions: Tensor,
max_distance: float = 100000000.0,
warp_mesh: Mesh | None = None,
return_warp_mesh: bool = False,
*,
implementation: Literal['warp'] | None = None,
) → tuple[Tensor, Tensor, Tensor, Tensor, Tensor] | tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Mesh]#

Intersect rays with a triangle mesh using Warp.

ray_mesh_intersect builds a Warp Mesh acceleration structure from triangle vertices and indices, casts each input ray against the mesh, and returns the closest hit within max_distance. Ray directions do not need to be normalized; the Warp implementation normalizes them before querying so returned hit distances are expressed in mesh-space length units.

Parameters:

  • mesh_vertices (torch.Tensor) – Mesh vertex positions with shape (num_vertices, 3).

  • mesh_indices (torch.Tensor) – Triangle connectivity with shape (num_faces, 3) or a flattened equivalent.

  • ray_origins (torch.Tensor) – Ray origins with shape (..., 3).

  • ray_directions (torch.Tensor) – Ray directions with the same shape as ray_origins.

  • max_distance (float, optional) – Maximum ray distance. Default is 1e8.

  • warp_mesh (wp.Mesh | None, optional) – Prepared Warp mesh returned by an earlier ray_mesh_intersect call with return_warp_mesh=True. If provided, the mesh tensors are not used to rebuild a Warp Mesh.

  • return_warp_mesh (bool, optional) – If True, append the Warp Mesh used for the query to the output tuple so it can be passed back through warp_mesh on later calls.

  • implementation (str, optional) – Explicit implementation name. Currently only "warp" is registered.

Returns:

By default, a tuple (hit_mask, hit_distance, hit_points, face_ids, hit_normals). If return_warp_mesh=True, the returned tuple is (hit_mask, hit_distance, hit_points, face_ids, hit_normals, warp_mesh). Missed rays have False in hit_mask, infinite hit_distance, zero hit_points and hit_normals, and -1 face_ids.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Notes

hit_normals are the mesh query normals returned by Warp. They preserve the mesh winding orientation and are not flipped to face the incoming ray. For repeated queries against a static mesh, call with return_warp_mesh=True once and pass the returned Warp mesh back through warp_mesh on later calls.

Visualization

This visualization shows a batch of rays intersecting a triangulated sphere, with hits, misses, hit points, and surface normals.

Ray mesh intersection overview with rays, hit points, and normals

Signed Distance Field#

physicsnemo.nn.functional.signed_distance_field(
mesh_vertices: Float[Tensor, 'num_vertices 3'],
mesh_indices: Tensor,
input_points: Float[Tensor, '... 3'],
max_dist: float = 100000000.0,
use_sign_winding_number: bool = False,
*,
implementation: Literal['warp'] | None = None,
) → tuple[Float[Tensor, '...'], Float[Tensor, '... 3'], Int[Tensor, '...']]#

Compute the signed distance field (SDF) for a mesh and query points.

The mesh must be a surface mesh consisting of triangles. This functional uses a Warp-backed implementation for accelerated execution.

Parameters:

  • mesh_vertices (torch.Tensor) – Coordinates of mesh vertices with shape (n_vertices, 3).

  • mesh_indices (torch.Tensor) – Triangle connectivity indexing into mesh_vertices. Expected shape is (n_faces, 3) or a flattened equivalent.

  • input_points (torch.Tensor) – Query points at which to evaluate the signed distance, with shape (..., 3).

  • max_dist (float, optional) – Maximum search distance for closest-point queries. Default is 1e8.

  • use_sign_winding_number (bool, optional) – Whether to use winding-number-based sign computation. Default is False. When False, the mesh should be watertight for reliable signs.

  • implementation (str, optional) – Explicit implementation name. Defaults to None, which uses normal dispatch (currently the Warp implementation).

Returns:

A tuple (sdf, hit_points, hit_faces) where: - sdf contains signed distances at each query point. - hit_points contains the closest point on the mesh for each query. - hit_faces contains the int64 index of the triangle holding each

closest point.

Queries with no triangle within max_dist return NaN for the distance and hit point and -1 for the hit face.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Examples

>>> mesh_vertices = torch.tensor(
...     [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)]
... )
>>> mesh_indices = torch.tensor([(0, 1, 2)])
>>> input_points = torch.tensor([(0.5, 0.5, 0.5)])
>>> sdf, hit_points, hit_faces = signed_distance_field(
...     mesh_vertices, mesh_indices, input_points
... )

Visualization

This visualization shows signed-distance values on a 2D slice through the domain, with the zero level-set contour indicating the implicit surface. The animation shows a sweep plane through the mesh (left) and corresponding SDF slice image (right).

Signed distance field 2D slice visualization
Signed distance field z-slice sweep animation