Autodiff and Adjoint Background

rfx gives you the gradient of a scalar design objective, such as reflected energy or |S21|, with respect to thousands of design variables at the cost of about one extra simulation. This page explains where that gradient comes from, what it does and does not tell you, and how to check it before you optimize.

A normal RF design loop is: choose geometry and materials, simulate, extract a metric, change the design, repeat. With one or two variables, a sweep or a finite difference is enough. With a design region of hundreds of cells, finite differences need one or two runs per variable, which is too slow.

A gradient answers a narrower question: at the current design, which direction in design space improves a scalar loss fastest?

design variables -> FDTD fields -> probes / ports / far field -> scalar loss J
   eps_r(x)            E, H(t)       S11, S21, power             dJ/d eps_r

Solvers such as Meep run a forward simulation and then a separate adjoint simulation with special sources, and assemble the gradient from both fields.

rfx takes a different route. The whole time-stepping loop is written in JAX, so jax.grad differentiates the discrete program directly (reverse-mode automatic differentiation). Mathematically this is a discrete adjoint: both compute a vector-Jacobian product for a scalar objective. In practice there is no adjoint run object to manage. You write a loss function and call jax.grad(loss_fn)(params).

The ideas from adjoint design still apply:

1. The objective must be an explicit scalar. Reflected energy, transmitted energy, directivity at one angle, mismatch to a target impedance. The solver does not know you want a “good antenna”.
2. Design variables must be continuous. Permittivity, a material density, a smooth shape parameter. Cell counts, on/off geometry and CAD booleans are not differentiable.
3. The gradient belongs to the observable you wrote. A reflected-energy proxy is not a calibrated S11. Verify the final design with the matching S-parameter or far-field workflow.
4. Optimization still needs engineering. Filters, projection, minimum feature size, mesh convergence and an independent final check are still your job.

• sim.forward(...) is the differentiable entry point. It takes overrides such as eps_override, sigma_override or a design_box with design_eps_override, and returns time series and other observables as JAX arrays.
• The result is the exact gradient of the discrete simulation you set up, not proof that your loss is the right RF metric.
• By default forward uses checkpoint=True, which recomputes each step in the backward pass to limit memory. For long runs, segmented checkpointing (checkpoint_segments on a uniform mesh, checkpoint_every on a non-uniform one) cuts memory to roughly sqrt(n_steps). See Memory Reduction.
• Keep design variables out of the CPML. Gradients through absorber cells are not meaningful design information.

Typical choices:

The periodic multilayer slab example uses the public Simulation.freeze_mesh() accessor and prints a gradient_record_length_witness result for its reflection objective. Run its smallest workflow without optimization or plotting:

JAX_PLATFORMS=cpu python examples/inverse_design/multilayer_ar_coating.py --gradient-witness-only

A passing record-length witness compares gradients at two run lengths; it does not establish mesh convergence or agreement with an external solver.

Before you optimize a new setup, compare the autodiff gradient with a central finite difference at one cell. The model below is a small closed PEC box, so every grid index is a physical cell and nothing lands in an absorber. The mesh is coarse so the check runs in seconds.

import jax
import jax.numpy as jnp
from rfx import Box, GaussianPulse, Simulation

sim = Simulation(freq_max=4e9, domain=(0.04, 0.03, 0.02), dx=2e-3,
                 boundary="pec")
sim.add_material("slab", eps_r=4.0, sigma=0.01)
sim.add(Box((0.014, 0.010, 0.006), (0.026, 0.020, 0.014)), material="slab")
sim.add_source((0.008, 0.015, 0.010), "ez",
               waveform=GaussianPulse(f0=3e9, bandwidth=0.8),
               amplitude_kind="current")
sim.add_probe((0.032, 0.015, 0.010), "ez")

# Design variable: the full permittivity grid.
params = jnp.full(sim.run(n_steps=1).grid.shape, 4.0, dtype=jnp.float32)


def loss_fn(eps_r):
    result = sim.forward(eps_override=eps_r, n_steps=200)
    return jnp.sum(result.time_series ** 2)


ad_grad = jax.grad(loss_fn)(params)

idx = (10, 7, 5)          # a cell inside the slab
h = 1e-2                  # FD step in eps_r units
fd_grad = (loss_fn(params.at[idx].add(h))
           - loss_fn(params.at[idx].add(-h))) / (2 * h)

rel_err = abs(ad_grad[idx] - fd_grad) / max(abs(fd_grad), 1e-30)
print(f"loss = {float(loss_fn(params)):.3e}, AD/FD relative error = {float(rel_err):.2e}")

On this model the two agree to a few parts in a thousand. The loss is in arbitrary units here (a raw probe signal), so only the relative error matters.

rfx runs in float32, so a very small finite-difference step is dominated by round-off. Start with h = 1e-2 for permittivity-like variables and raise it toward 1e-1 when the loss is large (see Gradient Behavior).

rfx.gradient_check(sim, design_params, objective_fn, eps=...) runs the same comparison for you. It adds design_params to the model’s permittivity, calls objective_fn on the forward result, and perturbs every element of design_params in turn, so keep it small (a scalar checks a uniform permittivity shift). Its default step eps=1e-3 is small for float32; pass a larger one.

Agreement between autodiff and finite differences validates the autodiff machinery: the derivative of the loss you wrote is computed correctly. It does not show that the loss observes the physics you intended. Both methods differentiate the same loss over the same observation window. If that window is empty, for example because the reflection you are minimizing never reaches the probe during the run, or the monitor sits in the absorber, they still agree, on a gradient of numerical noise.

So also check the loss’s absolute value against a physical expectation. A time-domain reflected-energy loss for a design that reflects a real fraction of the pulse lands around 1e-2 to 1e-1. A loss many orders of magnitude smaller (say ~1e-7) usually means the reflection never landed in the window, and the gradient is meaningless however well it matches finite differences. Fix the setup (a longer run, a correct monitor or split placement) before you read the derivative. Inverse Design explains the round-trip condition behind minimize_reflected_energy, and Gradient Behavior shows how to check that a run is long enough.

• Meep adjoint solver tutorial: the forward/adjoint two-run structure and MaterialGrid design variables.
• A. M. Hammond et al., “High-performance hybrid time/frequency-domain topology optimization for large-scale photonics inverse design”, Optics Express 30, 4467–4491 (2022).
• M. B. Giles and N. A. Pierce, “An Introduction to the Adjoint Approach to Design”, Flow, Turbulence and Combustion 65, 393–415 (2000).
• S. Molesky et al., “Outlook for inverse design in nanophotonics”, Nature Photonics 12, 659–670 (2018).
• JAX automatic differentiation guide and Autodiff Cookbook.
• A. Taflove and S. C. Hagness, Computational Electrodynamics: The Finite-Difference Time-Domain Method, 3rd ed., Artech House (2005).

• Loss far below a physical scale (for example ~1e-7 for a reflection proxy): the observable window is probably empty. Lengthen the run or move the monitor before trusting the gradient.
• Design cells inside the CPML: the gradient there is not design information. Keep the design region in the physical domain.
• Out-of-memory in the backward pass: use segmented checkpointing and the memory estimators in Memory Reduction.
• A proxy loss optimized well, but the real metric did not improve: verify the final design with the calibrated workflow for that metric (see the support matrix).