Optimization

Gradients through CAD parameters, constrained descent, and optimizations with a simulation as the objective.

An Optimization minimizes a scalar over the free parameters of a scene: a pure function of the parameters, or a declared simulation study. It returns the optimized parameters, the per-step history and a replayable trajectory; the viewer, the CLI and a script run the same object. The optimizer is optax, a core dependency.

Example

The objective form minimizes a pure function of the free parameters of a piece of geometry, with no simulation:

import jax.numpy as jnp
from cadjoint.geometry.parameters import Vector
from cadjoint.optimize import Optimization
from cadjoint.sdf.primitives import Sphere
from cadjoint.sdf.transforms import Translate

p = Vector([1.0, 1.0, 0.0], free=True, name="p")
scene = Translate(Sphere(radius=0.5), offset=p)

def objective(params):
    return jnp.sum(params["p"] ** 2)

run = Optimization(name="center", objective=objective, of=scene,
                   steps=40, learning_rate=0.1).run()

print(run.initial, "->", run.parameters)
# {'p': [1.0, 1.0, 0.0]} -> {'p': [0.1008, 0.1008, 0.0]}
print(run.history[0]["objective"], "->", run.objective)  # 2.0 -> 0.0203

Gradients

Every SDF is a pure JAX function, so jax.grad differentiates through geometry. Parameters can be passed as plain arrays:

import jax
import jax.numpy as jnp
from cadjoint.sdf import Sphere, Translate, Union, volume

def two_spheres(p1, p2):
    return Union(Translate(Sphere(radius=1.0), offset=p1),
                 Translate(Sphere(radius=1.0), offset=p2))

def loss(p1, p2):
    return volume(two_spheres(p1, p2))

value, (g1, g2) = jax.value_and_grad(loss, argnums=(0, 1))(
    jnp.array([0.0, 0.0, 0.0]), jnp.array([1.0, 0.0, 0.0]))
print(value, g1, g2)  # 6.873 [-1.891 0 0] [1.976 0 0]

volume integrates the SDF’s soft occupancy over a sampling box. Moving the spheres apart reduces their overlap, so the union’s volume grows along ±x.

Explicit arguments scale poorly once a scene has many parameters spread over its tree. Tagging parameters free=True on the geometry lets extract_parameters collect them, and functionalize turns the tree into a function of the parameter dictionaries:

from cadjoint import extract_parameters, functionalize
from cadjoint.geometry import Vector

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")
scene = Union(Translate(Sphere(radius=1.0), offset=p1),
              Translate(Sphere(radius=1.0), offset=p2))

free, fixed, metadata = extract_parameters(scene)
sdf = functionalize(scene)

value, grads = jax.value_and_grad(lambda params: volume(sdf(params, fixed)))(free)
print(grads)  # {'p1': [-1.891, 0, 0], 'p2': [1.976, 0, 0]}

extract_parameters returns three dictionaries: free parameter values keyed by name, fixed values keyed by tree path (union_0.smoothness, sphere_2.radius, …), and metadata holding the Parameter object behind each free value. functionalize(scene)(free, fixed) is a callable point → distance; the tree structure is fixed when it is built and the parameter values vary. Both forms give the same gradient. With parameter extraction, any node can declare a parameter free without changing the loss or the optimizer. Modelling covers parameters, functionalize and functionalize_parametric in full.

Objective and study forms

An Optimization takes one of two mutually exclusive forms: the objective form of the example, and the study form, which makes a declared ThermalStudy or ElasticStudy the objective, differentiating through geometry → mesh → solve → scalar:

from cadjoint.fem import Dirichlet, HeatFlux, Nodes, SimMesh, ThermalStudy
from cadjoint.geometry.parameters import Vector
from cadjoint.optimize import Optimization
from cadjoint.sdf.primitives import Box

bar = Box(Vector([1.0, 0.15, 0.15], free=True, name="size"))
mesh = SimMesh(name="bar-mesh", resolution=(22, 5, 5),
               bounds=(-1.1, -0.25, -0.25), size=(2.2, 0.5, 0.5))
study = ThermalStudy(
    name="bar-conduction",
    conductivity=2.0,
    bcs=[Dirichlet(Nodes.side("+x"), 0.0), HeatFlux(Nodes.side("-x"), 4.0)],
    mesh=mesh,
)

def volume(params):
    s = params["size"]
    return 8.0 * s[0] * s[1] * s[2]

run = Optimization(
    name="cool-bar",
    study=study,
    metric="max",
    regularizer=volume,
    regularizer_weight=0.1,
    steps=3,
    learning_rate=0.01,
    remesh_every=0,
).run(scene=bar)

print(run.history[0]["objective"], "->", run.objective)  # 4.018 -> 2.964
print(run.parameters)  # {'size': [0.970, 0.180, 0.179]}

Three Adam steps thicken the bar and lower peak temperature from 4.018 to 2.964; the volume penalty keeps it from growing without limit.

study= also accepts the name of a study, as scene files use it. Names resolve only inside a capture_studies() context, which the playground opens when it executes a scene; a standalone script passes the study instance.

Mixing the forms is refused: metric, regularizer, remesh_every, and a non-direct gradient_path all belong to the study form, and objective / of are rejected alongside study.

Metrics

metric reduces the solved field to one number.

metric ObjectiveMetric Means
"mean" ObjectiveMetric.MEAN The mean of the field — temperature, or displacement magnitude
"max" ObjectiveMetric.MAX Its maximum
"compliance" ObjectiveMetric.COMPLIANCE The work of the applied loads, f·u (twice the strain energy)

Every fixed option on Optimization has an enum in cadjoint.enums (ObjectiveMetric, GradientPath, OptimizerMethod with "adam" / "sgd"), and the member and the string are interchangeable.

"compliance" requires an elastic study with at least one Traction boundary condition; without loads the declaration is rejected.

The starter scene’s heat sink minimizes "max": the die is the hot spot, and a mean-temperature objective rewards deleting hot material instead of cooling the chip.

regularizer is a function of the parameters, added as objective + regularizer_weight * regularizer(params) and differentiated in reverse mode with the objective. In the starter scene it is a smoothed material volume computed by evaluating the traced SDF on a lattice.

Constrained optimization

Every update is projected back onto the constraint manifold. When extract_parameters reports a non-empty residual for the scene’s parameters, the optimizer chains a Newton projection after the optimizer transform:

import optax
from cadjoint import extract_parameters
from cadjoint.constraints import DistanceConstraint
from cadjoint.constraints.solve import make_manifold_projection
from cadjoint.geometry.parameters import Vector
from cadjoint.sdf.primitives import Sphere
from cadjoint.sdf.transforms import Translate

p = Vector([2.0, 0.0, 0.0], free=True, name="p")
anchor = Vector([0.0, 0.0, 0.0], free=False, name="anchor")
DistanceConstraint(p, anchor, 2.0)
scene = Translate(Sphere(radius=0.5), offset=p)
free, fixed, metadata = extract_parameters(scene)

transform = optax.chain(optax.adam(0.1),
                        make_manifold_projection(metadata, steps=2))

Optimization does this automatically. Horizontals stay horizontal and pinned dimensions stay pinned; the optimizer moves only within the sketch’s remaining degrees of freedom.

The same descent can be written by hand. The block below minimizes ‖p − target‖² subject to |p| = 2 from p₀ = (2, 0, 0) with target (1, 1.5, 0), whose optimum is p* = target · 2/|target| ≈ (1.109, 1.664, 0) at loss 0.0389, using five strategies for 40 steps each:

import jax
import jax.numpy as jnp
import optax
from cadjoint import extract_parameters
from cadjoint.constraints import (
    DistanceConstraint, constraint_residuals, make_manifold_projection, null_space,
)
from cadjoint.geometry import Vector

anchor = Vector(jnp.array([0.0, 0.0, 0.0]))
p = Vector(jnp.array([2.0, 0.0, 0.0]), free=True, name="p")
constraint = DistanceConstraint(anchor, p, distance=2.0)
free, fixed, metadata = extract_parameters(constraint)

target = jnp.array([1.0, 1.5, 0.0])
p_star = target * (2.0 / jnp.linalg.norm(target))       # optimum: (1.109, 1.664, 0)

def objective(params):
    return jnp.sum((params["p"] - target) ** 2)

value_and_grad = jax.value_and_grad(objective)

N0 = null_space(free, metadata)                          # tangent basis at the start
def grad_fixed(params):                                  # 1. fixed null space
    value, g = jax.value_and_grad(lambda z: objective(N0 @ z))(params @ N0)
    return value, N0 @ g

def grad_relin(params):                                  # 2./4. current null space
    N = null_space(params, metadata)
    value, g = value_and_grad(params)
    return value, N @ (g @ N)

LAMBDA = 25.0
def grad_penalty(params):                                # 5. penalty
    penalty = lambda q: objective(q) + LAMBDA * jnp.sum(constraint_residuals(q, metadata) ** 2)
    return objective(params), jax.grad(penalty)(params)

strategies = {
    "fixed null-space": (grad_fixed, optax.adam(0.05)),
    "relinearized null-space": (grad_relin, optax.sgd(0.15)),
    "manifold projection": (value_and_grad,
                            optax.chain(optax.adam(0.05), make_manifold_projection(metadata))),
    "Riemannian GD": (grad_relin,
                      optax.chain(optax.sgd(0.15), make_manifold_projection(metadata))),
    "penalty": (grad_penalty, optax.adam(2.5 / (1.0 + 2.0 * LAMBDA))),
}

for label, (grad_fn, opt) in strategies.items():
    params, state = free, opt.init(free)
    for _ in range(40):
        value, g = grad_fn(params)
        updates, state = opt.update(g, state, params)    # params= activates the projection
        params = optax.apply_updates(params, updates)
    violation = jnp.linalg.norm(constraint_residuals(params, metadata))
    print(f"{label:24s} loss {float(objective(params)):.4f}  |r| {float(violation):.1e}")
fixed null-space         loss 1.0000  |r| 5.0e-01
relinearized null-space  loss 0.1097  |r| 1.3e-01
manifold projection      loss 0.0392  |r| 1.2e-07
Riemannian GD            loss 0.0389  |r| 0.0e+00
penalty                  loss 0.0450  |r| 8.4e-03
Strategy Step Cost per step Constraint after each step Use when
Fixed null-space Gradient projected onto N Nᵀ, N = null_space(...) computed once at x₀ One gradient; N reused Drifts: the step follows the tangent plane at x₀, not at the current point Linear constraints, where the tangent plane does not rotate
Relinearized null-space N recomputed at the current point, gradient projected onto the current tangent plane One gradient plus one null_space Drifts by O(lr²κ) per step, with no correction A cheap manifold-aware direction, followed by a projection
Manifold projection optax.chain(opt, make_manifold_projection(metadata)) One gradient plus a Newton correction Exact Any optax optimizer; what Optimization uses
Riemannian GD Relinearized direction, then make_manifold_projection One gradient, one null_space, a Newton correction Exact Tangent-plane direction with exact feasibility; reached the optimum here
Penalty Loss ‖p − t‖² + λ‖r(p)‖² One gradient Approximate, error shrinks as λ grows No projection available; soft constraints

make_manifold_projection(metadata) computes params + update from the chained optimizer, projects it onto the constraint manifold with a Newton correction, and returns projected − params as the effective update, so optax.apply_updates(params, updates) lands on the manifold. It needs the current parameters: opt.update(grads, state, params) activates it, and standard optax transforms ignore the extra argument.

Relinearized steps use plain SGD rather than Adam: as the null space rotates, components that were zero acquire non-zero projections, and Adam’s element-wise scaling amplifies them and pushes the step off the tangent plane.

The penalty’s Hessian eigenvalue in the normal direction is 2 + 2λ, so a stable learning rate is at most 1/(1 + 2λ).

Frozen topology

Meshing has a discrete half (which cells are inside, how they connect) and a continuous half (where the nodes sit). Only the continuous half is differentiable, so the optimizer freezes topology and differentiates through node positions: each step re-projects the boundary nodes onto the current design’s zero set, Newton-corrected against the true SDF, with total motion clamped to half a cell diagonal.

remesh_every (default 6; 0 disables it) re-extracts the mesh at the current design every that many steps. The objective jumps at those steps are the discretization changing: on the bracket the step-12 re-extraction raised the objective from 6.02 to 8.86 and the step-18 one lowered it from 7.99 to 5.53, because the web is about two cells thick and a one-cell change in the loaded wall changes the discrete compliance substantially.

The final reported objective is evaluated on a freshly extracted mesh, so it does not depend on the last frozen topology.

A frozen topology also fixes every array shape, so the objective is compiled once per topology with jax.jit and each step costs a solve rather than a re-trace. On the starter scene this takes a step from ~11 s to 0.5–0.8 s, with the objective identical to 13 digits (performance notes).

Gradient paths

gradient_path selects which stages run outside JAX.

Value GradientPath Chain
"direct" (default) GradientPath.DIRECT Everything in-process. Node positions re-projected onto the true SDF, solved by the chosen backend.
"tesseract" GradientPath.TESSERACT Lattice samples → mesher Tesseract → solver Tesseract. Meshes the trilinear interpolant, so the gradient carries only normal boundary motion.
"tesseract-dc" GradientPath.TESSERACT_DC Dual contouring stays in JAX on the true SDF; only TetGen is wrapped, by the tetfill Tesseract with its exact pass-through VJP. Tet meshes only.

"tesseract-dc" keeps dual contouring in JAX against the true field, which preserves fin creases; the starter scene declares it. Both Tesseract paths need the tesseract extra. Whichever path runs the descent, the final result is evaluated on the direct path.

The stages behind both paths are plugins resolved by kind (mesher, tetfill, thermal_solver, elastic_solver), so where each runs is configuration. The plugins page gives the cost of each boundary.

Run results

run() returns an OptimizationRun:

  • history — one {"step", "objective", "grad_norm"} record per step, at the parameters before that step’s update. A non-finite objective or gradient norm aborts rather than propagating NaNs.
  • trajectory — {"step", "objective", "parameters"}, subsampled to at most 100 entries with the first and last always kept. This is what the viewer’s scrubber replays.
  • initial and parameters — before and after.
  • objective — the final objective.
  • result — the concrete SimulationResult from the final evaluation (study form only).

run(callback=fn) receives each history record as it is produced; the playground streams progress this way. method is "adam" or "sgd"; without optax it degrades to plain gradient descent and reports "gradient-descent".

Study runs execute inside an x64 scope, promote parameters to float64, and restore the target’s original parameter values on exit — including on failure.

A finished cool-sink run in the playground: the Results window replays the trajectory with its objective sparkline and scrubber (step 2 of 8), beside the Optimize window’s step count and learning rate

Bracket example

examples/fem_bracket_optimization.py runs the full loop on the parametric L-bracket from scenes/bracket.py: three named CAD parameters (web_thickness, rib_height, plate_thickness) → bracket SDF → HEX8 mesh → linear elastic solve → compliance plus mass → Adam with box-bound projection.

uv run python examples/fem_bracket_optimization.py            # ~10 min, CPU
uv run python examples/fem_bracket_optimization.py --smoke    # 2 cheap steps
uv run python examples/fem_bracket_optimization.py --backend calculix

Over 30 steps the objective falls 17.44 → 4.69, a 73% reduction; compliance drops 81% while mass rises 27%. The adjoint is validated against central differences at the initial design to 2.3e-5, 2.0e-3, and 3.3e-7 relative on the three parameters. --backend calculix runs CalculiX (Fortran) inside the same jax.grad call without changing the objective. Convergence CSV, a figure, and before/after VTK files land in examples/output/. The run record has the measured run, the rationale for the box bounds, and a known CalculiX deck fragility.

API