# Autodiff and Adjoint Background

Channel: development; package: 2.0.0rc1; source: 2c32b7cf79b796fae825bfa08fc73d9914e61723

Page: https://remilab.ai/rfx/design/gradients/
Source: https://github.com/bk-squared/rfx/blob/2c32b7cf79b796fae825bfa08fc73d9914e61723/docs/public/design/gradients.mdx

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.

## Why gradients

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?

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

| Method | Cost | Use it for |
|---|---|---|
| Parameter sweep | one run per design point | 1–3 variables, trade studies ([Parametric Sweeps](https://remilab.ai/rfx/design/sweeps/)) |
| Finite difference | one or two runs per variable | checking a gradient on a few variables |
| Adjoint / reverse-mode autodiff | about one forward plus one backward run, for any number of variables | design regions and topology optimization |

## Autodiff versus a separate adjoint solve

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.

## What the gradient covers in rfx

- `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](https://remilab.ai/rfx/run/memory-and-speed/).
- Keep design variables out of the CPML. Gradients through absorber cells are
  not meaningful design information.

Typical choices:

| Task | Design variable | First objective | Verify with |
|---|---|---|---|
| Dielectric matching block | bounded `eps_r` in a design region | reflected energy, or `forward(port_s11_freqs=...)` | the port's S-parameter workflow |
| Waveguide taper | dielectric profile or smooth shape | transmitted energy | `compute_waveguide_s_matrix()` with a convergence check |
| Patch tuning | substrate or feed-region parameters | a resonance or probe proxy | resonance extraction on a converged mesh |
| Far-field shaping | dielectric density away from the CPML | NTFF directivity | the far-field workflow on a refined mesh |

## Check the gradient first

The [periodic multilayer slab example](https://github.com/bk-squared/rfx/blob/2c32b7cf79b796fae825bfa08fc73d9914e61723/examples/inverse_design/multilayer_ar_coating.py)
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:

```bash
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.

```python
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](https://remilab.ai/rfx/guide/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.

### A passing finite-difference check is necessary, not sufficient

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](https://remilab.ai/rfx/design/optimization/) explains the
round-trip condition behind `minimize_reflected_energy`, and
[Gradient Behavior](https://remilab.ai/rfx/guide/gradient-behavior/) shows how to check that a
run is long enough.

## Further reading

- [Meep adjoint solver tutorial](https://meep.readthedocs.io/en/latest/Python_Tutorials/Adjoint_Solver/):
  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"](https://doi.org/10.1364/OE.442074),
  *Optics Express* 30, 4467–4491 (2022).
- M. B. Giles and N. A. Pierce, ["An Introduction to the Adjoint Approach to
  Design"](https://doi.org/10.1023/A:1011430410075), *Flow, Turbulence and
  Combustion* 65, 393–415 (2000).
- S. Molesky et al., ["Outlook for inverse design in nanophotonics"](https://doi.org/10.1038/s41566-018-0246-9),
  *Nature Photonics* 12, 659–670 (2018).
- JAX [automatic differentiation guide](https://docs.jax.dev/en/latest/automatic-differentiation.html)
  and [Autodiff Cookbook](https://docs.jax.dev/en/latest/notebooks/autodiff_cookbook.html).
- A. Taflove and S. C. Hagness, *Computational Electrodynamics: The
  Finite-Difference Time-Domain Method*, 3rd ed., Artech House (2005).

## Limits

- **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](https://remilab.ai/rfx/run/memory-and-speed/).
- **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](https://github.com/bk-squared/rfx/blob/2c32b7cf79b796fae825bfa08fc73d9914e61723/docs/guides/support_matrix.md)).