Authors: Azka Maula Iskandar Muda & Uğur Teğin, Koç University
This example builds a three-stage photonic classifier with six optical channels and four class outputs. A balanced 1 × 8 splitter feeds the six central input branches, and terminators absorb the two unused outer branches. Two 6 × 6 optical transformations process the encoded fields before a 4 × 6 transformation directs the optical power to four class detectors.
The input phase is applied before every optical stage. We compare this repeated-encoding architecture with an otherwise identical classifier that applies the phase only once. This controlled comparison follows the method introduced in Deep Inverse-Designed Nanophotonic Processors with Structural Nonlinearity from Repeated Phase Encoding.
We train the complex target matrices, construct the routed circuit in PhotonForge, and inverse-design the three silicon devices in Tidy3D. The final analysis compares the trained optical model with the realized transmission matrices.
What this notebook demonstrates:
- Prepare six optical features for a four-class MNIST classifier.
- Apply the image phase either once or before each of three passive stages.
- Connect a balanced SiEPIC splitter to the optical network in PhotonForge.
- Train the three passive transformations before electromagnetic realization.
- Inverse-design the corresponding silicon devices using Tidy3D.
- Compare the trained classifier with the realized optical response.
Setup¶
First, we prepare the Python environment and load the libraries used for data processing, optical training, circuit layout, and electromagnetic simulation. Local data, layouts, figures, and the realization progress cache are stored under output/deep. The cloud controls remain off until they are enabled in the cost-estimation and inverse-design sections.
[ ]:
%pip install --quiet --upgrade \
"numpy==2.4.4" "scipy==1.17.0" "matplotlib==3.10.9" "scikit-learn==1.8.0" \
"jax==0.10.0" "optax==0.2.8" "autograd==1.8.0" "gdstk==1.0.0" "klayout==0.30.10" \
"photonforge==1.4.7" "siepic-forge==1.1.2" \
"siepic_ebeam_pdk==0.4.53" "tidy3d[design,extras]==2.12.0"
[ ]:
import gzip
import inspect
import io
import json
import pickle
import struct
import time
import urllib.request
from contextlib import contextmanager
from importlib import resources
from importlib.metadata import version
from pathlib import Path
import autograd.numpy as anp
import gdstk
import jax
import jax.numpy as jnp
import klayout.db as kdb
import matplotlib.pyplot as plt
import numpy as np
import optax
import photonforge as pf
import scipy.interpolate
import siepic_forge
import tidy3d as td
import tidy3d.plugins.invdes as tdi
from autograd.tracer import getval
from IPython.display import display
from matplotlib_inline.backend_inline import set_matplotlib_formats
from requests.exceptions import ConnectionError as RequestsConnectionError
from requests.exceptions import ReadTimeout
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from tidy3d import web
from tidy3d.exceptions import WebError
from tidy3d.plugins.mode import ModeSolver
from tidy3d.plugins.smatrix import ModalComponentModeler, Port
from tidy3d.plugins.smatrix.run import compose_modeler_data_from_batch_data
from tidy3d.web.api import states as web_states
from tidy3d.web.api.autograd import engine as autograd_engine
from tidy3d.web.api.autograd import hooks as autograd_hooks
from tidy3d.web.api.autograd.io_utils import (
get_cached_vjp_traced_fields,
get_vjp_traced_fields,
)
from urllib3.exceptions import NewConnectionError
jax.config.update("jax_enable_x64", True)
td.config.simulation.use_local_subpixel = True
set_matplotlib_formats("png")
plt.rcParams.update({"figure.dpi": 110, "font.size": 9, "axes.titlesize": 10})
04:22:40 KST WARNING: The material-library variant 'Palik_Lossless' is deprecated and maps to 'Palik_LowLoss' because it contains a tiny fitted loss despite its name. Use 'Palik_NoLoss' where available for a zero-loss Palik model.
[3]:
ESTIMATE_COST = False
RUN_CLOUD = False
COST_APPROVED = False
AUTOGRAD_SOURCE_CHUNK_SIZE = 1
QUEUE_TIMEOUT_SECONDS = 600.0
QUEUE_RETRY_DELAY_SECONDS = 30.0
QUEUE_POLL_SECONDS = 5.0
QUEUE_RETRY_POLICY = "native_source_queue600_retry_forever_v1"
SEED = 7
WAVELENGTH_UM = 1.55
FREQ0_HZ = td.C_0 / WAVELENGTH_UM
N_CLASSES, N_LAYERS, MNIST_CLASSES = 4, 3, 10
if N_CLASSES < 2 or N_CLASSES > MNIST_CLASSES:
raise ValueError(f"N_CLASSES must be between 2 and {MNIST_CLASSES}.")
def channels_for_classes(n_classes):
channels = (3 * int(n_classes) + 1) // 2
return channels + ((channels - int(n_classes)) % 2)
def next_power_of_two(value):
return 1 << (int(value) - 1).bit_length()
N_CHANNELS = channels_for_classes(N_CLASSES)
N_SPLITTER_LEAVES = next_power_of_two(N_CHANNELS)
FINAL_OUTPUT_INDICES = tuple(
range(
(N_CHANNELS - N_CLASSES) // 2,
(N_CHANNELS + N_CLASSES) // 2,
)
)
LAYER_INPUT_COUNTS = (N_CHANNELS,) * N_LAYERS
LAYER_OUTPUT_COUNTS = (N_CHANNELS,) * (N_LAYERS - 1) + (N_CLASSES,)
REALIZATION_TARGET_TAG = (
f"outputs{'-'.join(map(str, LAYER_OUTPUT_COUNTS))}_centered_final_rows_v1"
)
if N_CHANNELS > 784:
raise ValueError("N_CHANNELS exceeds the 784-dimensional PCA limit.")
if channels_for_classes(4) != 6:
raise ValueError("Class-to-channel parity rule changed unexpectedly.")
CLASS_LABELS = tuple(range(N_CLASSES))
MIN_SELECTION_ACCURACY = 1 / N_CLASSES
FEATURE_ORDER = tuple(range(N_CHANNELS))
DESIGN_UM_PER_CHANNEL = 2.1
DESIGN_EDGE_UM = DESIGN_UM_PER_CHANNEL * N_CHANNELS
FANOUT_RULE = "balanced_next_power_of_two_centered_active_v1"
TERMINATOR_MODEL = "ebeam_terminator_te1550"
PREPROCESSING_TAG = "trainonly_pca_scaler_v1"
ROBUST_EPOCHS, ROBUST_REPLICAS, ROBUST_BATCH_SIZE, ROBUST_LR = 300, 16, 1_000, 0.001
MATRIX_TRAINING_SEED = 17
ROBUST_TARGET_ERROR, ROBUST_TARGET_MEAN_ACCURACY = 0.10, 0.92
ROBUST_WINDOW_MIN, ROBUST_WINDOW_MAX = 0.08, 0.12
ROBUST_MARGIN_TARGET, ROBUST_MARGIN_WEIGHT = 0.10, 1.0
ROBUST_VALIDATION_SAMPLES, ROBUST_AUDIT_SAMPLES = 256, 512
ROBUST_VALIDATION_INTERVAL = 10
PHASE_SCALE, GAMMA_TARGET, GAMMA_WEIGHT = 0.5, 240.0, 2.0
SIGMA_MIN, SIGMA_MAX = 10 ** (-6 / 20), 10 ** (-3 / 20)
TIDY3D_LR_MAX, TIDY3D_LR_WARMUP_UPDATES = 0.05, 5
TIDY3D_BETA1, TIDY3D_BETA2, TIDY3D_EPS = 0.9, 0.999, 1e-8
FABRICATION_MIN_FEATURE_UM = 0.2
FABRICATION_WEIGHT_MAX = 0.8
FABRICATION_START_UPDATE, FABRICATION_FULL_UPDATE = 6, 15
FABRICATION_BETA, FABRICATION_ETA0, FABRICATION_DELTA_ETA = 100.0, 0.5, 0.01
SIMULATION_ARCHITECTURE = "native_modal_component_smatrix_v2"
CLADDING_MODEL = "silica_pdk_n14447002763"
OPTIMIZER_DIRECTION = "gradient_ascent_on_negative_constrained_loss"
ACTIVE_MASK_RULE = "square_filter_field_physical_silica_corner_mask_v3"
REFLECTION_WEIGHT = 1.0
REALIZATION_OBJECTIVE = "frobenius"
REALIZATION_MODE = "synchronized" # "synchronized" or "sequential"
COST_ESTIMATE_MODES = ("synchronized", "sequential")
REALIZATION_STEPS = 40
CALIBRATION_STEPS, CALIBRATION_LR = 150, 0.002
def tidy3d_learning_rate(update_number):
update_number = int(update_number)
if update_number < 1:
raise ValueError("Adam update numbers are one-based.")
return TIDY3D_LR_MAX * min(update_number / TIDY3D_LR_WARMUP_UPDATES, 1.0)
def fabrication_penalty_weight(update_number):
update_number = int(update_number)
if update_number < 1:
raise ValueError("Fabrication update numbers are one-based.")
if update_number < FABRICATION_START_UPDATE:
return 0.0
if update_number >= FABRICATION_FULL_UPDATE:
return FABRICATION_WEIGHT_MAX
ramp_updates = FABRICATION_FULL_UPDATE - (FABRICATION_START_UPDATE - 1)
return (
FABRICATION_WEIGHT_MAX
* (update_number - (FABRICATION_START_UPDATE - 1))
/ ramp_updates
)
[ ]:
REPOSITORY_ROOT = Path.cwd().resolve()
EXPECTED_NOTEBOOK = (
REPOSITORY_ROOT / "deep_inverse_designed_photonic_neural_network.ipynb"
)
if (
not (REPOSITORY_ROOT / "pyproject.toml").is_file()
or not EXPECTED_NOTEBOOK.is_file()
):
raise RuntimeError(
"Start Jupyter from the repository root before running this notebook."
)
OUTPUT_DIR = REPOSITORY_ROOT / "output" / "deep"
DATA_DIR = OUTPUT_DIR / "data"
FIGURE_DIR = OUTPUT_DIR / "figures"
CHECKPOINT_DIR = OUTPUT_DIR / "checkpoints"
LAYOUT_DIR = OUTPUT_DIR / "layouts"
MNIST_DIR = DATA_DIR / "mnist"
for directory in (MNIST_DIR, FIGURE_DIR, CHECKPOINT_DIR, LAYOUT_DIR):
directory.mkdir(parents=True, exist_ok=True)
PROGRESS_PATH = OUTPUT_DIR / "progress.npz"
optimizer_state_paths = {
layer: CHECKPOINT_DIR / f"progress_layer_{layer + 1}.hdf5"
for layer in range(N_LAYERS)
}
AUTOGRAD_CACHE_DIR = CHECKPOINT_DIR / "autograd_source_cache"
AUTOGRAD_CACHE_DIR.mkdir(parents=True, exist_ok=True)
QUEUE_RETRY_LOG_PATH = CHECKPOINT_DIR / "queue_retries.json"
REPEATED_CHECKPOINT_PATH = CHECKPOINT_DIR / "repeated_encoding__matrix_training.pkl"
MATRIX_TRAINING_CONFIGURATION = {
"n_channels": N_CHANNELS,
"n_splitter_leaves": N_SPLITTER_LEAVES,
"class_labels": CLASS_LABELS,
"feature_order": FEATURE_ORDER,
"preprocessing": PREPROCESSING_TAG,
"fanout_rule": FANOUT_RULE,
"terminator_model": TERMINATOR_MODEL,
"physical_layer_inputs": LAYER_INPUT_COUNTS,
"physical_layer_outputs": LAYER_OUTPUT_COUNTS,
"physical_target": REALIZATION_TARGET_TAG,
}
PACKAGE_VERSIONS = {
"numpy": np.__version__,
"jax": jax.__version__,
"optax": version("optax"),
"scikit-learn": version("scikit-learn"),
"photonforge": pf.__version__,
"tidy3d": td.__version__,
"klayout": version("klayout"),
}
print(
"versions:",
", ".join(f"{name}={value}" for name, value in PACKAGE_VERSIONS.items()),
)
print(
f"configuration: {N_CLASSES} classes, {N_CHANNELS} channels; design={DESIGN_EDGE_UM:g} um square"
)
print(f"output folder: {OUTPUT_DIR.relative_to(REPOSITORY_ROOT)}")
print(
f"cloud flags: ESTIMATE_COST={ESTIMATE_COST}, "
f"COST_APPROVED={COST_APPROVED}, RUN_CLOUD={RUN_CLOUD}"
)
versions: numpy=2.4.4, jax=0.10.0, optax=0.2.8, scikit-learn=1.8.0, photonforge=1.4.7, tidy3d=2.12.0, klayout=0.30.10 configuration: 4 classes, 6 channels; design=12.6 um square output folder: output/deep cloud flags: ESTIMATE_COST=False, COST_APPROVED=False, RUN_CLOUD=False
Photonic Classifier Concept¶
The classifier alternates image-dependent phase encoding with passive optical transformations. This approach follows the repeated-data principle demonstrated experimentally by Yildirim et al., where multiple linear optical interactions produce a nonlinear mapping of the encoded data. If $D(z)=\operatorname{diag}[\exp(i\alpha z)]$ denotes the phase encoder, the three-stage optical mapping is
$$ a_{\mathrm{out}}=M_3D(z)M_2D(z)M_1D(z)a_{\mathrm{in}}. $$
The matrices $M_1$ and $M_2$ are 6 × 6, while $M_3$ maps the six internal channels to four class outputs. Square-law detection converts the final optical fields into class scores.
[5]:
from matplotlib.patches import FancyBboxPatch
fig, ax = plt.subplots(figsize=(11.5, 2.4), constrained_layout=True)
ax.set_xlim(-0.5, 7.5)
ax.set_ylim(-0.8, 0.8)
ax.axis("off")
blocks = [
(0, "8-way splitter\n6 active leaves", "#DDEBF7"),
(1, r"phase $z$", "#FFF2CC"),
(2, "device 1\n$6\\times6$", "#E2F0D9"),
(3, r"phase $z$", "#FFF2CC"),
(4, "device 2\n$6\\times6$", "#E2F0D9"),
(5, r"phase $z$", "#FFF2CC"),
(6, "device 3\n$4\\times6$", "#E2F0D9"),
(7, "4 class\ndetectors", "#FCE4D6"),
]
for x, label, color in blocks:
width = 0.72 if "phase" in label else 0.94
patch = FancyBboxPatch(
(x - width / 2, -0.33),
width,
0.66,
boxstyle="round,pad=0.04",
facecolor=color,
edgecolor="0.25",
)
ax.add_patch(patch)
ax.text(x, 0, label, ha="center", va="center", fontsize=9)
for left, right in zip(range(7), range(1, 8)):
ax.annotate(
"",
xy=(right - 0.48, 0),
xytext=(left + 0.48, 0),
arrowprops={"arrowstyle": "->", "lw": 1.3, "color": "0.25"},
)
ax.set_title("Repeated phase encoding across three passive optical stages")
plt.show()
Generate and Preprocess MNIST¶
First, we select the images labeled 0–3 and create a stratified validation set. We flatten each image and reduce its 784 pixel values to six features using principal component analysis. Both the PCA basis and the subsequent standardization are fitted on the training partition only.
We then apply the same transformations to the validation and test images and inspect the resulting feature coordinates.
[6]:
MNIST_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"
MNIST_FILES = {
"train_images": "train-images-idx3-ubyte.gz",
"train_labels": "train-labels-idx1-ubyte.gz",
"test_images": "t10k-images-idx3-ubyte.gz",
"test_labels": "t10k-labels-idx1-ubyte.gz",
}
archives = {}
for name, filename in MNIST_FILES.items():
path = MNIST_DIR / filename
if not path.exists():
urllib.request.urlretrieve(MNIST_URL + filename, path)
archives[name] = path
def read_idx(path):
raw = gzip.decompress(path.read_bytes())
magic = struct.unpack(">I", raw[:4])[0]
dimensions = magic & 0xFF
shape = struct.unpack(">" + "I" * dimensions, raw[4 : 4 + 4 * dimensions])
offset = 4 + 4 * dimensions
return np.frombuffer(raw[offset:], dtype=np.uint8).reshape(shape).copy()
train_images = read_idx(archives["train_images"])
train_labels_all = read_idx(archives["train_labels"])
test_images = read_idx(archives["test_images"])
test_labels = read_idx(archives["test_labels"])
train_mask = np.isin(train_labels_all, CLASS_LABELS)
test_mask = np.isin(test_labels, CLASS_LABELS)
train_images, train_labels_all = train_images[train_mask], train_labels_all[train_mask]
test_images, test_labels = test_images[test_mask], test_labels[test_mask]
print("active classes:", CLASS_LABELS)
print(train_images.shape, train_labels_all.shape, test_images.shape, test_labels.shape)
active classes: (0, 1, 2, 3) (24754, 28, 28) (24754,) (4157, 28, 28) (4157,)
[7]:
fit_images, val_images, train_labels, val_labels = train_test_split(
train_images,
train_labels_all,
test_size=5_000,
random_state=SEED,
stratify=train_labels_all,
)
train_flat = fit_images.reshape(len(fit_images), -1).astype(np.float32) / 255.0
val_flat = val_images.reshape(len(val_images), -1).astype(np.float32) / 255.0
test_flat = test_images.reshape(len(test_images), -1).astype(np.float32) / 255.0
pca = PCA(n_components=N_CHANNELS, svd_solver="covariance_eigh")
scaler = StandardScaler()
train_scores = pca.fit_transform(train_flat)
val_scores = pca.transform(val_flat)
test_scores = pca.transform(test_flat)
train_scaled = scaler.fit_transform(train_scores)
val_scaled = scaler.transform(val_scores)
test_scaled = scaler.transform(test_scores)
train_features = train_scaled[:, FEATURE_ORDER].astype(np.float32)
val_features = val_scaled[:, FEATURE_ORDER].astype(np.float32)
test_features = test_scaled[:, FEATURE_ORDER].astype(np.float32)
print(train_features.shape, val_features.shape, test_features.shape)
print(f"preprocessing: {PREPROCESSING_TAG}; PCA/scaler fit on training split only")
print(f"explained variance: {pca.explained_variance_ratio_.sum():.2%}")
fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.0), constrained_layout=True)
axes[0].imshow(train_images[0], cmap="gray", vmin=0, vmax=255)
axes[0].set(title=f"MNIST label {train_labels_all[0]}", xticks=[], yticks=[])
axes[1].bar(np.arange(N_CHANNELS), pca.explained_variance_ratio_)
axes[1].set(
xlabel="principal component", ylabel="variance fraction", title="Train-fitted PCA"
)
plt.show()
(19754, 6) (5000, 6) (4157, 6) preprocessing: trainonly_pca_scaler_v1; PCA/scaler fit on training split only explained variance: 46.04%
Build the PhotonForge Circuit¶
Next, we assemble the complete optical circuit in PhotonForge using components from the SiEPIC library. A balanced eight-way splitter supplies the six active channels, and packaged terminators close the two unused branches. We then connect the three passive stages and place 18 phase encoders along the optical paths.
At the output, four waveguides fan out from the detector pitch to grating couplers spaced by 127 µm. The electrical pads use a 150 µm pitch, and noncrossing S-bends separate the optical routes. Finally, we calculate the routed circuit response and check the layout for unintended optical, metal, and pad overlaps.
[8]:
PDK_ROOT = resources.files("siepic_ebeam_pdk")
def pdk_file(relative_path):
node = PDK_ROOT
for part in relative_path.split("/"):
node = node.joinpath(part)
return node
waveguide_node = pdk_file("CML/EBeam/source_data/strip_waveguide.txt")
with resources.as_file(waveguide_node) as waveguide_path:
waveguide_table = np.genfromtxt(waveguide_path, names=True, delimiter=",")
names = waveguide_table.dtype.names
neff_real = next(
name for name in names if "neff" in name.lower() and "real" in name.lower()
)
neff_imag = next(
(name for name in names if "neff" in name.lower() and "imag" in name.lower()), None
)
frequency_index = int(np.argmin(np.abs(waveguide_table["f"] - FREQ0_HZ)))
PDK_NEFF = complex(
waveguide_table[neff_real][frequency_index],
waveguide_table[neff_imag][frequency_index] if neff_imag else 1e-10,
)
print(f"SiEPIC 500-nm strip effective index at 1.55 um: {PDK_NEFF}")
def load_sparameters(relative_path, ports):
node = pdk_file(relative_path)
with resources.as_file(node) as path, np.load(path) as archive:
f = np.asarray(archive["f" if "f" in archive.files else "frequencies"])
s = np.asarray(archive["s" if "s" in archive.files else "s_array"])
return (s[None] if s.ndim == 2 else s), f, list(ports)
PDK_MODELS = {
"grating": load_sparameters(
"pymacros/opics_ebeam/data/gc_source/GC_TE1550_thickness=220 deltaw=0.npz",
("P0", "P1"),
),
"splitter": load_sparameters(
"pymacros/opics_ebeam/data/y_branch_source/Ybranch_Thickness =220 width=500.npz",
("P0", "P1", "P2"),
),
"terminator": load_sparameters(
"pymacros/opics_ebeam/data/ebeam_terminator_te1550/ebeam_terminator_te1550.npz",
("P0",),
),
}
pdk_model_checks = {}
for name, (s, f, ports) in PDK_MODELS.items():
maximum_singular_value = float(np.max(np.linalg.svd(s, compute_uv=False)))
pdk_model_checks[name] = {
"shape": s.shape,
"ports": ports,
"finite": bool(np.isfinite(s).all()),
"frequency in band": bool(f.min() <= FREQ0_HZ <= f.max()),
"passive": maximum_singular_value <= 1 + 1e-6,
"maximum singular value": maximum_singular_value,
"forward/reverse magnitude residual": float(
np.max(np.abs(np.abs(s) - np.abs(s.swapaxes(-1, -2))))
),
"complex phase-reference residual": float(
np.max(np.abs(s - s.swapaxes(-1, -2)))
),
}
terminator_s, terminator_f, _ = PDK_MODELS["terminator"]
terminator_frequency_index = int(np.argmin(np.abs(terminator_f - FREQ0_HZ)))
TERMINATOR_REFLECTION = complex(terminator_s[terminator_frequency_index, 0, 0])
TERMINATOR_RETURN_LOSS_DB = float(
-20 * np.log10(max(abs(TERMINATOR_REFLECTION), 1e-12))
)
pdk_model_checks["terminator"]["return loss dB at 1.55 um"] = TERMINATOR_RETURN_LOSS_DB
print(
"PDK checks:",
", ".join(f"{name}={value}" for name, value in pdk_model_checks.items()),
)
SiEPIC-EBeam-PDK Python module: siepic_ebeam_pdk, KLayout technology: EBeam
KLayout SiEPIC-Tools version 0.5.31
SiEPIC-Tools is up to date (0.5.31 vs 0.5.31).
Version check, time: 0.07045769691467285 seconds
SiEPIC-EBeam-PDK Python module: pymacros, v0.4.53
Libraries associated with Technology EBeam: ['EBeam', 'EBeam-ANT', 'EBeam-Dream', 'EBeam-SiN', 'EBeam_Beta']
SiEPIC 500-nm strip effective index at 1.55 um: (2.446947831912+1.280455805859e-09j)
PDK checks: grating={'shape': (100, 2, 2), 'ports': ['P0', 'P1'], 'finite': True, 'frequency in band': True, 'passive': True, 'maximum singular value': 0.7884114662993124, 'forward/reverse magnitude residual': 0.09337928989945549, 'complex phase-reference residual': 0.12567391014347995}, splitter={'shape': (51, 3, 3), 'ports': ['P0', 'P1', 'P2'], 'finite': True, 'frequency in band': True, 'passive': True, 'maximum singular value': 0.9956460637429584, 'forward/reverse magnitude residual': 0.005365000000000064, 'complex phase-reference residual': 0.0054007801290508126}, terminator={'shape': (101, 1, 1), 'ports': ['P0'], 'finite': True, 'frequency in band': True, 'passive': True, 'maximum singular value': 0.04896530000000001, 'forward/reverse magnitude residual': 0.0, 'complex phase-reference residual': 0.0, 'return loss dB at 1.55 um': 27.556735372370103}
[9]:
CHANNEL_PITCH, GRATING_PITCH = 250.0, 250.0
GRATING_LENGTH, SPLITTER_LENGTH = 392.0, 91.8
SPLITTER_GAP, MODULATOR_LENGTH, BLOCK_GAP = 80.0, 260.0, 1_500.0
ID_TAPER_LENGTH, ID_TAPER_MOUTH, ID_TAPER_OVERLAP = 3.1, 1.15, 0.0
ID_PORT_PITCH, OUTPUT_FIBER_PITCH, ROUTE_RADIUS = 1.75, 127.0, 5.0
BOND_PAD_PITCH, BOND_PAD_SIZE = 150.0, 100.0
BOND_PAD_OPEN_SIZE, BOND_PAD_EDGE_CLEARANCE = 95.0, 100.0
BOND_PAD_BANK_INNER_OFFSET = 225.0
N_SPLITTER_STAGES = int(np.log2(N_SPLITTER_LEAVES))
ID_LENGTH = DESIGN_EDGE_UM + 2 * ID_TAPER_LENGTH
CHIP_WIDTH = max(4_900.0, (N_SPLITTER_LEAVES - 1) * CHANNEL_PITCH + 1_150.0)
BOND_PAD_Y = CHIP_WIDTH / 2 - BOND_PAD_SIZE / 2 - BOND_PAD_EDGE_CLEARANCE
SPLITTER_LEAF_Y = (
np.arange(N_SPLITTER_LEAVES) - (N_SPLITTER_LEAVES - 1) / 2
) * CHANNEL_PITCH
active_leaf_start = (N_SPLITTER_LEAVES - N_CHANNELS) // 2
ACTIVE_LEAF_INDICES = tuple(range(active_leaf_start, active_leaf_start + N_CHANNELS))
DUMP_LEAF_INDICES = tuple(
leaf for leaf in range(N_SPLITTER_LEAVES) if leaf not in ACTIVE_LEAF_INDICES
)
LEAF_Y = SPLITTER_LEAF_Y[np.asarray(ACTIVE_LEAF_INDICES)]
IDEAL_ACTIVE_POWER_FRACTION = N_CHANNELS / N_SPLITTER_LEAVES
SPLITTER_X0 = GRATING_LENGTH + 320.0
SPLITTER_PITCH = SPLITTER_LENGTH + SPLITTER_GAP
PHASE_X0 = (
SPLITTER_X0 + (N_SPLITTER_STAGES - 1) * SPLITTER_PITCH + SPLITTER_LENGTH + 100.0
)
BANK_PITCH = MODULATOR_LENGTH + ID_LENGTH + 2 * BLOCK_GAP
PHASE_X = tuple(PHASE_X0 + layer * BANK_PITCH for layer in range(N_LAYERS))
ID_X = tuple(x + MODULATOR_LENGTH + BLOCK_GAP for x in PHASE_X)
OUTPUT_X = ID_X[-1] + ID_LENGTH + 320.0
CHIP_LENGTH = OUTPUT_X + 1_370.0
SPLITTER_Y = {}
for stage in range(N_SPLITTER_STAGES):
group_size = N_SPLITTER_LEAVES // 2**stage
SPLITTER_Y[stage] = np.asarray(
[
np.mean(SPLITTER_LEAF_Y[node * group_size : (node + 1) * group_size])
for node in range(2**stage)
]
)
print(
f"fanout: {N_SPLITTER_LEAVES} leaves, active={ACTIVE_LEAF_INDICES}, "
f"terminated={DUMP_LEAF_INDICES}, active power={IDEAL_ACTIVE_POWER_FRACTION:.1%}"
)
fanout: 8 leaves, active=(1, 2, 3, 4, 5, 6), terminated=(0, 7), active power=75.0%
[10]:
def pdk_port_spec(technology):
return next(
technology.ports[name]
for name in ("TE_1550_500", "strip_1550nm", "Strip")
if name in technology.ports
)
def data_model(model_name, ports=None):
s, f, default_ports = PDK_MODELS[model_name]
return pf.DataModel(
s_array=s,
frequencies=f,
ports=list(ports or default_ports),
interpolation_coords="real_imag",
)
def pdk_component(cell_name, model_name, aliases, technology):
base = siepic_forge.component(cell_name, technology=technology).copy()
for registered_model in tuple(base.models):
base.remove_model(registered_model)
component = pf.Component(f"SiEPIC {cell_name}", technology=technology)
reference = component.add_reference(base)
for alias, original in aliases.items():
component.add_port(reference[original], port_name=alias)
component.add_model(
data_model(model_name, aliases), f"Packaged {model_name}", set_active=True
)
return component
def splitter_component(technology):
aliases = {"P0": "P0", "P_LO": "P1", "P_HI": "P2"}
return pdk_component("ebeam_y_1550", "splitter", aliases, technology)
def terminator_component(technology):
base = siepic_forge.component(TERMINATOR_MODEL, technology=technology).copy()
for registered_model in tuple(base.models):
base.remove_model(registered_model)
component = pf.Component(
f"SiEPIC {TERMINATOR_MODEL} east-facing", technology=technology
)
reference = component.add_reference(pf.Reference(base, rotation=180))
component.add_port(reference["P0"], port_name="P0")
component.add_model(
data_model("terminator"), "Packaged terminator", set_active=True
)
return component
def phase_component(technology, name):
component = pf.Component(name, technology=technology)
spec = pdk_port_spec(technology)
component.add("Si", pf.Rectangle(corner1=(-5, -0.25), corner2=(265, 0.25)))
component.add("Si slab", pf.Rectangle(corner1=(0, -12), corner2=(260, 12)))
component.add("Si N", pf.Rectangle(corner1=(0, 0.5), corner2=(260, 2.25)))
component.add("Si N++", pf.Rectangle(corner1=(0, -2.25), corner2=(260, -0.5)))
component.add("M1_heater", pf.Rectangle(corner1=(0, 5.75), corner2=(260, 9.75)))
component.add("M1_heater", pf.Rectangle(corner1=(0, -9.75), corner2=(260, -5.75)))
signal_contact, signal_escape = (120, 7.75), (120, 17.75)
ground_contact, ground_escape = (140, -7.75), (-30, -37)
component.add(
"M1_heater",
pf.Path(signal_contact, 6).segment(signal_escape, join_limit="round"),
)
component.add("M1_heater", pf.Rectangle(center=ground_contact, size=(12, 12)))
component.add("VC", pf.Rectangle(center=ground_contact, size=(6, 6)))
component.add(
"M2_router",
pf.Path(ground_contact, 6).segment(
((ground_contact[0], ground_escape[1]), ground_escape), join_limit=-1.0
),
)
component.add_terminal(
pf.Terminal("M1_heater", pf.Rectangle(center=signal_escape, size=(6, 6))), "SIG"
)
component.add_terminal(
pf.Terminal("M2_router", pf.Rectangle(center=ground_escape, size=(6, 6))), "GND"
)
component.add_port(pf.Port((0, 0), 0, spec), "P0")
component.add_port(pf.Port((260, 0), 180, spec), "P1")
component.add_model(
pf.WaveguideModel(n_complex=PDK_NEFF, length=260, verbose=False),
"PDK waveguide",
set_active=True,
)
return component
def taper_polygon(x0, x1, y, w0, w1):
xs = np.linspace(x0, x1, 101)
widths = np.linspace(w0, w1, 101)
lower = [(x, y - width / 2) for x, width in zip(xs, widths)]
upper = [(x, y + width / 2) for x, width in zip(xs[::-1], widths[::-1])]
return pf.Polygon(lower + upper)
def matched_transfer_component(response, technology, name):
response = np.asarray(response)
n_outputs, n_inputs = response.shape
input_y = (np.arange(n_inputs) - (n_inputs - 1) / 2) * ID_PORT_PITCH
output_y = (np.arange(n_outputs) - (n_outputs - 1) / 2) * ID_PORT_PITCH
component = pf.Component(name, technology=technology)
spec = pdk_port_spec(technology)
for channel, y in enumerate(input_y):
component.add(
"Si",
taper_polygon(
0, ID_TAPER_LENGTH + ID_TAPER_OVERLAP, y, 0.5, ID_TAPER_MOUTH
),
)
component.add_port(pf.Port((0, float(y)), 0, spec), f"W{channel:02d}")
for channel, y in enumerate(output_y):
component.add(
"Si",
taper_polygon(
ID_TAPER_LENGTH + DESIGN_EDGE_UM - ID_TAPER_OVERLAP,
ID_LENGTH,
y,
ID_TAPER_MOUTH,
0.5,
),
)
component.add_port(pf.Port((ID_LENGTH, float(y)), 180, spec), f"E{channel:02d}")
component.add(
"Si",
pf.Rectangle(
corner1=(ID_TAPER_LENGTH, -DESIGN_EDGE_UM / 2),
corner2=(ID_TAPER_LENGTH + DESIGN_EDGE_UM, DESIGN_EDGE_UM / 2),
),
)
ports = [f"W{i:02d}" for i in range(n_inputs)] + [
f"E{i:02d}" for i in range(n_outputs)
]
s = np.zeros((1, n_inputs + n_outputs, n_inputs + n_outputs), dtype=complex)
s[0, n_inputs:, :n_inputs] = response
s[0, :n_inputs, n_inputs:] = response.T
component.add_model(
pf.DataModel(
s_array=s,
frequencies=[FREQ0_HZ],
ports=ports,
interpolation_coords="real_imag",
),
"matched route-calibration response",
set_active=True,
)
return component
def route_probe_component(technology, name, n_inputs, n_outputs):
component = pf.Component(name, technology=technology)
spec = pdk_port_spec(technology)
input_y = (np.arange(n_inputs) - (n_inputs - 1) / 2) * ID_PORT_PITCH
output_y = (np.arange(n_outputs) - (n_outputs - 1) / 2) * ID_PORT_PITCH
ports, pairs = [], []
for side, x, link_angle, probe_angle, positions in (
("W", 0, 0, 180, input_y),
("E", ID_LENGTH, 180, 0, output_y),
):
for channel, y in enumerate(positions):
link, probe = f"{side}{channel:02d}", f"{side}_PROBE{channel:02d}"
component.add_port(pf.Port((x, float(y)), link_angle, spec), link)
component.add_port(pf.Port((x, float(y)), probe_angle, spec), probe)
pairs.append((len(ports), len(ports) + 1))
ports += [link, probe]
s = np.zeros((1, len(ports), len(ports)), dtype=complex)
for link, probe in pairs:
s[0, link, probe] = s[0, probe, link] = 1
component.add_model(
pf.DataModel(
s_array=s,
frequencies=[FREQ0_HZ],
ports=ports,
interpolation_coords="real_imag",
),
"matched route probe",
set_active=True,
)
return component
def noncrossing_reference_length_route(port1=None, port2=None, model=None):
reference = pf.parametric.route(
port1=port1, port2=port2, radius=ROUTE_RADIUS, model=model
)
reference_length = pf.route_length(reference)
length_matched_model = pf.WaveguideModel(
n_complex=PDK_NEFF, length=reference_length, verbose=False
)
return pf.parametric.route_s_bend(
port1=port1, port2=port2, model=length_matched_model
)
def optical_route(first, second, waypoints=None, s_bend=False):
kwargs = {"model": pf.WaveguideModel(n_complex=PDK_NEFF, verbose=False)}
route_function = (
noncrossing_reference_length_route if s_bend else pf.parametric.route
)
if not s_bend:
kwargs["radius"] = ROUTE_RADIUS
if waypoints:
kwargs["waypoints"] = waypoints
return first, second, route_function, kwargs
[11]:
def add_optical_instances(instances, responses, technology, route_calibration=False):
grating = pdk_component(
"ebeam_gc_te1550", "grating", {"P0": "P0", "Fiber": "P1"}, technology
)
splitter = splitter_component(technology)
terminator = terminator_component(technology)
instances["gc_in"] = {"component": grating, "origin": (0, 0)}
for stage, y_values in SPLITTER_Y.items():
for node, y in enumerate(y_values):
instances[f"split_{stage}_{node}"] = {
"component": splitter,
"origin": (SPLITTER_X0 + stage * SPLITTER_PITCH, float(y)),
}
for leaf in DUMP_LEAF_INDICES:
instances[f"dump_{leaf}"] = {
"component": terminator,
"origin": (PHASE_X0, float(SPLITTER_LEAF_Y[leaf])),
}
for layer in range(N_LAYERS):
phase = phase_component(technology, f"phase cell {layer + 1}")
inverse = (
route_probe_component(
technology,
f"route probe {layer + 1}",
LAYER_INPUT_COUNTS[layer],
LAYER_OUTPUT_COUNTS[layer],
)
if route_calibration
else matched_transfer_component(
responses[layer], technology, f"inverse region {layer + 1}"
)
)
instances[f"id_{layer}"] = {"component": inverse, "origin": (ID_X[layer], 0)}
for channel, y in enumerate(LEAF_Y):
instances[f"phase_{layer}_{channel}"] = {
"component": phase,
"origin": (PHASE_X[layer], float(y)),
}
final_output_y = (np.arange(N_CLASSES) - (N_CLASSES - 1) / 2) * OUTPUT_FIBER_PITCH
for output, y in enumerate(final_output_y):
instances[f"gc_out_{output}"] = {
"component": grating,
"origin": (OUTPUT_X, float(y)),
"rotation": 180,
}
return grating
def m1_bond_pad(technology, bond_pad):
component = pf.Component("M1-to-M2 bond pad", technology=technology)
component.add_reference(bond_pad)
component.add("M1_heater", pf.Rectangle(center=(0, 0), size=(20, 20)))
component.add("VC", pf.Rectangle(center=(0, 0), size=(6, 6)))
component.add_terminal(
pf.Terminal("M1_heater", pf.Rectangle(center=(0, 0), size=(6, 6))), "T0"
)
return component
def bond_pad_x_positions(x):
channel = np.arange(N_CHANNELS)
signal = (
x + BOND_PAD_BANK_INNER_OFFSET + (N_CHANNELS - 1 - channel) * BOND_PAD_PITCH
)
ground = x - BOND_PAD_BANK_INNER_OFFSET - channel * BOND_PAD_PITCH
return signal, ground
def add_edge_pads(instances, technology):
ground_pad = siepic_forge.component("ebeam_BondPad", technology=technology)
signal_pad = m1_bond_pad(technology, ground_pad)
for layer, x in enumerate(PHASE_X):
signal_pad_x, ground_pad_x = bond_pad_x_positions(x)
for channel in range(N_CHANNELS):
instances[f"sig_pad_{layer}_{channel}"] = {
"component": signal_pad,
"origin": (signal_pad_x[channel], BOND_PAD_Y),
}
instances[f"gnd_pad_{layer}_{channel}"] = {
"component": ground_pad,
"origin": (ground_pad_x[channel], -BOND_PAD_Y),
}
def optical_routes():
routes = [optical_route(("gc_in", "P0"), ("split_0_0", "P0"))]
for stage in range(N_SPLITTER_STAGES - 1):
for parent in range(len(SPLITTER_Y[stage])):
for branch, child in enumerate((2 * parent, 2 * parent + 1)):
routes.append(
optical_route(
(f"split_{stage}_{parent}", ("P_LO", "P_HI")[branch]),
(f"split_{stage + 1}_{child}", "P0"),
)
)
leaf_stage = N_SPLITTER_STAGES - 1
for channel, leaf in enumerate(ACTIVE_LEAF_INDICES):
routes.append(
optical_route(
(f"split_{leaf_stage}_{leaf // 2}", ("P_LO", "P_HI")[leaf % 2]),
(f"phase_0_{channel}", "P0"),
)
)
for layer in range(N_LAYERS):
routes.append(
optical_route(
(f"phase_{layer}_{channel}", "P1"),
(f"id_{layer}", f"W{channel:02d}"),
s_bend=True,
)
)
if layer + 1 < N_LAYERS:
routes.append(
optical_route(
(f"id_{layer}", f"E{channel:02d}"),
(f"phase_{layer + 1}_{channel}", "P0"),
s_bend=True,
)
)
elif channel < N_CLASSES:
routes.append(
optical_route(
(f"id_{layer}", f"E{channel:02d}"),
(f"gc_out_{channel}", "P0"),
s_bend=True,
)
)
for leaf in DUMP_LEAF_INDICES:
routes.append(
optical_route(
(f"split_{leaf_stage}_{leaf // 2}", ("P_LO", "P_HI")[leaf % 2]),
(f"dump_{leaf}", "P0"),
)
)
return routes
def electrical_routes():
routes, terminals = [], []
for layer, x in enumerate(PHASE_X):
signal_pad_x, ground_pad_x = bond_pad_x_positions(x)
for channel, y in enumerate(LEAF_Y):
signal_waypoints = [(signal_pad_x[channel], y + 17.75)]
ground_waypoints = [(ground_pad_x[channel], y - 37)]
routes.append(
(
(f"phase_{layer}_{channel}", "SIG"),
(f"sig_pad_{layer}_{channel}", "T0"),
pf.parametric.route_manhattan,
{
"waypoints": signal_waypoints,
"width": 6.0,
"join_limit": "round",
},
)
)
routes.append(
(
(f"phase_{layer}_{channel}", "GND"),
(f"gnd_pad_{layer}_{channel}", "T0"),
pf.parametric.route_manhattan,
{"waypoints": ground_waypoints, "width": 6.0, "join_limit": -1.0},
)
)
terminals += [
(f"sig_pad_{layer}_{channel}", "T0", f"signal_{layer}_{channel}"),
(f"gnd_pad_{layer}_{channel}", "T0", f"ground_{layer}_{channel}"),
]
return routes, terminals
[50]:
def build_chip(responses, tag, route_calibration=False):
technology = siepic_forge.ebeam()
pf.config.default_technology = technology
instances = {}
add_optical_instances(instances, responses, technology, route_calibration)
add_edge_pads(instances, technology)
terminal_routes, terminals = electrical_routes()
ports = [("gc_in", "Fiber", "fiber_in")]
ports += [(f"gc_out_{i}", "Fiber", f"fiber_out_{i:02d}") for i in range(N_CLASSES)]
if route_calibration:
for layer in range(N_LAYERS):
ports += [
(f"id_{layer}", f"W_PROBE{i:02d}", f"id{layer}_in_{i:02d}")
for i in range(LAYER_INPUT_COUNTS[layer])
]
ports += [
(f"id_{layer}", f"E_PROBE{i:02d}", f"id{layer}_out_{i:02d}")
for i in range(LAYER_OUTPUT_COUNTS[layer])
]
chip = pf.component_from_netlist(
{
"name": f"deep_inverse_designed_pnn_{tag}",
"instances": instances,
"routes": optical_routes(),
"terminal routes": terminal_routes,
"terminals": terminals,
"ports": ports,
"models": [(pf.CircuitModel(verbose=False), "Circuit")],
"active models": {"optical": "Circuit"},
}
)
# chip.add(
# "FloorPlan",
# pf.Rectangle(
# corner1=(-500, -CHIP_WIDTH / 2), corner2=(CHIP_LENGTH - 500, CHIP_WIDTH / 2)
# ),
# )
pf.set_unique_names(chip)
chip.write_gds(LAYOUT_DIR / f"deep_inverse_designed_pnn__{tag}.gds")
chip.write_oas(LAYOUT_DIR / f"deep_inverse_designed_pnn__{tag}.oas")
return chip
identity_responses = [
np.eye(N_CHANNELS, dtype=complex),
np.eye(N_CHANNELS, dtype=complex),
np.eye(N_CHANNELS, dtype=complex)[np.asarray(FINAL_OUTPUT_INDICES)],
]
full_chip = build_chip(identity_responses, "full_chip")
route_calibration_chip = build_chip(
identity_responses, "route_calibration", route_calibration=True
)
display(full_chip)
[13]:
def nested_model_types(component, seen=None):
seen = set() if seen is None else seen
if id(component) in seen:
return []
seen.add(id(component))
model_types = [type(model).__name__ for model in component.models.values()]
for reference in component.references:
model_types.extend(nested_model_types(reference.component, seen))
return model_types
registered_model_types = nested_model_types(full_chip)
netlist = full_chip.get_netlist()
physical_netlist = full_chip.get_netlist(include_virtual_connections=False)
virtual_links = netlist.get(
"virtual connections", netlist.get("virtual_connections", ())
)
layout_counts = {
"channels": N_CHANNELS,
"classes": N_CLASSES,
"splitter leaves": N_SPLITTER_LEAVES,
"active splitter leaves": list(ACTIVE_LEAF_INDICES),
"terminated splitter leaves": list(DUMP_LEAF_INDICES),
"ideal active fanout power fraction": IDEAL_ACTIVE_POWER_FRACTION,
"input gratings": 1,
"output gratings": N_CLASSES,
"layer input counts": list(LAYER_INPUT_COUNTS),
"layer output counts": list(LAYER_OUTPUT_COUNTS),
"splitters": N_SPLITTER_LEAVES - 1,
"terminators": len(DUMP_LEAF_INDICES),
"terminator model": TERMINATOR_MODEL,
"terminator return loss dB at 1.55 um": TERMINATOR_RETURN_LOSS_DB,
"phase cells": N_LAYERS * N_CHANNELS,
"inverse-designed regions": N_LAYERS,
"optical routes": len(optical_routes()),
"electrical routes": len(electrical_routes()[0]),
"bond pads": len(electrical_routes()[1]),
"virtual links": len(virtual_links),
"netlist instances": len(netlist.get("instances", {})),
"physical connections": len(physical_netlist.get("connections", ())),
"top-level optical ports": len(full_chip.ports),
"top-level electrical terminals": len(full_chip.terminals),
"component model types": sorted(set(registered_model_types)),
"component Tidy3D models": registered_model_types.count("Tidy3DModel"),
"route radius um": ROUTE_RADIUS,
"design edge um": DESIGN_EDGE_UM,
}
print(
f"circuit: {layout_counts['splitters']} splitters, "
f"{layout_counts['phase cells']} phase cells, "
f"{layout_counts['inverse-designed regions']} inverse-designed regions, "
f"{layout_counts['output gratings']} output fibers"
)
def s_value(s_matrix, source, destination):
candidates = (
(f"{source}@0", f"{destination}@0"),
(destination, source),
(f"{destination}@0", f"{source}@0"),
)
key = next(key for key in candidates if key in s_matrix.elements)
return complex(np.asarray(s_matrix.elements[key]).reshape(-1)[0])
chip_s = route_calibration_chip.s_matrix([FREQ0_HZ], show_progress=False)
route_inputs = np.asarray(
[
[
s_value(
chip_s,
"fiber_in" if layer == 0 else f"id{layer - 1}_out_{channel:02d}",
f"id{layer}_in_{channel:02d}",
)
for channel in range(N_CHANNELS)
]
for layer in range(N_LAYERS)
]
)
route_output = np.asarray(
[
s_value(chip_s, f"id{N_LAYERS - 1}_out_{output:02d}", f"fiber_out_{output:02d}")
for output in range(N_CLASSES)
]
)
route_loss_db = -20 * np.log10(np.maximum(np.abs(route_inputs), 1e-12))
ROUTE_POWER_REFERENCE = float(
np.prod(np.mean(np.abs(route_inputs) ** 2, axis=1))
* np.mean(np.abs(route_output) ** 2)
)
print("route response shapes:", route_inputs.shape, route_output.shape)
print(
"finite circuit response:",
bool(np.isfinite(route_inputs).all() and np.isfinite(route_output).all()),
)
print(f"fixed routed power reference: {ROUTE_POWER_REFERENCE:.6f}")
print(f"output fibers: {N_CLASSES} on {OUTPUT_FIBER_PITCH:.1f} µm pitch")
circuit: 7 splitters, 18 phase cells, 3 inverse-designed regions, 4 output fibers route response shapes: (3, 6) (4,) finite circuit response: True fixed routed power reference: 0.034648 output fibers: 4 on 127.0 µm pitch
[14]:
route_gds = LAYOUT_DIR / "deep_inverse_designed_pnn__full_chip.gds"
layout = kdb.Layout()
layout.read(str(route_gds))
top_cell = layout.top_cell()
def layout_region(layer, datatype=0):
return kdb.Region(top_cell.begin_shapes_rec(layout.layer(layer, datatype)))
dbu = layout.dbu
silicon_region = layout_region(1).merged()
m1_region, m2_region = layout_region(11).merged(), layout_region(12).merged()
mlopen_region = layout_region(13).merged()
floorplan_region = layout_region(99).merged()
optical_route_regions = []
for instance in top_cell.each_inst():
if instance.cell.name.startswith("route_s_bend__") or instance.cell.name.startswith(
"route__"
):
region = kdb.Region(instance.cell.begin_shapes_rec(layout.layer(1, 0)))
region.transform(instance.cplx_trans)
optical_route_regions.append(region)
optical_route_intersections = sum(
not (first & second).is_empty()
for index, first in enumerate(optical_route_regions)
for second in optical_route_regions[index + 1 :]
)
tolerance = 0.002
rule_checks = {
"merged Si width markers < 60 nm": silicon_region.width_check(
round((0.060 - tolerance) / dbu), False, kdb.Metrics.Euclidian, 80
).size(),
"merged Si spacing markers < 70 nm": silicon_region.space_check(
round((0.070 - tolerance) / dbu), False, kdb.Metrics.Euclidian, 80
).size(),
"Si outside floor plan": (silicon_region - floorplan_region).count(),
"M1 width markers < 3 um": m1_region.width_check(
round((3.0 - tolerance) / dbu), False, kdb.Metrics.Euclidian, 70
).size(),
"M1 spacing markers < 3 um": m1_region.space_check(
round((3.0 - tolerance) / dbu)
).size(),
"M2 width markers < 5 um": m2_region.width_check(
round((5.0 - tolerance) / dbu), False, kdb.Metrics.Euclidian, 70
).size(),
"M2 spacing markers < 8 um": m2_region.space_check(
round((8.0 - tolerance) / dbu)
).size(),
"MLOpen width markers < 10 um": mlopen_region.width_check(
round((10.0 - tolerance) / dbu)
).size(),
"MLOpen spacing markers < 10 um": mlopen_region.space_check(
round((10.0 - tolerance) / dbu)
).size(),
}
expected_metal_components = 2 * N_LAYERS * N_CHANNELS
interface_checks = {
"taper-mouth clearance um": ID_PORT_PITCH - ID_TAPER_MOUTH,
"taper-to-region overlap um": ID_TAPER_OVERLAP,
"pad M2 edge clearance um": BOND_PAD_PITCH - BOND_PAD_SIZE,
"pad opening edge clearance um": BOND_PAD_PITCH - BOND_PAD_OPEN_SIZE,
"pad-to-chip-edge clearance um": (CHIP_WIDTH / 2 - BOND_PAD_Y - BOND_PAD_SIZE / 2),
"M1 connected components": m1_region.count(),
"M2 connected components": m2_region.count(),
"MLOpen connected components": mlopen_region.count(),
"optical route-route intersections": optical_route_intersections,
"electrical nets remain isolated": bool(
m1_region.count() == expected_metal_components
and m2_region.count() == expected_metal_components
and mlopen_region.count() == expected_metal_components
),
"optical routes remain separate": bool(optical_route_intersections == 0),
"taper mouths do not touch": bool(ID_PORT_PITCH - ID_TAPER_MOUTH > 0),
"tapers meet the fixed rail without intrusion": bool(ID_TAPER_OVERLAP == 0),
}
electrical_terminal_names = [terminal[2] for terminal in electrical_routes()[1]]
layout_verification = {
"KLayout rule screen": rule_checks,
"custom interface and layer checks": interface_checks,
"virtual links": len(virtual_links),
"electrical terminals": len(electrical_terminal_names),
"unique electrical terminals": len(set(electrical_terminal_names)),
"active splitter leaves": list(ACTIVE_LEAF_INDICES),
"terminated splitter leaves": list(DUMP_LEAF_INDICES),
"all splitter leaves assigned": len(ACTIVE_LEAF_INDICES) + len(DUMP_LEAF_INDICES)
== N_SPLITTER_LEAVES,
"terminator return loss above 20 dB": TERMINATOR_RETURN_LOSS_DB > 20.0,
"packaged DRC deck": "siepic_ebeam_pdk/drc/SiEPIC_EBeam_DRC.lydrc",
}
print(
"layout checks:",
f"rule markers={sum(rule_checks.values())}, "
f"route overlaps={optical_route_intersections}, "
f"electrical isolation={interface_checks['electrical nets remain isolated']}",
)
layout checks: rule markers=0, route overlaps=0, electrical isolation=True
[15]:
fig, axes = plt.subplots(1, 3, figsize=(11.2, 3.0), constrained_layout=True)
for layer in range(N_LAYERS):
axes[0].plot(
np.arange(N_CHANNELS),
np.abs(route_inputs[layer]),
marker=".",
label=f"path {layer + 1}",
)
axes[1].plot(
np.arange(N_CHANNELS), np.unwrap(np.angle(route_inputs[layer])), marker="."
)
centered = route_loss_db[layer] - route_loss_db[layer].mean()
axes[2].plot(np.arange(N_CHANNELS), centered, marker=".")
axes[0].set(xlabel="channel", ylabel="amplitude", title="Routed amplitude")
axes[1].set(xlabel="channel", ylabel="phase (rad)", title="Routed phase")
axes[2].set(
xlabel="channel",
ylabel="loss relative to mean (dB)",
title="Propagation-loss spread",
)
axes[0].legend(frameon=False)
fig.suptitle(
f"{N_SPLITTER_LEAVES}-leaf fanout: active {ACTIVE_LEAF_INDICES}, "
f"terminated {DUMP_LEAF_INDICES}, ideal active power {IDEAL_ACTIVE_POWER_FRACTION:.0%}"
)
plt.show()
Train the Optical Target¶
Before optimizing the device geometry, we train the three optical transformations directly as complex matrices. Each matrix is written as
$$ M_\ell=U_\ell\,\operatorname{diag}(\sigma_\ell)\,V_\ell^\dagger, $$
where $U_\ell$ and $V_\ell$ are unitary matrices. We constrain each singular amplitude to
$$ 10^{-6/20}\leq \sigma_{\ell,j}\leq10^{-3/20}, $$
which keeps every transformation passive and gives each singular channel an insertion loss between 3 and 6 dB.
The detected powers are $p=|a_{\mathrm{out}}|^2$, and the class logits are $\gamma p/P_{\mathrm{ref}}$, where $\gamma$ is a learned positive scale. We minimize cross-entropy together with penalties on the optical margin, long-range channel coupling, and excessive $\gamma$.
During training, 16 perturbed copies of the three matrices are generated with 8–12% relative transmission error. This teaches the classifier to tolerate the difference between the target matrices and their later electromagnetic realizations.
[16]:
def unitary(raw):
return jax.scipy.linalg.expm(raw - jnp.conj(raw.T))
def passive_matrix(left_raw, right_raw, singular_raw):
left, right = unitary(left_raw), unitary(right_raw)
singular = SIGMA_MIN + (SIGMA_MAX - SIGMA_MIN) * jax.nn.sigmoid(singular_raw)
return left @ jnp.diag(singular.astype(jnp.complex64)) @ jnp.conj(right.T)
batched_passive_matrix = jax.vmap(passive_matrix)
def layer_matrices(params):
left = params["left_re"] + 1j * params["left_im"]
right = params["right_re"] + 1j * params["right_im"]
return batched_passive_matrix(left, right, params["singular_raw"])
route_inputs_jax = jnp.asarray(route_inputs)
route_output_jax = jnp.asarray(route_output)
readout_slice = slice((N_CHANNELS - N_CLASSES) // 2, (N_CHANNELS + N_CLASSES) // 2)
channel_rows, channel_columns = jnp.indices((N_CHANNELS, N_CHANNELS))
normalized_channel_distance = jnp.abs(channel_rows - channel_columns) / (N_CHANNELS - 1)
def optical_forward(matrices, features):
phase = jnp.exp(1j * PHASE_SCALE * features)
fields = jnp.ones_like(features, dtype=jnp.complex64)
for layer in range(N_LAYERS):
fields = fields * route_inputs_jax[layer] * phase
fields = fields @ matrices[layer].T
return fields[:, readout_slice] * route_output_jax
def routing_transport(matrices):
power = jnp.abs(matrices) ** 2
layer_power = jnp.sum(power, axis=(1, 2))
layer_cost = jnp.sum(
power * normalized_channel_distance**2, axis=(1, 2)
) / jnp.maximum(layer_power, 1e-12)
return jnp.mean(layer_cost)
def routing_distance_statistics(matrices):
power = np.abs(np.asarray(matrices)) ** 2
distance = np.abs(np.arange(N_CHANNELS)[:, None] - np.arange(N_CHANNELS)[None, :])
layer_power = power.sum(axis=(1, 2))
mean_distance = (power * distance).sum(axis=(1, 2)) / layer_power
rms_distance = np.sqrt((power * distance**2).sum(axis=(1, 2)) / layer_power)
percentile_90 = []
for layer_power_map in power:
weights = np.bincount(
distance.ravel(), weights=layer_power_map.ravel(), minlength=N_CHANNELS
)
percentile_90.append(np.searchsorted(np.cumsum(weights) / weights.sum(), 0.9))
return mean_distance, rms_distance, np.asarray(percentile_90)
def logits_and_scores(params, features):
matrices = layer_matrices(params)
intensities = jnp.abs(optical_forward(matrices, features)) ** 2
scores = intensities
gamma = jnp.exp(params["log_gamma"])
return gamma * scores / ROUTE_POWER_REFERENCE, scores, matrices, gamma
def initialize_parameters(key):
keys = jax.random.split(key, 4)
shape = (N_LAYERS, N_CHANNELS, N_CHANNELS)
return {
"left_re": 1e-2 * jax.random.normal(keys[0], shape),
"left_im": 1e-2 * jax.random.normal(keys[1], shape),
"right_re": 1e-2 * jax.random.normal(keys[2], shape),
"right_im": 1e-2 * jax.random.normal(keys[3], shape),
"singular_raw": jnp.full((N_LAYERS, N_CHANNELS), 3.0),
"log_gamma": jnp.asarray(np.log(10.0), dtype=jnp.float32),
}
def host_tree(tree):
return jax.tree.map(lambda value: np.asarray(value), tree)
@jax.jit
def evaluate(params, features, labels):
logits, _, matrices, gamma = logits_and_scores(params, features)
loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == labels)
return loss, accuracy, routing_transport(matrices), gamma
[17]:
def adjacent_unitary(angles, phases):
unitary_matrix = jnp.eye(N_CHANNELS, dtype=jnp.complex128)
for offset in (0, 1):
for channel in range(offset, N_CHANNELS - 1, 2):
first_row = unitary_matrix[channel]
second_row = unitary_matrix[channel + 1]
cosine, sine = jnp.cos(angles[channel]), jnp.sin(angles[channel])
phase = jnp.exp(1j * phases[channel])
unitary_matrix = unitary_matrix.at[channel].set(
cosine * first_row + phase * sine * second_row
)
unitary_matrix = unitary_matrix.at[channel + 1].set(
-jnp.conj(phase) * sine * first_row + cosine * second_row
)
return unitary_matrix
def sample_layer_perturbation(key):
keys = jax.random.split(key, 8)
return {
"input_loss": jax.random.normal(keys[0], (N_CHANNELS,)),
"output_loss": jax.random.normal(keys[1], (N_CHANNELS,)),
"input_phase": jax.random.normal(keys[2], (N_CHANNELS,)),
"output_phase": jax.random.normal(keys[3], (N_CHANNELS,)),
"input_angles": jax.random.normal(keys[4], (N_CHANNELS - 1,)),
"output_angles": jax.random.normal(keys[5], (N_CHANNELS - 1,)),
"input_mixing_phase": jax.random.normal(keys[6], (N_CHANNELS - 1,)),
"output_mixing_phase": jax.random.normal(keys[7], (N_CHANNELS - 1,)),
}
def structured_perturb_layer(matrix, perturbation, target_error, pair_sign):
input_loss = jax.nn.sigmoid(pair_sign * perturbation["input_loss"])
output_loss = jax.nn.sigmoid(pair_sign * perturbation["output_loss"])
input_phase = 0.15 * pair_sign * perturbation["input_phase"]
output_phase = 0.15 * pair_sign * perturbation["output_phase"]
input_phase = input_phase - jnp.mean(input_phase)
output_phase = output_phase - jnp.mean(output_phase)
input_angles = 0.05 * pair_sign * perturbation["input_angles"]
output_angles = 0.05 * pair_sign * perturbation["output_angles"]
input_mixing_phase = jnp.pi * jnp.tanh(
pair_sign * perturbation["input_mixing_phase"]
)
output_mixing_phase = jnp.pi * jnp.tanh(
pair_sign * perturbation["output_mixing_phase"]
)
input_diagonal = 10 ** (-input_loss / 20) * jnp.exp(1j * input_phase)
output_diagonal = 10 ** (-output_loss / 20) * jnp.exp(1j * output_phase)
input_mixing = adjacent_unitary(input_angles, input_mixing_phase)
output_mixing = adjacent_unitary(output_angles, output_mixing_phase)
full_matrix = output_diagonal[:, None] * (output_mixing @ matrix @ input_mixing)
full_matrix = full_matrix * input_diagonal[None, :]
difference = full_matrix - matrix
matrix_norm = jnp.maximum(jnp.linalg.norm(matrix), 1e-12)
full_error = jnp.linalg.norm(difference) / matrix_norm
interpolation = jnp.minimum(1.0, target_error / jnp.maximum(full_error, 1e-12))
perturbed = matrix + interpolation * difference
actual_error = jnp.linalg.norm(perturbed - matrix) / matrix_norm
return perturbed, actual_error
def sample_chip_perturbation(key):
return jax.vmap(sample_layer_perturbation)(jax.random.split(key, N_LAYERS))
def perturb_chip(matrices, perturbations, target_errors, pair_sign):
return jax.vmap(structured_perturb_layer, in_axes=(0, 0, 0, None))(
matrices, perturbations, target_errors, pair_sign
)
def perturb_replicas(matrices, perturbations, target_errors, pair_signs):
return jax.vmap(perturb_chip, in_axes=(None, 0, 0, 0))(
matrices, perturbations, target_errors, pair_signs
)
def optical_scores(matrices, features):
return jnp.abs(optical_forward(matrices, features)) ** 2
def normalized_optical_margin(scores, labels):
label_mask = jax.nn.one_hot(labels, N_CLASSES, dtype=scores.dtype)
label_mask = label_mask.reshape((1,) * (scores.ndim - 2) + label_mask.shape)
correct = jnp.sum(scores * label_mask, axis=-1)
competitor = jnp.max(jnp.where(label_mask.astype(bool), -jnp.inf, scores), axis=-1)
return (correct - competitor) / (jnp.sum(scores, axis=-1) + 1e-12)
[18]:
def robust_objective(params, key, features, labels):
matrices = layer_matrices(params)
clean_scores = optical_scores(matrices, features)
gamma = jnp.exp(params["log_gamma"])
clean_logits = gamma * clean_scores / ROUTE_POWER_REFERENCE
clean_loss = optax.softmax_cross_entropy_with_integer_labels(
clean_logits, labels
).mean()
structure_key, error_key = jax.random.split(key)
perturbations = jax.vmap(sample_chip_perturbation)(
jax.random.split(structure_key, ROBUST_REPLICAS)
)
pair_signs = jnp.ones((ROBUST_REPLICAS,))
target_errors = jax.random.uniform(
error_key,
(ROBUST_REPLICAS, N_LAYERS),
minval=ROBUST_WINDOW_MIN,
maxval=ROBUST_WINDOW_MAX,
)
perturbed_matrices, actual_errors = perturb_replicas(
matrices, perturbations, target_errors, pair_signs
)
perturbed_scores = jax.vmap(lambda values: optical_scores(values, features))(
perturbed_matrices
)
perturbed_logits = gamma * perturbed_scores / ROUTE_POWER_REFERENCE
replica_labels = jnp.broadcast_to(labels, (ROBUST_REPLICAS, labels.shape[0]))
perturbed_losses = optax.softmax_cross_entropy_with_integer_labels(
perturbed_logits, replica_labels
).mean(axis=1)
perturbed_accuracies = jnp.mean(
jnp.argmax(perturbed_scores, axis=-1) == replica_labels, axis=1
)
margins = normalized_optical_margin(perturbed_scores, labels)
margin_penalty = jnp.maximum(ROBUST_MARGIN_TARGET - margins, 0).mean()
gamma_term = GAMMA_WEIGHT * jnp.square(
jnp.maximum(gamma - GAMMA_TARGET, 0.0) / GAMMA_TARGET
)
mean_perturbed_loss = jnp.mean(perturbed_losses)
loss = (
0.35 * clean_loss
+ 0.65 * mean_perturbed_loss
+ ROBUST_MARGIN_WEIGHT * margin_penalty
+ routing_transport(matrices)
+ gamma_term
)
clean_accuracy = jnp.mean(jnp.argmax(clean_scores, axis=-1) == labels)
metrics = jnp.asarray(
(
clean_loss,
mean_perturbed_loss,
margin_penalty,
clean_accuracy,
jnp.mean(perturbed_accuracies),
jnp.min(perturbed_accuracies),
jnp.mean(actual_errors),
)
)
return loss, metrics
robust_steps = ROBUST_EPOCHS * int(np.ceil(len(train_features) / ROBUST_BATCH_SIZE))
robust_schedule = optax.cosine_decay_schedule(ROBUST_LR, robust_steps, alpha=0.1)
robust_optimizer = optax.adam(robust_schedule)
@jax.jit
def robust_update(params, optimizer_state, key, features, labels):
(loss, metrics), gradients = jax.value_and_grad(robust_objective, has_aux=True)(
params, key, features, labels
)
updates, optimizer_state = robust_optimizer.update(
gradients, optimizer_state, params
)
return optax.apply_updates(params, updates), optimizer_state, loss, metrics
[19]:
def perturbation_batch_metrics(params, features, labels, keys, target_errors):
matrices = layer_matrices(params)
perturbations = jax.vmap(sample_chip_perturbation)(keys)
pair_signs = jnp.ones((len(keys),))
perturbed_matrices, actual_errors = perturb_replicas(
matrices, perturbations, target_errors, pair_signs
)
scores = jax.vmap(lambda values: optical_scores(values, features))(
perturbed_matrices
)
replica_labels = jnp.broadcast_to(labels, (len(keys), labels.shape[0]))
accuracies = jnp.mean(jnp.argmax(scores, axis=-1) == replica_labels, axis=1)
mean_margins = jnp.mean(normalized_optical_margin(scores, labels), axis=1)
maximum_singular = jnp.max(
jnp.linalg.svd(perturbed_matrices, compute_uv=False), axis=(1, 2)
)
return accuracies, mean_margins, jnp.mean(actual_errors, axis=1), maximum_singular
def evaluate_perturbation_bank(
params, features, labels, keys, severities, chunk_size=8
):
feature_array, label_array = jnp.asarray(features), jnp.asarray(labels)
severities = np.asarray(severities, dtype=float)
target_errors = (
np.repeat(severities[:, None], N_LAYERS, axis=1)
if severities.ndim == 1
else severities
)
collected = [[], [], [], []]
for start in range(0, len(keys), chunk_size):
stop = min(start + chunk_size, len(keys))
values = perturbation_batch_metrics(
params,
feature_array,
label_array,
jnp.asarray(keys[start:stop]),
jnp.asarray(target_errors[start:stop]),
)
for destination, value in zip(collected, values):
destination.append(np.asarray(value))
accuracy, margin, actual_error, maximum_singular = (
np.concatenate(values) for values in collected
)
return {
"accuracy": accuracy,
"margin": margin,
"actual_error": actual_error,
"maximum_singular": maximum_singular,
}
def fixed_error_bank(seed, size, minimum, maximum):
rng = np.random.default_rng(seed)
severities = rng.uniform(minimum, maximum, (size, N_LAYERS))
keys = np.asarray(jax.random.split(jax.random.PRNGKey(seed), size))
return keys, severities
validation_keys, validation_severities = fixed_error_bank(
SEED + 1_000, ROBUST_VALIDATION_SAMPLES, ROBUST_WINDOW_MIN, ROBUST_WINDOW_MAX
)
audit_keys, audit_severities = fixed_error_bank(
SEED + 2_000, ROBUST_AUDIT_SAMPLES, ROBUST_WINDOW_MIN, ROBUST_WINDOW_MAX
)
exact_keys, exact_severities = fixed_error_bank(
SEED + 3_000, ROBUST_AUDIT_SAMPLES, ROBUST_TARGET_ERROR, ROBUST_TARGET_ERROR
)
def save_repeated_checkpoint(state):
temporary_path = REPEATED_CHECKPOINT_PATH.with_suffix(".tmp")
with temporary_path.open("wb") as handle:
pickle.dump(state, handle, protocol=pickle.HIGHEST_PROTOCOL)
temporary_path.replace(REPEATED_CHECKPOINT_PATH)
[20]:
def validate_repeated_checkpoint(state):
expected_shape = (N_LAYERS, N_CHANNELS, N_CHANNELS)
actual_shape = np.asarray(state["params"]["left_re"]).shape
if actual_shape != expected_shape:
raise ValueError(
f"Surrogate checkpoint shape {actual_shape} does not match {expected_shape}."
)
configuration = state.get("configuration")
if configuration is None:
raise ValueError("Checkpoint is missing preprocessing metadata.")
if configuration != MATRIX_TRAINING_CONFIGURATION:
raise ValueError(
f"Surrogate checkpoint configuration {configuration} does not match "
f"{MATRIX_TRAINING_CONFIGURATION}."
)
training_rng = np.random.default_rng(MATRIX_TRAINING_SEED)
if REPEATED_CHECKPOINT_PATH.exists():
with REPEATED_CHECKPOINT_PATH.open("rb") as handle:
run = pickle.load(handle)
validate_repeated_checkpoint(run)
training_rng.bit_generator.state = run["rng_state"]
print(f"loaded N={N_CHANNELS} trained matrices at epoch {run['completed_epoch']}")
else:
initial_params = host_tree(
initialize_parameters(jax.random.PRNGKey(MATRIX_TRAINING_SEED))
)
run = {
"params": initial_params,
"optimizer_state": host_tree(robust_optimizer.init(initial_params)),
"jax_key": np.asarray(jax.random.PRNGKey(MATRIX_TRAINING_SEED + 10_000)),
"rng_state": training_rng.bit_generator.state,
"history": [],
"best_score": -np.inf,
"best_clean_accuracy": -np.inf,
"best_params": initial_params,
"best_epoch": 0,
"completed_epoch": 0,
"seed": MATRIX_TRAINING_SEED,
"configuration": MATRIX_TRAINING_CONFIGURATION,
}
training_params = run["params"]
training_optimizer_state = run["optimizer_state"]
training_key = jnp.asarray(run["jax_key"])
training_history = list(run["history"])
best_score = run["best_score"]
best_clean_accuracy = run["best_clean_accuracy"]
best_params = run["best_params"]
best_epoch = run["best_epoch"]
completed_epoch = run["completed_epoch"]
for epoch in range(completed_epoch + 1, ROBUST_EPOCHS + 1):
epoch_metrics, epoch_sizes = [], []
epoch_indices = training_rng.permutation(len(train_features))
for start in range(0, len(train_features), ROBUST_BATCH_SIZE):
indices = epoch_indices[start : start + ROBUST_BATCH_SIZE]
training_key, batch_key = jax.random.split(training_key)
training_params, training_optimizer_state, batch_loss, batch_metrics = (
robust_update(
training_params,
training_optimizer_state,
batch_key,
jnp.asarray(train_features[indices]),
jnp.asarray(train_labels[indices]),
)
)
epoch_metrics.append(
np.concatenate(([float(batch_loss)], np.asarray(batch_metrics)))
)
epoch_sizes.append(len(indices))
training_metrics = np.average(np.vstack(epoch_metrics), axis=0, weights=epoch_sizes)
validation_metrics = np.full(4, np.nan)
if epoch % ROBUST_VALIDATION_INTERVAL == 0 or epoch == ROBUST_EPOCHS:
_, clean_accuracy, _, _ = evaluate(
training_params, jnp.asarray(val_features), jnp.asarray(val_labels)
)
perturbed_validation = evaluate_perturbation_bank(
training_params,
val_features,
val_labels,
validation_keys,
validation_severities,
)
mean_accuracy = float(np.mean(perturbed_validation["accuracy"]))
fifth_accuracy = float(np.quantile(perturbed_validation["accuracy"], 0.05))
mean_margin = float(np.mean(perturbed_validation["margin"]))
validation_metrics = np.asarray(
(float(clean_accuracy), mean_accuracy, fifth_accuracy, mean_margin)
)
if (
float(clean_accuracy) >= MIN_SELECTION_ACCURACY
and mean_accuracy > best_score
):
best_score = mean_accuracy
best_clean_accuracy = float(clean_accuracy)
best_params = host_tree(training_params)
best_epoch = epoch
print(
f"epoch {epoch:4d}: clean={float(clean_accuracy):.3%}, "
f"8-12% mean={mean_accuracy:.3%}"
)
training_history.append(
tuple(np.concatenate((training_metrics, validation_metrics)))
)
completed_epoch = epoch
run = {
"params": host_tree(training_params),
"optimizer_state": host_tree(training_optimizer_state),
"jax_key": np.asarray(training_key),
"rng_state": training_rng.bit_generator.state,
"history": training_history,
"best_score": best_score,
"best_clean_accuracy": best_clean_accuracy,
"best_params": host_tree(best_params),
"best_epoch": best_epoch,
"completed_epoch": completed_epoch,
"seed": MATRIX_TRAINING_SEED,
"configuration": MATRIX_TRAINING_CONFIGURATION,
}
save_repeated_checkpoint(run)
params = best_params
trained_matrices = np.asarray(layer_matrices(params))
print(
f"selected trained matrices: N={N_CHANNELS}, seed {MATRIX_TRAINING_SEED}, epoch {best_epoch}; "
f"8-12% mean validation accuracy {best_score:.3%}"
)
REALIZATION_TARGETS = (
trained_matrices[0],
trained_matrices[1],
trained_matrices[2, np.asarray(FINAL_OUTPUT_INDICES), :],
)
print("physical target shapes:", [target.shape for target in REALIZATION_TARGETS])
def physical_forward(matrices, features):
phase = jnp.exp(1j * PHASE_SCALE * features)
fields = jnp.ones_like(features, dtype=jnp.complex64)
for layer in range(N_LAYERS):
fields = fields * route_inputs_jax[layer] * phase
fields = fields @ jnp.asarray(matrices[layer]).T
return fields * route_output_jax
square_target_fields = optical_forward(
jnp.asarray(trained_matrices), jnp.asarray(test_features)
)
rectangular_target_fields = physical_forward(
REALIZATION_TARGETS, jnp.asarray(test_features)
)
rectangular_target_error = float(
np.max(
np.abs(np.asarray(square_target_fields) - np.asarray(rectangular_target_fields))
)
)
if rectangular_target_error > 1e-10:
raise ValueError("The 6-input/4-output target changed the trained detector fields.")
print("maximum rectangular-target field difference:", rectangular_target_error)
loaded N=6 trained matrices at epoch 300 selected trained matrices: N=6, seed 17, epoch 290; 8-12% mean validation accuracy 89.636% physical target shapes: [(6, 6), (6, 6), (4, 6)] maximum rectangular-target field difference: 0.0
[21]:
test_loss, test_accuracy, test_transport, trained_gamma = evaluate(
params, jnp.asarray(test_features), jnp.asarray(test_labels)
)
singular_amplitudes = np.linalg.svd(trained_matrices, compute_uv=False)
mean_distance, rms_distance, percentile_90_distance = routing_distance_statistics(
trained_matrices
)
selected_validation = evaluate_perturbation_bank(
params, val_features, val_labels, validation_keys, validation_severities
)
selected_audit = evaluate_perturbation_bank(
params, test_features, test_labels, audit_keys, audit_severities
)
selected_exact = evaluate_perturbation_bank(
params, test_features, test_labels, exact_keys, exact_severities
)
selected_validation_mean = float(np.mean(selected_validation["accuracy"]))
selected_validation_fifth = float(np.quantile(selected_validation["accuracy"], 0.05))
selected_audit_mean = float(np.mean(selected_audit["accuracy"]))
selected_audit_fifth = float(np.quantile(selected_audit["accuracy"], 0.05))
selected_exact_mean = float(np.mean(selected_exact["accuracy"]))
selected_exact_fifth = float(np.quantile(selected_exact["accuracy"], 0.05))
matrix_training_metrics = {
"test accuracy": float(test_accuracy),
"test cross-entropy": float(test_loss),
"routing transport cost": float(test_transport),
"mean coupling distance by layer": mean_distance.tolist(),
"rms coupling distance by layer": rms_distance.tolist(),
"90th-percentile coupling distance by layer": percentile_90_distance.tolist(),
"training-only inverse temperature": float(trained_gamma),
"minimum singular amplitude": float(singular_amplitudes.min()),
"maximum singular amplitude": float(singular_amplitudes.max()),
"validation mean over 8-12%": selected_validation_mean,
"validation fifth percentile over 8-12%": selected_validation_fifth,
"test mean over 8-12%": selected_audit_mean,
"test fifth percentile over 8-12%": selected_audit_fifth,
"test mean at 10%": selected_exact_mean,
"test fifth percentile at 10%": selected_exact_fifth,
"selected seed": MATRIX_TRAINING_SEED,
"selected epoch": best_epoch,
"selection clean-accuracy floor": MIN_SELECTION_ACCURACY,
"92% mean-accuracy target reached": selected_exact_mean
>= ROBUST_TARGET_MEAN_ACCURACY,
}
performance_ready = (
best_epoch > 0
and float(test_accuracy) > 1 / N_CLASSES
and selected_exact_mean > 1 / N_CLASSES
)
matrix_training_ready = (
performance_ready
and np.isfinite(float(test_transport))
and singular_amplitudes.min() >= SIGMA_MIN - 1e-5
and singular_amplitudes.max() <= SIGMA_MAX + 1e-5
)
np.savez_compressed(
CHECKPOINT_DIR / "repeated_encoding__trained_matrices.npz",
matrices=trained_matrices,
history=np.asarray(training_history),
)
print(
f"trained target: accuracy={test_accuracy:.2%}, "
f"cross-entropy={test_loss:.6f}, "
f"10% perturbation mean={selected_exact_mean:.2%}, "
f"fifth percentile={selected_exact_fifth:.2%}"
)
print(
"matrix-training requirement for a cost estimate:",
"satisfied" if matrix_training_ready else "not satisfied",
)
history_array = np.asarray(training_history, dtype=float)
epochs = np.arange(1, len(history_array) + 1)
validation_mask = np.isfinite(history_array[:, 8])
fig, axes = plt.subplots(2, 2, figsize=(9.0, 6.0), constrained_layout=True)
axes[0, 0].plot(epochs, history_array[:, 0])
axes[0, 1].plot(epochs, 100 * history_array[:, 4], label="clean training")
axes[0, 1].plot(epochs, 100 * history_array[:, 5], label="perturbed mean")
axes[1, 0].plot(
epochs[validation_mask], 100 * history_array[validation_mask, 9], label="mean"
)
axes[1, 0].plot(
epochs[validation_mask],
100 * history_array[validation_mask, 10],
label="fifth percentile",
)
axes[1, 1].plot(epochs, 100 * history_array[:, 7])
axes[0, 0].set(xlabel="epoch", ylabel="objective", title="Robust objective")
axes[0, 1].set(xlabel="epoch", ylabel="accuracy (%)", title="Online accuracy")
axes[1, 0].set(
xlabel="epoch", ylabel="accuracy (%)", title="Validation over 8-12% error"
)
axes[1, 1].set(
xlabel="epoch", ylabel="relative error (%)", title="Sampled transmission error"
)
axes[0, 1].legend(frameon=False)
axes[1, 0].legend(frameon=False)
fig.savefig(FIGURE_DIR / "repeated_encoding_training.png", dpi=180, bbox_inches="tight")
plt.show()
trained target: accuracy=90.52%, cross-entropy=0.270203, 10% perturbation mean=89.12%, fifth percentile=87.65% matrix-training requirement for a cost estimate: satisfied
[22]:
robustness_levels = np.asarray((0.00, 0.03, 0.06, 0.09, 0.12, 0.15, 0.20))
robustness_keys = np.asarray(jax.random.split(jax.random.PRNGKey(SEED + 2_000), 100))
robustness_curve = []
for error_level in robustness_levels:
severities = np.full(len(robustness_keys), error_level)
robustness_curve.append(
evaluate_perturbation_bank(
params,
test_features,
test_labels,
robustness_keys,
severities,
chunk_size=10,
)
)
actual_error = np.asarray(
[np.median(values["actual_error"]) for values in robustness_curve]
)
median_accuracy = np.asarray(
[np.median(values["accuracy"]) for values in robustness_curve]
)
lower_accuracy = np.asarray(
[np.quantile(values["accuracy"], 0.05) for values in robustness_curve]
)
upper_accuracy = np.asarray(
[np.quantile(values["accuracy"], 0.95) for values in robustness_curve]
)
fig, axes = plt.subplots(1, 2, figsize=(9.4, 3.8), constrained_layout=True)
axes[0].plot(100 * actual_error, 100 * median_accuracy, marker="o")
axes[0].fill_between(
100 * actual_error, 100 * lower_accuracy, 100 * upper_accuracy, alpha=0.18
)
axes[0].set(
xlabel="actual relative transmission error (%)",
ylabel="test accuracy (%)",
title="Accuracy under matrix error",
)
axes[1].boxplot(
(robustness_curve[0]["margin"], robustness_curve[5]["margin"]),
tick_labels=("nominal", "15% error"),
showfliers=False,
)
axes[1].axhline(
ROBUST_MARGIN_TARGET, color="black", linestyle=":", label="training target"
)
axes[1].set(ylabel="mean normalized optical margin", title="Optical decision margin")
axes[1].legend(frameon=False)
fig.savefig(
FIGURE_DIR / "repeated_encoding_matrix_error.png", dpi=180, bbox_inches="tight"
)
plt.show()
fig, ax = plt.subplots(figsize=(6.2, 3.0), constrained_layout=True)
for layer in range(N_LAYERS):
ax.plot(
np.arange(1, N_CHANNELS + 1),
singular_amplitudes[layer],
marker=".",
label=f"layer {layer + 1}",
)
ax.axhline(SIGMA_MIN, color="black", linestyle=":", label="6 dB bound")
ax.axhline(SIGMA_MAX, color="black", linestyle="--", label="3 dB bound")
ax.set(
xlabel="singular channel",
ylabel="singular amplitude",
title="Trained target matrices",
)
ax.legend(frameon=False, ncol=2)
fig.savefig(
FIGURE_DIR / "repeated_encoding_singular_amplitudes.png",
dpi=180,
bbox_inches="tight",
)
plt.show()
Compare Single and Repeated Encoding¶
To isolate the effect of re-encoding, we train a second classifier that applies the image phase only at the input. Its optical mapping is
$$ a_{\mathrm{single}} = M_3M_2M_1D(z)a_{\mathrm{in}}, $$
whereas repeated encoding gives
$$ a_{\mathrm{repeated}} = M_3D(z)M_2D(z)M_1D(z)a_{\mathrm{in}}. $$
Both classifiers use three passive matrices, the same data split, initialization seeds, readout, perturbation levels, and training procedure. The comparison therefore changes only the position and number of phase encoders.
We compare their clean and perturbed classification results. The matrices from the repeated-encoding model are then used as fixed transmission targets for three silicon photonic devices optimized using Tidy3D.
[23]:
def single_encoding_forward(matrices, features):
fields = jnp.ones_like(features, dtype=jnp.complex64)
image_phase = jnp.exp(1j * PHASE_SCALE * features)
for layer in range(N_LAYERS):
fields = fields * route_inputs_jax[layer]
if layer == 0:
fields = fields * image_phase
fields = fields @ matrices[layer].T
return fields[:, readout_slice] * route_output_jax
def single_encoding_scores(matrices, features):
return jnp.abs(single_encoding_forward(matrices, features)) ** 2
def single_encoding_objective(parameters, key, features, labels):
matrices = layer_matrices(parameters)
clean_scores = single_encoding_scores(matrices, features)
gamma = jnp.exp(parameters["log_gamma"])
clean_logits = gamma * clean_scores / ROUTE_POWER_REFERENCE
clean_loss = optax.softmax_cross_entropy_with_integer_labels(
clean_logits, labels
).mean()
structure_key, error_key = jax.random.split(key)
perturbations = jax.vmap(sample_chip_perturbation)(
jax.random.split(structure_key, ROBUST_REPLICAS)
)
target_errors = jax.random.uniform(
error_key,
(ROBUST_REPLICAS, N_LAYERS),
minval=ROBUST_WINDOW_MIN,
maxval=ROBUST_WINDOW_MAX,
)
perturbed_matrices, actual_errors = perturb_replicas(
matrices, perturbations, target_errors, jnp.ones((ROBUST_REPLICAS,))
)
perturbed_scores = jax.vmap(
lambda values: single_encoding_scores(values, features)
)(perturbed_matrices)
perturbed_logits = gamma * perturbed_scores / ROUTE_POWER_REFERENCE
replica_labels = jnp.broadcast_to(labels, (ROBUST_REPLICAS, labels.shape[0]))
perturbed_losses = optax.softmax_cross_entropy_with_integer_labels(
perturbed_logits, replica_labels
).mean(axis=1)
perturbed_accuracies = jnp.mean(
jnp.argmax(perturbed_scores, axis=-1) == replica_labels, axis=1
)
margin_penalty = jnp.maximum(
ROBUST_MARGIN_TARGET - normalized_optical_margin(perturbed_scores, labels), 0
).mean()
gamma_term = GAMMA_WEIGHT * jnp.square(
jnp.maximum(gamma - GAMMA_TARGET, 0.0) / GAMMA_TARGET
)
loss = (
0.35 * clean_loss
+ 0.65 * jnp.mean(perturbed_losses)
+ ROBUST_MARGIN_WEIGHT * margin_penalty
+ routing_transport(matrices)
+ gamma_term
)
metrics = jnp.asarray(
(
clean_loss,
jnp.mean(perturbed_losses),
margin_penalty,
jnp.mean(jnp.argmax(clean_scores, axis=-1) == labels),
jnp.mean(perturbed_accuracies),
jnp.min(perturbed_accuracies),
jnp.mean(actual_errors),
)
)
return loss, metrics
@jax.jit
def single_encoding_update(parameters, optimizer_state, key, features, labels):
(loss, metrics), gradients = jax.value_and_grad(
single_encoding_objective, has_aux=True
)(parameters, key, features, labels)
updates, optimizer_state = robust_optimizer.update(
gradients, optimizer_state, parameters
)
return optax.apply_updates(parameters, updates), optimizer_state, loss, metrics
@jax.jit
def single_encoding_clean_metrics(parameters, features, labels):
matrices = layer_matrices(parameters)
scores = single_encoding_scores(matrices, features)
logits = jnp.exp(parameters["log_gamma"]) * scores / ROUTE_POWER_REFERENCE
loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
accuracy = jnp.mean(jnp.argmax(scores, axis=-1) == labels)
return loss, accuracy
@jax.jit
def single_encoding_perturbation_metrics(
parameters, features, labels, keys, target_errors
):
matrices = layer_matrices(parameters)
perturbations = jax.vmap(sample_chip_perturbation)(keys)
perturbed, actual_errors = perturb_replicas(
matrices, perturbations, target_errors, jnp.ones((len(keys),))
)
scores = jax.vmap(lambda values: single_encoding_scores(values, features))(
perturbed
)
replica_labels = jnp.broadcast_to(labels, (len(keys), labels.shape[0]))
accuracies = jnp.mean(jnp.argmax(scores, axis=-1) == replica_labels, axis=1)
margins = jnp.mean(normalized_optical_margin(scores, labels), axis=1)
return accuracies, margins, jnp.mean(actual_errors, axis=1)
def evaluate_single_encoding_bank(
parameters, features, labels, keys, severities, chunk_size=8
):
feature_array, label_array = jnp.asarray(features), jnp.asarray(labels)
target_errors = np.asarray(severities, dtype=float)
if target_errors.ndim == 1:
target_errors = np.repeat(target_errors[:, None], N_LAYERS, axis=1)
collected = [[], [], []]
for start in range(0, len(keys), chunk_size):
stop = min(start + chunk_size, len(keys))
values = single_encoding_perturbation_metrics(
parameters,
feature_array,
label_array,
jnp.asarray(keys[start:stop]),
jnp.asarray(target_errors[start:stop]),
)
for destination, value in zip(collected, values):
destination.append(np.asarray(value))
accuracy, margin, actual_error = (np.concatenate(values) for values in collected)
return {"accuracy": accuracy, "margin": margin, "actual_error": actual_error}
[24]:
single_encoding_configuration = {
**MATRIX_TRAINING_CONFIGURATION,
"encoding": "input phase before first stage only",
"training seed": MATRIX_TRAINING_SEED,
"epochs": ROBUST_EPOCHS,
"batch size": ROBUST_BATCH_SIZE,
"perturbation replicas": ROBUST_REPLICAS,
"learning rate": ROBUST_LR,
"robust error window": (ROBUST_WINDOW_MIN, ROBUST_WINDOW_MAX),
"margin target": ROBUST_MARGIN_TARGET,
"margin weight": ROBUST_MARGIN_WEIGHT,
"selection interval": ROBUST_VALIDATION_INTERVAL,
}
single_encoding_checkpoint = CHECKPOINT_DIR / ("single_encoding__matrix_training.pkl")
if single_encoding_checkpoint.exists():
with single_encoding_checkpoint.open("rb") as handle:
single_run = pickle.load(handle)
if single_run.get("configuration") != single_encoding_configuration:
raise ValueError(
"The single-encoding checkpoint uses a different configuration."
)
else:
single_initial = host_tree(
initialize_parameters(jax.random.PRNGKey(MATRIX_TRAINING_SEED))
)
single_rng = np.random.default_rng(MATRIX_TRAINING_SEED)
single_run = {
"params": single_initial,
"optimizer_state": host_tree(robust_optimizer.init(single_initial)),
"jax_key": np.asarray(jax.random.PRNGKey(MATRIX_TRAINING_SEED + 10_000)),
"rng_state": single_rng.bit_generator.state,
"best_score": -np.inf,
"best_clean_accuracy": -np.inf,
"best_params": single_initial,
"best_epoch": 0,
"completed_epoch": 0,
"configuration": single_encoding_configuration,
}
single_rng = np.random.default_rng(MATRIX_TRAINING_SEED)
single_rng.bit_generator.state = single_run["rng_state"]
single_params = single_run["params"]
single_optimizer_state = single_run["optimizer_state"]
single_key = jnp.asarray(single_run["jax_key"])
single_best_score = single_run["best_score"]
single_best_clean = single_run["best_clean_accuracy"]
single_best_params = single_run["best_params"]
single_best_epoch = single_run["best_epoch"]
[25]:
for epoch in range(single_run["completed_epoch"] + 1, ROBUST_EPOCHS + 1):
epoch_indices = single_rng.permutation(len(train_features))
for start in range(0, len(train_features), ROBUST_BATCH_SIZE):
indices = epoch_indices[start : start + ROBUST_BATCH_SIZE]
single_key, batch_key = jax.random.split(single_key)
single_params, single_optimizer_state, _, _ = single_encoding_update(
single_params,
single_optimizer_state,
batch_key,
jnp.asarray(train_features[indices]),
jnp.asarray(train_labels[indices]),
)
if epoch % ROBUST_VALIDATION_INTERVAL == 0 or epoch == ROBUST_EPOCHS:
_, clean_accuracy = single_encoding_clean_metrics(
single_params, jnp.asarray(val_features), jnp.asarray(val_labels)
)
validation = evaluate_single_encoding_bank(
single_params,
val_features,
val_labels,
validation_keys,
validation_severities,
)
mean_accuracy = float(np.mean(validation["accuracy"]))
if (
float(clean_accuracy) >= MIN_SELECTION_ACCURACY
and mean_accuracy > single_best_score
):
single_best_score = mean_accuracy
single_best_clean = float(clean_accuracy)
single_best_params = host_tree(single_params)
single_best_epoch = epoch
single_run = {
"params": host_tree(single_params),
"optimizer_state": host_tree(single_optimizer_state),
"jax_key": np.asarray(single_key),
"rng_state": single_rng.bit_generator.state,
"best_score": single_best_score,
"best_clean_accuracy": single_best_clean,
"best_params": host_tree(single_best_params),
"best_epoch": single_best_epoch,
"completed_epoch": epoch,
"configuration": single_encoding_configuration,
}
temporary_path = single_encoding_checkpoint.with_suffix(".tmp")
with temporary_path.open("wb") as handle:
pickle.dump(single_run, handle, protocol=pickle.HIGHEST_PROTOCOL)
temporary_path.replace(single_encoding_checkpoint)
[26]:
single_params = single_best_params
single_test_loss, single_test_accuracy = single_encoding_clean_metrics(
single_params, jnp.asarray(test_features), jnp.asarray(test_labels)
)
single_validation = evaluate_single_encoding_bank(
single_params, val_features, val_labels, validation_keys, validation_severities
)
single_validation_mean = float(np.mean(single_validation["accuracy"]))
single_validation_fifth = float(np.quantile(single_validation["accuracy"], 0.05))
encoding_comparison = {
"controlled difference": "image phase once or before every stage",
"single encoding": {
"selected epoch": int(single_best_epoch),
"clean test accuracy": float(single_test_accuracy),
"validation mean over 8-12%": single_validation_mean,
"validation fifth percentile over 8-12%": single_validation_fifth,
},
"repeated encoding": {
"selected epoch": int(best_epoch),
"clean test accuracy": float(test_accuracy),
"validation mean over 8-12%": selected_validation_mean,
"validation fifth percentile over 8-12%": selected_validation_fifth,
},
}
matrix_training_metrics["encoding comparison"] = encoding_comparison
print(
f"single encoding: clean={single_test_accuracy:.2%}, "
f"8-12% mean={single_validation_mean:.2%}; "
f"repeated encoding: clean={test_accuracy:.2%}, "
f"8-12% mean={selected_validation_mean:.2%}"
)
comparison_labels = ("clean test", "8-12% validation mean", "8-12% validation fifth")
single_values = 100 * np.asarray(
(float(single_test_accuracy), single_validation_mean, single_validation_fifth)
)
repeated_values = 100 * np.asarray(
(float(test_accuracy), selected_validation_mean, selected_validation_fifth)
)
positions = np.arange(len(comparison_labels))
fig, ax = plt.subplots(figsize=(7.4, 3.8), constrained_layout=True)
width = 0.36
ax.bar(positions - width / 2, single_values, width, label="encode once")
ax.bar(positions + width / 2, repeated_values, width, label="encode every stage")
for x, value in zip(positions - width / 2, single_values):
ax.text(x, value + 0.35, f"{value:.1f}", ha="center", va="bottom", fontsize=8)
for x, value in zip(positions + width / 2, repeated_values):
ax.text(x, value + 0.35, f"{value:.1f}", ha="center", va="bottom", fontsize=8)
ax.set_xticks(positions, comparison_labels)
ax.set(ylabel="accuracy (%)", title="Effect of repeated phase encoding")
ax.set_ylim(max(0, min(single_values.min(), repeated_values.min()) - 5), 100)
ax.legend(frameon=False)
fig.savefig(FIGURE_DIR / "encoding_comparison.png", dpi=180, bbox_inches="tight")
plt.show()
single encoding: clean=86.89%, 8-12% mean=85.79%; repeated encoding: clean=90.52%, 8-12% mean=89.64%
[27]:
target_amplitude = np.abs(trained_matrices)
target_phase = np.angle(trained_matrices)
channel_ticks = np.unique(np.linspace(0, N_CHANNELS - 1, min(5, N_CHANNELS), dtype=int))
fig, axes = plt.subplots(
2, N_LAYERS, figsize=(3.2 * N_LAYERS, 6.0), constrained_layout=True
)
for layer in range(N_LAYERS):
amplitude_image = axes[0, layer].imshow(
target_amplitude[layer], cmap="magma", vmin=0, vmax=target_amplitude.max()
)
phase_image = axes[1, layer].imshow(
target_phase[layer], cmap="twilight_shifted", vmin=-np.pi, vmax=np.pi
)
axes[0, layer].set(
title=f"Layer {layer + 1}", xticks=channel_ticks, yticks=channel_ticks
)
axes[1, layer].set(
xticks=channel_ticks, yticks=channel_ticks, xlabel="input channel"
)
axes[0, 0].set_ylabel("output channel")
axes[1, 0].set_ylabel("output channel")
fig.colorbar(amplitude_image, ax=axes[0, :], label="amplitude", shrink=0.85)
fig.colorbar(
phase_image,
ax=axes[1, :],
label="phase (rad)",
shrink=0.85,
ticks=(-np.pi, 0, np.pi),
format=plt.FuncFormatter(
lambda value, _: {-np.pi: r"$-\pi$", 0: "0", np.pi: r"$\pi$"}.get(value, "")
),
)
fig.savefig(
FIGURE_DIR / "repeated_encoding_target_matrices.png", dpi=180, bbox_inches="tight"
)
plt.show()
Set Up the Silicon Photonic Devices¶
We now convert the three target matrices into three inverse-design models. To reduce the computational cost, the silicon slab is represented by a two-dimensional effective-index model at 1.55 µm. Separate effective indices describe the 150 nm and 220 nm silicon thicknesses, while the surrounding region retains the SiEPIC silica index.
All three devices use a 12.6 µm square design region with 50 nm pixels and a 200 nm spatial filter. Six waveguides enter each device on a 1.75 µm pitch. The first two devices also have six outputs, while the final device has four outputs aligned with the class detectors.
We excite each input separately to recover the complete transmission and reflection matrices. Before running the inverse design, the following cells inspect the geometry, solve the quasi-TE port mode, and validate the 18 source simulations.
[28]:
N_CLAD = 1.4447002763
N_EFF_TE_150_SLAB, N_EFF_TE_220_SLAB = 2.5443521652, 2.8493681982
PDK_MODE_NEFF = float(np.real(PDK_NEFF))
EIM_GRID_STEPS_UM = (0.02, 0.01, 0.005)
SI_3D = td.material_library["cSi"]["Palik_NoLoss"]
SIO2_3D = td.material_library["SiO2"]["Palik_NoLoss"]
td.config.simulation.use_local_subpixel = True
def vertical_profile_simulation(grid_step):
profiles = (
td.Structure(
geometry=td.Box(center=(-1, 0, 0), size=(0.8, td.inf, 0.22)), medium=SI_3D
),
td.Structure(
geometry=td.Box(center=(0, 0, -0.035), size=(0.8, td.inf, 0.15)),
medium=SI_3D,
),
)
return td.Simulation(
center=(0, 0, 0),
size=(3, 3, 3),
medium=SIO2_3D,
structures=profiles,
sources=[],
monitors=[],
run_time=1e-12,
boundary_spec=td.BoundarySpec.all_sides(td.Periodic()),
grid_spec=td.GridSpec.uniform(dl=grid_step),
)
def variational_permittivity(simulation, point, reference=(-1, 0)):
vertical_simulation = simulation.updated_copy(
center=(reference[0], reference[1], 0),
size=(0, 3, 3),
sources=[],
monitors=[],
symmetry=(0, 0, 0),
boundary_spec=simulation.boundary_spec.updated_copy(x=td.Boundary.periodic()),
)
solver = ModeSolver(
simulation=vertical_simulation,
plane=td.Box(center=vertical_simulation.center, size=(3, 0, 3)),
mode_spec=td.ModeSpec(num_modes=1, target_neff=N_EFF_TE_220_SLAB),
freqs=[FREQ0_HZ],
direction="+",
)
mode_data = solver.solve()
reference_index = float(np.real(np.asarray(mode_data.n_eff).reshape(-1)[0]))
if point == reference:
return reference_index**2, reference_index
reference_eps = simulation.epsilon(
box=td.Box(center=(*reference, 0), size=(0, 0, td.inf)), freq=FREQ0_HZ
)
point_eps = simulation.epsilon(
box=td.Box(center=(*point, 0), size=(0, 0, td.inf)), freq=FREQ0_HZ
)
z = np.asarray(reference_eps.z)
delta_eps = np.squeeze(point_eps.values) - np.squeeze(reference_eps.values)
field_interpolator = scipy.interpolate.interp1d(
np.asarray(mode_data.Ex.z),
np.abs(np.squeeze(mode_data.Ex.values)) ** 2,
bounds_error=False,
fill_value=0.0,
)
field_intensity = field_interpolator(z)
correction = np.trapezoid(delta_eps * field_intensity, x=z) / np.trapezoid(
field_intensity, x=z
)
return max(1.0, float(np.real(reference_index**2 + correction))), reference_index
previous_log_level = td.config.logging.level
td.config.logging.level = "ERROR"
eim_convergence = []
for grid_step in EIM_GRID_STEPS_UM:
profile_simulation = vertical_profile_simulation(grid_step)
eps_220, reference_index = variational_permittivity(profile_simulation, (-1, 0))
eps_150, _ = variational_permittivity(profile_simulation, (0, 0))
eps_background = N_CLAD**2
eim_convergence.append(
{
"grid nm": 1e3 * grid_step,
"220 nm variational index": float(np.sqrt(eps_220)),
"150 nm variational index": float(np.sqrt(eps_150)),
"silica background index": float(np.sqrt(eps_background)),
}
)
td.config.logging.level = previous_log_level
print(
"effective-index calibration:",
f"grid={eim_convergence[-1]['grid nm']} nm, "
f"n150={eim_convergence[-1]['150 nm variational index']:.6f}, "
f"n220={eim_convergence[-1]['220 nm variational index']:.6f}",
)
eim_calibration = eim_convergence[-1]
N_EFF_TE_150 = eim_calibration["150 nm variational index"]
N_EFF_TE_220 = eim_calibration["220 nm variational index"]
N_EFF_TE_BACKGROUND = eim_calibration["silica background index"]
EPS_CLAD = N_EFF_TE_BACKGROUND**2
EPS_CORE = N_EFF_TE_220**2
EPS_DESIGN = (N_EFF_TE_150**2, N_EFF_TE_220**2)
grid_nm = np.asarray([record["grid nm"] for record in eim_convergence])
fig, axes = plt.subplots(1, 2, figsize=(7.4, 2.8), constrained_layout=True)
axes[0].plot(
grid_nm,
[record["220 nm variational index"] for record in eim_convergence],
"o-",
label="220 nm",
)
axes[0].plot(
grid_nm,
[record["150 nm variational index"] for record in eim_convergence],
"o-",
label="150 nm",
)
axes[0].invert_xaxis()
axes[0].set(xlabel="calibration grid (nm)", ylabel="variational index")
axes[0].legend()
axes[1].plot(
grid_nm,
[record["220 nm variational index"] - N_EFF_TE_220 for record in eim_convergence],
"o-",
label="220 nm",
)
axes[1].plot(
grid_nm,
[record["150 nm variational index"] - N_EFF_TE_150 for record in eim_convergence],
"o-",
label="150 nm",
)
axes[1].axhline(0, color="black", lw=0.8, ls="--")
axes[1].invert_xaxis()
axes[1].set(
xlabel="calibration grid (nm)", ylabel=r"$n_{\mathrm{2D}}-n_{\mathrm{finest}}$"
)
axes[1].legend()
plt.show()
CORE_THICKNESS, PORT_PITCH, WG_WIDTH, MODE_WINDOW = 0.22, ID_PORT_PITCH, 0.5, 1.75
DESIGN_SIZE, PIXEL, FILTER_RADIUS, FILTER_BETA = (
(DESIGN_EDGE_UM, DESIGN_EDGE_UM, CORE_THICKNESS),
0.05,
0.2,
50.0,
)
TAPER_LENGTH, TAPER_MOUTH, TAPER_OVERLAP = (
ID_TAPER_LENGTH,
ID_TAPER_MOUTH,
ID_TAPER_OVERLAP,
)
CORNER_RADIUS, RAIL_THICKNESS, RAIL_INNER_OVERLAP = 0.8, 0.30, 0.0
PML_EXTENSION = 2.0
DESIGN_HALF = DESIGN_EDGE_UM / 2
DOMAIN_X = (-DESIGN_HALF - TAPER_LENGTH - 4.65, DESIGN_HALF + TAPER_LENGTH + 4.65)
DOMAIN_Y = (-DESIGN_HALF - 1.55, DESIGN_HALF + 1.55)
INPUT_MONITOR_X = -DESIGN_HALF - TAPER_LENGTH - 1.55
OUTPUT_MONITOR_X = DOMAIN_X[1] - 1.55
INPUT_PORT_Y = (np.arange(N_CHANNELS) - (N_CHANNELS - 1) / 2) * PORT_PITCH
OUTPUT_PORT_Y = {
n_outputs: (np.arange(n_outputs) - (n_outputs - 1) / 2) * PORT_PITCH
for n_outputs in set(LAYER_OUTPUT_COUNTS)
}
CLAD = td.Medium(permittivity=EPS_CLAD)
CORE = td.Medium(permittivity=EPS_CORE)
MODE_SPEC = td.ModeSpec(num_modes=1, target_neff=float(np.real(PDK_MODE_NEFF)))
def vertices(xs, widths, y):
xs, widths = np.asarray(xs), np.asarray(widths)
lower = np.c_[xs, y - widths / 2]
upper = np.c_[xs[::-1], (y + widths / 2)[::-1]]
return np.r_[lower, upper]
def rounded_box(length, width, radius, layer):
polygon = gdstk.rectangle(
(-length / 2, -width / 2), (length / 2, width / 2), layer=layer
)
polygon.fillet(radius)
return polygon
def geometry_from_gds(polygons, layer):
library = gdstk.Library()
cell = library.new_cell(f"GEOMETRY_{layer}")
cell.add(*polygons)
return td.Geometry.from_gds(
cell,
axis=2,
slab_bounds=(-CORE_THICKNESS / 2, CORE_THICKNESS / 2),
gds_layer=layer,
gds_dtype=0,
reference_plane="middle",
)
effective-index calibration: grid=5.0 nm, n150=2.436998, n220=2.848820
[29]:
def fixed_polygons(n_outputs):
layer, polygons = 3, []
left, right = -DESIGN_HALF, DESIGN_HALF
taper_left, taper_right = left - TAPER_LENGTH, right + TAPER_LENGTH
output_y = OUTPUT_PORT_Y[n_outputs]
for y in INPUT_PORT_Y:
polygons.extend(
[
gdstk.rectangle(
(DOMAIN_X[0] - PML_EXTENSION, y - WG_WIDTH / 2),
(taper_left, y + WG_WIDTH / 2),
layer=layer,
),
gdstk.Polygon(
vertices(
np.linspace(taper_left, left, 101),
np.linspace(WG_WIDTH, TAPER_MOUTH, 101),
y,
),
layer=layer,
),
]
)
for y in output_y:
polygons.extend(
[
gdstk.Polygon(
vertices(
np.linspace(right, taper_right, 101),
np.linspace(TAPER_MOUTH, WG_WIDTH, 101),
y,
),
layer=layer,
),
gdstk.rectangle(
(taper_right, y - WG_WIDTH / 2),
(DOMAIN_X[1] + PML_EXTENSION, y + WG_WIDTH / 2),
layer=layer,
),
]
)
rail = RAIL_THICKNESS
outer = rounded_box(
DESIGN_EDGE_UM + 2 * rail,
DESIGN_EDGE_UM + 2 * rail,
CORNER_RADIUS + rail,
layer,
)
inner = rounded_box(DESIGN_EDGE_UM, DESIGN_EDGE_UM, CORNER_RADIUS, layer)
ring = gdstk.boolean([outer], [inner], "not", layer=layer, precision=1e-9) or []
mask_bound = DESIGN_HALF + 0.5
side_masks, bridges = [], []
for x0, x1, positions, bridge_x0, bridge_x1 in (
(-mask_bound, 0, INPUT_PORT_Y, -mask_bound, left),
(0, mask_bound, output_y, right, mask_bound),
):
lower = positions[0] - TAPER_MOUTH / 2
upper = positions[-1] + TAPER_MOUTH / 2
side_masks.extend(
[
gdstk.rectangle((x0, -mask_bound), (x1, lower), layer=layer),
gdstk.rectangle((x0, upper), (x1, mask_bound), layer=layer),
]
)
bridges.extend(
[
gdstk.rectangle(
(bridge_x0, lower - rail), (bridge_x1, positions[0]), layer=layer
),
gdstk.rectangle(
(bridge_x0, positions[-1]), (bridge_x1, upper + rail), layer=layer
),
]
)
polygons.extend(
gdstk.boolean(ring, side_masks, "and", layer=layer, precision=1e-9) or []
)
polygons.extend(
gdstk.boolean(ring, bridges, "and", layer=layer, precision=1e-9) or []
)
return gdstk.boolean(polygons, [], "or", layer=layer, precision=1e-9) or []
def corner_mask_polygons():
rectangle = gdstk.rectangle(
(-DESIGN_HALF, -DESIGN_HALF), (DESIGN_HALF, DESIGN_HALF), layer=4
)
rounded = rounded_box(DESIGN_EDGE_UM, DESIGN_EDGE_UM, CORNER_RADIUS, 4)
return gdstk.boolean([rectangle], [rounded], "not", layer=4, precision=1e-9) or []
FIXED_GEOMETRY = {
n_outputs: geometry_from_gds(fixed_polygons(n_outputs), 3)
for n_outputs in set(LAYER_OUTPUT_COUNTS)
}
CORNER_MASK_GEOMETRY = geometry_from_gds(corner_mask_polygons(), 4)
def fabrication_penalty(weight=1.0):
return tdi.ErosionDilationPenalty(
weight=weight,
length_scale=FABRICATION_MIN_FEATURE_UM,
beta=FABRICATION_BETA,
eta0=FABRICATION_ETA0,
delta_eta=FABRICATION_DELTA_ETA,
)
def design_region(fabrication_weight=0.0):
penalties = (
(fabrication_penalty(fabrication_weight),) if fabrication_weight > 0 else ()
)
return tdi.TopologyDesignRegion(
size=DESIGN_SIZE,
center=(0, 0, 0),
eps_bounds=EPS_DESIGN,
transformations=(
tdi.FilterProject(radius=FILTER_RADIUS, beta=FILTER_BETA, eta=0.5),
),
penalties=penalties,
pixel_size=PIXEL,
uniform=(False, False, True),
priority=0,
)
print("square filter grid:", design_region().params_shape)
INPUT_NAMES = tuple(f"input_{index:02d}" for index in range(N_CHANNELS))
def output_names(n_outputs):
return tuple(f"output_{index:02d}" for index in range(n_outputs))
def modal_ports(n_outputs):
return tuple(
Port(
center=(INPUT_MONITOR_X, float(y), 0),
size=(0, MODE_WINDOW, td.inf),
name=name,
direction="+",
mode_spec=MODE_SPEC,
)
for name, y in zip(INPUT_NAMES, INPUT_PORT_Y)
) + tuple(
Port(
center=(OUTPUT_MONITOR_X, float(y), 0),
size=(0, MODE_WINDOW, td.inf),
name=name,
direction="-",
mode_spec=MODE_SPEC,
)
for name, y in zip(output_names(n_outputs), OUTPUT_PORT_Y[n_outputs])
)
def device_base_simulation(n_outputs):
return td.Simulation(
center=(0, 0, 0),
size=(np.ptp(DOMAIN_X), np.ptp(DOMAIN_Y), 0),
medium=CLAD,
structures=[
td.Structure(
geometry=FIXED_GEOMETRY[n_outputs],
medium=CORE,
name="fixed_rails_tapers",
priority=2,
),
td.Structure(
geometry=CORNER_MASK_GEOMETRY,
medium=CLAD,
name="rounded_corner_mask",
priority=1,
),
],
sources=[],
monitors=[],
run_time=15e-12,
shutoff=1e-4,
grid_spec=td.GridSpec.auto(wavelength=WAVELENGTH_UM, min_steps_per_wvl=13),
boundary_spec=td.BoundarySpec(
x=td.Boundary.pml(), y=td.Boundary.pml(), z=td.Boundary.periodic()
),
subpixel=True,
)
def native_modal_modeler(layer):
n_outputs = LAYER_OUTPUT_COUNTS[layer]
return ModalComponentModeler(
name=f"deep_pnn_modal_layer_{layer + 1}",
simulation=device_base_simulation(n_outputs),
ports=modal_ports(n_outputs),
freqs=(FREQ0_HZ,),
run_only=tuple((name, 0) for name in INPUT_NAMES),
)
class InverseDesignMulti(tdi.InverseDesignMulti):
def run_async(self, simulations, **kwargs):
from tidy3d.web import run_async
simulation_data = {}
kwargs.setdefault("path_dir", AUTOGRAD_CACHE_DIR)
simulation_items = list(simulations.items())
for start in range(0, len(simulation_items), AUTOGRAD_SOURCE_CHUNK_SIZE):
chunk = dict(simulation_items[start : start + AUTOGRAD_SOURCE_CHUNK_SIZE])
batch_data = run_async(chunk, verbose=self.verbose, **kwargs)
for task_name in chunk:
simulation_data[task_name] = batch_data[task_name]
return simulation_data
NATIVE_MODELERS = {layer: native_modal_modeler(layer) for layer in range(N_LAYERS)}
NATIVE_TASK_NAMES = {
layer: tuple(modeler.sim_dict) for layer, modeler in NATIVE_MODELERS.items()
}
NATIVE_SIMULATIONS = {
layer: tuple(modeler.sim_dict.values())
for layer, modeler in NATIVE_MODELERS.items()
}
NATIVE_MONITOR_NAMES = {
layer: tuple(
tuple(monitor.name for monitor in simulation.monitors)
for simulation in NATIVE_SIMULATIONS[layer]
)
for layer in range(N_LAYERS)
}
def make_design(layer, target):
return InverseDesignMulti(
design_region=design_region(),
task_name=f"deep_id_pnn_layer_{layer + 1}",
simulations=NATIVE_SIMULATIONS[layer],
output_monitor_names=NATIVE_MONITOR_NAMES[layer],
verbose=False,
), np.asarray(target)
designs, targets = {}, {}
for layer in range(N_LAYERS):
designs[layer], targets[layer] = make_design(layer, REALIZATION_TARGETS[layer])
print(
f"layer {layer + 1}: {LAYER_INPUT_COUNTS[layer]} inputs, "
f"{LAYER_OUTPUT_COUNTS[layer]} outputs, "
f"{len(designs[layer].simulations)} source simulations"
)
square filter grid: (253, 253, 1)
04:23:41 KST WARNING: 'ModalComponentModeler' was refactored (tidy3d 'v2.10.0'). Existing functionality is available differently. Please consult the migration documentation: https://docs.flexcompute.com/projects/tidy3d/en/latest/api/microwav e/microwave_migration.html
layer 1: 6 inputs, 6 outputs, 6 source simulations layer 2: 6 inputs, 6 outputs, 6 source simulations layer 3: 6 inputs, 4 outputs, 6 source simulations
[30]:
initial_parameters = designs[N_LAYERS - 1].design_region.initial_parameters
preview_batch = designs[N_LAYERS - 1].to_simulation(initial_parameters)
preview_simulation = next(iter(preview_batch.values()))
fig, axes = plt.subplots(1, 2, figsize=(12.0, 4.8), constrained_layout=True)
preview_simulation.plot(z=0, source_alpha=0.9, monitor_alpha=0.65, ax=axes[0])
preview_simulation.plot_pml(z=0, ax=axes[0])
preview_simulation.plot_eps(
z=0,
freq=FREQ0_HZ,
source_alpha=0,
monitor_alpha=0,
eps_lim=(EPS_CLAD, EPS_CORE),
ax=axes[1],
)
axes[0].set(
xlim=DOMAIN_X,
ylim=DOMAIN_Y,
title=f"Device 3: {N_CHANNELS} inputs, {N_CLASSES} outputs",
)
axes[1].set(
xlim=DOMAIN_X, ylim=DOMAIN_Y, title=r"Relative permittivity $\varepsilon_r(x,y)$"
)
fig.savefig(
FIGURE_DIR / "tidy3d_preflight_domain_permittivity.png",
dpi=180,
bbox_inches="tight",
)
plt.show()
port_mode_indices, port_te_fractions = [], []
for port_y in INPUT_PORT_Y:
port_plane = td.Box(
center=(INPUT_MONITOR_X, float(port_y), 0), size=(0, MODE_WINDOW, td.inf)
)
port_mode = ModeSolver(
simulation=preview_simulation,
plane=port_plane,
mode_spec=MODE_SPEC,
freqs=[FREQ0_HZ],
direction="+",
).solve()
port_mode_indices.append(float(np.real(np.asarray(port_mode.n_eff).reshape(-1)[0])))
port_te_fractions.append(float(np.asarray(port_mode.TE_fraction).reshape(-1)[0]))
port_mode_indices = np.asarray(port_mode_indices)
port_te_fractions = np.asarray(port_te_fractions)
port_mode_report = {
"polarization": "fundamental quasi-TE",
"ports checked": len(INPUT_PORT_Y),
"minimum TE fraction": float(port_te_fractions.min()),
"maximum TE fraction": float(port_te_fractions.max()),
"minimum 2D rail effective index": float(port_mode_indices.min()),
"maximum 2D rail effective index": float(port_mode_indices.max()),
"PDK strip effective index": PDK_MODE_NEFF,
"maximum absolute effective-index difference": float(
np.max(np.abs(port_mode_indices - PDK_MODE_NEFF))
),
}
print(
f"port modes: {port_mode_report['ports checked']} checked, "
f"minimum TE fraction={port_mode_report['minimum TE fraction']:.4f}, "
f"n_eff range={port_mode_report['minimum 2D rail effective index']:.6f}-"
f"{port_mode_report['maximum 2D rail effective index']:.6f}"
)
fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.0), constrained_layout=True)
channels = np.arange(1, N_CHANNELS + 1)
axes[0].plot(channels, port_mode_indices, marker="o", label="2D port mode")
axes[0].axhline(PDK_MODE_NEFF, color="0.25", linestyle="--", label="SiEPIC PDK value")
axes[1].plot(channels, port_te_fractions, marker="o")
axes[0].set(xlabel="port", ylabel=r"$n_{\mathrm{eff}}$", title="Port effective index")
axes[1].set(
xlabel="port", ylabel="TE fraction", title="Polarization selection", ylim=(0, 1.02)
)
axes[0].legend(frameon=False)
plt.show()
preflight = {}
for layer, design in designs.items():
finalized = design.to_simulation(design.design_region.initial_parameters)
task_names = []
for task_name, simulation in finalized.items():
simulation.validate_pre_upload()
task_names.append(task_name)
native_source_x = float(next(iter(finalized.values())).sources[0].center[0])
preflight[layer] = {
"simulations": len(finalized),
"task names": task_names,
"sources per simulation": 1,
"input ports": LAYER_INPUT_COUNTS[layer],
"output ports": LAYER_OUTPUT_COUNTS[layer],
"monitors per simulation": (
LAYER_INPUT_COUNTS[layer] + LAYER_OUTPUT_COUNTS[layer]
),
"polarization": "fundamental quasi-TE",
"minimum port TE fraction": port_mode_report["minimum TE fraction"],
"port effective-index range": [
port_mode_report["minimum 2D rail effective index"],
port_mode_report["maximum 2D rail effective index"],
],
"variational TE index, 150 nm": N_EFF_TE_150,
"variational TE index, 220 nm": N_EFF_TE_220,
"silica background index": N_EFF_TE_BACKGROUND,
"variational calibration grid nm": eim_calibration["grid nm"],
"simulation architecture": "native ModalComponentModeler S matrix",
"incident normalization": True,
"mode window um": MODE_WINDOW,
"source-to-input-monitor um": INPUT_MONITOR_X - native_source_x,
"input-monitor-to-taper um": (-DESIGN_HALF - TAPER_LENGTH) - INPUT_MONITOR_X,
"taper mouth clearance um": PORT_PITCH - TAPER_MOUTH,
"taper rail overlap um": TAPER_OVERLAP,
}
for layer, report in preflight.items():
print(
f"layer {layer + 1}: {report['simulations']} local simulations, "
f"{report['input ports']} inputs, {report['output ports']} outputs"
)
port modes: 6 checked, minimum TE fraction=1.0000, n_eff range=2.483866-2.487522
layer 1: 6 local simulations, 6 inputs, 6 outputs layer 2: 6 local simulations, 6 inputs, 6 outputs layer 3: 6 local simulations, 6 inputs, 4 outputs
Define the Realization Objective¶
For each device, the six source simulations are combined into an incident-normalized transmission matrix $T$ and reflection matrix $R$. The first two transmission matrices are 6 × 6, and the final transmission matrix is 4 × 6. Reflection remains 6 × 6 for all three devices.
$$ L_k= \left\|T-T_{\mathrm{target}}\right\|_F +\left\|R\right\|_F +w_kL_{\mathrm{fab}}. $$
The first term matches both the amplitude and phase of the target transmission matrix. The second suppresses light reflected toward the input ports. The fabrication term penalizes solid and void features below 200 nm.
The fabrication weight $w_k$ is zero for updates 1–5, increases from 0.08 at update 6 to 0.8 at update 15, and remains at 0.8 afterward.
[31]:
def smatrix_element(smatrix, port_out, port_in):
value = (
smatrix.sel(
port_out=port_out,
mode_index_out=0,
port_in=port_in,
mode_index_in=0,
)
.sel(f=FREQ0_HZ, method="nearest")
.data
)
return value.item() if np.ndim(value) == 0 else anp.squeeze(value)
def response_from_batch(batch_data, layer):
modeler = NATIVE_MODELERS[layer]
native_batch = {
native_name: batch_data[actual_name]
for native_name, actual_name in zip(NATIVE_TASK_NAMES[layer], batch_data)
}
modeler_data = compose_modeler_data_from_batch_data(
modeler=modeler, batch_data=native_batch
)
smatrix = modeler_data.smatrix()
transmissions, reflections = [], []
for input_name in INPUT_NAMES:
port = modeler.get_port_by_name(port_name=input_name)
task_name = modeler.get_task_name(port=port, mode_index=0)
incident = modeler._normalization_factor(port, modeler_data.data[task_name])
incident_value = np.asarray(getval(incident))
if not np.all(np.isfinite(incident_value)) or np.any(
np.abs(incident_value) == 0
):
raise ValueError(
f"Invalid incident amplitude for {input_name}: {incident_value}"
)
transmissions.append(
anp.stack(
[
smatrix_element(smatrix, output_name, input_name)
for output_name in output_names(LAYER_OUTPUT_COUNTS[layer])
]
)
)
reflections.append(
anp.stack(
[
smatrix_element(smatrix, reflected_name, input_name)
for reflected_name in INPUT_NAMES
]
)
)
return anp.stack(transmissions, axis=1), anp.stack(reflections, axis=1)
def frobenius_norm(values):
return anp.sqrt(anp.sum(anp.real(values * anp.conj(values))) + 1e-30)
def postprocess(layer, target):
target = anp.asarray(target)
def realization_score(batch_data):
transmission, reflection = response_from_batch(batch_data, layer)
transmission_error = frobenius_norm(transmission - target)
reflection_error = frobenius_norm(reflection)
constrained_optical_loss = (
transmission_error + REFLECTION_WEIGHT * reflection_error
)
return -constrained_optical_loss
return realization_score
Estimate Cloud Cost¶
Before starting the inverse design, we estimate the cost of all six source simulations for each device. The cell uploads each simulation for pricing, records its estimated forward cost, and then deletes the estimate-only task.
An optimization update requires both forward and adjoint simulations, so the projected cost is twice the forward cost multiplied by the requested number of updates. The notebook reports the estimate for each layer and for the complete realization schedule.
Set ESTIMATE_COST=True in a fresh kernel to calculate the estimate. Running the optimization also requires explicit approval through COST_APPROVED=True and RUN_CLOUD=True.
[32]:
cost_report = None
if ESTIMATE_COST and matrix_training_ready:
layers = []
for layer, design in designs.items():
simulations = design.to_simulation(design.design_region.initial_parameters)
forward_costs = []
for task_name, simulation in sorted(simulations.items()):
task_id = web.upload(
simulation,
task_name=f"layer_{layer + 1}_{task_name}_cost",
folder_name="deep_inverse_designed_pnn_cost",
)
forward_costs.append(float(web.estimate_cost(task_id)))
web.delete(task_id)
forward_batch = sum(forward_costs)
layers.append(
{
"layer": layer + 1,
"sources": len(simulations),
"forward FlexCredits": forward_batch,
"projected FlexCredits": 2 * forward_batch * REALIZATION_STEPS,
}
)
projected_per_mode = sum(layer["projected FlexCredits"] for layer in layers)
cost_report = {
"mode": REALIZATION_MODE,
"validation modes": COST_ESTIMATE_MODES,
"steps per layer": REALIZATION_STEPS,
"layers": layers,
"projected FlexCredits per mode": projected_per_mode,
"projected FlexCredits": len(COST_ESTIMATE_MODES) * projected_per_mode,
}
cost_lines = [
f"Layer {layer['layer']}: {layer['forward FlexCredits']:.4f} forward, "
f"{layer['projected FlexCredits']:.4f} projected FlexCredits"
for layer in layers
]
cost_lines.append(f"Projected per mode: {projected_per_mode:.4f} FlexCredits")
cost_lines.append(
f"Projected {len(COST_ESTIMATE_MODES)}-mode total: "
f"{cost_report['projected FlexCredits']:.4f} FlexCredits"
)
print("\n".join(cost_lines))
elif ESTIMATE_COST:
print(
"Cost estimation is closed until the trained matrices pass the required checks."
)
else:
print("Cost estimation is disabled.")
Cost estimation is disabled.
Run the Inverse Design¶
Adam updates the continuous material parameters by maximizing $-L_k$. The learning rate follows
$$ \eta_k=0.05\min\left(\frac{k}{5},1\right), $$
so it increases from 0.01 to 0.05 over the first five updates. After each update, the parameters are clipped to $[0,1]$ before the next filtered device is constructed.
Each forward or adjoint batch submits one source simulation at a time. If a source remains queued for 600 seconds, the task is aborted and retried after 30 seconds without advancing Adam.
The notebook saves the current parameters, transmission and reflection matrices, and Adam state after every update. This allows the optimization to continue from the same point after interruption. The displayed analysis uses a completed 40-update response for each of the three devices.
[33]:
CLOUD_RUN_LABEL = "deep_classifier"
optimizer_designs = {
layer: designs[layer].updated_copy(task_name=f"{CLOUD_RUN_LABEL}_layer_{layer + 1}")
for layer in range(N_LAYERS)
}
EXPECTED_TIDY3D_RETRY_VERSION = "2.12.0"
ACTIVE_QUEUE_CONTEXT = {"layer": None, "update": None, "job": None}
queue_retry_records = []
if QUEUE_RETRY_LOG_PATH.exists():
saved_retry_log = json.loads(QUEUE_RETRY_LOG_PATH.read_text())
if saved_retry_log.get("policy") != QUEUE_RETRY_POLICY:
raise ValueError("Saved queue-retry policy does not match this notebook.")
queue_retry_records = list(saved_retry_log.get("records", []))
class SourceQueueTimeout(RuntimeError):
def __init__(self, queued_seconds):
self.queued_seconds = float(queued_seconds)
super().__init__(f"source remained queued for {queued_seconds:.1f} seconds")
def validate_native_retry_hooks():
if td.__version__ != EXPECTED_TIDY3D_RETRY_VERSION:
raise RuntimeError(
f"Per-source retry requires Tidy3D {EXPECTED_TIDY3D_RETRY_VERSION}; "
f"found {td.__version__}."
)
for name in ("_run_async_tidy3d", "_run_async_tidy3d_bwd"):
function = getattr(autograd_hooks, name, None)
parameters = inspect.signature(function).parameters if function else {}
accepts_keywords = any(
parameter.kind == inspect.Parameter.VAR_KEYWORD
for parameter in parameters.values()
)
if "simulations" not in parameters or not accepts_keywords:
raise RuntimeError(
f"Tidy3D autograd hook {name} has an unsupported signature."
)
def save_queue_retry_log():
payload = {
"policy": QUEUE_RETRY_POLICY,
"source chunk size": AUTOGRAD_SOURCE_CHUNK_SIZE,
"queue timeout seconds": QUEUE_TIMEOUT_SECONDS,
"retry delay seconds": QUEUE_RETRY_DELAY_SECONDS,
"records": queue_retry_records,
}
temporary_path = QUEUE_RETRY_LOG_PATH.with_suffix(".tmp")
temporary_path.write_text(json.dumps(payload, indent=2))
temporary_path.replace(QUEUE_RETRY_LOG_PATH)
def record_queue_wait(phase, source, queued_seconds, retried):
identity = {
"layer": int(ACTIVE_QUEUE_CONTEXT["layer"]),
"update": int(ACTIVE_QUEUE_CONTEXT["update"]),
"source": str(source),
"phase": str(phase),
}
record = next(
(
item
for item in queue_retry_records
if all(item.get(key) == value for key, value in identity.items())
),
None,
)
if record is None:
record = {**identity, "retry count": 0, "accumulated queue seconds": 0.0}
queue_retry_records.append(record)
record["accumulated queue seconds"] += float(queued_seconds)
if retried:
record["retry count"] += 1
save_queue_retry_log()
def monitor_single_source(job, clock=time.monotonic, sleep=time.sleep):
queued_started = None
queued_seconds = 0.0
while True:
status = str(job.get_info().status).lower()
now = float(clock())
if status in web_states.QUEUED_STATES:
if queued_started is None:
queued_started = now
current_queue_seconds = queued_seconds + now - queued_started
if current_queue_seconds >= QUEUE_TIMEOUT_SECONDS:
raise SourceQueueTimeout(current_queue_seconds)
elif queued_started is not None:
queued_seconds += now - queued_started
queued_started = None
if status in web_states.SUCCESS_STATES:
return queued_seconds
if status in web_states.ERROR_STATES or status in web_states.DIVERGED_STATES:
raise RuntimeError(f"Tidy3D source task ended with status {status!r}.")
sleep(QUEUE_POLL_SECONDS)
def transient_web_error(error):
text = str(error).lower()
return any(
fragment in text
for fragment in (
"connection",
"timeout",
"timed out",
"bad gateway",
"service unavailable",
"gateway timeout",
"502",
"503",
"504",
)
)
def abort_source_job(job, wait):
task_id = getattr(job, "task_id", None)
if task_id is None:
return
while True:
try:
web.abort(task_id)
break
except (RequestsConnectionError, ReadTimeout, NewConnectionError):
if not wait:
return
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
except WebError as error:
if not transient_web_error(error):
raise
if not wait:
return
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
if wait:
while True:
try:
if str(job.get_info().status).lower() in web_states.END_STATES:
return
time.sleep(QUEUE_POLL_SECONDS)
except (RequestsConnectionError, ReadTimeout, NewConnectionError):
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
except WebError as error:
if not transient_web_error(error):
raise
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
def run_one_source_with_queue_retry(
phase, source, make_batch, finish_batch, start_kwargs
):
while True:
batch = None
job = None
try:
batch = make_batch()
job = next(iter(batch.jobs.values()))
ACTIVE_QUEUE_CONTEXT["job"] = job
if not job.load_if_cached:
batch.start(**start_kwargs)
queued_seconds = monitor_single_source(job)
else:
queued_seconds = 0.0
result = finish_batch(batch, job)
record_queue_wait(phase, source, queued_seconds, retried=False)
ACTIVE_QUEUE_CONTEXT["job"] = None
return result
except SourceQueueTimeout as error:
abort_source_job(job, wait=True)
record_queue_wait(phase, source, error.queued_seconds, retried=True)
print(
f"{phase} source {source} exceeded {QUEUE_TIMEOUT_SECONDS:g} queued "
f"seconds; retrying only this source after {QUEUE_RETRY_DELAY_SECONDS:g} seconds."
)
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
except (RequestsConnectionError, ReadTimeout, NewConnectionError) as error:
abort_source_job(job, wait=True)
record_queue_wait(phase, source, 0.0, retried=True)
print(f"Temporary connection failure for {phase} source {source}: {error}")
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
except WebError as error:
if not transient_web_error(error):
raise
abort_source_job(job, wait=True)
record_queue_wait(phase, source, 0.0, retried=True)
print(f"Temporary web failure for {phase} source {source}: {error}")
time.sleep(QUEUE_RETRY_DELAY_SECONDS)
except KeyboardInterrupt:
abort_source_job(job, wait=False)
raise
finally:
ACTIVE_QUEUE_CONTEXT["job"] = None
[34]:
def queue_aware_forward_hook(simulations, **run_kwargs):
if len(simulations) != 1:
raise RuntimeError("Per-source retry received more than one forward source.")
source = next(iter(simulations))
base_kwargs = dict(run_kwargs)
disable_result_cache = base_kwargs.pop("disable_result_cache", False)
path_dir = Path(base_kwargs.pop("path_dir", "."))
path_dir.mkdir(parents=True, exist_ok=True)
priority = base_kwargs.get("priority")
vgpu_allocation = base_kwargs.get("vgpu_allocation")
ignore_memory_limit = base_kwargs.get("ignore_memory_limit")
start_kwargs = {
"priority": priority,
"vgpu_allocation": vgpu_allocation,
"ignore_memory_limit": ignore_memory_limit,
}
def make_batch():
batch_kwargs = autograd_engine.parse_run_kwargs(**base_kwargs)
batch = autograd_engine._build_batch(
simulations=simulations,
num_workers=base_kwargs.get("num_workers"),
**batch_kwargs,
)
if disable_result_cache:
batch = autograd_engine._with_result_cache_disabled(batch)
if batch.simulation_type == "autograd_fwd":
prepared = {
name: simulation.updated_copy(
simulation_type="autograd_fwd", deep=False
)
for name, simulation in batch.simulations.items()
}
batch = batch.updated_copy(simulations=prepared)
artifacts = {
name: autograd_engine._sim_fields_keys_artifacts(keys)
for name, keys in base_kwargs["sim_fields_keys_dict"].items()
}
batch._upload_jobs(_sidecar_artifacts_by_task=artifacts)
else:
batch.upload()
return batch
def finish_batch(batch, job):
batch.download(path_dir=path_dir)
batch_data = batch.load(path_dir=path_dir, skip_download=True)
task_ids = getattr(batch_data, "task_ids", None)
return batch_data, (
dict(task_ids) if task_ids is not None else {source: job.task_id}
)
return run_one_source_with_queue_retry(
"forward",
source,
make_batch,
finish_batch,
start_kwargs,
)
def queue_aware_backward_hook(simulations, **run_kwargs):
verbose = run_kwargs.get("verbose", True)
cached = {
name: value
for name, simulation in simulations.items()
if (value := get_cached_vjp_traced_fields(simulation, verbose=verbose))
is not None
}
remaining = {
name: simulation
for name, simulation in simulations.items()
if name not in cached
}
if not remaining:
return cached
if len(remaining) != 1:
raise RuntimeError("Per-source retry received more than one adjoint source.")
source = next(iter(remaining))
base_kwargs = dict(run_kwargs)
base_kwargs.pop("path_dir", None)
priority = base_kwargs.get("priority")
vgpu_allocation = base_kwargs.get("vgpu_allocation")
ignore_memory_limit = base_kwargs.get("ignore_memory_limit")
start_kwargs = {
"priority": priority,
"vgpu_allocation": vgpu_allocation,
"ignore_memory_limit": ignore_memory_limit,
}
def make_batch():
batch_kwargs = autograd_engine.parse_run_kwargs(**base_kwargs)
batch = autograd_engine._build_batch(
simulations=remaining,
num_workers=base_kwargs.get("num_workers"),
**batch_kwargs,
)
return autograd_engine._with_result_cache_disabled(batch)
def finish_batch(batch, job):
return {
**cached,
source: get_vjp_traced_fields(
task_id_adj=job.task_id,
verbose=batch.verbose,
cache_simulation=job.simulation,
),
}
return run_one_source_with_queue_retry(
"adjoint",
source,
make_batch,
finish_batch,
start_kwargs,
)
@contextmanager
def native_source_queue_retry(layer, update_number):
validate_native_retry_hooks()
original_forward = autograd_hooks._run_async_tidy3d
original_backward = autograd_hooks._run_async_tidy3d_bwd
ACTIVE_QUEUE_CONTEXT.update(
layer=int(layer) + 1,
update=int(update_number),
job=None,
)
autograd_hooks._run_async_tidy3d = queue_aware_forward_hook
autograd_hooks._run_async_tidy3d_bwd = queue_aware_backward_hook
try:
yield
finally:
active_job = ACTIVE_QUEUE_CONTEXT.get("job")
if active_job is not None:
abort_source_job(active_job, wait=False)
autograd_hooks._run_async_tidy3d = original_forward
autograd_hooks._run_async_tidy3d_bwd = original_backward
ACTIVE_QUEUE_CONTEXT.update(layer=None, update=None, job=None)
if RUN_CLOUD:
validate_native_retry_hooks()
[35]:
class ClippedAdamAscentOptimizer(tdi.AdamOptimizer):
def initial_state(self, parameters):
return super().initial_state(anp.clip(parameters, 0.0, 1.0))
def update(self, parameters, gradient, state=None):
parameters = anp.clip(parameters, 0.0, 1.0)
updated, new_state = super().update(parameters, gradient, state)
return anp.clip(updated, 0.0, 1.0), new_state
def optimizer_for_update(layer, update_number):
penalty_weight = fabrication_penalty_weight(update_number)
penalties = (fabrication_penalty(penalty_weight),) if penalty_weight > 0 else ()
scheduled_region = optimizer_designs[layer].design_region.updated_copy(
penalties=penalties
)
scheduled_design = optimizer_designs[layer].updated_copy(
design_region=scheduled_region
)
return ClippedAdamAscentOptimizer(
design=scheduled_design,
learning_rate=tidy3d_learning_rate(update_number),
maximize=True,
beta1=TIDY3D_BETA1,
beta2=TIDY3D_BETA2,
eps=TIDY3D_EPS,
num_steps=REALIZATION_STEPS,
store_full_results=True,
results_cache_fname=None,
)
optimizers = {layer: optimizer_for_update(layer, 1) for layer in range(N_LAYERS)}
ascent_probe = np.full(designs[0].design_region.params_shape, 0.5)
ascent_gradient = np.ones_like(ascent_probe)
ascent_state = optimizers[0].initial_state(ascent_probe)
ascent_updated, ascent_state = optimizers[0].update(
ascent_probe, gradient=-ascent_gradient, state=ascent_state
)
ascent_alignment = float(
np.vdot(ascent_gradient.ravel(), (ascent_updated - ascent_probe).ravel()).real
)
if ascent_alignment <= 0:
raise ValueError("Adam update is not aligned with score-gradient ascent.")
print("optimizer direction:", OPTIMIZER_DIRECTION, "alignment:", ascent_alignment)
unit_fabrication_penalty = fabrication_penalty(weight=1.0)
max_realization_steps = max(REALIZATION_STEPS, 1)
matrix_history_shape = (
max_realization_steps,
N_LAYERS,
max(LAYER_OUTPUT_COUNTS),
N_CHANNELS,
)
metric_history_shape = (max_realization_steps, N_LAYERS)
optimizer direction: gradient_ascent_on_negative_constrained_loss alignment: 640.0899935992227
[36]:
transmission_history = np.full(matrix_history_shape, np.nan + 0j)
reflection_history = np.full(matrix_history_shape, np.nan + 0j)
transmission_error_history = np.full(metric_history_shape, np.nan)
reflection_error_history = np.full(metric_history_shape, np.nan)
realization_loss_history = np.full(metric_history_shape, np.nan)
learning_rate_history = np.full(metric_history_shape, np.nan)
fabrication_weight_history = np.full(metric_history_shape, np.nan)
fabrication_penalty_history = np.full(metric_history_shape, np.nan)
weighted_fabrication_penalty_history = np.full(metric_history_shape, np.nan)
scheduled_constrained_loss_history = np.full(metric_history_shape, np.nan)
fixed_reference_loss_history = np.full(metric_history_shape, np.nan)
test_accuracy_history = np.full(metric_history_shape, np.nan)
latest_parameters = np.stack(
[
np.asarray(designs[layer].design_region.initial_parameters)
for layer in range(N_LAYERS)
]
)
latest_adam_m = np.zeros_like(latest_parameters)
latest_adam_v = np.zeros_like(latest_parameters)
latest_adam_t = np.zeros(N_LAYERS, dtype=int)
response_steps = np.full(N_LAYERS, -1, dtype=int)
completed_steps = np.zeros(N_LAYERS, dtype=int)
optimization_run_summary = {}
CACHE_CAN_CONTINUE = True
history_names = (
"transmission_history",
"reflection_history",
"transmission_error_history",
"reflection_error_history",
"realization_loss_history",
"learning_rate_history",
"fabrication_weight_history",
"fabrication_penalty_history",
"weighted_fabrication_penalty_history",
"scheduled_constrained_loss_history",
"fixed_reference_loss_history",
"test_accuracy_history",
)
if PROGRESS_PATH.exists():
with np.load(PROGRESS_PATH, allow_pickle=False) as progress:
required = set(history_names) | {
"latest_parameters",
"latest_adam_m",
"latest_adam_v",
"latest_adam_t",
"response_steps",
"completed_steps",
"n_channels",
"n_splitter_leaves",
"class_labels",
"preprocessing_tag",
"realization_target_tag",
"layer_input_counts",
"layer_output_counts",
}
missing = required.difference(progress.files)
if missing:
raise ValueError(f"Progress cache is missing {sorted(missing)}.")
identity_checks = (
int(progress["n_channels"]) == N_CHANNELS,
int(progress["n_splitter_leaves"]) == N_SPLITTER_LEAVES,
np.array_equal(progress["class_labels"], CLASS_LABELS),
str(progress["preprocessing_tag"]) == PREPROCESSING_TAG,
str(progress["realization_target_tag"]) == REALIZATION_TARGET_TAG,
np.array_equal(progress["layer_input_counts"], LAYER_INPUT_COUNTS),
np.array_equal(progress["layer_output_counts"], LAYER_OUTPUT_COUNTS),
)
if not all(identity_checks):
raise ValueError("Progress cache does not match this classifier.")
for name in history_names:
globals()[name] = np.asarray(progress[name])
latest_parameters = np.asarray(progress["latest_parameters"])
latest_adam_m = np.asarray(progress["latest_adam_m"])
latest_adam_v = np.asarray(progress["latest_adam_v"])
latest_adam_t = np.asarray(progress["latest_adam_t"], dtype=int)
response_steps = np.asarray(progress["response_steps"], dtype=int)
completed_steps = np.asarray(progress["completed_steps"], dtype=int)
CACHE_CAN_CONTINUE = (
bool(progress["continuation_allowed"])
if "continuation_allowed" in progress.files
else True
)
if completed_steps.shape != (N_LAYERS,) or not np.array_equal(
completed_steps, latest_adam_t
):
raise ValueError("Cached Adam steps do not match completed updates.")
if (
RUN_CLOUD
and np.any(completed_steps < REALIZATION_STEPS)
and not CACHE_CAN_CONTINUE
):
raise ValueError(
"This completed-data cache is analysis-only and cannot be resumed."
)
print(f"loaded cache: {completed_steps.tolist()} updates")
history_extension = max(0, REALIZATION_STEPS - len(test_accuracy_history))
if history_extension:
matrix_padding = ((0, history_extension), (0, 0), (0, 0), (0, 0))
metric_padding = ((0, history_extension), (0, 0))
transmission_history = np.pad(
transmission_history, matrix_padding, constant_values=np.nan
)
reflection_history = np.pad(
reflection_history, matrix_padding, constant_values=np.nan
)
for name in history_names[2:-1]:
globals()[name] = np.pad(
globals()[name], metric_padding, constant_values=np.nan
)
test_accuracy_history = np.pad(
test_accuracy_history,
((0, history_extension), (0, 0)),
constant_values=np.nan,
)
max_realization_steps = len(test_accuracy_history)
loaded cache: [40, 40, 40] updates
[37]:
def result_from_progress(layer):
base = optimizers[layer].initialize_result()
steps = int(latest_adam_t[layer])
parameters = np.asarray(latest_parameters[layer])
state = {
"m": np.asarray(latest_adam_m[layer]),
"v": np.asarray(latest_adam_v[layer]),
"t": steps,
}
history = {key: list(value) for key, value in base.history.items()}
history["params"] = [base.get_last("params")] + [parameters] * steps
history["opt_state"] = [base.get_last("opt_state")] + [state] * steps
history["grad"] = [np.zeros_like(parameters)] * steps
history["objective_fn_val"] = list(
-scheduled_constrained_loss_history[:steps, layer]
)
history["penalty"] = list(weighted_fabrication_penalty_history[:steps, layer])
history["post_process_val"] = list(-realization_loss_history[:steps, layer])
return tdi.InverseDesignResult(design=base.design, **history)
cloud_results = {}
for layer in range(N_LAYERS):
if optimizer_state_paths[layer].exists():
try:
cloud_results[layer] = tdi.InverseDesignResult.from_file(
str(optimizer_state_paths[layer])
)
except OSError:
if latest_adam_t[layer] == 0:
raise
cloud_results[layer] = result_from_progress(layer)
elif latest_adam_t[layer] > 0:
cloud_results[layer] = result_from_progress(layer)
else:
cloud_results[layer] = optimizers[layer].initialize_result()
completed_steps[layer] = len(cloud_results[layer].objective_fn_val)
adam_update = int(cloud_results[layer].get_last("opt_state")["t"])
if adam_update != completed_steps[layer]:
raise ValueError(
f"Layer {layer + 1} Adam state t={adam_update} does not match "
f"its {completed_steps[layer]} completed objective evaluations."
)
optimizer_forward = {
layer: {
"transmission": transmission_history[
response_steps[layer], layer, : LAYER_OUTPUT_COUNTS[layer]
],
"reflection": reflection_history[response_steps[layer], layer],
"params": latest_parameters[layer],
"step": int(response_steps[layer]),
}
for layer in range(N_LAYERS)
if response_steps[layer] >= 0
}
[38]:
def save_realization_progress():
buffer = io.BytesIO()
np.savez_compressed(
buffer,
transmission_history=transmission_history,
reflection_history=reflection_history,
transmission_error_history=transmission_error_history,
reflection_error_history=reflection_error_history,
realization_loss_history=realization_loss_history,
learning_rate_history=learning_rate_history,
fabrication_weight_history=fabrication_weight_history,
fabrication_penalty_history=fabrication_penalty_history,
weighted_fabrication_penalty_history=weighted_fabrication_penalty_history,
scheduled_constrained_loss_history=scheduled_constrained_loss_history,
fixed_reference_loss_history=fixed_reference_loss_history,
test_accuracy_history=test_accuracy_history,
latest_parameters=latest_parameters,
latest_adam_m=latest_adam_m,
latest_adam_v=latest_adam_v,
latest_adam_t=latest_adam_t,
response_steps=response_steps,
completed_steps=completed_steps,
realization_objective=np.asarray(REALIZATION_OBJECTIVE),
realization_mode=np.asarray(REALIZATION_MODE),
requested_steps_per_layer=np.asarray(REALIZATION_STEPS),
n_channels=np.asarray(N_CHANNELS),
n_splitter_leaves=np.asarray(N_SPLITTER_LEAVES),
active_leaf_indices=np.asarray(ACTIVE_LEAF_INDICES),
dump_leaf_indices=np.asarray(DUMP_LEAF_INDICES),
class_labels=np.asarray(CLASS_LABELS),
preprocessing_tag=np.asarray(PREPROCESSING_TAG),
fanout_rule=np.asarray(FANOUT_RULE),
terminator_model=np.asarray(TERMINATOR_MODEL),
ideal_active_power_fraction=np.asarray(IDEAL_ACTIVE_POWER_FRACTION),
terminator_reflection=np.asarray(TERMINATOR_REFLECTION),
design_um_per_channel=np.asarray(DESIGN_UM_PER_CHANNEL),
design_edge_um=np.asarray(DESIGN_EDGE_UM),
port_pitch_um=np.asarray(PORT_PITCH),
mode_window_um=np.asarray(MODE_WINDOW),
rail_thickness_um=np.asarray(RAIL_THICKNESS),
rail_inner_overlap_um=np.asarray(RAIL_INNER_OVERLAP),
taper_overlap_um=np.asarray(TAPER_OVERLAP),
simulation_architecture=np.asarray(SIMULATION_ARCHITECTURE),
cladding_model=np.asarray(CLADDING_MODEL),
background_index=np.asarray(N_EFF_TE_BACKGROUND),
optimizer_direction=np.asarray(OPTIMIZER_DIRECTION),
active_mask_rule=np.asarray(ACTIVE_MASK_RULE),
realization_target_tag=np.asarray(REALIZATION_TARGET_TAG),
layer_input_counts=np.asarray(LAYER_INPUT_COUNTS),
layer_output_counts=np.asarray(LAYER_OUTPUT_COUNTS),
queue_retry_policy=np.asarray(QUEUE_RETRY_POLICY),
autograd_source_chunk_size=np.asarray(AUTOGRAD_SOURCE_CHUNK_SIZE),
queue_timeout_seconds=np.asarray(QUEUE_TIMEOUT_SECONDS),
queue_retry_delay_seconds=np.asarray(QUEUE_RETRY_DELAY_SECONDS),
tidy3d_lr_max=np.asarray(TIDY3D_LR_MAX),
tidy3d_lr_warmup_updates=np.asarray(TIDY3D_LR_WARMUP_UPDATES),
fabrication_min_feature_um=np.asarray(FABRICATION_MIN_FEATURE_UM),
fabrication_weight_max=np.asarray(FABRICATION_WEIGHT_MAX),
fabrication_start_update=np.asarray(FABRICATION_START_UPDATE),
fabrication_full_update=np.asarray(FABRICATION_FULL_UPDATE),
fabrication_beta=np.asarray(FABRICATION_BETA),
fabrication_eta0=np.asarray(FABRICATION_ETA0),
fabrication_delta_eta=np.asarray(FABRICATION_DELTA_ETA),
continuation_allowed=np.asarray(True),
)
temporary_path = PROGRESS_PATH.with_suffix(".tmp")
temporary_path.write_bytes(buffer.getvalue())
temporary_path.replace(PROGRESS_PATH)
[39]:
def keep_optimizer_forward(layer):
def callback(result, step_index, aux_data):
transmission, reflection = response_from_batch(aux_data["sim_data"], layer)
transmission, reflection = np.asarray(transmission), np.asarray(reflection)
step_index = int(step_index)
update_number = step_index + 1
transmission_history[step_index, layer] = np.nan + 0j
transmission_history[step_index, layer, : LAYER_OUTPUT_COUNTS[layer]] = (
transmission
)
reflection_history[step_index, layer] = reflection
transmission_error_history[step_index, layer] = np.linalg.norm(
transmission - targets[layer]
)
reflection_error_history[step_index, layer] = np.linalg.norm(reflection)
realization_loss_history[step_index, layer] = (
transmission_error_history[step_index, layer]
+ REFLECTION_WEIGHT * reflection_error_history[step_index, layer]
)
learning_rate_history[step_index, layer] = tidy3d_learning_rate(update_number)
fabrication_weight_history[step_index, layer] = fabrication_penalty_weight(
update_number
)
latest_parameters[layer] = np.asarray(getval(aux_data["params"]))
material_density = np.asarray(
designs[layer].design_region.material_density(latest_parameters[layer])
)
fabrication_penalty_history[step_index, layer] = float(
unit_fabrication_penalty.evaluate(material_density, PIXEL)
)
weighted_fabrication_penalty_history[step_index, layer] = float(
result.penalty[-1]
)
scheduled_constrained_loss_history[step_index, layer] = (
realization_loss_history[step_index, layer]
+ weighted_fabrication_penalty_history[step_index, layer]
)
fixed_reference_loss_history[step_index, layer] = (
realization_loss_history[step_index, layer]
+ FABRICATION_WEIGHT_MAX * fabrication_penalty_history[step_index, layer]
)
response_steps[layer] = step_index
optimizer_forward[layer] = {
"transmission": transmission,
"reflection": reflection,
"params": latest_parameters[layer],
"step": step_index,
"learning_rate": learning_rate_history[step_index, layer],
"fabrication_weight": fabrication_weight_history[step_index, layer],
}
completed_steps[layer] = len(result.objective_fn_val)
latest_state = result.get_last("opt_state")
latest_adam_m[layer] = np.asarray(latest_state["m"])
latest_adam_v[layer] = np.asarray(latest_state["v"])
latest_adam_t[layer] = int(latest_state["t"])
save_realization_progress()
checkpoint_path = optimizer_state_paths[layer]
temporary_checkpoint = checkpoint_path.with_name(
checkpoint_path.stem + ".tmp.hdf5"
)
result.to_file(str(temporary_checkpoint))
temporary_checkpoint.replace(checkpoint_path)
return callback
def uncalibrated_test_accuracy(matrices):
fields = physical_forward(matrices, jnp.asarray(test_features))
scores = jnp.abs(fields) ** 2
return float(jnp.mean(jnp.argmax(scores, axis=-1) == jnp.asarray(test_labels)))
def latest_network_matrices():
return tuple(
optimizer_forward[layer]["transmission"]
if layer in optimizer_forward
else targets[layer]
for layer in range(N_LAYERS)
)
def advance_layer(layer, local_step):
done_steps = len(cloud_results[layer].objective_fn_val)
if done_steps <= local_step and done_steps < REALIZATION_STEPS:
update_number = done_steps + 1
step_optimizer = optimizer_for_update(layer, update_number)
cloud_results[layer] = cloud_results[layer].updated_copy(
design=step_optimizer.design
)
with native_source_queue_retry(layer, update_number):
cloud_results[layer] = step_optimizer.continue_run(
result=cloud_results[layer],
num_steps=1,
post_process_fn=postprocess(layer, targets[layer]),
callback=keep_optimizer_forward(layer),
)
completed_steps[layer] = len(cloud_results[layer].objective_fn_val)
adam_update = int(cloud_results[layer].get_last("opt_state")["t"])
if adam_update != completed_steps[layer]:
raise ValueError(
f"Layer {layer + 1} Adam state reset at update {update_number}."
)
save_realization_progress()
[40]:
needs_cloud_updates = bool(np.any(completed_steps < REALIZATION_STEPS))
if (
RUN_CLOUD
and needs_cloud_updates
and COST_APPROVED
and matrix_training_ready
and cost_report is not None
):
if REALIZATION_MODE == "sequential":
for layer in range(N_LAYERS):
for local_step in range(REALIZATION_STEPS):
advance_layer(layer, local_step)
if response_steps[layer] >= local_step and not np.isfinite(
test_accuracy_history[local_step, layer]
):
test_accuracy_history[local_step, layer] = (
uncalibrated_test_accuracy(latest_network_matrices())
)
print(
f"layer {layer + 1}, iteration {local_step + 1:2d}: "
f"optical loss={realization_loss_history[local_step, layer]:.6e}, "
f"fabrication={fabrication_penalty_history[local_step, layer]:.6e}, "
f"constrained={scheduled_constrained_loss_history[local_step, layer]:.6e}, "
f"test accuracy={test_accuracy_history[local_step, layer]:.3%}"
)
save_realization_progress()
else:
for global_step in range(REALIZATION_STEPS):
for layer in range(N_LAYERS):
advance_layer(layer, global_step)
available = np.all(np.isfinite(transmission_error_history[global_step]))
if available and not np.isfinite(test_accuracy_history[global_step, 0]):
synchronized_matrices = tuple(
transmission_history[
global_step, layer, : LAYER_OUTPUT_COUNTS[layer]
]
for layer in range(N_LAYERS)
)
test_accuracy_history[global_step, 0] = uncalibrated_test_accuracy(
synchronized_matrices
)
print(
f"iteration {global_step + 1:2d}: optical loss="
f"{np.mean(realization_loss_history[global_step]):.6e}, "
f"fabrication="
f"{np.mean(fabrication_penalty_history[global_step]):.6e}, "
f"constrained="
f"{np.mean(scheduled_constrained_loss_history[global_step]):.6e}, "
f"test accuracy={test_accuracy_history[global_step, 0]:.3%}"
)
save_realization_progress()
matched_tasks = [
task
for task in web.get_tasks(folder="default")
if CLOUD_RUN_LABEL in task.get("task_name", task.get("taskName", ""))
]
completed_costs = [
float(web.real_cost(task.get("task_id", task.get("taskId")), verbose=False))
for task in matched_tasks
if task.get("status", "") in ("success", "completed")
]
optimization_run_summary = {
"mode": REALIZATION_MODE,
"task count": len(matched_tasks),
"completed task count": len(completed_costs),
"actual FlexCredits": float(sum(completed_costs)),
}
save_realization_progress()
elif RUN_CLOUD and not needs_cloud_updates:
print("The cached inverse design already contains all requested updates.")
elif RUN_CLOUD:
print(
"Cloud optimization is closed: estimate cost in this kernel and "
"enable explicit approval."
)
else:
print("Cloud optimization is disabled.")
Cloud optimization is disabled.
[41]:
realization_complete = bool(np.all(completed_steps >= REALIZATION_STEPS))
realization_progress = {
"mode": REALIZATION_MODE,
"iterations": np.arange(1, max_realization_steps + 1),
"transmission_history": transmission_history,
"reflection_history": reflection_history,
"transmission_error_history": transmission_error_history,
"reflection_error_history": reflection_error_history,
"realization_loss_history": realization_loss_history,
"learning_rate_history": learning_rate_history,
"fabrication_weight_history": fabrication_weight_history,
"fabrication_penalty_history": fabrication_penalty_history,
"weighted_fabrication_penalty_history": weighted_fabrication_penalty_history,
"scheduled_constrained_loss_history": scheduled_constrained_loss_history,
"fixed_reference_loss_history": fixed_reference_loss_history,
"test_accuracy_history": test_accuracy_history,
}
print("device update order:", REALIZATION_MODE)
print("completed device-optimization steps:", completed_steps.tolist())
print("device optimization complete:", realization_complete)
device update order: synchronized completed device-optimization steps: [40, 40, 40] device optimization complete: True
[42]:
def padded_layer_matrices(matrices):
padded = np.full((N_LAYERS, max(LAYER_OUTPUT_COUNTS), N_CHANNELS), np.nan + 0j)
for layer, matrix in enumerate(matrices):
padded[layer, : LAYER_OUTPUT_COUNTS[layer]] = matrix
return padded
def cloud_payload():
transmissions, reflections, parameters, densities = [], [], [], []
for layer in range(N_LAYERS):
transmissions.append(np.asarray(optimizer_forward[layer]["transmission"]))
reflections.append(np.asarray(optimizer_forward[layer]["reflection"]))
layer_parameters = np.asarray(optimizer_forward[layer]["params"])
parameters.append(layer_parameters)
densities.append(
np.asarray(designs[layer].design_region.material_density(layer_parameters))
)
return {
"realization_objective": REALIZATION_OBJECTIVE,
"realization_target_tag": REALIZATION_TARGET_TAG,
"layer_input_counts": np.asarray(LAYER_INPUT_COUNTS),
"layer_output_counts": np.asarray(LAYER_OUTPUT_COUNTS),
"n_splitter_leaves": N_SPLITTER_LEAVES,
"active_leaf_indices": np.asarray(ACTIVE_LEAF_INDICES),
"dump_leaf_indices": np.asarray(DUMP_LEAF_INDICES),
"fanout_rule": FANOUT_RULE,
"terminator_model": TERMINATOR_MODEL,
"terminator_reflection": TERMINATOR_REFLECTION,
"cladding_model": CLADDING_MODEL,
"background_index": N_EFF_TE_BACKGROUND,
"optimizer_direction": OPTIMIZER_DIRECTION,
"active_mask_rule": ACTIVE_MASK_RULE,
"realization_mode": REALIZATION_MODE,
"transmission": padded_layer_matrices(transmissions),
"reflection": np.stack(reflections),
"target": padded_layer_matrices(targets.values()),
"parameters": np.stack(parameters),
"density": np.stack(densities),
"iterations": realization_progress["iterations"],
"transmission_history": transmission_history,
"reflection_history": reflection_history,
"transmission_error_history": transmission_error_history,
"reflection_error_history": reflection_error_history,
"realization_loss_history": realization_loss_history,
"learning_rate_history": learning_rate_history,
"fabrication_weight_history": fabrication_weight_history,
"fabrication_penalty_history": fabrication_penalty_history,
"weighted_fabrication_penalty_history": weighted_fabrication_penalty_history,
"scheduled_constrained_loss_history": scheduled_constrained_loss_history,
"fixed_reference_loss_history": fixed_reference_loss_history,
"test_accuracy_history": test_accuracy_history,
"completed_steps": np.asarray(completed_steps),
"actual_cloud_cost_flexcredits": np.asarray(
optimization_run_summary.get("actual FlexCredits", np.nan)
),
}
result_data = (
cloud_payload()
if realization_complete and len(optimizer_forward) == N_LAYERS
else None
)
print(
"analysis results:",
"available" if result_data is not None else "awaiting saved device simulations",
)
analysis results: available
Analyze the Realized Classifier¶
After optimization, each layer provides a complex transmission matrix $T_\ell$, a 6 × 6 reflection matrix $R_\ell$, and a material-density pattern. We first compare the electromagnetic response with the trained target:
$$ E_{T,\ell}=\|T_\ell-T_{\ell,\mathrm{target}}\|_F, \qquad P_{\mathrm{refl},\ell j}=\sum_i|R_{\ell,ij}|^2. $$
The transmission error includes both amplitude and phase differences. The reflected power measures how much of each incident channel returns toward the six input ports. We also calculate the insertion loss of the realized singular channels from $-20\log_{10}(\sigma)$.
Next, we replace the trained matrices with the realized matrices:
$$ a_{\mathrm{out}}= T_3D(z)T_2D(z)T_1D(z)a_{\mathrm{in}}, \qquad p=|a_{\mathrm{out}}|^2. $$
The propagation includes the routed PhotonForge input and output responses. We use the training partition to refit one positive detector scale $\gamma$ and one common output phase. The phase leaves $p$ unchanged, while $\gamma$ adjusts the confidence of the class logits. The held-out images are used only for the final comparison.
The figures show the optimization histories, target and realized matrices, reflection, insertion loss, device patterns, confusion matrices, and channel fields for representative images.
[43]:
def fixed_logits(matrices, features, final_phase, log_gamma):
fields = physical_forward(matrices, features)
fields = fields * jnp.exp(1j * final_phase)
scores = jnp.abs(fields) ** 2
return jnp.exp(log_gamma) * scores / ROUTE_POWER_REFERENCE, scores
def metric_pair(logits, labels):
loss = optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == labels)
return float(loss), float(accuracy)
reconstruction = None
if result_data is not None:
realized_matrices = tuple(
np.asarray(result_data["transmission"])[layer, : LAYER_OUTPUT_COUNTS[layer]]
for layer in range(N_LAYERS)
)
realized_reflections = np.asarray(result_data["reflection"])
calibration = {
"final_phase": jnp.asarray(0.0),
"log_gamma": params["log_gamma"],
}
calibration_optimizer = optax.adam(CALIBRATION_LR)
calibration_state = calibration_optimizer.init(calibration)
if result_data is not None:
@jax.jit
def calibrate(calibration, state):
def loss_fn(values):
logits, _ = fixed_logits(
realized_matrices,
jnp.asarray(train_features),
values["final_phase"],
values["log_gamma"],
)
return optax.softmax_cross_entropy_with_integer_labels(
logits, jnp.asarray(train_labels)
).mean()
loss, gradient = jax.value_and_grad(loss_fn)(calibration)
updates, state = calibration_optimizer.update(gradient, state, calibration)
return optax.apply_updates(calibration, updates), state, loss
calibration_history = []
for _ in range(CALIBRATION_STEPS):
calibration, calibration_state, loss = calibrate(calibration, calibration_state)
calibration_history.append(float(loss))
if result_data is not None:
repeated_logits, _, _, _ = logits_and_scores(params, jnp.asarray(test_features))
uncalibrated_logits, _ = fixed_logits(
realized_matrices,
jnp.asarray(test_features),
jnp.asarray(0.0),
params["log_gamma"],
)
calibrated_logits, calibrated_scores = fixed_logits(
realized_matrices,
jnp.asarray(test_features),
calibration["final_phase"],
calibration["log_gamma"],
)
reconstruction = {
"single-encoding ideal matrix": (
float(single_test_loss),
float(single_test_accuracy),
),
"repeated-encoding ideal matrix": metric_pair(repeated_logits, test_labels),
"repeated-encoding realized, uncalibrated": metric_pair(
uncalibrated_logits, test_labels
),
"repeated-encoding realized, calibrated": metric_pair(
calibrated_logits, test_labels
),
}
for name, (loss, accuracy) in reconstruction.items():
print(f"{name}: cross-entropy={loss:.6f}, accuracy={accuracy:.2%}")
single-encoding ideal matrix: cross-entropy=0.333241, accuracy=86.89% repeated-encoding ideal matrix: cross-entropy=0.270203, accuracy=90.52% repeated-encoding realized, uncalibrated: cross-entropy=0.324838, accuracy=89.05% repeated-encoding realized, calibrated: cross-entropy=0.319369, accuracy=89.05%
[44]:
if result_data is not None:
density_layers = [
np.squeeze(np.asarray(result_data["density"])[layer])
for layer in range(N_LAYERS)
]
fig, axes = plt.subplots(
1, N_LAYERS, figsize=(3.75 * N_LAYERS, 3.2), constrained_layout=True
)
for layer, density in enumerate(density_layers):
axes[layer].imshow(
density.T,
origin="lower",
extent=(-DESIGN_HALF, DESIGN_HALF, -DESIGN_HALF, DESIGN_HALF),
cmap="gray_r",
vmin=0,
vmax=1,
interpolation="nearest",
)
axes[layer].set(
xlabel=r"$x$ ($\mu$m)", ylabel=r"$y$ ($\mu$m)", title=f"Layer {layer + 1}"
)
contours = gdstk.contour(density.T, 0.5, PIXEL, PIXEL / 10, layer=3)
x_span, y_span = (density.shape[0] - 1) * PIXEL, (density.shape[1] - 1) * PIXEL
for polygon in contours:
polygon.translate(-x_span / 2, -y_span / 2)
silicon = (
gdstk.boolean(
fixed_polygons(LAYER_OUTPUT_COUNTS[layer]) + contours,
[],
"or",
layer=3,
precision=0.001,
)
or []
)
library = gdstk.Library(unit=1e-6, precision=1e-9)
cell = library.new_cell(f"INVERSE_DESIGNED_LAYER_{layer + 1}")
cell.add(*silicon)
library.write_gds(LAYOUT_DIR / f"inverse_designed_layer_{layer + 1}.gds")
library.write_oas(LAYOUT_DIR / f"inverse_designed_layer_{layer + 1}.oas")
fig.savefig(
FIGURE_DIR / "realized_device_densities.png", dpi=180, bbox_inches="tight"
)
plt.show()
progress_source = (
result_data
if result_data is not None and "test_accuracy_history" in result_data
else realization_progress
)
progress_accuracy = np.asarray(progress_source["test_accuracy_history"], dtype=float)
accuracy_recalculated_from_matrices = False
if not np.any(np.isfinite(progress_accuracy)):
transmission_values = np.asarray(
progress_source["transmission_history"], dtype=complex
)
progress_accuracy = np.full(transmission_values.shape[:2], np.nan)
for step in range(len(transmission_values)):
step_matrices = tuple(
transmission_values[step, layer, : LAYER_OUTPUT_COUNTS[layer]]
for layer in range(N_LAYERS)
)
if all(np.isfinite(matrix).all() for matrix in step_matrices):
progress_accuracy[step, 0] = uncalibrated_test_accuracy(step_matrices)
accuracy_recalculated_from_matrices = np.any(np.isfinite(progress_accuracy))
has_realization_history = np.any(
np.isfinite(np.asarray(progress_source["transmission_error_history"], dtype=float))
)
if has_realization_history:
iterations = np.asarray(progress_source["iterations"], dtype=int)
fig, axes = plt.subplots(3, 2, figsize=(9.2, 9.2), constrained_layout=True)
metric_panels = (
("transmission_error_history", "Frobenius norm", "Transmission target error"),
("reflection_error_history", "Frobenius norm", "Reflection"),
("realization_loss_history", "optical loss", "Transmission + reflection error"),
("fabrication_penalty_history", "unit-weight penalty", "Fabrication penalty"),
("fixed_reference_loss_history", "reference loss", "Optical + 0.8 fabrication"),
)
colors = plt.cm.tab10(np.arange(N_LAYERS))
for axis, (key, ylabel, title) in zip(axes.flat[:5], metric_panels):
values = np.asarray(progress_source[key], dtype=float)
valid_rows = np.any(np.isfinite(values), axis=1)
for layer in range(N_LAYERS):
valid = np.isfinite(values[:, layer])
axis.plot(
iterations[valid],
values[valid, layer],
color=colors[layer],
linewidth=1.1,
alpha=0.75,
label=f"layer {layer + 1}",
)
mean_values = np.full(len(iterations), np.nan)
mean_values[valid_rows] = np.nanmean(values[valid_rows], axis=1)
axis.plot(
iterations[valid_rows],
mean_values[valid_rows],
color="black",
linewidth=2.2,
label="mean",
)
axis.set(
xlabel="device-optimization step", ylabel=ylabel, title=title, yscale="log"
)
axes[0, 0].legend(frameon=False, ncol=2)
progress_mode = str(
np.asarray(
progress_source["mode"]
if "mode" in progress_source
else progress_source["realization_mode"]
).item()
)
accuracy_axis = axes[2, 1]
if progress_mode == "sequential" and not accuracy_recalculated_from_matrices:
for layer in range(N_LAYERS):
valid = np.isfinite(progress_accuracy[:, layer])
accuracy_axis.plot(
iterations[valid],
100 * progress_accuracy[valid, layer],
color=colors[layer],
linewidth=1.8,
marker="o",
markersize=3,
label=f"after layer {layer + 1} update",
)
else:
valid = np.isfinite(progress_accuracy[:, 0])
accuracy_axis.plot(
iterations[valid],
100 * progress_accuracy[valid, 0],
color="#1f77b4",
linewidth=2.0,
marker="o",
markersize=3,
label=(
"matched saved matrices"
if accuracy_recalculated_from_matrices
else "reconstructed stack"
),
)
accuracy_axis.axhline(
100 * test_accuracy,
color="black",
linestyle="--",
linewidth=1.2,
label="trained target matrices",
)
accuracy_axis.set(
xlabel="device-optimization step",
ylabel="test accuracy (%)",
title="Uncalibrated network accuracy",
)
accuracy_axis.legend(frameon=False)
fig.savefig(FIGURE_DIR / "realization_progress.png", dpi=180, bbox_inches="tight")
plt.show()
if result_data is not None:
fig, axes = plt.subplots(
N_LAYERS, 4, figsize=(11.2, 2.7 * N_LAYERS), constrained_layout=True
)
for layer in range(N_LAYERS):
images = (
np.abs(result_data["target"][layer, : LAYER_OUTPUT_COUNTS[layer]]),
np.abs(realized_matrices[layer]),
np.angle(result_data["target"][layer, : LAYER_OUTPUT_COUNTS[layer]]),
np.angle(realized_matrices[layer]),
)
titles = (
"target amplitude",
"realized amplitude",
"target phase",
"realized phase",
)
for column, (image, title) in enumerate(zip(images, titles)):
cmap = "viridis" if column < 2 else "twilight"
limits = {} if column < 2 else {"vmin": -np.pi, "vmax": np.pi}
axes[layer, column].imshow(image, cmap=cmap, aspect="equal", **limits)
axes[layer, column].set(
title=f"Layer {layer + 1}: {title}", xlabel="input", ylabel="output"
)
fig.savefig(
FIGURE_DIR / "realized_target_response.png", dpi=180, bbox_inches="tight"
)
plt.show()
[45]:
reflection_summary = None
if result_data is not None:
reflected_power = np.sum(np.abs(realized_reflections) ** 2, axis=1)
reflected_power_db = 10 * np.log10(np.maximum(reflected_power, 1e-12))
reflection_summary = {
f"layer {layer + 1}": {
"mean reflected power dB": float(
10 * np.log10(max(np.mean(reflected_power[layer]), 1e-12))
),
"median reflected power dB": float(
10 * np.log10(max(np.median(reflected_power[layer]), 1e-12))
),
"worst-case reflected power dB": float(
10 * np.log10(max(np.max(reflected_power[layer]), 1e-12))
),
}
for layer in range(N_LAYERS)
}
for name, values in reflection_summary.items():
print(
f"{name}: mean={values['mean reflected power dB']:.2f} dB, "
f"worst={values['worst-case reflected power dB']:.2f} dB"
)
insertion_loss_db = [
-20 * np.log10(np.maximum(np.linalg.svd(matrix, compute_uv=False), 1e-12))
for matrix in realized_matrices
]
fig, axes = plt.subplots(1, 2, figsize=(8.4, 3.2), constrained_layout=True)
axes[0].boxplot(
[reflected_power_db[layer] for layer in range(N_LAYERS)],
tick_labels=[str(layer + 1) for layer in range(N_LAYERS)],
)
for layer in range(N_LAYERS):
singular_losses = np.sort(insertion_loss_db[layer])
axes[1].plot(
np.arange(1, len(singular_losses) + 1),
singular_losses,
marker=".",
label=f"layer {layer + 1}",
)
axes[0].set(
xlabel="layer", ylabel="total reflected power (dB)", title="Input reflection"
)
axes[1].set(
xlabel="singular channel",
ylabel="insertion loss (dB)",
title="Realized transmission",
)
axes[1].legend(frameon=False)
fig.savefig(
FIGURE_DIR / "reflection_insertion_loss.png", dpi=180, bbox_inches="tight"
)
plt.show()
def confusion(labels, predictions):
result = np.zeros((N_CLASSES, N_CLASSES), dtype=int)
np.add.at(result, (np.asarray(labels), np.asarray(predictions)), 1)
return result
def normalized_confusion(labels, predictions):
counts = confusion(labels, predictions)
return counts / np.maximum(counts.sum(axis=1, keepdims=True), 1)
def repeated_field_trace(matrices, features):
fields = jnp.ones_like(features, dtype=jnp.complex64)
phase = jnp.exp(1j * PHASE_SCALE * features)
trace = []
for layer in range(N_LAYERS):
fields = fields * route_inputs_jax[layer] * phase
fields = fields @ jnp.asarray(matrices[layer]).T
trace.append(fields * route_output_jax if layer == N_LAYERS - 1 else fields)
return tuple(trace)
def highest_margin_correct_sample(scores, labels, predictions, class_label):
labels = np.asarray(labels)
predictions = np.asarray(predictions)
scores = np.asarray(scores)
candidates = np.flatnonzero((labels == class_label) & (predictions == class_label))
if len(candidates) == 0:
return int(np.flatnonzero(labels == class_label)[0])
class_scores = scores[candidates]
correct = class_scores[:, class_label]
competitors = np.max(np.delete(class_scores, class_label, axis=1), axis=1)
margins = (correct - competitors) / np.maximum(class_scores.sum(axis=1), 1e-12)
return int(candidates[np.argmax(margins)])
result_figure_summary = {
"repeated-encoding training": "repeated_encoding_training.png",
"matrix-error response": "repeated_encoding_matrix_error.png",
"singular amplitudes": "repeated_encoding_singular_amplitudes.png",
"encoding comparison": "encoding_comparison.png",
"target matrices": "repeated_encoding_target_matrices.png",
"local simulation domains": "tidy3d_preflight_domain_permittivity.png",
}
if result_data is not None:
single_scores = np.asarray(
single_encoding_scores(
layer_matrices(single_params), jnp.asarray(test_features)
)
)
repeated_scores = np.asarray(
jnp.abs(
optical_forward(jnp.asarray(trained_matrices), jnp.asarray(test_features))
)
** 2
)
realized_scores = np.asarray(calibrated_scores)
prediction_sets = (
np.argmax(single_scores, axis=1),
np.argmax(repeated_scores, axis=1),
np.argmax(realized_scores, axis=1),
)
confusion_titles = (
"Encode once",
"Repeat encoding: ideal matrices",
"Repeat encoding: realized devices",
)
fig, axes = plt.subplots(1, 3, figsize=(10.4, 3.25), constrained_layout=True)
for axis, predictions, title in zip(axes, prediction_sets, confusion_titles):
matrix = normalized_confusion(test_labels, predictions)
image = axis.imshow(matrix, cmap="Blues", vmin=0, vmax=1)
for row in range(N_CLASSES):
for column in range(N_CLASSES):
color = "white" if matrix[row, column] > 0.55 else "black"
axis.text(
column,
row,
f"{100 * matrix[row, column]:.1f}",
ha="center",
va="center",
color=color,
fontsize=8,
)
accuracy = np.mean(predictions == np.asarray(test_labels))
axis.set(
xticks=np.arange(N_CLASSES),
yticks=np.arange(N_CLASSES),
xlabel="predicted digit",
ylabel="true digit",
title=f"{title}\naccuracy {100 * accuracy:.2f}%",
)
fig.colorbar(image, ax=axes, label="fraction of true class", shrink=0.82)
fig.savefig(
FIGURE_DIR / "encoding_confusion_comparison.png", dpi=180, bbox_inches="tight"
)
plt.show()
sample_indices = tuple(
highest_margin_correct_sample(
realized_scores, test_labels, prediction_sets[2], class_label
)
for class_label in CLASS_LABELS
)
sample_traces = tuple(
tuple(
np.asarray(values[0])
for values in repeated_field_trace(
realized_matrices, jnp.asarray(test_features[index : index + 1])
)
)
for index in sample_indices
)
phase_map = plt.get_cmap("twilight")
fig, axes = plt.subplots(
N_CLASSES,
N_LAYERS + 1,
figsize=(3.0 * (N_LAYERS + 1), 2.15 * N_CLASSES),
constrained_layout=True,
)
axes = np.atleast_2d(axes)
for row, (sample_index, trace) in enumerate(zip(sample_indices, sample_traces)):
axes[row, 0].imshow(test_images[sample_index], cmap="gray", vmin=0, vmax=255)
axes[row, 0].set(
title=f"digit {test_labels[sample_index]}\nsample {sample_index}",
xticks=[],
yticks=[],
)
common_limit = 1.08 * max(np.max(np.abs(values)) for values in trace)
for layer, values in enumerate(trace):
amplitudes = np.abs(values)
phases = np.angle(values)
channel_index = np.arange(len(values))
axes[row, layer + 1].bar(
channel_index,
amplitudes,
color=phase_map((phases + np.pi) / (2 * np.pi)),
edgecolor="black",
linewidth=0.25,
)
axes[row, layer + 1].set(
xticks=channel_index,
xlabel="channel" if row == N_CLASSES - 1 else None,
ylabel="amplitude",
ylim=(0, common_limit),
title=(
f"after layer {layer + 1}"
if layer < N_LAYERS - 1
else f"detector fields\npredicted {prediction_sets[2][sample_index]}"
),
)
phase_scale = plt.cm.ScalarMappable(
norm=plt.Normalize(-np.pi, np.pi), cmap=phase_map
)
phase_scale.set_array([])
fig.colorbar(
phase_scale,
ax=axes[:, 1:],
location="bottom",
shrink=0.55,
pad=0.04,
label="optical phase (rad)",
)
fig.savefig(
FIGURE_DIR / "sample_channel_field_evolution.png", dpi=180, bbox_inches="tight"
)
plt.show()
result_figure_summary.update(
{
"optimization evolution": "realization_progress.png",
"device densities": "realized_device_densities.png",
"target and realized matrices": "realized_target_response.png",
"reflection and insertion loss": "reflection_insertion_loss.png",
"confusion matrices": "encoding_confusion_comparison.png",
"sample channel fields": "sample_channel_field_evolution.png",
"field sample indices": list(sample_indices),
"field sample selection": "largest normalized correct-class margin",
}
)
layer 1: mean=-49.99 dB, worst=-48.36 dB layer 2: mean=-49.28 dB, worst=-44.90 dB layer 3: mean=-50.58 dB, worst=-49.26 dB
Summary¶
In this notebook, we trained a six-channel classifier for MNIST digits 0–3 and compared two ways of applying phase encoding. The repeated model applies the image phase before all three passive transformations, while the baseline applies it only once. We transferred the trained transformations to three inverse-designed silicon devices and included the routed PhotonForge circuit in the final classifier.
Repeated encoding performs better than the single-encoding baseline. Most of this advantage remains when the trained matrices are replaced by the Tidy3D transmission matrices. The matrix, reflection, and insertion-loss plots show where the realized devices depart from their targets.
The device optimization uses a two-dimensional effective-index model and therefore misses out-of-plane scattering. Before fabrication, the devices should be optimized in three dimensions or fine-tuned from the two-dimensional designs.
[46]:
accuracy_series = (
test_accuracy_history[:, 0]
if REALIZATION_MODE == "synchronized"
else test_accuracy_history.T.reshape(-1)
)
valid_progress = np.flatnonzero(np.isfinite(accuracy_series))
latest_uncalibrated_accuracy = (
float(reconstruction["repeated-encoding realized, uncalibrated"][1])
if reconstruction is not None
else float(accuracy_series[valid_progress[-1]])
if len(valid_progress)
else None
)
print(
f"classifier: {N_CLASSES} classes, {N_CHANNELS} channels, "
f"{DESIGN_EDGE_UM:g} µm devices"
)
print(
f"realization: mode={REALIZATION_MODE}, "
f"completed={completed_steps.tolist()}, complete={realization_complete}"
)
if latest_uncalibrated_accuracy is not None:
print(f"latest uncalibrated test accuracy: {latest_uncalibrated_accuracy:.2%}")
if reflection_summary is not None:
worst_reflection = max(
values["worst-case reflected power dB"]
for values in reflection_summary.values()
)
print(f"worst reflected power: {worst_reflection:.2f} dB")
print(f"progress cache: {PROGRESS_PATH.relative_to(REPOSITORY_ROOT)}")
saved_figures = [
value for value in result_figure_summary.values() if isinstance(value, str)
]
print("saved figures:", ", ".join(saved_figures))
classifier: 4 classes, 6 channels, 12.6 µm devices realization: mode=synchronized, completed=[40, 40, 40], complete=True latest uncalibrated test accuracy: 89.05% worst reflected power: -44.90 dB progress cache: output/deep/progress.npz saved figures: repeated_encoding_training.png, repeated_encoding_matrix_error.png, repeated_encoding_singular_amplitudes.png, encoding_comparison.png, repeated_encoding_target_matrices.png, tidy3d_preflight_domain_permittivity.png, realization_progress.png, realized_device_densities.png, realized_target_response.png, reflection_insertion_loss.png, encoding_confusion_comparison.png, sample_channel_field_evolution.png, largest normalized correct-class margin
References¶
-
A. M. I. Muda and U. Teğin, “Deep inverse-designed nanophotonic processors with structural nonlinearity from repeated phase encoding,” arXiv:2608.02094 [physics.optics] (2026). arXiv:2608.02094
-
M. Yildirim, N. U. Dinc, I. Oguz, D. Psaltis, and C. Moser, “Nonlinear processing with linear optics,” Nature Photonics 18, 1076–1082 (2024). doi:10.1038/s41566-024-01494-z
-
Y. LeCun, L. Bottou, Y. Bengio, and P. Haffner, “Gradient-based learning applied to document recognition,” Proceedings of the IEEE 86, 2278–2324 (1998). doi:10.1109/5.726791
-
M. Hammer and O. V. Ivanova, “Effective index approximations of photonic crystal slabs: a 2-to-1-D assessment,” Optical and Quantum Electronics 41, 267–283 (2009). doi:10.1007/s11082-009-9349-3
-
Flexcompute, Effective index approximation in Tidy3D.
-
V. Nikkhah et al., “Inverse-designed low-index-contrast structures on a silicon photonics platform for vector–matrix multiplication,” Nature Photonics 18, 501–508 (2024). doi:10.1038/s41566-024-01394-2
-
SiEPIC, SiEPIC EBeam PDK.
TERMS & CONDITIONS
How the process works:
Try this real URL to see an example: