# Inverse design optimization of a waveguide taper#

In this notebook, we will show how to use the adjoint plugin to compute the gradient with respect to the boundaries of a structure defined using a JaxPolySlab.

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, jax, tidy3d and its adjoint plugin.

[1]:

import matplotlib.pylab as plt
import numpy as np
import jax
import jax.numpy as jnp

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_eps = 4
wg_medium = td.Medium(permittivity=wg_eps)
wg_jax_medium = tda.JaxMedium(permittivity=wg_eps)

wg_length = 1 * wavelength
taper_length = 10.

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 JaxPolySlab for use in the JaxSimulation.

[5]:

def make_taper(ys) -> tda.JaxPolySlab:
"""Create a JaxPolySlab for the taper based on the supplied y values."""

# note, jax doesnt work well with concatenating, so we just construct these vertices for Tidy3D in a loop.

vertices = []
for x, y in zip(xs, ys):
vertices.append((x, y))
for x, y in zip(xs[::-1], ys[::-1]):
vertices.append((x, -y))

return tda.JaxPolySlab(
vertices=vertices,
slab_bounds=(-1, 1),
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 JaxSimulation 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) -> tda.JaxSimulation:
"""Make a JaxSimulation for 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 = tda.JaxStructure(geometry=taper_geo, medium=wg_jax_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 tda.JaxSimulation(
size=(Lx, Ly, 0),
structures=[wg_in, wg_out],
input_structures=[taper],
output_monitors=[mode_monitor],
monitors=[field_monitor] if include_field_mnt else [],
sources=[mode_source],
run_time=100/freq0,
grid_spec=td.GridSpec.auto(min_steps_per_wvl=30),
boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=False),
)

[8]:

sim = make_sim(ys0)
ax = sim.plot(z=0)

[15:29:13] WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.


Note: we get a warning because the polyslab edges in x intersect with the waveguides. We can ignore this warning because we don’t actually care about the gradient w.r.t these edges.

## Defining Objective#

Now that we have a function to create our JaxSimulation, 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 JaxSimulationData to give this result.

[9]:

def measure_transmission(sim_data: tda.JaxSimulationData) -> float:
"""Measure the first order transmission."""
amp_data = sim_data['mode'].amps
amp = amp_data.sel(f=freq0, direction="+", mode_index=0)
return 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 = (jnp.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 JaxSimulation given some parameters.

[11]:

def make_sim_params(parameters: np.ndarray, include_field_mnt: bool=False) -> tda.JaxSimulation:
"""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, we implement 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 150nm.

We also include some tunable parameters to adjust the characteristics of this penalty.

The 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]:

def smooth_penalty(parameters: np.ndarray, min_radius: float = .150, alpha: float = 1.0, kappa: float = 10.0) -> float:
"""Penalty between 0 and alpha based on radius of curvature."""

"""Quadratic bezier fit (and derivatives) for three points.
(x(t), y(t)) = P(t) = P0*t^2 + P1*2*t*(1-t) + P2*(1-t)^2
t in [0, 1]
"""

# ensure curve goes through (x1, y1) at t=0.5
p1 = 2 * pc - p0/2 - p2/2

def p(t):
"""Bezier curve parameterization."""
term0 = (1-t)**2 * (p0 - p1)
term1 = p1
term2 = t**2 * (p2 - p1)
return term0 + term1 + term2

def d_p(t):
"""First derivative function."""
d_term0 = 2 * (1-t) * (p1 - p0)
d_term2 = 2 * t * (p2 - p1)
return d_term0 + d_term2

def d2_p(t):
"""Second derivative function."""
d2_term0 = 2 * p0
d2_term1 = -4 * p1
d2_term2 = 2 * p2
return d2_term0 + d2_term1 +  d2_term2

return p, d_p, d2_p

def get_fit_vals(xs, ys):
""" Get the values of the bezier curve and its derivatives at t=0.5 along the taper."""

ps = jnp.stack((xs, ys), axis=1)
p0 = ps[:-2]
pc = ps[1:-1]
p2 = ps[2:]

p, d_p, d_2p = quad_fit(p0, pc, p2)

ps = p(0.5)
dps = d_p(0.5)
d2ps = d_2p(0.5)
return ps.T, dps.T, d2ps.T

"""Get the radii of curvature at each (internal) point along the taper."""
ps, dps, d2ps = get_fit_vals(xs, ys)
xp, yp = dps
xp2, yp2 = d2ps
num = (xp**2 + yp**2) ** (3.0/2.0)
den = abs(xp * yp2 - yp * xp2) + 1e-2
return num / den

"""Get the penalty for a given radius."""
return alpha * ((1 + jnp.exp(arg)) ** (-1))

ys = get_ys(parameters)

# return the average penalty over the points
return jnp.sum(penalty_fn(rs)) / len(rs)



### 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)

[15:29:14] WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.


Lets get the penalty value corresponding to this design, it should be relatively low, but the random variations could introduce a bit of penalty.

[14]:

# test penalty
penalty = smooth_penalty(params0)
print(f'starting penalty = {float(penalty):.2e}')

starting penalty = 8.13e-04


Finally, let’s run this simulation to get a feeling for the initial device performance.

[15]:

sim_data = web.run(sim.to_simulation()[0], task_name='taper fields')

[15:29:15] Created task 'taper fields' with task_id 'fdve-8ab49af6-f4e2-4add-b616-46afa53f635cv1'.    webapi.py:186

           View task using web UI at                                                                  webapi.py:188
cv1'.

[15:29:17] status = queued                                                                            webapi.py:321

[15:29:19] status = preprocess                                                                        webapi.py:315

[15:29:26] Maximum FlexCredit cost: 0.025. Use 'web.real_cost(task_id)' to get the billed FlexCredit  webapi.py:338
cost after a simulation run.

           starting up solver                                                                         webapi.py:342

           running solver                                                                             webapi.py:352

[15:29:33] status = postprocess                                                                       webapi.py:383

[15:29:52] loading SimulationData from simulation_data.hdf5                                           webapi.py:568

[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()



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)
return measure_transmission(sim_data) - smooth_penalty(parameters)


To test our our objective, we will use jax to make and run a function that returns the objective value and its gradient.

[18]:

grad_fn = jax.value_and_grad(objective, argnums=(0,))

[19]:

val, grad = grad_fn(params0, verbose=True)

[15:29:55] WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

           WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

[15:29:55] Created task 'adjoint_taper' with task_id 'fdve-002c4abb-07d5-4b70-9c38-6af5185e0235v1'.   webapi.py:186

           View task using web UI at                                                                  webapi.py:188
5v1'.

[15:29:58] status = queued                                                                            webapi.py:321

[15:30:00] status = preprocess                                                                        webapi.py:315

[15:30:04] Maximum FlexCredit cost: 0.025. Use 'web.real_cost(task_id)' to get the billed FlexCredit  webapi.py:338
cost after a simulation run.

           starting up solver                                                                         webapi.py:342

[15:30:05] running solver                                                                             webapi.py:352

[15:30:11] early shutoff detected, exiting.                                                           webapi.py:366

           status = postprocess                                                                       webapi.py:383

[15:30:31] status = success                                                                           webapi.py:390

[15:30:33] loading SimulationData from simulation_data.hdf5                                           webapi.py:568

[15:30:33] WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

           WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

[15:30:34] WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

           WARNING: 'JaxPolySlab'-containing 'JaxSimulation.input_structures[0]' intersects with  simulation.py:173
'JaxSimulation.structures[0]'. Note that in this version of the adjoint plugin, there
may be errors in the gradient when 'JaxPolySlab' intersects with background
structures. Skipping the rest of the structures.

[15:30:34] Created task 'adjoint_taper_adj' with task_id                                              webapi.py:186
'fdve-1b71acbc-8cc5-4850-8a88-827a82eedbeev1'.

           View task using web UI at                                                                  webapi.py:188
ev1'.

[15:30:36] status = queued                                                                            webapi.py:321

[15:30:38] status = preprocess                                                                        webapi.py:315

[15:30:42] Maximum FlexCredit cost: 0.025. Use 'web.real_cost(task_id)' to get the billed FlexCredit  webapi.py:338
cost after a simulation run.

           starting up solver                                                                         webapi.py:342

[15:30:43] running solver                                                                             webapi.py:352

[15:30:48] early shutoff detected, exiting.                                                           webapi.py:366

           status = postprocess                                                                       webapi.py:383

[15:30:55] status = success                                                                           webapi.py:390

[20]:

print(f'objective = {val:.2e}')

objective = 7.16e-01
gradient = [ 2.76337285e-03  1.08436763e-01  3.29497084e-02 -1.46241695e-01
-2.71953315e-01 -3.51883262e-01 -4.11713511e-01 -3.05958450e-01
-2.12514460e-01 -5.29730171e-02  1.60660297e-01  1.29041135e-01
6.37797937e-02  1.15576386e-03  9.54667032e-02  2.89692819e-01
3.85698915e-01  3.61970603e-01  5.04512310e-01  3.32683325e-01
3.44295323e-01  1.95844412e-01  7.33737573e-02  1.57067314e-01
4.89773415e-02  9.43381563e-02 -5.66829443e-02 -5.72257861e-03
-5.61871156e-02 -9.01155919e-03 -6.34363368e-02 -1.46888215e-02
7.19396323e-02  3.99515666e-02  4.63337824e-02  3.01078297e-02
4.63940464e-02  8.56711119e-02 -1.78261306e-02 -2.95523144e-02
6.39195368e-02 -7.65922666e-02 -4.26236540e-02 -4.13464420e-02
-4.45509255e-02 -8.30242559e-02  3.82876806e-02 -1.15348011e-01
1.51078645e-02 -2.50359382e-02 -2.52359938e-02  8.92452151e-03
-1.20067805e-01  3.94128934e-02 -1.29042342e-01 -9.70354304e-03
-4.90549393e-02 -5.39230034e-02 -4.24545780e-02 -4.39608693e-02
-4.92368639e-03 -1.15904994e-01 -9.26655158e-03 -4.22562584e-02
-4.55936752e-02 -4.65928465e-02 -4.83092293e-03 -1.13954321e-01
3.25321034e-02 -1.04063943e-01 -3.24416161e-03 -3.34401205e-02
-3.40586975e-02 -3.11472379e-02  1.12128817e-03 -8.74605924e-02
3.37448195e-02 -8.11599344e-02  3.27244774e-02 -7.41031319e-02
6.74696267e-03  7.73959234e-03 -6.26113936e-02  3.15839760e-02
-5.60581014e-02  3.00975665e-02 -4.85852212e-02  2.77132913e-02
-4.07127440e-02  2.47099474e-02 -3.37789580e-02  2.19077840e-02
-2.69655362e-02  8.06999672e-03 -2.25553755e-03 -1.10493414e-03
5.23142284e-03 -9.95822344e-03  6.25229673e-03 -2.56414595e-03
3.64503503e-05]


Now we can run our typical Adam optimization loop with this value_and_grad() function. See tutorial 3 for more details on the implementation.

[21]:

objective_history = []
param_history = []

def optimize(
parameters,
step_size=.01,
num_steps=41,
beta1=0.9,
beta2=0.999,
epsilon=1e-8,
):

parameters = parameters.copy()

# add starting parameter to the history
param_history.append(parameters.copy())

mt = np.zeros_like(parameters)
vt = np.zeros_like(parameters)

for i in range(num_steps):

t = i + 1
print(f"working on step = {t}")

mt = beta1 * mt + (1 - beta1) * grad
vt = beta2 * vt + (1 - beta2) * grad**2

mt_hat = mt / (1 - beta1**t)
vt_hat = vt / (1 - beta2**t)

update = step_size * (mt_hat / (np.sqrt(vt_hat) + epsilon))

objective_history.append(val)
print(f"\tJ = {val:.4e}")

parameters += update
param_history.append(parameters.copy())

# add the final objective value to the history
objective_history.append(val)
return parameters


[22]:

# turn off warnings to reduce verbosity
td.config.logging_level='ERROR'
params_best = optimize(params0)

working on step = 1
J = 7.1601e-01
working on step = 2
J = 5.3330e-01
working on step = 3
J = 6.5026e-01
working on step = 4
J = 6.7317e-01
working on step = 5
J = 6.8268e-01
working on step = 6
J = 7.5234e-01
working on step = 7
J = 7.5970e-01
working on step = 8
J = 7.6914e-01
working on step = 9
J = 7.8249e-01
working on step = 10
J = 8.3684e-01
working on step = 11
J = 8.1982e-01
working on step = 12
J = 8.2771e-01
working on step = 13
J = 8.4098e-01
working on step = 14
J = 8.6267e-01
working on step = 15
J = 8.6746e-01
working on step = 16
J = 8.7501e-01
working on step = 17
J = 8.8936e-01
working on step = 18
J = 8.9435e-01
working on step = 19
J = 9.0456e-01
working on step = 20
J = 9.0376e-01
working on step = 21
J = 9.1942e-01
working on step = 22
J = 9.2938e-01
working on step = 23
J = 9.2713e-01
working on step = 24
J = 9.3594e-01
working on step = 25
J = 9.3641e-01
working on step = 26
J = 9.4146e-01
working on step = 27
J = 9.4642e-01
working on step = 28
J = 9.5448e-01
working on step = 29
J = 9.5408e-01
working on step = 30
J = 9.6132e-01
working on step = 31
J = 9.6286e-01
working on step = 32
J = 9.6624e-01
working on step = 33
J = 9.6376e-01
working on step = 34
J = 9.6769e-01
working on step = 35
J = 9.6965e-01
working on step = 36
J = 9.7187e-01
working on step = 37
J = 9.7227e-01
working on step = 38
J = 9.7473e-01
working on step = 39
J = 9.7451e-01
working on step = 40
J = 9.7666e-01
working on step = 41
J = 9.7552e-01


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.to_simulation()[0], task_name='taper final')


[16:13:56] Created task 'taper final' with task_id 'fdve-3d3689de-4134-4ae5-b930-bba4477d7f0ev1'.     webapi.py:186

           View task using web UI at                                                                  webapi.py:188
ev1'.

[16:13:58] status = queued                                                                            webapi.py:321

[16:14:01] status = preprocess                                                                        webapi.py:315

[16:14:08] Maximum FlexCredit cost: 0.025. Use 'web.real_cost(task_id)' to get the billed FlexCredit  webapi.py:338
cost after a simulation run.

           starting up solver                                                                         webapi.py:342

           running solver                                                                             webapi.py:352

[16:14:16] early shutoff detected, exiting.                                                           webapi.py:366

           status = postprocess                                                                       webapi.py:383

[16:14:27] loading SimulationData from simulation_data.hdf5                                           webapi.py:568


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 71.68% to 99.58%