Skip to content

Gradients through the simulation

Every number in a StackDevice (the vertices of every polygon, the heights, thicknesses and sidewall angles of every level) is differentiable through the whole simulation. The example differentiates the phase a 90° bend adds to the light with respect to the bend's radius with jax.value_and_grad, and compares the result with a central finite difference of the same simulation. Run it on a GPU.

import fdtdx
import jax
import jax.numpy as jnp
import gdsfactory as gf
import bifrost as bf
from gdsfactory.technology import LayerLevel, LayerStack

gf.gpdk.PDK.activate()

RADIUS = 2.0  # um: the parameter we differentiate with respect to
RESOLUTION = 0.04  # um, the grid of the bend tutorial

layer_stack = LayerStack(
    layers={"core": LayerLevel(layer=gf.gpdk.LAYER.WG, zmin=0.0, thickness=0.22, material="core", mesh_order=1)}
)
materials = {"core": fdtdx.Material(permittivity=3.48**2), "clad": fdtdx.Material(permittivity=1.444**2)}

# a gdsfactory bend (any differentiable cell works); the ports are extended 2 um, and the box is 6 x 6 x 3.4 um.
# 2 um is tighter than the strip's 3.5 um minimum radius, so that check is off.
component = gf.components.bend_circular(radius=RADIUS, angle=90, allow_min_radius_violation=True)
stack = bf.StackDevice.from_layer_stack(
    component, layer_stack, materials, background="clad", zrange_layers=["core"],
    extend_ports=(("o1", 2.0), ("o2", 2.0)), boundary_offset=(1.75, 0, 0, -1.75, 1.59, -1.59), name="stack",
)

config = stack.get_simulation_config(
    time=400e-15, grid=fdtdx.core.grid.UniformGrid(spacing=RESOLUTION * 1e-6),
    # how the backward pass trades memory for time: it re-simulates from 60 stored states
    # instead of keeping every step, so the memory stays small
    gradient_config=fdtdx.GradientConfig(method="checkpointed", num_checkpoints=60),
)
wave = fdtdx.WaveCharacter(wavelength=1.55e-6)
pulse = fdtdx.GaussianPulseProfile(
    center_wave=wave, spectral_width=fdtdx.WaveCharacter(frequency=0.05 * wave.get_frequency())
)
te0 = dict(mode_index=0, filter_pol="te")
volume, constraints = stack.get_simulation_volume()
objects = [volume, stack]
for obj, placing in (
    stack.create_mode_plane_source_at(
        "o1", (0.5, 0.3), wave_character=wave, temporal_profile=pulse, port_offset=-0.5, name="source", **te0,
    ),
    stack.create_mode_overlap_detector_at(
        "o1", (0.5, 0.3), direction="in", wave_characters=(wave,), scaling_mode="pulse", name="in", **te0,
    ),
    # the output monitor sits 0.5 um out in the straight extension (port_offset < 0 moves a plane out of the
    # component), so that it stays in the straight lead for every radius we try; on the port itself it would
    # sit inside the bend as soon as the radius grows
    stack.create_mode_overlap_detector_at(
        "o2", (0.5, 0.3), wave_characters=(wave,), scaling_mode="pulse", port_offset=-0.5, name="out", **te0,
    ),
):
    objects.append(obj)
    constraints += placing
boundaries, placing = fdtdx.boundary_objects_from_config(
    fdtdx.BoundaryConfig.from_uniform_bound(thickness=round(1.0 / RESOLUTION)), volume
)
objects += boundaries.values()
constraints += placing

key = jax.random.PRNGKey(0)
objects, arrays, params, config, _ = fdtdx.place_objects(objects, config, constraints, key)
params = {**params, **bf.init_stack_params(objects)}
placed = objects["stack"]  # the device as placed: it can swap in the vertices of another component


def ratio(radius):
    """The complex ratio of the TE0 amplitude on the output monitor to that on the input monitor."""
    bend = gf.components.bend_circular(radius=radius, angle=90, allow_min_radius_violation=True)  # new radius
    p = {**params, "stack": {**params["stack"], **placed.pack(bend)}}  # its vertices into the stack's parameters
    arrays2, objects2, _ = bf.apply_stack(arrays, objects, p, key)
    _, arrays2 = fdtdx.run_fdtd(arrays=arrays2, objects=objects2, config=config, key=key, show_progress=False)
    a_in = objects2["in"].compute_overlap(arrays2.detector_states["in"])
    a_out = objects2["out"].compute_overlap(arrays2.detector_states["out"])
    return (a_out / a_in)[0]


# The overall sign of a mode is arbitrary in each mode solve, so the phase of the ratio is only defined
# modulo pi (ratio**2 is not affected). We compare it with the phase at the base radius, which keeps the
# result away from the branch cut: it is the extra phase (rad) that a bend of this radius adds.
reference = jnp.angle(ratio(RADIUS) ** 2) / 2


def phase_shift(radius):
    """The phase (rad, modulo pi) of the bend of this radius relative to the base radius."""
    return jnp.angle(ratio(radius) ** 2 * jnp.exp(-2j * reference)) / 2


value, gradient = jax.value_and_grad(phase_shift)(RADIUS)  # the simulation runs forward and backward
h = 0.05  # um: a central difference over +-h
finite = (phase_shift(RADIUS + h) - phase_shift(RADIUS - h)) / (2 * h)
print(f"phase shift at {RADIUS} um: {float(value):+.3f} rad (zero by construction)")
print(f"d(phase)/dR: autodiff {float(gradient):+.2f} rad per um, finite difference {float(finite):+.2f} rad per um")

Output:

phase shift at 2.0 um: +0.000 rad (zero by construction)
d(phase)/dR: autodiff -5.64 rad per um, finite difference -5.64 rad per um

On an NVIDIA L40S the script took 133 s: four simulations, one of them also run backward. In a scan of the phase shift at radii from 1.95 to 2.1 µm, the finite difference was -5.39, -5.58 and -5.64 rad per µm for steps of ±0.01, ±0.02 and ±0.05 µm: the smallest steps are the noisiest, because the grid's cells are 0.04 µm wide.

How the script works:

  • stack.pack(component) returns the stack's vertex parameters for a new component of the same polygon topology. Rebuilding the bend with a new radius inside the objective and merging its vertices into the parameters ({**params["stack"], **placed.pack(bend)}) makes the result a function of the radius.
  • bf.apply_stack writes the stack into the arrays, fdtdx.run_fdtd solves, and JAX differentiates through all of it. fdtdx.GradientConfig(method="checkpointed", ...) bounds the memory of the backward pass: without a gradient configuration the run stores every time step and runs out of memory quickly. Simulations with lossy materials need the checkpointed method, because fdtdx's reversible gradients amplify round-off in lossy voxels; StackDevice warns about that combination.
  • This objective runs eagerly: bend_circular picks its number of points from the radius, which needs a concrete value. A cell function whose structure does not depend on the parameters (the same polygons and number of vertices, no if on a parameter) can be traced, and then the whole objective can be wrapped in bf.jit (jax.jit with an XLA option for gradients through fdtdx's time loop), as Inverse design does; discrete choices and KLayout operations work only eagerly (see the gdsfactoryx README). So does apply_stack in dispersive simulations with mode or plane-wave sources (see dispersive materials).

Worth copying into your own objectives:

  • Choose an objective with a signal. The transmission of a good bend is nearly flat in the radius (it barely moves around 2 µm), and the grid's cell-scale noise is as large as its slope, so its derivative tells you little. The phase changes by several radians per µm. A bend tight enough to lose light is the opposite case: its transmission has a signal, and Inverse design optimizes it.
  • Keep monitors in the straight extension for every parameter value. The output monitor sits 0.5 µm out in the port extension. On the port itself, a larger radius would put the end of the arc past the monitor.
  • A mode's overall sign is arbitrary in each mode solve. The ratio of two mode amplitudes is defined only up to a sign, so its phase jumps by π between some parameter values, which would ruin a finite difference. The script compares ratio**2, which does not depend on the sign, with a reference at the base radius.

The stack's parameters

The stack's parameters live in params[name], in stack units (µm and degrees):

params["stack"] = {
    "vertices/<layer>": ...,       # (E, 2), one entry per source layer with polygons
    "zmin/<level>": ...,           # ()
    "thickness/<level>": ...,      # ()
    "sidewall_angle/<level>": ..., # ()
}

Thicknesses, heights and sidewall angles are ordinary entries and can be differentiated directly. Every level's zmin is independent: making a film thicker does not lift the levels deposited on it. To keep them stacked, derive their zmin from the level below inside the objective, e.g. for a cladding deposited on the core:

def stacked(t_core):
    return {**params["stack"], "thickness/core": t_core, "zmin/clad": params["stack"]["zmin/core"] + t_core}

Every level in the box is computed over the whole box height, so zmin, thickness and roughness can move it anywhere inside the box. Levels outside the box are not part of the device: level_margin= (µm) in from_layer_stack includes the levels within that distance of the box, so they can be moved into it.

To differentiate with respect to a cell parameter, rebuild the component inside the objective and swap in its vertices with stack.pack, as the script does. pack requires the polygon topology the device was built from: the same number of polygons, and of vertices per polygon, on every layer.

What the derivatives are

They are the derivatives of the discretized simulation, not of the physical device; through a whole fdtdx solve they agree with its finite differences to 1e-3. Averaging puts an interface's material into the voxel it crosses, so a response is piecewise smooth in an interface's position, with a kink where the interface crosses a sample voxel's side. At high contrast its slope inside one voxel can be far from the physical one, while its values and its slope over a voxel follow them. In the anisotropic slab test (benchmarks/subpixel/em_slab.py --grad, contrast about 9), d|t|^2/d(thickness) is 0.13 at 20 nm and 0.59 at 10 nm, against 0.82 from the transfer matrix, while |t|^2 is within 0.6 %. Refine the grid where the optimum depends on a gradient with respect to an interface's position.

Where a surface lies exactly on a voxel face (a thickness or zmin that is a whole number of cells from the box's bottom, at some Yee position), the voxel fractions have a kink: the derivative there is the mean of the two one-sided ones, half of the surface's movement in each of the two voxels, with every material's fraction changing consistently.

On a non-uniform grid the grid is part of the compiled program, like a uniform one: a fixed grid compiles once, and the derivatives are those of that grid's discretization (none with respect to the cell edges). The kinks above sit on that grid's own voxel faces, so they are closer together in fine cells; refine where an interface whose position you optimize can move.

Polygon topology

The device takes a few things from the vertices it is built with, once, and indexes later vertices with them by polygon and vertex: which corners lie inside another polygon's edge (T-junctions, such as a taper on the side of an MMI), which edges are seams between abutting polygons (left out of the distance), the order of the polygons, which way each one runs, and where each port extension meets the polygons on the port face. A parameter change that breaks one of these would silently give a wrong geometry: a taper moved off an MMI side leaves a bump on that side, two abutting polygons moved apart keep the gap between them filled, and a component that lists two polygons of the same size in the other order hands each one the other's topology.

So whenever the component is concrete (not traced), pack calls stack.check_topology(component), which raises ValueError naming each layer and what changed in it. Traced components, inside jax.grad or jax.jit, are never checked, so the check costs nothing in the objective. Round-off, layout snapping and moves that keep the topology, such as a taper sliding along an MMI side or abutting polygons moving together, pass. pack(component, check=False) skips the check.

All of this uses one distance, the device's topology_tol: a corner this close to another polygon's edge lies on it, vertices this close to each other coincide, and a vertex this close to a port face is on it. The default is 2 database units of the layout (gf.kcl.dbu, so 2 nm), because snapping to the 1 nm grid can move a corner of a rotated component up to about 1.4 nm off the slanted edge it lies on; float32 round-off of the coordinates (2 epsilons times the largest one, about 1.2 nm at 5000 µm from the origin) is added, whatever the device's dtype. A corner within it bends the other polygon's edge through itself, by at most that distance, instead of leaving a slit between the polygons. The check, a new device (from_layer_stack) and retopologize all use the same tolerance, so whatever the check accepts, a device built for the new vertices would also have. Moves smaller than it (a taper 1 nm off an MMI side, a waveguide end 1 nm short of an extended port's face) count as keeping the topology. Abutting edges whose corners are up to that far apart but not at the same point (which snapping does not produce from corners that coincide) are a seam too: a voxel whose center falls in the gap between them, if there is one, is filled as if it were far outside. Pass topology_tol=... (µm) to from_layer_stack to change it.

The check also reports polygons on a stack layer that had none when the device was built: the levels read such a layer as empty, and a level that needs it (e.g. a slab level of a device built from a strip waveguide) was not drawn at all. Only StackDevice.from_layer_stack builds a device for them; retopologize raises.

When the new topology is intended, stack.retopologize(component) returns the same device (box, levels, materials, placement) with the topology of the new component, whose polygons may also change in number and size. Swap it into the placed objects and refresh its parameters:

objects = objects.copy()
objects["stack"] = objects["stack"].retopologize(component)
params = {**params, "stack": {**params["stack"], **objects["stack"].pack(component)}}

check_topology and retopologize also take a parameter dict (e.g. params["stack"]) instead of a component, for vertices optimized directly; port extensions and new layers are then not checked, and the extensions are not found again.