I recently started my sabbatical at Nvidia with the team that develops the warp library/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 computes a cute fact about triangle meshes.
If you have a solid enclosed by a triangle mesh, you can compute its volume via Gauss theorem: 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 a bunch of 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 our degrees of freedom then the gradient of volume will define a normal direction per-vertex.
If we consider a single vertex and think of those signed-volume tetrahedra, the only tetrahedra that matter are the ones for the triangles that contain 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 that in a round-about way 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(...).