Optimize geometry with PyTorch

Recover mesh vertex heights with torch.optim.Adam through RayD's PyTorch backend, with the scene refit in place every iteration.

This tutorial solves the same problem as Fit a mesh to target hit distances, but in PyTorch: recover the four corner heights of a quad from the distances at which 256 rays hit it. You will use torch.optim.Adam and ordinary autograd; RayD records the intersection as an autograd node.

Install the PyTorch backend first, as described in Install RayD.

The complete script#

import torch
import rayd.torch as rt
 
dev = torch.device("cuda")
xy = torch.tensor([[-1.0, -1.0], [1.0, -1.0], [1.0, 1.0], [-1.0, 1.0]], device=dev)
faces = torch.tensor([[0, 1, 2], [0, 2, 3]], device=dev, dtype=torch.int32)
target_z = torch.tensor([0.2, -0.1, 0.3, 0.0], device=dev)
 
# 256 rays on a 16 x 16 grid, shooting down from z = 2.
n = 16
g = -0.9 + 1.8 * (torch.arange(n, device=dev) + 0.5) / n
gx, gy = torch.meshgrid(g, g, indexing="ij")
origins = torch.stack([gx.flatten(), gy.flatten(), torch.full((n * n,), 2.0, device=dev)], dim=1)
directions = torch.tensor([[0.0, 0.0, -1.0]], device=dev).expand(n * n, 3).contiguous()
ray = rt.Ray(origins, directions)
 
def vertices(z):
    return torch.cat([xy, z[:, None]], dim=1).contiguous()
 
scene = rt.Scene()
mesh_id = scene.add_mesh(rt.Mesh(vertices(target_z), faces), dynamic=True)
scene.build()
 
def hit_distance():
    hit = scene.intersect(ray)
    return torch.where(hit.is_valid(), hit.t, torch.zeros_like(hit.t))
 
with torch.no_grad():
    target_t = hit_distance()
 
z = torch.zeros(4, device=dev, requires_grad=True)
opt = torch.optim.Adam([z], lr=0.02)
for it in range(201):
    scene.update_mesh_vertices(mesh_id, vertices(z))
    scene.sync()
    loss = (hit_distance() - target_t).square().mean()
    opt.zero_grad()
    loss.backward()
    opt.step()
    if it % 40 == 0:
        print(f"iter {it:3d}  loss {loss.item():.2e}  z {[round(v, 3) for v in z.tolist()]}")

Output on an RTX 5080 with RayD 0.8.0 and PyTorch 2.10 (z is printed after the step):

iter   0  loss 3.05e-02  z [0.02, 0.02, 0.02, 0.02]
iter  40  loss 1.10e-04  z [0.206, -0.128, 0.306, -0.032]
iter  80  loss 4.89e-06  z [0.2, -0.094, 0.299, 0.006]
iter 120  loss 4.56e-08  z [0.2, -0.1, 0.3, -0.0]
iter 160  loss 8.44e-10  z [0.2, -0.1, 0.3, -0.0]
iter 200  loss 1.43e-12  z [0.2, -0.1, 0.3, -0.0]

How it works#

  1. Tensors in, tensors out

    Every input is a contiguous CUDA tensor: vertices (V, 3) float32, faces (F, 3) int32, ray origins and directions (N, 3) float32. .contiguous() matters for the expanded direction tensor and the concatenated vertices.

  2. One dynamic mesh

    The mesh is added with dynamic=True, so update_mesh_vertices can replace its vertices each iteration. sync() refits the acceleration structure without a rebuild.

  3. Autograd through the hit

    The vertex tensor depends on z, which requires a gradient, so scene.intersect records an autograd node. loss.backward() fills z.grad with the fixed-winner derivative of the hit distances.

  4. A standard optimizer

    Any torch.optim optimizer works. The target was recorded under torch.no_grad(), so it carries no graph.

Tips#

  • torch.where(hit.is_valid(), hit.t, 0) keeps missed rays (t = inf) out of the loss.
  • The PyTorch backend supports torch.compile and forward-mode AD through torch.autograd.forward_ad. Some operations are forward-only; check rt.api_manifest()["derivatives"] or Differentiation before you optimize through reflection accumulation or diffraction paths.
  • For large batches, Scene(devices=[0, 1]) splits the rays across GPUs and returns gradients on the first device; see Select devices and use several GPUs.

Next steps#