Adjoint optimization of an integrated bandpass filter#
In this notebook, we will use broadband optimization to design a bandpass filter in an integrated photonic device.
In short, this device takes in broadband light through a waveguide and encourages transmission to the output waveguide for a certain frequency range, while suppressing transmission outside of this range.
If you are unfamiliar with inverse design, we also recommend our intro to inverse design tutorials and our primer on automatic differentiation with tidy3d.
[1]:
import autograd as ag
import autograd.numpy as np
import matplotlib.pylab as plt
import tidy3d as td
import tidy3d.web as web
Setup#
First, we define some parameters that set up our optimization, such as the frequency points we are interested in as well as the geometric, material, and source parameters.
[2]:
wvl_min, wvl_max = 0.8, 1.2
fmin, fmax = td.C_0 / wvl_max, td.C_0 / wvl_min
num_freqs = 51
freqs = np.linspace(fmin, fmax, num_freqs)
freq0 = np.mean(freqs)
lambda0 = td.C_0 / freq0
fwidth = freq0 / 20
run_time = 50 / fwidth
wvl0 = td.C_0 / freq0
We will use a photonic integrated circuit with a 2D design region, so we set up some parameters to define that region.
[3]:
size_design_x = 10 * wvl0
size_design_y = 5 * wvl0
width_wg = 0.4 * wvl0
length_wg = 2 * wvl0
size_buffer = 1.5 * wvl0
And then some positions and length scales to help us keep track of things later on.
[4]:
size_sim_x = length_wg + size_design_x + length_wg
size_sim_y = size_buffer + size_design_y + size_buffer
center_wg_in_x = -size_sim_x / 2 + length_wg / 2
center_wg_in_y = 0
center_wg_out_x = +size_sim_x / 2 - length_wg / 2
center_wg_out_y = 0
center_design_x = 0
center_design_y = 0
Our material will just be a simple dielectric with relative permittivity of 4.
[5]:
eps_mat = 4.0
mat = td.Medium(permittivity=eps_mat)
Next, we define some discretization parameters for our simulation and design region.
[6]:
min_steps_per_wvl = 25
dl_design_region = 0.015
[7]:
nx = int(np.ceil(size_design_x / dl_design_region))
ny = int(np.ceil(size_design_y / dl_design_region))
Static Components#
In the next step, we will define the “static” tidy3d simulation components that do not change over the course of optimization. For example, the sources, monitors, and many of the structures.
We’ll put these into a base simulation, which we can easily update with our optimizing components later.
[8]:
waveguide_in = td.Structure(
geometry=td.Box(
center=(center_wg_in_x - length_wg, center_wg_in_y, 0),
size=(3 * length_wg, width_wg, td.inf),
),
medium=mat,
)
waveguide_out = td.Structure(
geometry=td.Box(
center=(center_wg_out_x + length_wg, center_wg_out_y, 0),
size=(3 * length_wg, width_wg, td.inf),
),
medium=mat,
)
design_region_static = td.Structure(
geometry=td.Box(
size=(size_design_x, size_design_y, td.inf),
center=(center_design_x, center_design_y, 0),
),
medium=td.Medium(permittivity=np.mean([1, eps_mat])),
)
[9]:
mode_plane_size = 5 * width_wg
num_modes = 1
mode_spec = td.ModeSpec(num_modes=num_modes)
mode_index = 0
src_mode = td.ModeSource(
size=(0, mode_plane_size, td.inf),
center=(center_wg_in_x, center_wg_in_y, 0),
direction="+",
mode_spec=mode_spec,
mode_index=mode_index,
num_freqs=10,
source_time=td.GaussianPulse(
freq0=freq0,
fwidth=fwidth,
),
)
mnt_mode = td.ModeMonitor(
size=(0, mode_plane_size, td.inf),
center=(center_wg_out_x, center_wg_out_y, 0),
freqs=freqs,
mode_spec=mode_spec,
name="mode",
)
mnt_field_0 = td.FieldMonitor(
size=(td.inf, td.inf, 0), center=(0, 0, 0), freqs=[freq0], name="field_single"
)
mnt_field_all = mnt_field_0.updated_copy(freqs=freqs)
[10]:
# mesh override structure over the design region helps keep the resolution constant across this space
override_structure_design = td.MeshOverrideStructure(
geometry=design_region_static.geometry,
dl=3 * [dl_design_region],
)
sim_base = td.Simulation(
size=(size_sim_x, size_sim_y, 0),
structures=[waveguide_in, waveguide_out, design_region_static],
sources=[src_mode],
monitors=[mnt_mode],
grid_spec=td.GridSpec.auto(
wavelength=wvl0,
min_steps_per_wvl=min_steps_per_wvl,
override_structures=[override_structure_design],
),
boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=False),
run_time=run_time,
symmetry=(0,-1,0),
)
[11]:
_, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 4))
sim_base.plot(z=0, ax=ax1)
sim_base.plot_eps(z=0, ax=ax2)
plt.show()
Design Region Parameterization#
Next, we define the functions that allow us to generate our design region using topology optimization.
We’ll apply several functions that enable us to filter and project our optimization parameters into a permittivity grid, and then turn this permittivity grid into a structure that can be added to our base simnulation as the optimization progresses.
[12]:
from tidy3d.plugins.autograd import rescale, make_filter_and_project
# radius of the circular filter (um) and the threshold strength
radius = 0.120 # <= larger radius = bigger feature sizes
beta = 50 # <= higher beta = more binarized
beta_penalty = 10
filter_project = make_filter_and_project(radius, dl_design_region)
def get_density(params: np.ndarray, beta: float) -> np.ndarray:
"""Get the material density values (0, 1) array as a function of the parameters (0, 1)"""
return filter_project(params, beta)
def get_eps(params: np.ndarray, beta: float) -> np.ndarray:
"""Get the permittivity values (1, eps_mat) array as a function of the parameters (0, 1)"""
processed_params = get_density(params, beta)
eps = rescale(processed_params, 1, eps_mat)
return eps
def get_design_region(params, beta) -> td.Structure:
"""Get the design region evaluted at a set of parameters."""
xmin = center_design_x - size_design_x / 2.0
xmax = center_design_x + size_design_x / 2.0
ymin = center_design_y - size_design_y / 2.0
ymax = center_design_y + size_design_y / 2.0
coord_bounds_x = np.linspace(xmin, xmax, nx + 1)
coord_bounds_y = np.linspace(ymin, ymax, ny + 1)
coord_centers_x = 0.5 * (coord_bounds_x[1:] + coord_bounds_x[:-1])
coord_centers_y = 0.5 * (coord_bounds_y[1:] + coord_bounds_y[:-1])
coords = dict(x=coord_centers_x, y=coord_centers_y, z=[0])
eps = get_eps(params, beta=beta).reshape((nx, ny, 1))
permittivity = td.SpatialDataArray(eps, coords=coords)
medium = td.CustomMedium(permittivity=permittivity)
return design_region_static.updated_copy(medium=medium)
def get_sim(params, beta) -> td.Simulation:
"""Get the design region evaluted at a set of parameters."""
design_region = get_design_region(params=params, beta=beta)
new_structures = list(sim_base.structures)
new_structures[-1] = design_region
return sim_base.updated_copy(structures=new_structures)
[13]:
# get simulation at a set of starting parameters and try it out
params0 = np.random.random((nx, ny))
sim0 = get_sim(params0, beta=beta)
[14]:
_, (ax1, ax2) = plt.subplots(1, 2, tight_layout=True, figsize=(10, 4))
sim0.plot(z=0, ax=ax1)
sim0.plot_eps(z=0, ax=ax2)
plt.show()
Objective Function#
Next, we will write the code that will define our metric for optimization. Specifically, we will analyze the transmission spectrum of our device, and compute the weighted sum with a set of weights representing our ideal bandpass filter operation.
We’ll also include a penalty to discourage small feature sizes, which could be difficult to fabricate.
[15]:
def get_transmission(
params, beta, step_num: int, verbose: bool, use_broadband: bool
) -> np.ndarray:
"""Compute transmission amplitudes as function of frequency."""
sim = get_sim(params=params, beta=beta)
task_name = "bend_transmission_mf"
if step_num is not None:
task_name += "_step_" + str(step_num)
data = web.run(sim, task_name=task_name, verbose=verbose)
amps = data[mnt_mode.name].amps.sel(direction="+", mode_index=mode_index)
if not use_broadband:
amps = amps.interp(f=freq0)
return np.abs(amps) ** 2
[16]:
bandwidth_um = 0.1
import xarray as xr
def target(freq: float) -> xr.DataArray:
"""Target transmission as a function of frequency."""
wvl_min = lambda0 - bandwidth_um / 2.0
wvl_max = lambda0 + bandwidth_um / 2.0
wvl = td.C_0 / freq
values = 1.0 * (np.logical_and(wvl > wvl_min, wvl < wvl_max))
return xr.DataArray(values, coords=dict(f=freq))
target_spectrum = target(freqs)
target_spectrum.plot.scatter()
plt.ylabel('transmission')
plt.show()
[17]:
def get_metric(transmission: xr.DataArray) -> float:
in_band = target_spectrum
out_band = 1 - target_spectrum
avg_T_in_band = np.sum(in_band * transmission).values / np.sum(in_band).values
avg_T_out_band = np.sum(out_band * transmission).values / np.sum(out_band).values
# ideally, (avg_T_in_band = 1 and avg_T_out_band = 0) -> metric = 1
# at worst, (avg_T_in_band = 0 and avg_T_out_band = 1) -> metric = -1
return avg_T_in_band - avg_T_out_band
Fabrication Penalty#
Next, we introduce a builtin fabrication penalty function, that evaluates our design region to reduce our objective function if any small feature sizes are measured compared to the radius defined earlier.
[18]:
from tidy3d.plugins.autograd import make_erosion_dilation_penalty
penalty = make_erosion_dilation_penalty(radius, dl_design_region, beta=beta_penalty)
[19]:
# for debugging, can be helpful to be able to change our objective function behavior
use_metric = True
use_penalty = True
use_broadband = True
def objective(
params: np.ndarray,
beta: float,
step_num: int = 0,
verbose: bool = False,
use_metric: bool = use_metric,
use_penalty: bool = use_penalty,
use_broadband: bool = use_broadband,
) -> float:
metric = 0.0
penalty_val = 0.0
if use_metric:
transmission = get_transmission(
params=params,
beta=beta,
step_num=step_num,
verbose=verbose,
use_broadband=use_broadband,
)
metric = get_metric(transmission)
if use_penalty:
penalty_weight = 1.0#np.minimum(1, beta / 25)
density = get_density(params, beta)
penalty_val = penalty_weight * penalty(density)
return metric - penalty_val
[20]:
grad_fn = ag.value_and_grad(objective)
[21]:
val, grad = grad_fn(params0, beta=beta, verbose=True)
print(grad.shape)
13:20:14 EDT Created task 'bend_transmission_mf_step_0' with task_id 'fdve-248484fb-a0f2-4a87-b096-381a9cd4a037' and task_type 'FDTD'.
View task using web UI at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-248484fb-a0f 2-4a87-b096-381a9cd4a037'.
13:20:17 EDT status = queued
To cancel the simulation, use 'web.abort(task_id)' or 'web.delete(task_id)' or abort/delete the task in the web UI. Terminating the Python script will not stop the job running on the cloud.
13:20:22 EDT status = preprocess
13:20:25 EDT Maximum FlexCredit cost: 0.038. Use 'web.real_cost(task_id)' to get the billed FlexCredit cost after a simulation run.
starting up solver
running solver
13:20:31 EDT early shutoff detected at 16%, exiting.
status = postprocess
13:20:58 EDT status = success
View simulation result at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-248484fb-a0f 2-4a87-b096-381a9cd4a037'.
13:21:12 EDT loading simulation from simulation_data.hdf5
Created task 'bend_transmission_mf_step_0_adjoint' with task_id 'fdve-67f26474-b9b4-4b23-b10c-f83117cd67fb' and task_type 'FDTD'.
View task using web UI at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-67f26474-b9b 4-4b23-b10c-f83117cd67fb'.
13:21:15 EDT status = queued
To cancel the simulation, use 'web.abort(task_id)' or 'web.delete(task_id)' or abort/delete the task in the web UI. Terminating the Python script will not stop the job running on the cloud.
13:21:19 EDT status = preprocess
13:21:22 EDT Maximum FlexCredit cost: 0.029. Use 'web.real_cost(task_id)' to get the billed FlexCredit cost after a simulation run.
starting up solver
running solver
13:21:30 EDT early shutoff detected at 16%, exiting.
status = postprocess
13:21:44 EDT status = success
View simulation result at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-67f26474-b9b 4-4b23-b10c-f83117cd67fb'.
13:21:58 EDT loading simulation from simulation_data.hdf5
(640, 320)
[22]:
print(np.linalg.norm(grad))
val
0.06403837219017432
[22]:
-0.7156905024842533
Optimization#
Now that we have a function to compute our gradient, we can use an open source gradient-based optimizer to perform inverse design of our device.
In this case, we use the optax package to implement the “adam” optimization algorithm, which performs gradient-ascent with some extra terms to improve the performance and convergence.
[23]:
# assert False
import optax
# hyperparameters
num_steps = 35
learning_rate = 0.1
# initialize adam optimizer with starting parameters
params = np.ones_like(params0) * 0.5
optimizer = optax.adam(learning_rate=learning_rate)
opt_state = optimizer.init(params)
# store history
objective_history = []
params_history = [params]
beta_history = []
# gradually increase the binarization strength
beta0 = 1
beta_max = 50
for i in range(num_steps):
# compute gradient and current objective funciton value
density = get_density(params, beta)
plt.subplots(1, 1, figsize=(2, 2))
plt.imshow(np.flipud(1 - density.T), cmap="gray", vmin=0, vmax=1)
plt.axis("off")
plt.show()
perc_done = (i / (num_steps - 1))
beta = beta0 + perc_done * (beta_max - beta0)
value, gradient = grad_fn(
np.array(params),
step_num=i + 1,
beta=beta,
)
# outputs
print(f"step = {i + 1}")
print(f"\tbeta = {beta:.4e}")
print(f"\tJ = {value:.4e}")
print(f"\tgrad_norm = {np.linalg.norm(gradient):.4e}")
# compute and apply updates to the optimizer based on gradient (-1 sign to maximize obj_fn)
updates, opt_state = optimizer.update(-gradient, opt_state, params)
params = optax.apply_updates(params, updates)
# cap the parameters
params = np.clip(params, 0.0, 1.0)
# save history
objective_history.append(value)
params_history.append(params)
beta_history.append(beta)
power = objective(params_history[-1], beta=beta)
objective_history.append(power)
step = 1
beta = 1.0000e+00
J = -1.1031e+00
grad_norm = 2.6938e-02
step = 2
beta = 2.4412e+00
J = -9.3950e-01
grad_norm = 3.2958e-02
step = 3
beta = 3.8824e+00
J = -9.2553e-01
grad_norm = 3.6470e-02
step = 4
beta = 5.3235e+00
J = -5.6456e-01
grad_norm = 4.7396e-02
step = 5
beta = 6.7647e+00
J = -4.2836e-01
grad_norm = 3.2926e-02
step = 6
beta = 8.2059e+00
J = -3.7891e-01
grad_norm = 1.7939e-02
step = 7
beta = 9.6471e+00
J = -2.1445e-01
grad_norm = 2.2958e-02
step = 8
beta = 1.1088e+01
J = -8.9444e-02
grad_norm = 2.2555e-02
step = 9
beta = 1.2529e+01
J = 1.1362e-02
grad_norm = 2.2093e-02
step = 10
beta = 1.3971e+01
J = 1.1210e-01
grad_norm = 2.0336e-02
step = 11
beta = 1.5412e+01
J = 1.9396e-01
grad_norm = 2.1514e-02
step = 12
beta = 1.6853e+01
J = 2.7025e-01
grad_norm = 2.1704e-02
step = 13
beta = 1.8294e+01
J = 3.4755e-01
grad_norm = 1.8862e-02
step = 14
beta = 1.9735e+01
J = 4.1219e-01
grad_norm = 1.5809e-02
step = 15
beta = 2.1176e+01
J = 4.6086e-01
grad_norm = 1.4479e-02
step = 16
beta = 2.2618e+01
J = 5.0012e-01
grad_norm = 1.4009e-02
step = 17
beta = 2.4059e+01
J = 5.3230e-01
grad_norm = 1.3394e-02
step = 18
beta = 2.5500e+01
J = 5.5752e-01
grad_norm = 1.2881e-02
step = 19
beta = 2.6941e+01
J = 5.7865e-01
grad_norm = 1.1554e-02
step = 20
beta = 2.8382e+01
J = 5.9717e-01
grad_norm = 1.0243e-02
step = 21
beta = 2.9824e+01
J = 6.1391e-01
grad_norm = 9.7453e-03
step = 22
beta = 3.1265e+01
J = 6.3077e-01
grad_norm = 9.4796e-03
step = 23
beta = 3.2706e+01
J = 6.4800e-01
grad_norm = 8.1805e-03
step = 24
beta = 3.4147e+01
J = 6.6236e-01
grad_norm = 7.1684e-03
step = 25
beta = 3.5588e+01
J = 6.7200e-01
grad_norm = 7.3689e-03
step = 26
beta = 3.7029e+01
J = 6.8189e-01
grad_norm = 7.4940e-03
step = 27
beta = 3.8471e+01
J = 6.9442e-01
grad_norm = 7.2038e-03
step = 28
beta = 3.9912e+01
J = 7.0622e-01
grad_norm = 6.1630e-03
step = 29
beta = 4.1353e+01
J = 7.1699e-01
grad_norm = 5.6555e-03
step = 30
beta = 4.2794e+01
J = 7.2542e-01
grad_norm = 5.0796e-03
step = 31
beta = 4.4235e+01
J = 7.3167e-01
grad_norm = 4.8476e-03
step = 32
beta = 4.5676e+01
J = 7.3623e-01
grad_norm = 5.0041e-03
step = 33
beta = 4.7118e+01
J = 7.4070e-01
grad_norm = 4.6467e-03
step = 34
beta = 4.8559e+01
J = 7.4490e-01
grad_norm = 4.7943e-03
step = 35
beta = 5.0000e+01
J = 7.5133e-01
grad_norm = 4.1301e-03
[24]:
params_final = params_history[-1]
Analyzing Results#
Now we can take a look at our optimization results and make some figures.
We’ll first plot the objective function value history, and then look at the transmission spectrum and field profiles of our final device.
[25]:
plt.plot(objective_history)
plt.xlabel("iterations")
plt.ylabel("objective function")
plt.show()
[26]:
mnt_fld = td.FieldMonitor(
size=(td.inf, td.inf, 0), center=(0, 0, 0), freqs=[fmin, freq0, fmax], name="fields"
)
sim_final = get_sim(params_final, beta=beta)
sim_final = sim_final.updated_copy(monitors=list(sim_final.monitors) + [mnt_fld])
[27]:
sim_final.plot_eps(z=0, monitor_alpha=0.0, source_alpha=0.0)
plt.show()
[28]:
transmission = get_transmission(
params=params_final, beta=beta, step_num=None, verbose=True, use_broadband=True
)
14:50:33 EDT Created task 'bend_transmission_mf' with task_id 'fdve-2c8321a0-a619-4955-be76-0413614c27ac' and task_type 'FDTD'.
View task using web UI at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-2c8321a0-a61 9-4955-be76-0413614c27ac'.
14:50:34 EDT status = success
14:50:35 EDT loading simulation from simulation_data.hdf5
[29]:
plt.scatter(freqs / 1e12, target(freqs), label="target")
plt.plot(freqs / 1e12, transmission, label="measured")
plt.xlabel("frequency (THz)")
plt.ylabel("transmission (P_out / P_in)")
plt.legend()
plt.show()
We see that the final transmission spectrum more or less matches our expected bandpass filter!
[30]:
sim_data_final = web.run(sim_final, task_name="final")
14:50:36 EDT Created task 'final' with task_id 'fdve-21e1d529-3cd7-49cb-a52e-94c75a2da034' and task_type 'FDTD'.
View task using web UI at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-21e1d529-3cd 7-49cb-a52e-94c75a2da034'.
14:50:38 EDT status = queued
To cancel the simulation, use 'web.abort(task_id)' or 'web.delete(task_id)' or abort/delete the task in the web UI. Terminating the Python script will not stop the job running on the cloud.
14:50:42 EDT status = preprocess
14:50:46 EDT Maximum FlexCredit cost: 0.025. Use 'web.real_cost(task_id)' to get the billed FlexCredit cost after a simulation run.
starting up solver
running solver
14:50:52 EDT early shutoff detected at 16%, exiting.
status = postprocess
14:51:04 EDT status = success
View simulation result at 'https://tidy3d.simulation.cloud/workbench?taskId=fdve-21e1d529-3cd 7-49cb-a52e-94c75a2da034'.
14:51:07 EDT loading simulation from simulation_data.hdf5
[31]:
f, axes = plt.subplots(3, 1, tight_layout=True, figsize=(10, 6))
for ax, f in zip(axes, (mnt_fld.freqs)):
sim_data_final.plot_field(mnt_fld.name, field_name="E", val="abs^2", f=f, ax=ax)
ax.set_title(f"wvl={1000 * f/td.C_0:.0f} (nm)")
We see that the field patterns at the middle and edges of our spectrum behave as we would expect, respectively blocking and transmitting the light.
Using Inverse Design Plugin#
Tidy3D now supports a higher level syntax for defining inverse design problems in a simpler way.
Here we’ll show how to set up this same problem using that tool and run it for 2 iterations.
For a full tutorial, refer to this example.
[38]:
import tidy3d.plugins.invdes as tdi
[48]:
# transformations on the parameters that lead to the material density array (0,1)
filter_project = tdi.FilterProject(radius=radius, beta=beta_history[-1])
# penalties applied to the state of the material density, after these transformations are applied
penalty = tdi.ErosionDilationPenalty(weight=1.0, length_scale=radius)
design_region = tdi.TopologyDesignRegion(
size=design_region_static.geometry.size,
center=design_region_static.geometry.center,
eps_bounds=(1.0, eps_mat), # the minimum and maximum permittivity values in the final grid
transformations=[filter_project],
penalties=[penalty],
pixel_size=dl_design_region,
)
[49]:
design = tdi.InverseDesign(
simulation=sim_base,
design_region=design_region,
task_name="invdes",
)
[50]:
def post_process_fn(sim_data, **kwargs) -> float:
amps = sim_data[mnt_mode.name].amps.sel(direction="+", mode_index=mode_index)
transmission = np.abs(amps) ** 2
return get_metric(transmission)
[52]:
optimizer = tdi.AdamOptimizer(
design=design,
num_steps=1,
learning_rate=learning_rate,
results_cache_fname="data/invdes_history_filter.hdf5",
)
[53]:
result = optimizer.run(params0=params_history[-1].reshape((nx, ny, 1)), post_process_fn=post_process_fn)
step (1/1)
objective_fn_val = 6.785e-01
grad_norm = 6.173e-03
post_process_val = 8.023e-01
penalty = 1.238e-01
[55]:
sim_last = result.sim_last
ax = sim_last.plot_eps(z=0, monitor_alpha=0.0, source_alpha=0.0)
[ ]: