Closest point projection

[jorgensd]

Would be neat to have. Example use scipy. Will rewrite to hand-written newton.

import numpy as np
from mpi4py import MPI
import ufl
import basix.ufl
import dolfinx
from scipy.optimize import minimize
import numpy.typing as npt


def project_point_to_element(
    mesh: dolfinx.mesh.Mesh,
    cells: npt.NDArray[np.int64],
    target_points: npt.NDArray[np.float64 | np.float32],
    method=None,
    tol: float | None = None,
):
    """
    Projects a 3D point onto a cell in a potentially higher order mesh.

    Args:
        mesh: {py:class}`dolfinx.mesh.Mesh`, the mesh containing the cell.
        cells: {py:class}`numpy.ndarray`, the indices of the cells to project onto.
        target_point: (3,) numpy array, the 3D point to project.
        method: str, the optimization method to use.
        tol: float, the tolerance for the optimizer.

    Returns:
        dict: A dictionary containing the reference coordinates, closest 3D point, and distance.
    """

    # Extract scalar element of mesh
    element = mesh.ufl_domain().ufl_coordinate_element().sub_elements[0]

    # Get the coordinates of the nodes for the specified cell
    node_coords = mesh.geometry.x[mesh.geometry.dofmap[cells]][
        :, :, : mesh.geometry.dim
    ]
    target_points = target_points.reshape(-1, 3)
    # cmap = mesh.geometry.cmap

    # Constraints and Bounds
    cell_type = mesh.topology.cell_type
    if cell_type == dolfinx.mesh.CellType.triangle:
        method = method or "SLSQP"
        constraint = {"type": "ineq", "fun": lambda x: 1.0 - np.sum(x)}
    else:
        method = method or "L-BFGS-B"
        constraint = {}
    bounds = [(0.0, 1.0) for _ in range(mesh.topology.dim)]

    # Set initial guess and tolerance for solver
    initial_guess = np.full(
        mesh.topology.dim, 1 / (mesh.topology.dim + 1), dtype=np.float64
    )
    tol = np.sqrt(np.finfo(mesh.geometry.x.dtype).eps) if tol is None else tol
    closest_points = np.zeros((target_points.shape[0], 3), dtype=mesh.geometry.x.dtype)
    for i, (coord, target_point) in enumerate(zip(node_coords, target_points)):
        coord = coord.reshape(-1, mesh.geometry.dim)

        def S(x_ref):
            N_vals = element.tabulate(0, x_ref.reshape(1, mesh.topology.dim))[0, 0, :]
            return np.dot(N_vals, coord)
            # return cmap.push_forward(x_ref.reshape(1, mesh.topology.dim), coord)[0]

        def dSdx_ref(x_ref):
            """Evaluate jacobian (tangent vectors) at the given reference coordinates."""
            dN = element.tabulate(1, x_ref.reshape(1, mesh.topology.dim))[
                1 : mesh.topology.dim + 1, 0, :
            ]
            return np.dot(dN, coord)

        def objective(x_ref):
            surface_point = S(x_ref)
            diff = surface_point - target_point
            return 0.5 * np.sum(diff**2)

        def objective_grad(x_ref):
            diff = S(x_ref) - target_point
            tangents = dSdx_ref(x_ref)
            return np.dot(tangents, diff)

        res = minimize(
            objective,
            initial_guess,
            method=method,
            jac=objective_grad,
            bounds=bounds,
            constraints=constraint,
            tol=tol,
        )
        closest_points[i] = S(res.x)
        assert res.success, (
            f"Optimization failed for {cells[i]} and {target_point=}: {res.message}"
        )
    return closest_points


# --- Example Usage ---
if __name__ == "__main__":
    # Dummy curved quadratic triangle in 3D (6 nodes)
    curved_nodes = np.array(
        [
            [0.0, 0.0, 0.0],  # Node 0: Vertex
            [1.0, 0.0, 0.0],  # Node 1: Vertex
            [0.0, 1.0, 0.0],  # Node 2: Vertex
            [0.6, 0.6, 0.0],  # Node 4: Edge 1-2
            [0.1, 0.5, 0.2],  # Node 5: Edge 2-0 (curved upward in Z)
            [0.5, 0.1, 0.2],  # Node 3: Edge 0-1 (curved upward in Z)
        ]
    )
    cells = np.array(
        [[0, 1, 2, 3, 4, 5]], dtype=np.int64
    )  # Single curved triangle element
    c_el = ufl.Mesh(basix.ufl.element("Lagrange", "triangle", 2, shape=(3,)))
    mesh = dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells=cells, x=curved_nodes, e=c_el)

    point_to_project = np.array([0.3, 0.3, -0.2])  # A point floating above the element
    import time

    start = time.perf_counter()
    result = project_point_to_element(
        mesh, np.array([0], dtype=np.int32), point_to_project
    )
    end = time.perf_counter()
    print("Closest Point on Element:", result, f"Time taken: {end - start:.6e} seconds")

    with dolfinx.io.VTXWriter(mesh.comm, "closest_point.bp", mesh) as writer:
        writer.write(0.0)

    curved_nodes_tet = np.array(
        [
            [0.0, 0.0, 0.0],  # 0: Vertex
            [1.0, 0.0, 0.0],  # 1: Vertex
            [0.0, 1.0, 0.0],  # 2: Vertex
            [0.0, 0.0, 1.0],  # 3: Vertex
            [0.0, 0.5, 0.5],  # 4: Edge 2-3
            [0.5, 0.0, 0.5],  # 5: Edge 1-3
            [0.5, 0.5, 0.0],  # 6: Edge 1-2
            [0.0, 0.0, 0.5],  # 9: Edge 0-3
            [0.0, 0.5, 0.2],  # 8: Edge 0-2 (Curved)
            [0.5, 0.0, 0.2],  # 7: Edge 0-1 (Curved)
        ],
        dtype=np.float64,
    )
    cells_tet = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], dtype=np.int64)

    domain_tet = ufl.Mesh(basix.ufl.element("Lagrange", "tetrahedron", 2, shape=(3,)))
    mesh_tet = dolfinx.mesh.create_mesh(
        MPI.COMM_WORLD, cells=cells_tet, x=curved_nodes_tet, e=domain_tet
    )
    start = time.perf_counter()
    result = project_point_to_element(
        mesh_tet, np.array([0], dtype=np.int32), point_to_project
    )
    end = time.perf_counter()
    print("Closest Point on Element:", result, f"Time taken: {end - start:.6e} seconds")

    with dolfinx.io.VTXWriter(
        mesh_tet.comm, "closest_point_tet.bp", mesh_tet
    ) as writer:
        writer.write(0.0)

[jorgensd]

Current result:

Closest Point on Element: [[0.21761519 0.21761519 0.17579796]] Time taken: 1.352163e-03 seconds
Closest Point on Element: [[0.47210013 0.47210013 0.0421489 ]] Time taken: 4.841130e-04 seconds

[finsberg]

In other words: Given a point (or a number of points) and a mesh, find the point (or number of points) on the mesh that is closet to the point (number of points). I agree this would be useful. How does this handle cases when the target point is inside the mesh? Would it be an idea to first find collisions?

[jorgensd]

No, the function does this given a single cell (as it would otherwise be computationally infeasible). One should first use an other method to find the candidate cells to search through, by for instance:

• Finding the M cells that has the closest bounding boxes to the point(s)
• Finding the M cells whose convex hull is closest to the point(s).
  On these candidates, compute the projection onto the cell.

Here is an implementation using pure numpy arrays (which can be ported to C++)

import numpy as np
from mpi4py import MPI
import ufl
import basix.ufl
import dolfinx
from scipy.optimize import minimize
import numpy.typing as npt


def project_point_to_element(
    mesh: dolfinx.mesh.Mesh,
    cells: npt.NDArray[np.int64],
    target_points: npt.NDArray[np.float64 | np.float32],
    method=None,
    tol: float | None = None,
    max_iter: int = 25,
):
    """
    Projects a 3D point onto a cell in a potentially higher order mesh.

    Args:
        mesh: {py:class}`dolfinx.mesh.Mesh`, the mesh containing the cell.
        cells: {py:class}`numpy.ndarray`, the indices of the cells to project onto.
        target_point: (3,) numpy array, the 3D point to project.
        method: str, the optimization method to use.
        tol: float, the tolerance for the optimizer.

    Returns:
        dict: A dictionary containing the reference coordinates, closest 3D point, and distance.
    """

    # Extract scalar element of mesh
    element = mesh.ufl_domain().ufl_coordinate_element().sub_elements[0]
    tdim = mesh.topology.dim
    # Get the coordinates of the nodes for the specified cell
    node_coords = mesh.geometry.x[mesh.geometry.dofmap[cells]][
        :, :, : mesh.geometry.dim
    ]
    target_points = target_points.reshape(-1, 3)
    # cmap = mesh.geometry.cmap

    # Constraints and Bounds
    cell_type = mesh.topology.cell_type
    if cell_type == dolfinx.mesh.CellType.triangle:
        method = method or "SLSQP"
        constraint = {"type": "ineq", "fun": lambda x: 1.0 - np.sum(x)}
    else:
        method = method or "L-BFGS-B"
        constraint = {}
    bounds = [(0.0, 1.0) for _ in range(mesh.topology.dim)]

    # Set initial guess and tolerance for solver
    initial_guess = np.full(
        mesh.topology.dim, 1 / (mesh.topology.dim + 1), dtype=np.float64
    )
    tol = np.sqrt(np.finfo(mesh.geometry.x.dtype).eps) if tol is None else tol
    closest_points = np.zeros((target_points.shape[0], 3), dtype=mesh.geometry.x.dtype)

    if cell_type not in [
        dolfinx.mesh.CellType.triangle,
        dolfinx.mesh.CellType.tetrahedron,
    ]:

        def clip(x_k):
            return np.clip(x_k, 0.0, 1.0)
    else:

        def clip(x_k):
            x_k = np.maximum(x_k, 0.0)  # Clamp negative values
            s = np.sum(x_k)
            if s > 1.0:
                x_k /= s  # Pull back to valid simplex domain
            return x_k

    c1 = 1e-5  # Standard Armijo constant
    for i, (coord, target_point) in enumerate(zip(node_coords, target_points)):
        coord = coord.reshape(-1, mesh.geometry.dim)

        def exact_hessian_and_grad(x_ref):
            tdim = mesh.topology.dim

            # Tabulate up to 2nd order derivatives
            # Shape for 2D: (6, 1, num_nodes, 1)
            # Shape for 3D: (10, 1, num_nodes, 1)
            tab = element.tabulate(2, x_ref.reshape(1, tdim))

            # --- 0th Derivative (Value) ---
            N = tab[0, 0, :]
            surface_point = np.dot(N, coord)
            diff = surface_point - target_point  # Residual vector 'r'

            # --- 1st Derivatives (Jacobian / Tangents) ---
            dN = tab[1 : tdim + 1, 0, :]
            tangents = np.dot(dN, coord)  # Shape: (tdim, 3)

            # Gradient: J^T * r (Works universally for 2D and 3D)
            g = np.dot(tangents, diff)

            # First-order term of Hessian (Gauss-Newton): J^T * J
            H_approx = np.dot(tangents, tangents.T)

            # --- The Exact Second-Order Term ---
            # Extract 2nd derivatives.
            d2N = tab[tdim + 1 :, 0, :]
            second_deriv_surface = np.dot(d2N, coord)  # Shape: (num_2nd_derivs, 3)

            # Dot the residual with the second derivatives of the surface mapping
            S_2_dot_r = np.dot(second_deriv_surface, diff)

            # Build the symmetric second-order Hessian matrix
            if tdim == 1:
                H_second_order = np.array([[S_2_dot_r[0]]])
            elif tdim == 2:
                H_second_order = np.array(
                    [[S_2_dot_r[0], S_2_dot_r[1]], [S_2_dot_r[1], S_2_dot_r[2]]]
                )
            elif tdim == 3:
                H_second_order = np.array(
                    [
                        [S_2_dot_r[0], S_2_dot_r[1], S_2_dot_r[3]],  # xx, xy, xz
                        [S_2_dot_r[1], S_2_dot_r[2], S_2_dot_r[4]],  # xy, yy, yz
                        [S_2_dot_r[3], S_2_dot_r[4], S_2_dot_r[5]],  # xz, yz, zz
                    ]
                )
            else:
                raise RuntimeError(
                    f"Unsupported dimension {tdim} for Hessian construction"
                )

            # Return the exact gradient, the hessian decomposed in its two components, and current physical point
            return g, H_approx, H_second_order, surface_point

        x_k = initial_guess.copy()
        I = np.eye(mesh.topology.dim) * 1e-14
        for k in range(max_iter):
            # Tabulate, basis + first and second order derivatives
            g, H_approx, H_second_order, surface_point = exact_hessian_and_grad(x_k)
            diff = surface_point - target_point
            current_dist_sq = np.sum(diff**2)

            # Solve the dense system: H * dx = -g
            H = H_approx + H_second_order
            g_orig = g.copy()  # Store original gradient for Armijo check

            # --- ACTIVE SET PROJECTION ---
            # If a coordinate is pinned to a 0.0 boundary and the unconstrained gradient
            # wants to push it further negative, we MUST remove it from the Newton system.
            # Otherwise, the blocked full-Hessian step corrupts the free coordinates!
            for d in range(tdim):
                if x_k[d] <= 1e-12 and g[d] > 0:
                    H[d, :] = 0.0
                    H[:, d] = 0.0
                    H[d, d] = 1.0  # Keep matrix invertible
                    g[d] = 0.0  # Zero out the gradient step

            # Solve the dense system: H * dx = -g
            dx = np.linalg.solve(H, -g)

            # --- SAFETY CHECK 1: Is H_exact indefinite? ---
            # If the exact step points uphill (dot product > 0), the Hessian is indefinite
            # due to large residuals. Fall back to Gauss-Newton (J^T * J).
            if np.dot(g_orig, dx) > 0:
                H_gn = H_approx + I
                dx = np.linalg.solve(H_gn, -g)
                # If extreme numerical noise STILL makes it point uphill, use steepest descent
                if np.dot(g_orig, dx) > 0:
                    dx = -g_orig

            # Linesearch
            alpha = 1.0
            for _ in range(25):
                x_new = x_k + alpha * dx
                x_new = clip(x_new)

                actual_step = x_new - x_k

                tab_new = element.tabulate(0, x_new.reshape(1, tdim))
                S_new = np.dot(tab_new[0, 0, :], coord)
                new_dist_sq = np.sum((S_new - target_point) ** 2)

                # Armijo Sufficient Decrease Check
                # Note: We multiply the directional derivative by 2 because
                # our dist_sq is 2 * f(x)
                expected_decrease = 2.0 * c1 * np.dot(g_orig, actual_step)

                if expected_decrease >= 0:
                    alpha *= 0.5
                    continue

                if new_dist_sq <= current_dist_sq + expected_decrease:
                    break  # Sufficient decrease achieved!

                # If not, cut the step size and try again
                alpha *= 0.5
            if np.linalg.norm(x_new - x_k) < tol:
                break
            x_k = x_new.copy()

        closest_points[i] = surface_point
    return closest_points


# --- Example Usage ---
if __name__ == "__main__":
    # Dummy curved quadratic triangle in 3D (6 nodes)
    curved_nodes = np.array(
        [
            [0.0, 0.0, 0.0],  # Node 0: Vertex
            [1.0, 0.0, 0.0],  # Node 1: Vertex
            [0.0, 1.0, 0.0],  # Node 2: Vertex
            [0.6, 0.6, 0.0],  # Node 4: Edge 1-2
            [0.1, 0.5, 0.2],  # Node 5: Edge 2-0 (curved upward in Z)
            [0.5, 0.1, 0.2],  # Node 3: Edge 0-1 (curved upward in Z)
        ]
    )
    cells = np.array(
        [[0, 1, 2, 3, 4, 5]], dtype=np.int64
    )  # Single curved triangle element
    c_el = ufl.Mesh(basix.ufl.element("Lagrange", "triangle", 2, shape=(3,)))
    mesh = dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells=cells, x=curved_nodes, e=c_el)

    point_to_project = np.array([0.3, 0.3, -0.2])  # A point floating above the element
    import time

    start = time.perf_counter()
    result = project_point_to_element(
        mesh, np.array([0], dtype=np.int32), point_to_project
    )
    end = time.perf_counter()
    print("Closest Point on Element:", result, f"Time taken: {end - start:.6e} seconds")

    with dolfinx.io.VTXWriter(mesh.comm, "closest_point.bp", mesh) as writer:
        writer.write(0.0)

    curved_nodes_tet = np.array(
        [
            [0.0, 0.0, 0.0],  # 0: Vertex
            [1.0, 0.0, 0.0],  # 1: Vertex
            [0.0, 1.0, 0.0],  # 2: Vertex
            [0.0, 0.0, 1.0],  # 3: Vertex
            [0.0, 0.5, 0.5],  # 4: Edge 2-3
            [0.5, 0.0, 0.5],  # 5: Edge 1-3
            [0.5, 0.5, 0.0],  # 6: Edge 1-2
            [0.0, 0.0, 0.5],  # 9: Edge 0-3
            [0.0, 0.5, 0.2],  # 8: Edge 0-2 (Curved)
            [0.5, 0.0, 0.2],  # 7: Edge 0-1 (Curved)
        ],
        dtype=np.float64,
    )
    cells_tet = np.array([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]], dtype=np.int64)

    domain_tet = ufl.Mesh(basix.ufl.element("Lagrange", "tetrahedron", 2, shape=(3,)))
    mesh_tet = dolfinx.mesh.create_mesh(
        MPI.COMM_WORLD, cells=cells_tet, x=curved_nodes_tet, e=domain_tet
    )
    start = time.perf_counter()
    result = project_point_to_element(
        mesh_tet, np.array([0], dtype=np.int32), point_to_project
    )
    end = time.perf_counter()
    print("Closest Point on Element:", result, f"Time taken: {end - start:.6e} seconds")

    with dolfinx.io.VTXWriter(
        mesh_tet.comm, "closest_point_tet.bp", mesh_tet
    ) as writer:
        writer.write(0.0)

with output:

Closest Point on Element: [[0.21763919 0.21763919 0.17580848]] Time taken: 5.299030e-04 seconds
Closest Point on Element: [[0.47210016 0.47210017 0.04214885]] Time taken: 5.117880e-04 seconds