Fit a mesh to target hit distances

In this tutorial you recover the shape of a mesh from what rays see. A quad has four unknown corner heights. You know only the distances at which 256 downward rays hit the true quad. Starting from a flat quad, gradient descent on the hit distances finds the four heights.

You will learn how to:

• keep one dynamic scene and update its vertices every iteration,
• get gradients of hit distances with respect to vertex positions,
• avoid mixing two traced scenes in one Dr.Jit kernel.

Before you start, install the Dr.Jit backend as described in Install RayD.

The complete script#

import drjit as dr
import rayd.drjit as rd
from drjit.cuda import Array3i
from drjit.cuda.ad import Array3f, Float
 
# A 2 x 2 quad whose four corner heights we want to recover.
x = Float([-1.0, 1.0, 1.0, -1.0])
y = Float([-1.0, -1.0, 1.0, 1.0])
faces = Array3i([0, 0], [1, 2], [2, 3])
target_z = Float([0.2, -0.1, 0.3, 0.0])
 
# 256 rays on a 16 x 16 grid, shooting down from z = 2.
n = 16
i = dr.arange(Float, n * n)
row = dr.floor(i / n)
gx = -0.9 + 1.8 * (row + 0.5) / n
gy = -0.9 + 1.8 * (i - n * row + 0.5) / n
ray = rd.RayAD(Array3f(gx, gy, 2.0), Array3f(0.0, 0.0, -1.0))
 
scene = rd.Scene()
mesh_id = scene.add_mesh(rd.Mesh(Array3f(x, y, target_z), faces), dynamic=True)
scene.build()
 
def hit_distance():
    hit = scene.intersect(ray)
    return dr.select(hit.is_valid(), hit.t, 0.0)
 
# Record the target distances, then start from a flat quad.
target_t = hit_distance()
dr.eval(target_t)
 
z = Float([0.0, 0.0, 0.0, 0.0])
lr = 2.0
for it in range(61):
    dr.enable_grad(z)
    scene.update_mesh_vertices(mesh_id, Array3f(x, y, z))
    scene.sync()
    t = hit_distance()
    loss = dr.mean(dr.square(t - target_t))
    dr.backward(loss)
    if it % 10 == 0:
        print(f"iter {it:3d}  loss {loss[0]:.2e}  z {z}")
    z = dr.detach(z - lr * dr.grad(z))

Output on an RTX 5080 with RayD 0.8.0:

iter   0  loss 3.05e-02  z [0, 0, 0, 0]
iter  10  loss 4.98e-05  z [0.19369, -0.0740094, 0.28933, 0.0216304]
iter  20  loss 1.81e-06  z [0.198458, -0.0953122, 0.298268, 0.00449763]
iter  30  loss 6.71e-08  z [0.199688, -0.0991099, 0.29968, 0.000881729]
iter  40  loss 2.50e-09  z [0.199939, -0.0998289, 0.299939, 0.000170729]
iter  50  loss 9.30e-11  z [0.199988, -0.0999669, 0.299988, 3.30306e-05]
iter  60  loss 3.45e-12  z [0.199998, -0.0999936, 0.299998, 6.41398e-06]

The heights converge to the target [0.2, -0.1, 0.3, 0.0].

How it works#

1. Build one dynamic scene

   The mesh is added with dynamic=True, so update_mesh_vertices can replace its vertices after build(). The scene is built once, with the true heights, to record the target.

2. Record the target and evaluate it

   target_t holds the true hit distances. dr.eval(target_t) runs that trace now. Without it, Dr.Jit would fuse the target trace and the first optimization trace into one kernel.

3. Update the vertices from the parameters

   Each iteration turns the AD array z into vertex positions and passes them to update_mesh_vertices, then calls sync() to refit the acceleration structure. Because the positions are AD arrays, the scene records how the triangles depend on z.

4. Differentiate the hit distances

   intersect with a RayAD returns an IntersectionAD. dr.select replaces the distance of a missed ray (inf) with 0 so that it does not poison the loss. dr.backward sends the loss gradient back through the hit distances to z.

5. Take a gradient step

   Plain gradient descent with a step of 2.0 is enough here. dr.detach starts the next iteration with a fresh AD variable.

Pitfalls#

• The gradient holds the winning triangle fixed. Each ray here keeps hitting the same quad, so the gradient is exact. If a perturbation makes rays switch triangles or miss, the gradient does not see that change; see Differentiation.
• Without the dr.select, a single missed ray makes the loss infinite.

Next steps#

• Do the same in PyTorch with torch.optim: Optimize geometry with PyTorch.
• Move meshes through transforms instead of vertices: Update geometry in place.