## `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`](https://nikolas-claussen.github.io/triangulax/src/interpolation.html#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`](https://nikolas-claussen.github.io/triangulax/src/interpolation.html#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 ``` python 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 ``` python 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.* ``` python 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) ``` ### Closest-face search ------------------------------------------------------------------------ source ### find_closest_faces ``` python 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. ``` python 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 ``` python 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. ``` python 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) ``` ``` python 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 ``` python 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. ``` python ``` 4.37 ms ± 118 μs per loop (mean ± std. dev. of 7 runs, 100 loops each) ``` python 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) ``` python 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)