Compact inverse-designed grating coupler

This notebook reconstructs the official Tidy3D Autograd6GratingCoupler.ipynb problem with FDTDX. It retains the 3D silicon/BOX stack, 10° tilted Gaussian beam, electric y symmetry, exact 216×108 half-domain controls at 20 nm, random Gaussian-smoothed cold start, double constant-padded 80 nm conic filter, erosion/dilation penalty, and Adam/β continuation.

The result is intentionally shown as a failed parity audit. The final design is fully binary and trustworthy as a forward simulation, but it does not match the published optimizer performance. Every geometry, field, spectrum, and comparison below is generated from the recorded FDTDX artifact or a fresh forward solve—not from an illustrative schematic alone.

Imports, accelerator, and benchmark record

import os
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")

import json
from pathlib import Path
from time import perf_counter
import warnings
import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, Rectangle
import numpy as np
import pandas as pd
import jax
import jax.numpy as jnp
import yaml
import fdtdx
from IPython import get_ipython

from benchmarks.invdes_grating_coupler_runner import (
    BETA_MAX, BORDER_BUFFER_UM, BOX_BOTTOM_Z_UM, BOX_THICKNESS_UM,
    CONTROL_SPACING_UM, DESIGN_CENTER_X_UM, DESIGN_SIZE_UM,
    DOMAIN_SIZE_UM, EPS_AIR, EPS_SILICA, EPS_SILICON,
    FDTD_DL_UM, FDTD_DZ_UM, FIBER_TILT_DEG, GRATING_LENGTH_UM,
    GRATING_WIDTH_UM, LEARNING_RATE, MIN_FEATURE_UM,
    PARAMETER_SHAPE, PUBLISHED_ITERATIONS, RUN_TIME_S,
    SLAB_BOTTOM_Z_UM, SLAB_CENTER_Z_UM, SLAB_TOP_Z_UM,
    SOURCE_Z_UM, SPOT_SIZE_UM, WAVELENGTH0_UM, WAVELENGTHS_UM,
    fabrication_penalty, hard_binary_design, initial_parameters,
    optimize_grating_coupler, pre_process, simulate_grating_coupler,
)

get_ipython().run_line_magic("matplotlib", "inline")
warnings.filterwarnings("ignore", category=FutureWarning, module=r"tidy3d\.components\.mode")
warnings.filterwarnings("ignore", message=r"GaussianModeOverlapDetector.*truncates")
plt.rcParams.update({"figure.figsize": (10, 4.8), "figure.dpi": 120})
meta = yaml.safe_load(Path("benchmarks/cases/invdes_grating_coupler/case.yaml").read_text())
result = json.loads(Path("progress.json").read_text())["cases"]["invdes_grating_coupler"]
print(f"JAX devices: {jax.devices()}")
print("Scientific result:", result["status"].upper())

JAX devices: [CudaDevice(id=0)]
Scientific result: FAIL

1. Published physical, fabrication, and optimizer parameters

pd.Series({
    "wavelength_um": WAVELENGTH0_UM,
    "spectrum_um": (WAVELENGTHS_UM[0], WAVELENGTHS_UM[-1], len(WAVELENGTHS_UM)),
    "domain_um": DOMAIN_SIZE_UM,
    "grating_region_um": (GRATING_LENGTH_UM, GRATING_WIDTH_UM, DESIGN_SIZE_UM[2]),
    "optimizer_controls_half_domain": PARAMETER_SHAPE,
    "control_pitch_um": CONTROL_SPACING_UM,
    "minimum_feature_radius_um": MIN_FEATURE_UM,
    "border_buffer_um": BORDER_BUFFER_UM,
    "silicon_epsilon": EPS_SILICON,
    "silica_epsilon": EPS_SILICA,
    "Gaussian_spot_diameter_um": SPOT_SIZE_UM,
    "fiber_tilt_deg": FIBER_TILT_DEG,
    "field_symmetry": (0, -1, 0),
    "published_updates": PUBLISHED_ITERATIONS,
    "allowed_updates": meta["invdes"]["max_iterations"],
    "Adam_learning_rate": LEARNING_RATE,
    "beta_schedule": (1.0, BETA_MAX),
    "FDTD_dx_dy_um": FDTD_DL_UM,
    "FDTD_dz_max_um": FDTD_DZ_UM,
    "automatic_window_fs": RUN_TIME_S / 1e-15,
}, name="value").to_frame()

	value
wavelength_um	1.55
spectrum_um	(1.52, 1.58, 61)
domain_um	(6.25, 6.18, 3.6799999999999997)
grating_region_um	(4.0, 4.0, 0.22)
optimizer_controls_half_domain	(216, 108)
control_pitch_um	0.02
minimum_feature_radius_um	0.08
border_buffer_um	0.16
silicon_epsilon	12.1104
silica_epsilon	2.0736
Gaussian_spot_diameter_um	2.5
fiber_tilt_deg	10.0
field_symmetry	(0, -1, 0)
published_updates	75
allowed_updates	150
Adam_learning_rate	0.2
beta_schedule	(1.0, 30.0)
FDTD_dx_dy_um	0.05
FDTD_dz_max_um	0.04
automatic_window_fs	688.709462

2. Random cold start and exact double-filter transform

The upstream notebook calls uniform, applies a one-dimensional Gaussian filter to the flattened vector, then reshapes it. We use seed 0 only to make that same random distribution reproducible. No optimized pixels or analytic grating seed are supplied. The forced silicon strip at the left interface is the notebook’s only material mask.

parameters0 = initial_parameters(seed=0)
processed0 = np.asarray(pre_process(parameters0, beta=1.0))
initial_binary = np.asarray(hard_binary_design(parameters0))
x_extent = (DESIGN_CENTER_X_UM - DESIGN_SIZE_UM[0] / 2,
            DESIGN_CENTER_X_UM + DESIGN_SIZE_UM[0] / 2)
y_half_extent = (0, DESIGN_SIZE_UM[1] / 2)

fig, axes = plt.subplots(1, 3, figsize=(15, 4.3), sharex=True, sharey=True)
for ax, data, title in zip(
    axes,
    (parameters0, processed0, initial_binary),
    ("random smoothed controls", "double filter, β=1", "fixed-threshold β=30"),
    strict=True,
):
    image = ax.imshow(data.T, origin="lower", extent=(*x_extent, *y_half_extent),
                      cmap="viridis", vmin=0, vmax=1, aspect="equal")
    ax.set(xlabel="x (µm)", title=title)
axes[0].set_ylabel("retained y (µm)")
fig.colorbar(image, ax=axes, label="silicon density", shrink=.78)
plt.show()
print(f"Initial morphology penalty: {float(fabrication_penalty(parameters0, 1.0)):.6f}")

Initial morphology penalty: 0.999455

3. Exact 3D stack, source, output port, and PML spacing

fig, axes = plt.subplots(1, 2, figsize=(14, 5.2))
ax = axes[0]
ax.add_patch(Rectangle((x_extent[0], -DESIGN_SIZE_UM[1] / 2), DESIGN_SIZE_UM[0], DESIGN_SIZE_UM[1],
                       facecolor="#55c2b5", alpha=.42, label="inverse-design silicon/air"))
guide_end = -DOMAIN_SIZE_UM[0] / 2 + 1.0
ax.add_patch(Rectangle((-DOMAIN_SIZE_UM[0] / 2, -.25), 1.0, .5,
                       facecolor="#087d83", label="500 nm Si output guide"))
ax.axvspan(-DOMAIN_SIZE_UM[0] / 2, -DOMAIN_SIZE_UM[0] / 2 + .6, color="#8221a8", alpha=.14)
ax.axvspan(DOMAIN_SIZE_UM[0] / 2 - .6, DOMAIN_SIZE_UM[0] / 2, color="#8221a8", alpha=.14)
ax.set(xlim=(-DOMAIN_SIZE_UM[0]/2, DOMAIN_SIZE_UM[0]/2),
       ylim=(-DOMAIN_SIZE_UM[1]/2, DOMAIN_SIZE_UM[1]/2), aspect="equal",
       xlabel="x (µm)", ylabel="y (µm)", title="Top view: exact physical extents")
ax.legend(loc="upper right")

ax = axes[1]
zmin, zmax = -DOMAIN_SIZE_UM[2]/2, DOMAIN_SIZE_UM[2]/2
ax.add_patch(Rectangle((-DOMAIN_SIZE_UM[0]/2, zmin), DOMAIN_SIZE_UM[0], BOX_BOTTOM_Z_UM-zmin,
                       color="#087d83", alpha=.7, label="Si substrate"))
ax.add_patch(Rectangle((-DOMAIN_SIZE_UM[0]/2, BOX_BOTTOM_Z_UM), DOMAIN_SIZE_UM[0], BOX_THICKNESS_UM,
                       color="#d6ecf5", alpha=.8, label="1.6 µm SiO₂ BOX"))
ax.add_patch(Rectangle((x_extent[0], SLAB_BOTTOM_Z_UM), DESIGN_SIZE_UM[0], DESIGN_SIZE_UM[2],
                       color="#55c2b5", alpha=.8, label="220 nm design layer"))
ax.add_patch(FancyArrowPatch((DESIGN_CENTER_X_UM + .75, SOURCE_Z_UM + .62),
                             (DESIGN_CENTER_X_UM - .15, SOURCE_Z_UM - .04),
                             arrowstyle="-|>", mutation_scale=16, color="#d56b2d", lw=2,
                             label="10° Gaussian beam"))
ax.axhline(SOURCE_Z_UM, color="#d56b2d", ls=":", lw=1)
ax.set(xlim=(-DOMAIN_SIZE_UM[0]/2, DOMAIN_SIZE_UM[0]/2), ylim=(zmin, zmax),
       xlabel="x (µm)", ylabel="z relative to simulation center (µm)",
       title="Vertical stack and downward tilted source")
ax.legend(loc="lower right")
plt.tight_layout(); plt.show()

4. Differentiable FDTD optimization

Set RUN_FULL_OPTIMIZATION=True to repeat the complete cold 150-update run. The retained notebook defaults to the benchmark record because that run takes about half an hour on the local RTX 3080. It never loads an optimizer checkpoint. Binary checkpoints at updates 75 and 150 are each scored by fresh broadband forward solves.

RUN_FULL_OPTIMIZATION = False
if RUN_FULL_OPTIMIZATION:
    result = optimize_grating_coupler(meta)
assert result["status"] in {"pass", "fail"}, result
display(pd.DataFrame(result["metrics"]))

	name	compare	got	ref	pass	error	atol	rtol	max_got	max_ref
0	best_fom_parity	fom_ge	0.577275	0.630556	False	0.053281	0.0	0.0	NaN	NaN
1	optimizer_improvement	fom_ge	0.572028	0.000100	True	-0.571928	0.0	0.0	NaN	NaN
2	passivity_peak_coupling	upper_bound	NaN	NaN	True	0.000000	0.0	0.0	0.577275	1.05
3	source_normalization_drift	upper_bound	NaN	NaN	True	0.000000	0.0	0.0	0.016601	0.10
4	final_binary_fraction	fom_ge	1.000000	1.000000	True	0.000000	0.0	0.0	NaN	NaN

5. Convergence, fabrication penalty, and binary checkpoints

extras = result["extras"]
physical = np.asarray(extras["fom_history"])
constrained = np.asarray(extras["constrained_history"])
penalty = np.asarray(extras["penalty_history"])
golden = np.load("benchmarks/goldens/invdes_grating_coupler.npz", allow_pickle=True)

fig, axes = plt.subplots(1, 2, figsize=(13, 4.7))
axes[0].plot(physical, color="#087d83", label="FDTDX physical center coupling")
axes[0].plot(constrained, color="#55c2b5", alpha=.8, label="coupling − morphology penalty")
axes[0].plot(np.asarray(golden["raw_objective_history"]), "--", color="#8221a8",
             label="executed Tidy3D objective")
axes[0].scatter(extras["checkpoint_iterations"], extras["checkpoint_binary_foms"],
                marker="s", color="#d56b2d", zorder=4, label="fresh hard-binary checkpoints")
axes[0].axhline(extras["threshold"], color="#d29b00", ls=":", label="95% gate")
axes[0].set(xlabel="objective evaluation", ylabel="figure of merit", title="Cold-start convergence")
axes[0].legend(fontsize=8); axes[0].grid(alpha=.25)

axes[1].plot(penalty, color="#8221a8", label="erosion/dilation penalty")
axes[1].plot(np.asarray(extras["source_power_history"]) / extras["source_power_history"][0] - 1,
             color="#087d83", label="source-loading change")
axes[1].set(xlabel="objective evaluation", ylabel="diagnostic", title="Fabrication and normalization")
axes[1].legend(); axes[1].grid(alpha=.25)
plt.tight_layout(); plt.show()

6. Final controls, processed density, and fully binary device

artifact = np.load(extras["best_design_artifact"], allow_pickle=True)
controls = np.asarray(artifact["selected_checkpoint_parameters"])
processed = np.asarray(artifact["processed_density"])
binary_half = np.asarray(artifact["binary_density"])
binary_full = np.concatenate((np.fliplr(binary_half), binary_half), axis=1)
full_extent = (*x_extent, -DESIGN_SIZE_UM[1]/2, DESIGN_SIZE_UM[1]/2)

fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))
for ax, data, extent, title in (
    (axes[0], controls, (*x_extent, *y_half_extent), "selected controls"),
    (axes[1], processed, (*x_extent, *y_half_extent), "double-filtered β=30"),
    (axes[2], binary_full, full_extent, "fully binary mirrored device"),
):
    image = ax.imshow(data.T, origin="lower", extent=extent, cmap="viridis", vmin=0, vmax=1, aspect="equal")
    ax.set(xlabel="x (µm)", ylabel="y (µm)", title=title)
fig.colorbar(image, ax=axes, label="silicon density", shrink=.76)
plt.show()
print("Unique final material values:", np.unique(binary_half))
print("Binary fraction:", extras["final_binary_fraction"])

Unique final material values: [0. 1.]
Binary fraction: 1.0

7. Erosion/dilation process-corner audit

def morphology(values, eta_outer, eta_inner):
    first = fdtdx.filter_and_project(values, radius=MIN_FEATURE_UM,
                                     spacing=CONTROL_SPACING_UM, beta=20, eta=eta_outer)
    return fdtdx.filter_and_project(first, radius=MIN_FEATURE_UM,
                                    spacing=CONTROL_SPACING_UM, beta=20, eta=eta_inner)

opened = np.asarray(morphology(processed, .99, .01))
closed = np.asarray(morphology(processed, .01, .99))
fig, axes = plt.subplots(1, 3, figsize=(15, 4.3), sharex=True, sharey=True)
for ax, data, title in zip(axes, (opened, closed, closed-opened),
                           ("opened corner", "closed corner", "close − open"), strict=True):
    image = ax.imshow(data.T, origin="lower", extent=(*x_extent, *y_half_extent),
                      cmap="magma", aspect="equal")
    ax.set(xlabel="x (µm)", title=title)
axes[0].set_ylabel("retained y (µm)")
fig.colorbar(image, ax=axes, shrink=.76)
plt.show()
print(f"Final RMS erosion/dilation discrepancy: {np.linalg.norm(closed-opened)/np.sqrt(opened.size):.6f}")

Final RMS erosion/dilation discrepancy: 0.019405

8. Fresh 61-wavelength binary-design FDTDX solve

started = perf_counter()
validation = simulate_grating_coupler(binary_half, record_fields=True, record_spectrum=True)
print(f"Fresh forward runtime: {perf_counter() - started:.1f} s")
print(f"Reduced grid: {validation['grid']}; steps: {validation['steps']:,}")
print(f"Peak FDTDX coupling: {validation['peak_fom']:.6f}")

Fresh forward runtime: 25.3 s
Reduced grid: (125, 63, 94); steps: 8,771
Peak FDTDX coupling: 0.577275

9. Actual Yee-grid material in device and vertical planes

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
xy = axes[0].imshow(validation["permittivity_xy"].T, origin="lower",
                    extent=validation["extent_xy_um"], cmap="viridis", aspect="equal")
axes[0].set(xlabel="x (µm)", ylabel="y (µm)", title="binary material on executed xy Yee plane")
xz = axes[1].imshow(validation["permittivity_xz"].T, origin="lower",
                    extent=validation["extent_xz_um"], cmap="viridis", aspect="auto")
axes[1].set(xlabel="x (µm)", ylabel="z (µm)", title="executed xz stack near symmetry plane")
fig.colorbar(xy, ax=axes[0], label="relative permittivity")
fig.colorbar(xz, ax=axes[1], label="relative permittivity")
plt.tight_layout(); plt.show()

10. Fresh central-wavelength electric-field profiles

fig, axes = plt.subplots(1, 2, figsize=(14, 5))
field_xy = validation["field_intensity_xy"]
field_xz = validation["field_intensity_xz"]
im0 = axes[0].imshow((field_xy/(field_xy.max()+1e-30)).T, origin="lower",
                     extent=validation["extent_xy_um"], cmap="magma", vmin=0, vmax=1, aspect="equal")
axes[0].contour(np.linspace(validation["extent_xy_um"][0], validation["extent_xy_um"][1], field_xy.shape[0]),
                np.linspace(validation["extent_xy_um"][2], validation["extent_xy_um"][3], field_xy.shape[1]),
                validation["permittivity_xy"].T, levels=[3], colors="cyan", linewidths=.35)
axes[0].set(xlabel="x (µm)", ylabel="y (µm)", title="device-plane normalized |E|²")
im1 = axes[1].imshow((field_xz/(field_xz.max()+1e-30)).T, origin="lower",
                     extent=validation["extent_xz_um"], cmap="magma", vmin=0, vmax=1, aspect="auto")
axes[1].contour(np.linspace(validation["extent_xz_um"][0], validation["extent_xz_um"][1], field_xz.shape[0]),
                np.linspace(validation["extent_xz_um"][2], validation["extent_xz_um"][3], field_xz.shape[1]),
                validation["permittivity_xz"].T, levels=[3], colors="cyan", linewidths=.35)
axes[1].set(xlabel="x (µm)", ylabel="z (µm)", title="vertical normalized |E|²")
fig.colorbar(im1, ax=axes, label="normalized phasor |E|²", shrink=.76)
plt.show()

11. Broadband coupling and independent Tidy3D validation

tidy = np.load("benchmarks/goldens/tidy3d_validation_invdes_grating_coupler.npz", allow_pickle=True)
tidy_wavelengths = np.asarray(tidy["wavelengths_um"])
tidy_coupling = np.asarray(tidy["coupling"])
order = np.argsort(tidy_wavelengths)

fig, axes = plt.subplots(1, 2, figsize=(13, 4.7))
axes[0].plot(validation["wavelengths_um"], validation["coupling"], color="#087d83", lw=2,
             label="FDTDX final binary")
axes[0].plot(tidy_wavelengths[order], tidy_coupling[order], "--", color="#8221a8", lw=2,
             label="Tidy3D, same FDTDX geometry")
axes[0].axhline(float(golden["ref_fom"][0]), color="#d29b00", ls=":",
                label="published Tidy3D peak")
axes[0].set(xlabel="wavelength (µm)", ylabel="fundamental-mode power",
            title="Independent cross-solver spectrum")
axes[0].legend(); axes[0].grid(alpha=.25)

axes[1].plot(validation["wavelengths_um"], 10*np.log10(np.maximum(validation["coupling"], 1e-12)),
             color="#087d83", label="FDTDX")
axes[1].plot(tidy_wavelengths[order], 10*np.log10(np.maximum(tidy_coupling[order], 1e-12)),
             "--", color="#8221a8", label="Tidy3D validation")
axes[1].axhline(float(golden["executed_peak_loss_db"][0]), color="#d29b00", ls=":",
                label="notebook −1.78 dB")
axes[1].set(xlabel="wavelength (µm)", ylabel="coupling (dB)", title="Insertion loss")
axes[1].legend(); axes[1].grid(alpha=.25)
plt.tight_layout(); plt.show()

fdtdx_peak = float(validation["peak_fom"])
tidy_same_geometry = float(tidy["tidy3d_fom"][0])
published = float(golden["ref_fom"][0])
pd.Series({
    "FDTDX final binary peak": fdtdx_peak,
    "Tidy3D same-geometry peak": tidy_same_geometry,
    "cross-solver relative difference": abs(fdtdx_peak-tidy_same_geometry)/tidy_same_geometry,
    "published Tidy3D peak": published,
    "actual parity in Tidy3D": tidy_same_geometry/published,
    "repository 95% threshold": extras["threshold"],
    "Tidy3D task id": json.loads(str(tidy["meta_json"]))["tidy3d_task_id"],
}, name="value").to_frame()

	value
FDTDX final binary peak	0.577275
Tidy3D same-geometry peak	0.533884
cross-solver relative difference	0.081275
published Tidy3D peak	0.663743
actual parity in Tidy3D	0.804353
repository 95% threshold	0.630556
Tidy3D task id	fdve-7de41b63-0858-4c0c-9a4b-aa1c27472e89

Interpretation and reproduction

The two solvers agree on this final geometry within 10%, so the forward result is credible. The geometry nevertheless reaches only about 80% of the executed notebook’s performance, even after twice its recorded iteration count. This is a real optimizer-parity failure, not a plotting or normalization success. The example remains public precisely so future optimizer changes have a complete, inspectable target.

print("uv run fdtdx-bench run --case invdes_grating_coupler")

uv run fdtdx-bench run --case invdes_grating_coupler