Differentiability in cadjoint

cadjoint is a fully differentiable CAD system: every SDF is a pure JAX function, so gradients flow through geometry just like through any other computation. This makes it straightforward to optimize shape parameters directly with respect to any loss.

import jax
import jax.numpy as jnp

from cadjoint import extract_parameters, functionalize
from cadjoint.geometry import Vector
from cadjoint.render import render_raymarched
from cadjoint.sdf import SDF, Sphere, Translate, Union, volume

Because SDFs are plain JAX functions, we can differentiate through them out of the box. As a running example, we define a scene builder that places two unit spheres at positions p1 and p2 and returns their union.

def double_sphere(p1: jax.Array, p2: jax.Array) -> SDF:
    radius = 1.0
    sphere1 = Translate(Sphere(radius=radius), offset=p1)
    sphere2 = Translate(Sphere(radius=radius), offset=p2)
    return jax.jit(Union(sphere1, sphere2))

p1 = jnp.array([0.0, 0.0, 0.0])
p2 = jnp.array([1.0, 0.0, 0.0])

render_raymarched(
    double_sphere(p1, p2),
    camera_pos=jnp.array([3.0, 2.0, 4.0]),
    look_at=jnp.array([0.5, 0.0, 0.0]),
    resolution=(300, 300),
    aa_samples=2,
)

cadjoint provides a volume function that integrates an SDF over a bounding volume. We can use it as a loss and differentiate with respect to the sphere positions using jax.value_and_grad.

def loss(p1: jax.Array, p2: jax.Array) -> jax.Array:
    return volume(double_sphere(p1, p2))

grad = jax.value_and_grad(loss, argnums=(0, 1))
value, grads = grad(p1, p2)
print("Loss value:", value)
print("Gradient w.r.t p1:", grads[0])
print("Gradient w.r.t p2:", grads[1])

Loss value: 6.8726716
Gradient w.r.t p1: [-1.8908138e+00  2.6635826e-07 -1.0745152e-06]
Gradient w.r.t p2: [1.9764323e+00 1.3652647e-06 9.5690950e-07]

Passing positions explicitly works, but scales poorly: complex scenes may have dozens of parameters scattered across the tree. cadjoint lets you tag parameters as free directly on the geometry objects, then extract them automatically as a flat dictionary that JAX can differentiate through.

p1 = Vector(jnp.array([0.0, 0.0, 0.0]), free=True, name="p1")
p2 = Vector(jnp.array([1.0, 0.0, 0.0]), free=True, name="p2")
sphere1 = Translate(Sphere(radius=1.0), offset=p1)
sphere2 = Translate(Sphere(radius=1.0), offset=p2)
double_sphere_ = Union(sphere1, sphere2)

extract_parameters walks the SDF tree and returns three dictionaries: free (differentiable) parameters keyed by name, fixed parameters keyed by tree path, and metadata carrying the Parameter objects behind the free values.

free_params, fixed_params, metadata = extract_parameters(double_sphere_)

print("Free parameters:", free_params)

Free parameters: {'p1': Array([0., 0., 0.], dtype=float32), 'p2': Array([1., 0., 0.], dtype=float32)}

functionalize compiles the SDF tree into a curried pure function: given parameter dictionaries it returns a callable point → distance. This separates the geometry structure (fixed at compile time) from the parameter values (free to vary).

double_sphere_functional = functionalize(double_sphere_)
double_sphere_functional(free_params, fixed_params)(jnp.array([0.5, 0, 0]))

Array(-0.6, dtype=float32)

We can now differentiate the loss through the parameter dictionary directly.

def loss(params):
    return volume(double_sphere_functional(params, fixed_params))

grad = jax.value_and_grad(loss, argnums=(0))
value, grads = grad(free_params)
print("Loss value:", value)
grads

Loss value: 6.8726716

{'p1': Array([-1.8908138e+00,  2.6635826e-07, -1.0745152e-06], dtype=float32),
 'p2': Array([1.9764323e+00, 1.3652647e-06, 9.5690950e-07], dtype=float32)}

For simple scenes the two approaches are equivalent. The parameter-extraction workflow becomes essential as scenes grow: any node in the tree can declare its parameters free, and the compiler collects them automatically - no changes needed to the loss or optimizer.