Inverse design optimization of a waveguide taper#
Note: native autograd support is an experimental feature in tidy3d 2.7. To see the original implementation of this notebook using jax and the adjoint plugin, refer to this notebook.
In this notebook, we will show how to use tidy3d to optimize a taper with respect to the boundaries of a structure defined using a PolySlab.
We will apply this capability to design a non-adiabatic waveguide taper between a narrow and wide waveguide, based loosely on Michaels, Andrew, and Eli Yablonovitch. "Leveraging continuous material averaging for inverse electromagnetic design." Optics express 26.24 (2018): 31717-31737.
We start by importing our typical python packages, plus tidy3d and autograd.
[1]:
import matplotlib.pylab as plt
import numpy as np
import autograd as ag
import autograd.numpy as anp
import tidy3d as td
import tidy3d.web as web
Set up#
Next we will define some basic parameters of the waveguide, such as the input and output waveguide dimensions, taper width, and taper length.
[2]:
wavelength = 1.0
freq0 = td.C_0 / wavelength
wg_width_in = 0.5 * wavelength
wg_width_out = 5.0 * wavelength
wg_medium = td.material_library["Si3N4"]["Philipp1973Sellmeier"]
wg_length = 1 * wavelength
taper_length = 10.0
spc_pml = 1.5 * wavelength
Lx = wg_length + taper_length + wg_length
Ly = spc_pml + max(wg_width_in, wg_width_out) + spc_pml
Our taper is defined as a set of num_points connected vertices in a polygon. We define the fixed x positions of each vertex and then construct the y positions for the starting device (linear taper).
[3]:
num_points = 101
x_start = -taper_length / 2
x_end = +taper_length / 2
xs = np.linspace(x_start, x_end, num_points)
ys0 = (wg_width_in + (wg_width_out - wg_width_in) * (xs - x_start) / (x_end - x_start)) / 2.0
Let’s plot these points to make sure they seem reasonable.
[4]:
plt.plot(xs, +ys0, "ko-")
plt.plot(xs, -ys0, "ko-")
plt.xlabel("x")
plt.ylabel("y")
plt.title("taper points")
plt.show()
Let’s wrap this in a function that constructs these points and then creates a td.PolySlab for use in the td.Simulation.
[5]:
def make_taper(ys) -> td.PolySlab:
"""Create a PolySlab for the taper based on the supplied y values."""
vertices = anp.concatenate(
[
anp.column_stack((xs, ys)),
anp.column_stack((xs[::-1], -ys[::-1])),
]
)
return td.PolySlab(vertices=vertices, slab_bounds=(-td.inf, +td.inf), axis=2)
Now we’ll call this function and plot the geometry for a sanity check.
[6]:
# sanity check to see the polyslab
taper_geo = make_taper(ys0)
ax = taper_geo.plot(z=0)
Next, let’s write a function that generates a td.Simulation given a set of y coordinates for the taper, including the monitors, sources, and waveguide geometries.
[7]:
def make_sim(ys, include_field_mnt: bool = False) -> td.Simulation:
"""Make a td.Simulation containing the taper."""
wg_in_box = td.Box.from_bounds(
rmin=(-Lx, -wg_width_in / 2, -td.inf),
rmax=(-Lx / 2 + wg_length + 0.01, +wg_width_in / 2, +td.inf),
)
wg_out_box = td.Box.from_bounds(
rmin=(+Lx / 2 - wg_length - 0.01, -wg_width_out / 2, -td.inf),
rmax=(+Lx, +wg_width_out / 2, +td.inf),
)
taper_geo = make_taper(ys)
wg_in = td.Structure(geometry=wg_in_box, medium=wg_medium)
wg_out = td.Structure(geometry=wg_out_box, medium=wg_medium)
taper = td.Structure(geometry=taper_geo, medium=wg_medium)
mode_source = td.ModeSource(
center=(-Lx / 2 + wg_length / 2, 0, 0),
size=(0, td.inf, td.inf),
source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10),
direction="+",
)
mode_monitor = td.ModeMonitor(
center=(+Lx / 2 - wg_length / 2, 0, 0),
size=(0, td.inf, td.inf),
freqs=[freq0],
mode_spec=td.ModeSpec(),
name="mode",
)
field_monitor = td.FieldMonitor(
center=(0, 0, 0),
size=(td.inf, td.inf, 0),
freqs=[freq0],
name="field",
)
return td.Simulation(
size=(Lx, Ly, 0),
structures=[taper, wg_in, wg_out],
monitors=[mode_monitor, field_monitor] if include_field_mnt else [mode_monitor],
sources=[mode_source],
run_time=100 / freq0,
grid_spec=td.GridSpec.auto(min_steps_per_wvl=20),
boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=False),
)
[8]:
sim = make_sim(ys0)
ax = sim.plot(z=0)
Defining Objective#
Now that we have a function to create our td.Simulation, we need to define our objective function.
We will try to optimize the power transmission into the fundamental mode on the output waveguide, so we write a function to postprocess a td.SimulationData to give this result.
[9]:
def measure_transmission(sim_data: td.SimulationData) -> float:
"""Measure the first order transmission."""
amp_data = sim_data["mode"].amps
amp = amp_data.sel(f=freq0, direction="+", mode_index=0).values
return anp.sum(abs(amp) ** 2)
Next, we will define a few convenience functions to generate our taper y values passed on our objective function parameters.
We define a set of parameters that can range from -infinity to +infinity, but project onto the range [wg_width_out and wg_width_in] through a tanh() function.
We do this to constrain the taper values within this range in a smooth and differentiable way.
We also write an inverse function to get the parameters given a set of desired y values and assert that this function works properly.
[10]:
def get_ys(parameters: np.ndarray) -> np.ndarray:
"""Convert arbitrary parameters to y values for the vertices (parameter (-inf, inf) -> wg width of (wg_width_out, wg_width_in)."""
params_between_0_1 = (anp.tanh(np.pi * parameters) + 1.0) / 2.0
params_scaled = params_between_0_1 * (wg_width_out - wg_width_in) / 2.0
ys = params_scaled + wg_width_in / 2
return ys
def get_params(ys: np.ndarray) -> np.ndarray:
"""inverse of above, get parameters from ys"""
params_scaled = ys - wg_width_in / 2
params_between_0_1 = 2 * params_scaled / (wg_width_out - wg_width_in)
tanh_pi_params = 2 * params_between_0_1 - 1
params = np.arctanh(tanh_pi_params) / np.pi
return params
# assert that the inverse function works properly
params_test = 2 * (np.random.random((10,)) - 0.5)
ys_test = get_ys(params_test)
assert np.allclose(get_params(ys_test), params_test)
We then make a function that simply wraps our previously defined functions to generate a Simulation given some parameters.
[11]:
def make_sim_params(parameters: np.ndarray, include_field_mnt: bool = False) -> td.Simulation:
"""Make the simulation out of raw parameters."""
ys = get_ys(parameters)
return make_sim(ys, include_field_mnt=include_field_mnt)
Smoothness Penalty#
It is important to ensure that the final device does not contain feature sizes below a minimum radius of curvature, otherwise there could be considerable difficulty in fabricating the device reliably.
For this, the autograd plugin introduces a penalty function that approximates the radius of curvature about each vertex and introduces a significant penalty to the objective function if that value is below a minimum radius of curvature.
The local radii are determined by fitting a quadratic Bezier curve to each set of 3 adjacent points in the taper and analytically computing the radius of curvature from that curve fit. The details of this calculation are described in the paper linked at the introduction of this notebook.
[12]:
from tidy3d.plugins.autograd.invdes import make_curvature_penalty
curvature_penalty = make_curvature_penalty(min_radius=0.15)
def penalty(params: np.ndarray) -> float:
"""Compute penalty for a set of parameters."""
ys = get_ys(params)
points = anp.array([xs, ys]).T
return curvature_penalty(points)
Initial Starting Design#
As our initial design, we take a linear taper between the two waveguides.
[13]:
# desired ys
ys0 = np.linspace(wg_width_in / 2 + 0.001, wg_width_out / 2 - 0.001, num_points)
# corresponding parameters
params0 = get_params(ys0)
# make the simulation corresponding to these parameters
sim = make_sim_params(params0, include_field_mnt=True)
ax = sim.plot(z=0)
Lets get the penalty value corresponding to this design, which is 0 because the edges are straight (infinite radius of curvature).
[14]:
penalty_value = penalty(params0)
print(f"starting penalty = {float(penalty_value):.2e}")
starting penalty = 0.00e+00
Finally, let’s run this simulation to get a feeling for the initial device performance.
[15]:
sim_data = web.run(sim, task_name="taper fields")
↓ simulation_data.hdf5.gz ━━━━━━━━━━━━━ 100.0% • 6.4/6.4 MB • 9.3 MB/s • 0:00:00
15:21:22 EDT loading simulation from simulation_data.hdf5
[16]:
f, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 3))
sim_data.plot_field(field_monitor_name="field", field_name="Ez", val="real", ax=ax1)
sim_data.plot_field(field_monitor_name="field", field_name="E", val="abs", ax=ax2)
plt.show()
Gradient-based Optimization#
Now that we have our design and post processing functions set up, we are finally ready to put everything together to start optimizing our device with inverse design.
We first set up an objective function that takes the parameters, sets up and runs the simulation, and returns the transmission minus the penalty of the parameters.
[17]:
def objective(parameters: np.ndarray, verbose: bool = False) -> float:
"""Construct simulation, run, and measure transmission."""
sim = make_sim_params(parameters, include_field_mnt=False)
sim_data = web.run(sim, task_name="autograd taper", verbose=verbose)
return measure_transmission(sim_data) - penalty(parameters)
To test our our objective, we will use autograd to make and run a function that returns the objective value and its gradient.
[18]:
grad_fn = ag.value_and_grad(objective)
[19]:
val, grad = grad_fn(params0, verbose=True)
↓ simulation_data.hdf5.gz ━━━━━━━━━━━━━ 100.0% • 3.6/3.6 MB • 8.0 MB/s • 0:00:00
15:21:39 EDT loading simulation from simulation_data.hdf5
[20]:
print(f"objective = {val:.2e}")
print(f"gradient = {np.nan_to_num(grad)}")
objective = 7.22e-01
gradient = [ 5.55552699e-03 2.07151937e-01 1.12387950e-01 -2.49679408e-01
-5.26063840e-01 -7.04767199e-01 -7.95133051e-01 -6.55073586e-01
-4.45962339e-01 -1.85182400e-01 2.03505444e-01 3.12891399e-01
1.06171217e-01 3.26962915e-02 1.87562472e-01 4.63636353e-01
6.73876756e-01 7.81965246e-01 8.76915755e-01 7.94489090e-01
6.24658704e-01 4.78218133e-01 2.63204648e-01 1.99069762e-01
2.16589987e-01 1.06222319e-01 2.06259666e-02 -5.05152940e-02
-1.04464768e-01 -9.21181150e-02 -6.68642870e-02 9.45449182e-03
6.54451648e-02 7.49362798e-02 8.56240758e-02 6.68995577e-02
8.04351403e-02 9.66085109e-02 5.49320537e-02 2.19974811e-02
-1.21026713e-02 -6.21566856e-02 -8.32930680e-02 -8.23885410e-02
-8.64080696e-02 -8.73485769e-02 -8.25397644e-02 -7.44381304e-02
-5.63182726e-02 -4.95586481e-02 -5.18839668e-02 -5.89320578e-02
-7.57361310e-02 -8.10764917e-02 -9.17143458e-02 -9.85566004e-02
-9.59841372e-02 -1.03919252e-01 -9.26992535e-02 -8.53246014e-02
-8.90348519e-02 -8.24240868e-02 -8.94583036e-02 -8.90254691e-02
-8.70291888e-02 -9.21943049e-02 -8.38742218e-02 -8.32291695e-02
-7.90023105e-02 -7.13772841e-02 -7.40488232e-02 -7.02907159e-02
-6.78238179e-02 -6.48513894e-02 -5.91945169e-02 -5.56949243e-02
-5.03000719e-02 -4.85293919e-02 -4.61946529e-02 -4.27182970e-02
-3.96951477e-02 -3.52495053e-02 -3.19343727e-02 -2.84388821e-02
-2.60513345e-02 -2.27829572e-02 -1.98223867e-02 -1.79678402e-02
-1.46587081e-02 -1.33549190e-02 -1.12190773e-02 -8.80083441e-03
-7.42354649e-03 -5.22990871e-03 -4.35521402e-03 -2.90697927e-03
-1.72076833e-03 -1.28009940e-03 -4.32479310e-04 -1.92935342e-04
-6.55643387e-06]
Now we can run optax’s Adam optimization loop with this value_and_grad() function. See tutorial 3 for more details on the implementation.
[21]:
import optax
# turn off warnings to reduce verbosity
td.config.logging_level = "ERROR"
# hyperparameters
num_steps = 21
learning_rate = 0.01
# initialize adam optimizer with starting parameters
params = np.array(params0).copy()
optimizer = optax.adam(learning_rate=learning_rate)
opt_state = optimizer.init(params)
# store history
objective_history = []
param_history = [params]
for i in range(num_steps):
# compute gradient and current objective funciton value
value, gradient = grad_fn(params, verbose=False)
# convert nan to 0 (infinite radius of curvature) and multiply all by -1 to maximize obj_fn.
# outputs
print(f"step = {i + 1}")
print(f"\tJ = {value:.4e}")
print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}")
# compute and apply updates to the optimizer based on gradient
updates, opt_state = optimizer.update(-gradient, opt_state, params)
params = optax.apply_updates(params, updates)
params = np.array(params)
# save history
objective_history.append(value)
param_history.append(params)
step = 1
J = 7.2186e-01
grad_norm = 2.4788e+00
step = 2
J = 7.6162e-01
grad_norm = 2.7910e+00
step = 3
J = 8.0478e-01
grad_norm = 4.2973e+00
step = 4
J = 8.3605e-01
grad_norm = 3.4016e+00
step = 5
J = 8.7346e-01
grad_norm = 3.9689e+00
step = 6
J = 8.7388e-01
grad_norm = 4.0717e+00
step = 7
J = 8.9084e-01
grad_norm = 4.5136e+00
step = 8
J = 9.1410e-01
grad_norm = 4.6112e+00
step = 9
J = 9.2536e-01
grad_norm = 4.0967e+00
step = 10
J = 9.3123e-01
grad_norm = 4.7938e+00
step = 11
J = 9.3927e-01
grad_norm = 4.9101e+00
step = 12
J = 9.3973e-01
grad_norm = 4.5086e+00
step = 13
J = 9.4017e-01
grad_norm = 5.2787e+00
step = 14
J = 9.4989e-01
grad_norm = 4.1704e+00
step = 15
J = 9.4899e-01
grad_norm = 3.8176e+00
step = 16
J = 9.4511e-01
grad_norm = 3.9490e+00
step = 17
J = 9.4500e-01
grad_norm = 4.2445e+00
step = 18
J = 9.5812e-01
grad_norm = 3.2506e+00
step = 19
J = 9.5528e-01
grad_norm = 3.8069e+00
step = 20
J = 9.5948e-01
grad_norm = 3.4437e+00
step = 21
J = 9.6042e-01
grad_norm = 3.2553e+00
[22]:
params_best = param_history[-1]
We see that our objective has increased steadily from a transmission of 56% to about 95%!
[23]:
ax = plt.plot(objective_history)
plt.xlabel("iteration number")
plt.ylabel("objective function")
plt.show()
Our final device displays smooth features and no sharp corners. Without our penalty this would have not been the case!
[24]:
sim_best = make_sim_params(param_history[-1], include_field_mnt=True)
ax = sim_best.plot(z=0.01)
[25]:
sim_data_best = td.web.run(sim_best, task_name="taper final")
↓ simulation_data.hdf5.gz ━━━━━━━━━━━━━ 100.0% • 6.4/6.4 MB • 7.5 MB/s • 0:00:00
15:35:43 EDT loading simulation from simulation_data.hdf5
Comparing the field patterns, we see that the optimized device gives a much more uniform field profile at the output waveguide, as desired. One can further check that this device and field pattern matches the referenced paper quite nicely!
[26]:
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, tight_layout=True, figsize=(11, 7))
# plot original
sim_data.plot_field(field_monitor_name='field', field_name='Ez', val='real', ax=ax1)
sim_data.plot_field(field_monitor_name='field', field_name='E', val='abs', ax=ax2)
# plot optimized
sim_data_best.plot_field(field_monitor_name="field", field_name="Ez", val="real", ax=ax3)
sim_data_best.plot_field(field_monitor_name="field", field_name="E", val="abs", ax=ax4)
plt.show()
[27]:
transmission_start = float(measure_transmission(sim_data))
transmission_end = float(measure_transmission(sim_data_best))
print(
f"Transmission improved from {(transmission_start * 100):.2f}% to {(transmission_end * 100):.2f}%"
)
Transmission improved from 72.19% to 98.79%