# Inverse Design

Channel: release; package: 1.8.0; source: 7d1e97331656dd74ad0f5471653ac0c5effdff70

Page: https://remilab.ai/rfx/versions/v1.8.0/guide/inverse-design/
Source: https://github.com/bk-squared/rfx/blob/7d1e97331656dd74ad0f5471653ac0c5effdff70/docs/public/guide/inverse-design.md

rfx runs a JAX-differentiable FDTD: `jax.grad` returns reverse-mode gradients of
a scalar loss through the FDTD time-stepping, so the same solver that computes
the fields also computes design sensitivities. If you are coming from Meep's
adjoint terminology, start with
[Autodiff and Adjoint Background](https://remilab.ai/rfx/versions/v1.8.0/guide/autodiff-adjoint/).

## How it works

JAX traces the solver and objective path as a computation graph. Gradient
checkpointing recomputes forward states during backpropagation to bound
reverse-mode memory (see [Memory Reduction](https://remilab.ai/rfx/versions/v1.8.0/guide/memory-reduction/)).

```text
Forward:  eps_r → FDTD steps → probes / NTFF → scalar loss
Backward: jax.grad(loss)(eps_r) → ∂loss/∂eps_r
```

## Manual gradient loop

Most designs should use the higher-level [`optimize()` driver](#design-region-api),
which builds this loop for you. Drop to a manual `jax.grad` loop only for custom
objectives or training loops. Stay on the differentiable `Simulation.forward()`
API instead of hand-writing FDTD updates: build the `Simulation` once, pass an
`eps_override` array shaped like the grid permittivity, and differentiate a
scalar loss:

```python
import jax
import jax.numpy as jnp
from rfx import GaussianPulse, Simulation

# Stand-in for your own configured Simulation: one soft source, one probe,
# open (CPML) boundaries, deliberately small so the snippet runs in seconds.
sim = Simulation(freq_max=10e9, domain=(0.03, 0.02, 0.02), dx=2e-3,
                 boundary="cpml", cpml_layers=6)
sim.add_source((0.010, 0.010, 0.010), "ez",
               waveform=GaussianPulse(f0=5e9, bandwidth=0.8),
               amplitude_kind="current")   # required from 1.9; see Sources & Ports
sim.add_probe((0.020, 0.010, 0.010), "ez")

# Discover the grid shape from a quick run (result.grid is always populated):
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=150, checkpoint=True)
    return -jnp.sum(result.time_series ** 2)  # example proxy loss

grad = jax.grad(objective)(eps0)
```

This differentiates the proxy objective only. It does not by itself validate the
final RF observable; re-run the relevant port, resonance, far-field, or
convergence check on the optimized design (see [Validation](https://remilab.ai/rfx/versions/v1.8.0/guide/validation/)).

You may wrap the whole loss in an outer `jax.jit` for performance-sensitive
training loops:

```python
jit_objective = jax.jit(objective)
jit_grad = jax.jit(jax.grad(objective))
```

The supported `forward(eps_override=...)` path is regression-locked so the
outer-jitted loss and gradient match the eager path. Keep shape-changing setup
choices outside the jitted function; pass continuous arrays or scalar design
variables in as JAX arguments.

## Built-in objectives: choose the right family

### 1) Post-processed S-parameter objectives

Each factory returns `objective(result) -> scalar` that reads `result.s_params`,
so it applies to a completed `run(compute_s_params=True)` result. These raise
`ValueError` inside the default `optimize()`/`forward()` loop, which does not
build the post-processed S-parameter matrix (see family 2). An objective value
is only as trustworthy as the S-parameter extraction behind it — confirm that
port's S-parameter path is validated (see [Validation](https://remilab.ai/rfx/versions/v1.8.0/guide/validation/))
before treating the optimized number as a physical result.

```python
from rfx import minimize_s11, maximize_s21, target_impedance, maximize_bandwidth

obj_s11 = minimize_s11(freqs=jnp.array([5e9]), target_db=-10)
obj_s21 = maximize_s21(freqs=jnp.linspace(4e9, 6e9, 20))
obj_z = target_impedance(freq=5e9, z_target=50.0)
obj_bw = maximize_bandwidth(f_center=5e9, f_bw=2e9, s11_threshold=-10)
```

### 2) Differentiable loop objectives for `optimize()`

The `optimize()`/`forward()` pass emits `result.time_series` but not the
post-processed `result.s_params` matrix, so the family-1 objectives cannot run
inside a gradient loop. Use these time-domain proxies instead — they read only
`result.time_series` and compose with both `optimize()` and `topology_optimize()`:

```python
from rfx import minimize_reflected_energy, maximize_transmitted_energy

obj_reflect = minimize_reflected_energy(port_probe_idx=0)
obj_transmit = maximize_transmitted_energy(output_probe_idx=-1)
```

`minimize_reflected_energy` is a late-time-reflection S11 proxy;
`maximize_transmitted_energy` is an output-power S21 proxy. Both index
`result.time_series` columns: `port_probe_idx` selects the probe co-located with
the excitation port, `output_probe_idx` the downstream probe.

**Precondition — the split window must contain the reflection.**
`minimize_reflected_energy` splits the probe time series at `late_fraction`
(default: the second half) and treats the late half as reflected energy. That
premise only holds if the round trip from the port to the reflecting feature and
back arrives *after* the split point and *before* the run ends. On a short
round-trip geometry (a thin substrate, a feature close to the port), the
reflection lands in the early "incident" half and the late window is nearly empty,
so the proxy collapses to `~0` and its gradient becomes numerical noise — with no
error raised. Sanity-check the loss magnitude (a meaningfully reflecting design
sits near `1e-2`–`1e-1`, not `~1e-7`), then fix the window by its failure mode:
enlarge `n_steps` if the reflection arrives after the run ends; **raise**
`late_fraction` (which moves the split earlier) if it arrives before the split; or,
if the incident pulse and reflection overlap in time on a very short round trip,
narrow the source bandwidth to compress the incident pulse — or use an
impedance-referenced port (`add_port(..., impedance=Z0)` with
`forward(port_s11_freqs=...)`), which separates incident and reflected waves
exactly. Note that an AD-vs-finite-difference check will *not* catch an empty
window: both differentiate the same window and agree (see
[Autodiff and Adjoint Background](https://remilab.ai/rfx/versions/v1.8.0/guide/autodiff-adjoint/#a-passing-finite-difference-check-is-necessary-not-sufficient)).

For NTFF/directivity optimization, `maximize_directivity(...)` defaults to
`log_ratio=True`, which is sign-correct for every degree of freedom, including
design variables that change total radiated power (conductors/PEC, lossy, or
magnitude-changing dielectric DoFs). 
`log_ratio=False` is the legacy mode kept for back-compatibility: it drops a
quotient-rule term and yields wrong-sign gradients for power-changing DoFs, so
it is correct only for shape-preserving, constant-radiated-power ones. The
default flipped in 1.6.7 (CHANGELOG `[1.6.7] - 2026-07-28`, GitHub #129); code
written against the older default should pass `log_ratio=False` explicitly if it
relied on it.

## Design-region API

`optimize()` is the high-level driver: give it a configured `Simulation`, a
`DesignRegion` box (physical `corner_lo`/`corner_hi` in metres plus an
`eps_range` the design is clamped to), and a proxy objective from family 2. It
runs an Adam gradient loop over the region's permittivity and returns an
`OptimizeResult` with `eps_design` (optimized permittivity in the box),
`loss_history`, and the final `latent` parameters. For port/probe setup on the
base simulation, see [Sources & Ports](https://remilab.ai/rfx/versions/v1.8.0/guide/sources-ports/).

```python
from rfx import Simulation, DesignRegion, optimize, minimize_reflected_energy

sim = Simulation(freq_max=10e9, domain=(0.1, 0.04, 0.02), dx=2e-3,
                 boundary="cpml", cpml_layers=6)
feed = (0.015, 0.02, 0.01)
sim.add_port(position=feed, component="ez", impedance=50.0)
sim.add_probe(feed, "ez")   # time_series column 0 -> port_probe_idx=0

region = DesignRegion(
    corner_lo=(0.03, 0.0, 0.0),
    corner_hi=(0.07, 0.04, 0.02),
    eps_range=(1.0, 6.0),
)

result = optimize(
    sim,
    region,
    objective=minimize_reflected_energy(port_probe_idx=0),
    n_iters=5,      # a real design runs tens to hundreds of iterations
    lr=0.01,
)
# result.eps_design, result.loss_history, result.latent
```

The region is clamped to the grid interior; `optimize()` raises `ValueError` if
it lies entirely inside the CPML absorber, so keep `corner_lo`/`corner_hi` within
the physical domain rather than the padding.

For difficult or multimodal proxy losses, `optimize()` also exposes default-off
robustness knobs:

| Knob | Use | Default behavior |
|---|---|---|
| `n_starts=N` | run `N` Adam restarts and return the best one | `1`, the legacy single run |
| `best_iterate=True` | return the lowest-loss visited iterate instead of the final iterate | `False`, final iterate |
| `step_clamp=value` | cap the L2 norm of each Adam latent update | `None`, unclamped |
| `seed=...` | make extra restart initializations reproducible | used only when `n_starts > 1` |

These knobs do not change the objective or the electromagnetic model. They help
only when the loss surface itself benefits from restarts, overshoot protection,
or reproducible restart sampling.

## Scalar lumped-RLC value design

A lumped R/L/C element registered with `add_lumped_rlc(...)` is a circuit element
inside the FDTD update — it produces no S-parameters by itself. To measure the
load, pair it with a co-located `add_port(..., impedance=Z0)`: the port supplies
both the excitation and the S11 accumulator that `forward(port_s11_freqs=...)`
reads. On the uniform single-device `forward(...)` lane the registered R/L/C
values affect the differentiable run, and scalar component values can enter the
AD tape through `rlc_values_override`:

```python
import jax
import jax.numpy as jnp
from rfx import Simulation

sim = Simulation(freq_max=10e9, domain=(0.02, 0.02, 0.02),
                 boundary="cpml", cpml_layers=6)
load_pos = (0.0093, 0.0093, 0.0093)
sim.add_port(position=load_pos, component="ez", impedance=50.0)
sim.add_lumped_rlc(
    position=load_pos,
    component="ez",
    R=50.0,
    C=0.2e-12,
    topology="series",
)

def load_loss(R):
    result = sim.forward(
        n_steps=800,
        port_s11_freqs=jnp.array([5e9]),
        rlc_values_override={0: {"R": R}},
    )
    # forward()'s lumped-port S11 is per-frequency, shape (n_freqs,)
    return jnp.abs(result.s_params[0]) ** 2

dloss_dR = jax.grad(load_loss)(50.0)
```

The mapping key is the 0-based registration order of `add_lumped_rlc(...)`
calls; missing keys fall back to the registered float value. This surface is
uniform single-device only. Note the division of labour in the example: the S11
in the loss comes from the co-located `add_port` (the reference-impedance
measurement), while `add_lumped_rlc` contributes the circuit element being
designed — the RLC element is **not** a port type and produces no `s_params` on
its own.

## Far-field objectives with NTFF data

`maximize_directivity` optimizes the directivity ratio toward a target direction.
It reads `result.ntff_data`, `result.ntff_box`, and `result.grid`, so the base
simulation must register a near-to-far-field box first via
`sim.add_ntff_box(corner_lo, corner_hi)`; without it the objective raises
`ValueError`. `theta_target`/`phi_target` are in radians.

```python
from rfx import maximize_directivity

objective = maximize_directivity(
    theta_target=0.0,  # radians
    phi_target=0.0,
    log_ratio=True,
)
```

Pass this to `optimize()` like any family-2 proxy. (`optimize()` also forwards
`result.ntff_box` as a keyword to any custom objective whose signature includes
an `ntff_box` parameter; the built-in objective does not need that — it reads the
box from the result.) NTFF passes cost more than probe-only losses, so iterate on
a coarse grid, then re-run the final design through the far-field validation path
(see [Far-field & RCS](https://remilab.ai/rfx/versions/v1.8.0/guide/farfield-rcs/)).

## Tips

- **Memory**: `checkpoint=True` is the default; for large `n_steps`, use `checkpoint_segments` (uniform) or `checkpoint_every` (non-uniform) to trade ~2x compute for ~√n_steps memory. See [Memory Reduction](https://remilab.ai/rfx/versions/v1.8.0/guide/memory-reduction/).
- **Start with small grids** for design iteration, then scale up for the final verification run.
- **Learning rate**: `0.01–0.1` is a reasonable first range for permittivity optimization; see [Gradient Behavior](https://remilab.ai/rfx/versions/v1.8.0/guide/gradient-behavior/).
- **Proxy objectives first**: start with `minimize_reflected_energy()` or `maximize_transmitted_energy()`.
- **NTFF objectives cost more** than probe-only losses; reserve them for radiation targets.
- **GPU acceleration** depends on the installed JAX/CUDA environment; verify device placement for performance-sensitive runs.