interp: Linear interpolation on triangular meshes

This module provides JAX-compatible linear interpolation of vertex-based data on triangular meshes in 2D and 3D, in two steps:

  1. For each query point, find the closest mesh triangle (find_closest_faces).
  2. Interpolate the values at that triangle’s vertices using the barycentric coordinates of the closest point on the triangle (interpolate_barycentric).

In 3D, query points off the surface are assigned the value at the closest surface point.

Everything is compatible with jax.jit and automatic differentiation. We use brute-force search to find the closest point, which less efficient than a spatial data structure (e.g. kD-tree or AABB tree). These are unfortunately is harder to implement efficiently in JAX, and currently out of scope.

Closest-point helpers

source

get_closest_point_on_triangle


def get_closest_point_on_triangle(
    point:Float[Array, 'dim'], # Query point.
    a:Float[Array, 'dim'], b:Float[Array, 'dim'], c:Float[Array, 'dim']
)->Float[Array, 'dim']: # Closest point on the triangle.

Closest point on triangle abc to a query point.

For non-degenerate triangles this first projects onto the triangle plane, then clamps to the triangle if the projection lies outside. Degenerate triangles fall back to the closest point on the three edges.

source

get_closest_point_on_segment


def get_closest_point_on_segment(
    point:Float[Array, 'dim'], # Query point.
    a:Float[Array, 'dim'], b:Float[Array, 'dim']
)->Float[Array, 'dim']: # Closest point on the segment from a to b.

Closest point on a line segment.

triangle2d = jnp.array([[0., 0.], [1., 0.], [0., 1.]])
assert jnp.allclose(get_closest_point_on_segment(jnp.array([2., 0.5]),
                                                 triangle2d[0],
                                                 triangle2d[1]),
                    jnp.array([1., 0.]))
assert jnp.allclose(get_closest_point_on_triangle(jnp.array([0.2, 0.3]), *triangle2d),
                    jnp.array([0.2, 0.3]), atol=1e-10)
assert jnp.allclose(get_closest_point_on_triangle(jnp.array([0.8, 0.8]), *triangle2d),
                    jnp.array([0.5, 0.5]), atol=1e-10)

triangle3d = jnp.array([[0., 0., 0.], [1., 0., 0.], [0., 1., 0.]])
assert jnp.allclose(get_closest_point_on_triangle(jnp.array([0.2, 0.3, 1.5]), *triangle3d),
                    jnp.array([0.2, 0.3, 0.]), atol=1e-10)

shifted_triangle3d = jnp.array([[1., 0., 0.], [1., 1., 0.], [1., 0., 1.]])
assert jnp.allclose(get_closest_point_on_triangle(jnp.array([2., 0.25, 0.25]), *shifted_triangle3d),
                    jnp.array([1., 0.25, 0.25]), atol=1e-10)

find_closest_faces


def find_closest_faces(
    points:Float[Array, 'n_points dim'], # Query points (dim = 2 or 3).
    vertices:Float[Array, 'n_vertices dim'], # Mesh vertices.
    faces:Int[Array, 'n_faces 3'], # Triangle indices into vertices.
)->tuple: # Closest face index and closest point on that face for each query point.

Find the closest triangle and closest point on the mesh for each query point.

Brute-force O(n_points * n_faces) search: cheap point-triangle squared distances (using per-mesh precomputed edge vectors and unit normals) select the closest face, then the exact closest point is computed on that face only. 2D meshes are zero-padded to 3D for the search. Compatible with jax.jit and automatic differentiation.

vertices = jnp.array([[0., 0.], [1., 0.], [0., 1.], [1., 1.]])
faces = jnp.array([[0, 1, 2], [1, 3, 2]])
values = jnp.array([0., 1., 2., 3.])
points = jnp.array([[0., 0.], [0.25, 0.25], [0.8, 0.8], [1.5, 0.5]])

closest_faces, closest_points = find_closest_faces(points, vertices, faces)
assert jnp.array_equal(closest_faces, jnp.array([0, 0, 1, 1]))
assert jnp.allclose(closest_points, jnp.array([[0., 0.], [0.25, 0.25],
                                               [0.8, 0.8], [1., 0.5]]), atol=1e-10)

triangle3d = jnp.array([[0., 0., 0.], [1., 0., 0.], [0., 1., 0.]])
_, closest3d = find_closest_faces(jnp.array([[0.2, 0.3, 1.5]]), triangle3d, jnp.array([[0, 1, 2]]))
assert jnp.allclose(closest3d[0], jnp.array([0.2, 0.3, 0.]), atol=1e-10)

Barycentric interpolation

source

interpolate_barycentric


def interpolate_barycentric(
    points:Float[Array, 'n_points dim'], # Query points.
    vertices:Float[Array, 'n_vertices dim'], # Mesh vertices in 2D or 3D.
    faces:Int[Array, 'n_faces 3'], # Mesh triangles, indexing vertices.
    values:Float[Array, 'n_vertices ...'], # Values defined at mesh vertices. Additional trailing axes are preserved.
    distance_threshold:float=inf, # Squared-distance threshold. Points farther from the mesh than this are
masked with NaN in the returned array.
)->Float[Array, 'n_points ...']: # Interpolated values at the query points.

Interpolate vertex data onto query points using barycentric coordinates.

Each query point is first mapped to its closest triangle on the mesh via find_closest_faces, then the barycentric coordinates of the closest point are used to linearly blend the values at that triangle’s vertices. Query points off the mesh (e.g. away from the surface in 3D) are assigned the value at the closest mesh point.

interpolated = interpolate_barycentric(points, vertices, faces, values)
reference_values, reference_faces, reference_points, _ = _interpolate_barycentric_libigl(points, vertices,
                                                                                         faces, values)
assert jnp.allclose(interpolate_barycentric(vertices, vertices, faces, values), values)
assert np.allclose(np.asarray(closest_points), reference_points[:, :2], atol=1e-8)
assert np.allclose(np.asarray(interpolated), reference_values, atol=1e-8, equal_nan=True)
key = jax.random.key(0)

vertices2d_full, faces2d = igl.read_triangle_mesh(str(mesh_dir / 'disk.obj'))
vertices2d = jnp.asarray(vertices2d_full[:, :2])
faces2d = jnp.asarray(faces2d)
key, subkey = jax.random.split(key)
box_min = vertices2d.min(axis=0) - 0.1
box_max = vertices2d.max(axis=0) + 0.1
points2d = box_min + jax.random.uniform(subkey, shape=(32, 2)) * (box_max - box_min)
key, subkey = jax.random.split(key)
values2d = jax.random.normal(subkey, shape=(vertices2d.shape[0], 3))

closest_faces2d, closest_points2d = find_closest_faces(points2d, vertices2d, faces2d)
interpolated2d = interpolate_barycentric(points2d, vertices2d, faces2d, values2d)
reference_values2d, reference_faces2d, reference_points2d, _ = _interpolate_barycentric_libigl(points2d,
                                                                                               vertices2d,
                                                                                               faces2d,
                                                                                               values2d)

assert np.array_equal(np.asarray(closest_faces2d), reference_faces2d)
assert np.allclose(np.asarray(closest_points2d), reference_points2d[:, :2], atol=1e-8)
assert np.allclose(np.asarray(interpolated2d), reference_values2d, atol=1e-8)
Warning: readOBJ() ignored non-comment line 3:
  o flat_tri_ecmc
key = jax.random.key(0)

vertices3d, faces3d = igl.read_triangle_mesh(str(mesh_dir / 'sphere_finer.obj'))
vertices3d = jnp.asarray(vertices3d)
faces3d = jnp.asarray(faces3d)
key, subkey = jax.random.split(key)

N_points = 2000
vertex_ids = jax.random.randint(subkey, shape=(N_points,), minval=0, maxval=vertices3d.shape[0])
key, subkey = jax.random.split(key)
points3d = vertices3d[vertex_ids] + 0.05 * jax.random.normal(subkey, shape=(N_points, 3))
key, subkey = jax.random.split(key)
values3d = jax.random.normal(subkey, shape=(vertices3d.shape[0], 2, 2))

closest_faces3d, closest_points3d = find_closest_faces(points3d, vertices3d, faces3d)
interpolated3d = interpolate_barycentric(points3d, vertices3d, faces3d, values3d)
reference_values3d, reference_faces3d, reference_points3d, _ = _interpolate_barycentric_libigl(points3d,
                                                                                               vertices3d,
                                                                                               faces3d,
                                                                                               values3d)

# face indices can differ from igl for points near-equidistant to several faces,
# but closest points and interpolated values must agree
assert np.allclose(np.asarray(closest_points3d), reference_points3d, rtol=1e-5)
assert np.allclose(np.asarray(interpolated3d), reference_values3d, rtol=1e-5)
Warning: readOBJ() ignored non-comment line 4:
  o Icosphere

Speed testing

Compare against the igl reference (CGAL AABB tree, not differentiable). The JAX-version is slower due to the brute-force search; casting inputs to float32 can reduce the runtime if reduced precision is acceptable.

4.37 ms ± 118 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
interpolate_barycentric_jit = jax.jit(interpolate_barycentric)
_ = interpolate_barycentric_jit(points3d, vertices3d, faces3d, values3d)
19.2 ms ± 245 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)
points3d_f32, vertices3d_f32 = points3d.astype(jnp.float32), vertices3d.astype(jnp.float32)
values3d_f32 = values3d.astype(jnp.float32)
_ = interpolate_barycentric_jit(points3d_f32, vertices3d_f32, faces3d, values3d_f32)
15.5 ms ± 132 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)