I recently started a sabbatical at Nvidia with the team that develops the warp programming language. Warp is a way to write python code that will run on the GPU and compute first-derivatives automatically. There are a bunch of great examples maintained in the library, but to get my hands dirty I wrote this small geometry processing example that demonstrates a cute fact about triangle meshes.
You can compute the solid volume contained inside a triangle mesh as a sum over the triangles. Volume is the integral of unit mass in the interior and Guass theorm let's you move the integral to the surface. You can work this out with math, but another way of arriving at the exact same thing is to consider the signed volume of a tetrahedron formed by each triangle and the origin. Adding these up you get the positive volume of the solid (everything spilling outside the shape cancels out).
Now, let's consider differentiating the volume with respect to the surface. Gradients are directions of maximum increase so increasing volume should move points in the outward normal direction. For a mesh, if vertices are the degrees of freedom then the gradient of volume will define per-vertex normal directions.
Consider those signed-volume tetrahedra. The only tetrahedra impacted by moving a particular vertex are the ones for triangles incident on that vertex. The gradient of the volume of a single tetrahedra will contain a part proportional to the area of the triangle and another part propertional to the "fake" triangles that connect each incident edge on that vertex to the origin, but those ones will get canceled out by the neighboring faces with the opposite orientation. So, when we sum up gradients, each vertex gets a contribution from each incident triangle that is proportional to the area of that triangle and points in the direction of the triangle's normal. This matches exactly the definition of an area-weighted per-vertex normal (after normalization which is just an average).
Here's how to compute per-vertex normals via volume differentiation in warp:
import warp as wp
import igl
import polyscope as ps
import numpy as np
import sys
# read in a mesh from argv or if not provided, use icosahedron
if len(sys.argv) > 1:
mesh_path = sys.argv[1]
V, F = igl.read_triangle_mesh(mesh_path)
else:
V, F = igl.icosahedron()
@wp.func
def signed_volume_with_origin(
p0: wp.vec3, p1: wp.vec3, p2: wp.vec3
) -> float:
return wp.dot(p0, wp.cross(p1, p2)) / 6.0
@wp.kernel(enable_backward=True)
def total_volume_kernel(
points: wp.array(dtype=wp.vec3),
indices: wp.array(dtype=int),
total_volume: wp.array(dtype=float),
):
tid = wp.tid()
v = signed_volume_with_origin(
points[ indices[3 * tid + 0] ],
points[ indices[3 * tid + 1] ],
points[ indices[3 * tid + 2] ])
wp.atomic_add(total_volume, 0, v)
wp.init()
points = wp.array(V, dtype=wp.vec3, requires_grad=True)
indices = wp.array(F.flatten(), dtype=int)
total_volume = wp.zeros(1, dtype=float, requires_grad=True)
tape = wp.Tape()
with tape:
wp.launch(total_volume_kernel, dim=indices.shape[0] // 3, inputs=[points, indices], outputs=[total_volume])
tape.backward(loss=total_volume)
# ∂volume/∂points gives area-weighted vertex normals (pointing outward)
vertex_normals = points.grad.numpy().copy()
vertex_normals /= np.linalg.norm(vertex_normals, axis=1)[:, np.newaxis]
# Render the result to an image
ps.init("openGL3_egl")
ps_mesh = ps.register_surface_mesh("mesh", V, F)
ps_mesh.add_vector_quantity("vertex normals", vertex_normals, enabled=True)
ps.set_ground_plane_mode("shadow_only")
#ps.show()
ps.screenshot("example_vertex_normals_via_autodiff.jpg")
This example shows the syntax for setting up data in warp (points = wp.array(...), defining a kernel (@wp.kernel),
defining a local function that can be used inside a kernel (@wp.func), and launching the kernel
in a differentiable way (with tape: wp.launch(...).