Gradient Behavior in rfx

rfx gradient workflows use JAX reverse-mode automatic differentiation through the implemented discrete FDTD calculation. This guide documents where gradients are usually well behaved, where they are numerically stiff or noisy, and how to check them before using them for RF design decisions. For the conceptual background, see Autodiff and Adjoint Background.

JAX traces the supported FDTD time-stepping and objective path as a computation graph. Per-step jax.checkpoint (the checkpoint=True default) or the segmented checkpoint knobs — checkpoint_segments on uniform grids, checkpoint_every on non-uniform grids — reduce reverse-mode memory by recomputing forward states during backpropagation. See the Memory-Reduction Path for how to size and pick these knobs.

Setup used by the snippets below:

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

sim = Simulation(
    freq_max=5e9,
    domain=(0.04, 0.03, 0.02),
    dx=2e-3,
    boundary="cpml",
    cpml_layers=6,
)
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_port(
    (0.008, 0.015, 0.010),
    "ez",
    impedance=50.0,
    waveform=GaussianPulse(f0=3e9, bandwidth=0.8),
)
sim.add_probe((0.032, 0.015, 0.010), "ez")

# Values the later snippets read.
s11 = sim.run(n_steps=600, compute_s_params=True).s_params[0, 0, :]
eps_min, eps_max = 1.0, 12.0
latent = jnp.zeros((8, 8, 4))

import jax
import jax.numpy as jnp

# sim is an already configured Simulation with sources/probes/ports.
# Discover the permittivity-array shape from a quick run, then build the
# starting array (result.grid is populated after every run() call):
eps0 = jnp.ones(sim.run(n_steps=1).grid.shape, dtype=jnp.float32)

def objective(eps_r):
    result = sim.forward(eps_override=eps_r, n_steps=400, checkpoint=True)
    return jnp.sum(result.time_series ** 2)

grad = jax.grad(objective)(eps0)  # discrete reverse-mode AD gradient

The gradient is a sensitivity of this discrete objective. It is not, by itself, a calibrated S-parameter, far-field, or resonance validation result.

Continuous changes in eps_r inside a fixed design region are the safest first case. Use bounded parameterizations, such as a sigmoid or projection from latent variables, so the optimizer cannot leave the material range you intend to test.

# Example bounded dielectric parameterization.
eps_design = eps_min + (eps_max - eps_min) * jax.nn.sigmoid(latent)

Objectives averaged over a band are often less sensitive to a single noisy frequency point:

loss = jnp.mean(jnp.abs(s11) ** 2)

This does not remove the need for a final port-family validation run. It only makes the optimization objective less dominated by one sample.

Some port paths carry a differentiable |S| channel: the lumped/wire path via forward(port_s11_freqs=...), and the compute_waveguide_s_matrix / compute_msl_s_matrix calculators via their eps_override argument. The uniform forward(...) lane can also differentiate scalar R/L/C values through rlc_values_override for elements registered with add_lumped_rlc(...). Each is differentiable only where its own support entry says so. Treat the AD path as a contract for the sensitivity calculation; the physical claim still follows the port-family evidence envelope. See Inverse Design for end-to-end optimization loops built on these paths.

Stairstepping creates discontinuities when a geometry parameter moves a PEC edge across a Yee cell boundary. The gradient can be large, change sign abruptly, or represent the rasterized discretization more than the intended CAD motion.

Mitigation: keep topology fixed during a gradient run, use smooth dielectric or relaxed occupancy variables where the workflow documents them, apply filters or minimum-feature constraints, and re-run a final discrete geometry validation. Do not rely on a larger learning rate to repair a discontinuous objective.

Long simulations can have weak gradient signal, large dynamic range, or more roundoff sensitivity. Checkpointing preserves the mathematical reverse-mode program for the chosen computation, but recomputation still occurs in finite precision.

Mitigation: use the minimum run length that captures the observable, record the time-window choice, and compare gradients against finite differences on a small set of cells or scalar parameters.

Near waveguide cutoff, beta changes rapidly with frequency and geometry. High-Q resonances can also make a single frequency sample highly sensitive to small mesh, material, or run-length changes.

Mitigation: keep optimization bands away from cutoffs unless that is the actual design target; use broadband or mode-tracking objectives; and verify the final design with a convergence or cross-reference check.

rfx runs these workflows in precision="float32" by default (complex64 field and DFT buffers). The uniform single-device lane can use precision="float64" after JAX x64 is enabled, or precision="mixed" for float16 field storage with float32 accumulators. Enabling x64 alone does not select float64 field storage. Mixed precision with CPML has a higher absorber residual floor and is not suitable for low-reflection or S-parameter observables near that floor. Non-uniform, distributed, and subgridded lanes reject both non-default precision modes instead of silently running float32. If the finite-difference step is too small, cancellation can make the finite-difference witness look worse than the AD path.

Mitigation: start finite-difference checks with a step that is meaningful for the variable scale, often around h = 1e-2 for permittivity-like variables, then adjust based on the loss magnitude and local sensitivity.

Gradients with respect to physical material variables inside CPML are not useful for RF design. CPML is an artificial absorber, not a device material. Exclude CPML cells from design regions.

Grid dimensions, CPML layer count, timestep count, and shape insertion/removal are discrete choices. Use them as fixed setup parameters, or run an outer design search that launches separate differentiable problems with fixed topology.

Subpixel averaging of material properties at geometry boundaries is precomputed once at setup and is not part of the JAX computation graph. Gradients do not flow through the subpixel weights — the material boundary is treated as fixed during reverse-mode differentiation. Exclude geometry boundary positions from design variables when subpixel smoothing is active.

If a source, port, mesh, or monitor combination is outside the documented support scope, finite AD output is not validation evidence. The correct outcome for an unsupported combination is an explicit support error or a clearly scoped local engineering check, not a public claim.

Always check a new objective against finite differences at a few representative cells or scalar design parameters before trusting a full optimization run:

def fd_check(objective, eps_r, cell=(10, 5, 5), h=1e-2):
    """Finite-difference gradient witness for one cell."""
    eps_p = eps_r.at[cell].add(h)
    eps_m = eps_r.at[cell].add(-h)
    fd = (objective(eps_p) - objective(eps_m)) / (2 * h)
    ad = jax.grad(objective)(eps_r)[cell]
    rel_err = abs(ad - fd) / max(abs(fd), 1e-30)
    print(f"AD={ad:.6e}, FD={fd:.6e}, err={rel_err:.2%}")
    return rel_err

Use the result as a witness, not as a universal tolerance. If the witness fails, inspect the objective scale, perturbation size, run length, monitor placement, and support envelope before changing optimizer settings.

A passing witness is also not sufficient on its own. Finite differences and AD differentiate the same objective through the same observation window, so both agree even when that window is empty — an objective that never sees the physics it is supposed to (reflection that never reaches the probe, a monitor inside an absorber) produces a self-consistent gradient of numerical noise. Guard against this by checking the loss magnitude against a physical expectation, not only the AD-vs-FD relative error: a reflected-energy proxy on a meaningfully reflecting design lands near 1e-2–1e-1, so a value like ~1e-7 signals an empty window, not a matched design.

1. Start with a small, supported setup before scaling the grid or objective.
2. Use bounded continuous variables for dielectric or occupancy design.
3. Exclude CPML and fixed metal from the design region unless the guide for that workflow explicitly documents a relaxed variable.
4. Check AD against finite differences on a small set of cells or scalar parameters.
5. Prefer broadband or averaged losses when a single frequency point is noisy or near a resonance null.
6. Use the segmented checkpoint knobs (checkpoint_segments on uniform grids, checkpoint_every on non-uniform grids) when reverse-mode memory is the limiting factor.
7. Validate the final RF observable through the relevant port, resonance, far-field, or convergence workflow.