Adjoint Plugin: 5 Boundary Gradients
Contents
Adjoint Plugin: 5 Boundary Gradients#
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
import tidy3d.plugins.adjoint as tda
from tidy3d.plugins.adjoint.web import run
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."""
def quad_fit(p0, pc, p2):
"""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
def get_radii_curvature(xs, ys):
"""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
def penalty_fn(radius):
"""Get the penalty for a given radius."""
arg = -kappa * (min_radius - radius)
return alpha * ((1 + jnp.exp(arg)) ** (-1))
ys = get_ys(parameters)
rs = get_radii_curvature(xs, ys)
# 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')
View task using web UI at webapi.py:188 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-8ab49af6-f4e2-4add-b616-46afa53f635 cv1'.
[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 = run(sim, task_name='adjoint_taper', path='data/simulation.hdf5', verbose=verbose)
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.
View task using web UI at webapi.py:188 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-002c4abb-07d5-4b70-9c38-6af5185e023 5v1'.
[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.
View task using web UI at webapi.py:188 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-1b71acbc-8cc5-4850-8a88-827a82eedbe ev1'.
[20]:
print(f'objective = {val:.2e}')
print(f'gradient = {np.nan_to_num(grad[0])}')
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}")
val, grad = grad_fn(parameters, verbose=False)
grad = np.nan_to_num(np.array(grad[0]).copy())
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}")
print(f"\tgrad_norm = {np.linalg.norm(grad):.4e}")
parameters += update
param_history.append(parameters.copy())
# add the final objective value to the history
val, _ = grad_fn(parameters, verbose=False)
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
grad_norm = 1.3131e+00
working on step = 2
J = 5.3330e-01
grad_norm = 2.4857e+00
working on step = 3
J = 6.5026e-01
grad_norm = 5.3623e+00
working on step = 4
J = 6.7317e-01
grad_norm = 7.2553e+00
working on step = 5
J = 6.8268e-01
grad_norm = 7.0144e+00
working on step = 6
J = 7.5234e-01
grad_norm = 6.8726e+00
working on step = 7
J = 7.5970e-01
grad_norm = 7.1501e+00
working on step = 8
J = 7.6914e-01
grad_norm = 6.4100e+00
working on step = 9
J = 7.8249e-01
grad_norm = 7.3678e+00
working on step = 10
J = 8.3684e-01
grad_norm = 5.1055e+00
working on step = 11
J = 8.1982e-01
grad_norm = 6.7719e+00
working on step = 12
J = 8.2771e-01
grad_norm = 5.7634e+00
working on step = 13
J = 8.4098e-01
grad_norm = 5.9014e+00
working on step = 14
J = 8.6267e-01
grad_norm = 5.8077e+00
working on step = 15
J = 8.6746e-01
grad_norm = 6.1517e+00
working on step = 16
J = 8.7501e-01
grad_norm = 6.0151e+00
working on step = 17
J = 8.8936e-01
grad_norm = 5.4251e+00
working on step = 18
J = 8.9435e-01
grad_norm = 5.7949e+00
working on step = 19
J = 9.0456e-01
grad_norm = 4.9875e+00
working on step = 20
J = 9.0376e-01
grad_norm = 5.9177e+00
working on step = 21
J = 9.1942e-01
grad_norm = 4.7830e+00
working on step = 22
J = 9.2938e-01
grad_norm = 4.5944e+00
working on step = 23
J = 9.2713e-01
grad_norm = 5.4676e+00
working on step = 24
J = 9.3594e-01
grad_norm = 4.4541e+00
working on step = 25
J = 9.3641e-01
grad_norm = 4.7951e+00
working on step = 26
J = 9.4146e-01
grad_norm = 4.2165e+00
working on step = 27
J = 9.4642e-01
grad_norm = 5.0692e+00
working on step = 28
J = 9.5448e-01
grad_norm = 4.1341e+00
working on step = 29
J = 9.5408e-01
grad_norm = 5.0659e+00
working on step = 30
J = 9.6132e-01
grad_norm = 4.1441e+00
working on step = 31
J = 9.6286e-01
grad_norm = 4.3558e+00
working on step = 32
J = 9.6624e-01
grad_norm = 3.5453e+00
working on step = 33
J = 9.6376e-01
grad_norm = 4.2934e+00
working on step = 34
J = 9.6769e-01
grad_norm = 3.4260e+00
working on step = 35
J = 9.6965e-01
grad_norm = 3.7677e+00
working on step = 36
J = 9.7187e-01
grad_norm = 3.2665e+00
working on step = 37
J = 9.7227e-01
grad_norm = 3.5137e+00
working on step = 38
J = 9.7473e-01
grad_norm = 2.9129e+00
working on step = 39
J = 9.7451e-01
grad_norm = 3.2449e+00
working on step = 40
J = 9.7666e-01
grad_norm = 2.3411e+00
working on step = 41
J = 9.7552e-01
grad_norm = 2.9780e+00
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')
View task using web UI at webapi.py:188 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-3d3689de-4134-4ae5-b930-bba4477d7f0 ev1'.
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%