Differentiability

How Maxwell computes gradients through FDTD with a native CUDA adjoint, which inputs are trainable, and what limits apply.

Maxwell is differentiable through PyTorch autograd. You build a scene from tensors with requires_grad=True, run the simulation, form a scalar loss from the result, and call loss.backward(). The gradient arrives in the .grad of your tensors. This page explains how that works and where its limits are.

The adjoint method#

Backpropagating through thousands of time steps by storing every intermediate field would need far too much memory. Maxwell instead uses the adjoint method: after the forward run, it runs the time-reversed (transposed) update equations, seeded by the derivative of your loss with respect to the monitor outputs. One backward run gives the gradient with respect to every material cell at once, at roughly the cost of one more simulation, no matter how many parameters you have.

Both the forward and the reverse updates are native CUDA kernels. Python handles checkpointing, replay and the link to autograd; there is no Python or finite-difference fallback.

Checkpointing#

The reverse run needs the forward fields in reverse order. Maxwell stores the full field state every adjoint_checkpoint_stride steps and recomputes the steps in between during the backward pass. The default stride is the square root of the number of time steps, which balances memory against recomputation. You can set it on FDTDConfig(adjoint_checkpoint_stride=...) when you build a Simulation directly.

What you can differentiate#

The gradient flows from a scalar loss through monitor outputs back to tensors that shape the compiled material, plus a set of RF inputs.

InputHow to make it trainable
Design densityMaterialRegion(name, geometry=Box(...), density=tensor, eps_bounds=(lo, hi)). The density in 0..1 interpolates permittivity (and permeability with mu_bounds) inside the box. Optional filter_radius (metres, a box filter) and projection_beta (tanh projection) give manufacturable designs.
GeometryTensor positions or sizes in geometry such as Box(position=tensor, ...). The gradient flows through subpixel averaging of the cut cells.
Material perturbationThe perturbation tensor of a PerturbationMedium.
RF inputsPort source amplitudes, lumped R, L and C values, circuit sources, and supported port and network objectives.

Supported loss inputs include point, plane and flux monitor values, port quantities (PortData, NetworkData) and their postprocessing (S, Z and Y parameters, gain, efficiency).

SceneModule is the natural way to organize this: keep parameters as nn.Parameter, build the scene in to_scene(), and hand the module to a PyTorch optimizer.

density = torch.full((20, 20, 4), 0.5, device="cuda", requires_grad=True)
scene.add_material_region(
    mw.MaterialRegion(
        name="design",
        geometry=mw.Box(position=(0.0, 0.0, 0.0), size=(0.4, 0.4, 0.08)),
        density=density,
        eps_bounds=(1.0, 4.0),
    )
)
result = mw.Simulation.fdtd(scene, frequencies=[f0]).run()
loss = -result.monitor("flux_up")["flux"]
loss.backward()          # density.grad now holds d(loss)/d(density)

Supported scene classes#

The native reverse kernels cover standard and CPML fields; electric conductivity; Debye, Drude and Lorentz dispersion (electric and magnetic); full electric anisotropy; chi2, chi3 and two-photon nonlinearity; complex Bloch fields; mixed Bloch and CPML axes; TFSF and grating TFSF; multiple sources; and single-device RF compositions with lumped ports, R/L/C elements, fixed single-mode wave ports and linear circuits. Multi-GPU FDTD has its own distributed adjoint.

Limits#

  • Maxwell checks every trainable scene at preparation time. If a feature combination has no implemented reverse path, prepare() or run() raises before the forward run starts. It never returns a silently wrong gradient.
  • A trainable input must reach the compiled material or a supported RF channel. Parameters that only move a source, for example a dipole position, are not part of the adjoint path.
  • A tensor passed as a plain material constant, such as Material(eps_r=tensor), is not tracked: the loss has no gradient, and inside a SceneModule the run raises NotImplementedError. To optimize permittivity, use a MaterialRegion with eps_bounds.
  • MaterialRegion geometry must be a Box, and basis must be "density".
  • Callable CustomSourceTime(fn) is not supported in the adjoint; give a sampled table instead.
  • The time step is fixed during differentiation (gradient_convention="fixed_dt" in result.stats()["time_step"]), so the gradient does not include the effect of a parameter on the automatic time step.
  • Batch execution supports a narrower set: density, dielectric geometry and named source-scale inputs. See Batch and multi-GPU runs.

Check a gradient#

Compare the adjoint gradient against a central finite difference on a small scene before you trust it in a long optimization:

def loss_at(z):
    box_z = torch.tensor(z, device="cuda")
    result = mw.Simulation.fdtd(build_scene(box_z), frequencies=[f0]).run()
    return result.monitor("probe")["data"].abs().item() ** 2
 
h = 1e-3
fd = (loss_at(h) - loss_at(-h)) / (2 * h)

Use a step large enough to move the subpixel fill of the cut cells, and keep the grid and time step identical between the runs.

Related#