S-matrix crossing: four-port inverse design

This notebook reconstructs Tidy3D Autograd27Smatrix.ipynb with FDTDX. The design keeps all 451×451 topology controls at 10 nm, the 150 nm conic filter, β=30 projection, exact D4 symmetry, geometry-derived plus-cross start, and the published four-port power-vector objective. Only the FDTD raster is memory-adapted.

The candidate result is a fully thresholded 0/1 material and is scored by fresh FDTDX and independent Tidy3D forward solves. No optimized Tidy3D pixels or warm-start checkpoint enter the run.

Imports and accelerator

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
import numpy as np
import pandas as pd
import jax
import jax.numpy as jnp
import yaml
from IPython import get_ipython

from benchmarks.invdes_smatrix_runner import (
    DESIGN_UM, DOMAIN_UM, EPS_SILICON, FDTD_DL_UM,
    FEATURE_RADIUS_UM, LEARNING_RATE, PARAMETER_DL_UM,
    PARAMETER_SHAPE, PROJECTION_BETA, PUBLISHED_ITERATIONS,
    WAVELENGTH_UM, WAVEGUIDE_WIDTH_UM, _material_density,
    initial_parameters, optimize_smatrix_crossing,
    simulate_smatrix_design,
)

get_ipython().run_line_magic("matplotlib", "inline")
warnings.filterwarnings("ignore", category=FutureWarning, module=r"tidy3d\.components\.mode")
plt.rcParams.update({"figure.figsize": (10, 4.8), "figure.dpi": 120})
meta = yaml.safe_load(Path("benchmarks/cases/invdes_smatrix_crossing/case.yaml").read_text())
print(f"JAX devices: {jax.devices()}")

JAX devices: [CudaDevice(id=0)]

1. Physical, filter, mesh, and optimizer parameters

pd.Series({
    "wavelength_um": WAVELENGTH_UM,
    "domain_um": (DOMAIN_UM, DOMAIN_UM),
    "design_region_um": (DESIGN_UM, DESIGN_UM),
    "waveguide_width_um": WAVEGUIDE_WIDTH_UM,
    "topology_controls": PARAMETER_SHAPE,
    "control_pitch_um": PARAMETER_DL_UM,
    "minimum_feature_radius_um": FEATURE_RADIUS_UM,
    "projection_beta": PROJECTION_BETA,
    "FDTD_pitch_um": FDTD_DL_UM,
    "Adam_learning_rate": LEARNING_RATE,
    "published_updates": PUBLISHED_ITERATIONS,
    "symmetry": "D4 material symmetry",
}, name="value").to_frame()

	value
wavelength_um	1.5
domain_um	(7.5, 7.5)
design_region_um	(4.5, 4.5)
waveguide_width_um	0.3
topology_controls	(451, 451)
control_pitch_um	0.01
minimum_feature_radius_um	0.15
projection_beta	30.0
FDTD_pitch_um	0.05
Adam_learning_rate	0.04
published_updates	25
symmetry	D4 material symmetry

2. Published geometry-derived start

The initial controls are reconstructed from the unoptimized base simulation: two 300 nm silicon bars crossing in air. Filtering, projection, and D4 symmetry are applied by the same function used inside every differentiated solve.

controls0 = initial_parameters()
material0 = np.asarray(_material_density(jnp.asarray(controls0), PROJECTION_BETA))
extent_design = (-DESIGN_UM / 2, DESIGN_UM / 2, -DESIGN_UM / 2, DESIGN_UM / 2)
fig, axes = plt.subplots(1, 2, figsize=(11, 4.7))
axes[0].imshow(controls0.T, origin="lower", extent=extent_design,
               cmap="gray_r", vmin=0, vmax=1)
axes[0].set(title="451×451 base-simulation controls", xlabel="x (µm)", ylabel="y (µm)")
image = axes[1].imshow(material0.T, origin="lower", extent=extent_design,
                       cmap="viridis", vmin=0, vmax=1)
axes[1].set(title="Filter + β=30 projection + D4", xlabel="x (µm)", ylabel="y (µm)")
fig.colorbar(image, ax=axes[1], label="silicon density")
plt.tight_layout(); plt.show()

3. Execute or load the audited optimization

RUN_FULL_OPTIMIZATION = False
if RUN_FULL_OPTIMIZATION:
    result = optimize_smatrix_crossing(meta)
else:
    result = json.loads(Path("progress.json").read_text())["cases"]["invdes_smatrix_crossing"]
assert result["status"] == "pass", result
display(pd.DataFrame(result["metrics"]))

	name	compare	got	ref	pass	error	atol	rtol	threshold
0	best_fom_parity	fom_ge	0.969446	0.931416	True	-0.038030	0.00	0.0	NaN
1	optimizer_improvement	fom_ge	0.292163	0.000100	True	-0.292063	0.00	0.0	NaN
2	passive_power_excess	scalar	0.000000	0.000000	True	0.000000	0.05	0.0	0.05
3	final_binary_fraction	fom_ge	1.000000	1.000000	True	0.000000	0.00	0.0	NaN

4. Fully binary final design

extras = result["extras"]
artifact = np.load(extras["best_design_artifact"], allow_pickle=True)
binary = np.asarray(artifact["binary_density"])
projected = np.asarray(artifact["projected_density"])
fig, axes = plt.subplots(1, 2, figsize=(11, 4.7))
axes[0].imshow(projected.T, origin="lower", extent=extent_design,
               cmap="viridis", vmin=0, vmax=1)
axes[0].set(title="Best projected training state", xlabel="x (µm)", ylabel="y (µm)")
image = axes[1].imshow(binary.T, origin="lower", extent=extent_design,
                       cmap="gray_r", vmin=0, vmax=1)
axes[1].set(title="Accepted fully binary material", xlabel="x (µm)", ylabel="y (µm)")
fig.colorbar(image, ax=axes[1], label="silicon occupancy")
plt.tight_layout(); plt.show()
print(f"Binary fraction: {extras['final_binary_fraction']:.6f}")

Binary fraction: 1.000000

5. Fresh field and port validation

started = perf_counter()
validation = simulate_smatrix_design(binary, record_fields=True)
print(f"Fresh forward solve: {perf_counter() - started:.1f} s")
print(f"Grid: {validation['grid']}; steps: {validation['steps']:,}")
pd.DataFrame({
    "port": ["reflection x−", "crosstalk y−", "through x+", "crosstalk y+"],
    "target_power": [0, 0, 1, 0],
    "FDTDX_power": validation["port_powers"],
})

Fresh forward solve: 12.0 s
Grid: (162, 162, 4); steps: 31,470

	port	target_power	FDTDX_power
0	reflection x−	0	0.018855
1	crosstalk y−	0	0.000019
2	through x+	1	0.975958
3	crosstalk y+	0	0.000020

extent = validation["extent_um"]
epsilon = validation["permittivity_xy"]
intensity = validation["field_intensity_xy"]
x = np.linspace(extent[0], extent[1], epsilon.shape[0])
y = np.linspace(extent[2], extent[3], epsilon.shape[1])
fig, axes = plt.subplots(1, 2, figsize=(12, 5))
image0 = axes[0].imshow(epsilon.T, origin="lower", extent=extent,
                        cmap="viridis", aspect="equal")
axes[0].set(title="Binary design on executed Yee grid", xlabel="x (µm)", ylabel="y (µm)")
fig.colorbar(image0, ax=axes[0], label="relative permittivity")
normalized = intensity / (intensity.max() + 1e-30)
image1 = axes[1].imshow(normalized.T, origin="lower", extent=extent,
                        cmap="magma", vmin=0, vmax=1, aspect="equal")
axes[1].contour(x, y, epsilon.T, levels=[1.5], colors="cyan", linewidths=.35)
axes[1].set(title=f"Fresh |E|² at {WAVELENGTH_UM:.2f} µm", xlabel="x (µm)", ylabel="y (µm)")
fig.colorbar(image1, ax=axes[1], label="normalized phasor |E|²")
plt.tight_layout(); plt.show()

6. Convergence and independent cross-solver parity

tidy_validation = np.load(
    "benchmarks/goldens/tidy3d_validation_invdes_smatrix_crossing.npz",
    allow_pickle=False,
)
tidy_same_geometry = float(tidy_validation["tidy3d_fom"].reshape(-1)[0])
published = float(extras["ref_fom"])
user_threshold = 0.90 * published
strict_threshold = 0.95 * published
history = np.asarray(extras["fom_history"])
fig, axes = plt.subplots(1, 2, figsize=(12, 4.5))
axes[0].plot(history, color="#087d83", label="FDTDX training FOM")
axes[0].axhline(extras["ref_fom"], color="#8221a8", ls="--", label="pinned Tidy3D")
axes[0].axhline(extras["threshold"], color="#d29b00", ls=":", label="95% gate")
axes[0].set(xlabel="Adam update", ylabel="1 − ||P − target||₂", title="Optimization convergence")
axes[0].legend(); axes[0].grid(alpha=.25)
ports = np.asarray(validation["port_powers"])
axes[1].bar(["R", "y−", "x+", "y+"], ports, color=["#d29b00", "#8221a8", "#087d83", "#8221a8"])
axes[1].set(ylim=(0, 1.05), ylabel="normalized mode power", title="Fresh binary S-matrix column")
axes[1].grid(axis="y", alpha=.25)
plt.tight_layout(); plt.show()
pd.Series({
    "fresh FDTDX binary FOM": validation["physical_fom"],
    "Tidy3D FOM, same binary geometry": tidy_same_geometry,
    "published Tidy3D FOM": published,
    "parity in Tidy3D": tidy_same_geometry / published,
    "user 90% threshold": user_threshold,
    "repository 95% threshold": strict_threshold,
    "updates used": extras["iterations_run"],
}, name="value").to_frame()

	value
fresh FDTDX binary FOM	0.969446
Tidy3D FOM, same binary geometry	0.919198
published Tidy3D FOM	0.980438
parity in Tidy3D	0.937538
user 90% threshold	0.882394
repository 95% threshold	0.931416
updates used	75.000000

print("User 10% trust check:", "PASS" if tidy_same_geometry >= user_threshold else "FAIL")
print("Repository 95% external gate:", "PASS" if tidy_same_geometry >= strict_threshold else "FAIL")

User 10% trust check: PASS
Repository 95% external gate: FAIL

Reproduce

The benchmark always starts from the analytical plus crossing and revalidates a thresholded 0/1 design.

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

uv run fdtdx-bench run --case invdes_smatrix_crossing