This notebook will get users up and running with a very simple inverse design optimization with tidy3d. Inverse design uses the "adjoint method" to compute gradients of a figure of merit with respect to design parameters using only 2 simulations no matter how many design parameters are present. This gradient is then used to do high dimensional, gradient-based optimization of the system.
The setup we'll demonstrate here involves a point dipole source and a point field monitor on either side of a dielectric box. Using the adjoint plugin in tidy3d, we use gradient-based optimization to maximize the intensity enhancement at the measurement spot with respect to the box size in all 3 dimensions.
For more detailed notebooks, see these
[1]:
# To install other packages needed, uncomment lines below.
# !pip install optax
[2]:
import tidy3d as td
from tidy3d.web import run
import matplotlib.pylab as plt
import autograd as ag
import autograd.numpy as anp
import optax
Setup¶
First, we set up some basic parameters and "static" components of our simulation.
[3]:
# wavelength and frequency
wavelength = 1.55
freq0 = td.C_0 / wavelength
# permittivity of box
eps_box = 2
# size of sim in x,y,z
L = 10 * wavelength
# spc between sources, monitors, and PML / box
buffer = 1.0 * wavelength
[4]:
# create a source to the left of sim
source = td.PointDipole(
center=(-L / 2 + buffer, 0, 0),
source_time=td.GaussianPulse(freq0=freq0, fwidth=freq0 / 10.0),
polarization="Ez",
)
[5]:
# create a monitor to right of sim for measuring intensity
monitor = td.FieldMonitor(
center=(+L / 2 - buffer, 0, 0),
size=(0.0, 0.0, 0.0),
freqs=[freq0],
name="point",
)
[6]:
# create "base" simulation (the box will be added inside of the objective function later)
sim = td.Simulation(
size=(L, L, L),
grid_spec=td.GridSpec.auto(min_steps_per_wvl=25),
structures=[],
sources=[source],
monitors=[monitor],
run_time=120 / freq0,
)
Define objective function¶
Now we construct our objective function out of some helper functions. Our objective function measures the intensity enhancement at the measurement point as a function of a design parameter that controls the box size.
[7]:
# function to get box size (um) as a function of the design parameter (-inf, inf)
size_min = 0
size_max = L - 4 * buffer
def get_size(param: float):
"""Size of box as function of parameter, smoothly maps (-inf, inf) to (size_min, size_max)."""
param_01 = 0.5 * (anp.tanh(param) + 1)
return (size_max * param_01) + (size_min * (1 - param_01))
[8]:
# function to construct the simulation as a function of the design parameter
def make_sim(param: float) -> float:
"""Make simulation with a Box added, as given by the design parameter."""
# for normalization, ignore any structures and return base sim
if param is None:
return sim.copy()
# make a Box with the side length set by the parameter
size_box = get_size(param)
box = td.Structure(
geometry=td.Box(center=(0, 0, 0), size=(size_box, size_box, size_box)),
medium=td.Medium(permittivity=eps_box),
)
# add the box to the simulation
return sim.updated_copy(structures=[box])
[9]:
# function to compute and measure intensity as function of the design paramater
def measure_intensity(sim_data: td.SimulationData) -> float:
"""get intensity from SimulationData."""
return anp.sum(sim_data.get_intensity(monitor.name).values)
def intensity(param: float) -> float:
"""Intensity measured at monitor as function of parameter."""
# make the sim using the parameter value
sim_with_square = make_sim(param)
# run sim through tidy3d web API
data = run(sim_with_square, task_name="adjoint_quickstart", verbose=False)
# evaluate the intensity at the measurement position
return measure_intensity(data)
[10]:
# get the intensity with no box, for normalization (care about enhancement, not abs value)
intensity_norm = intensity(param=None)
print(f"With no box, intensity = {intensity_norm:.4f}.")
print("This value will be used for normalization of the objective function.")
With no box, intensity = 95.8381. This value will be used for normalization of the objective function.
[11]:
def objective_fn(param: float) -> float:
"""Intensity at measurement point, normalized by intensity with no box."""
return intensity(param) / intensity_norm
Optimization Loop¶
Next, we use autograd to construct a function that returns the gradient of our objective function and use this to run our gradient-based optimization in a for loop.
[12]:
# use autograd to get function that returns objective function and its gradient
val_and_grad_fn = ag.value_and_grad(objective_fn)
[13]:
# hyperparameters
num_steps = 9
learning_rate = 0.05
# initialize adam optimizer with starting parameter
param = -0.5
optimizer = optax.adam(learning_rate=learning_rate)
opt_state = optimizer.init(param)
# store history
objective_history = [1.0] # the normalized objective function with no box
param_history = [-100, param] # -100 is approximately "no box" (size=0)
for i in range(num_steps):
print(f"step = {i + 1}")
print(f"\tparam = {param:.4f}")
print(f"\tsize = {get_size(param):.4f} um")
# compute gradient and current objective funciton value
value, gradient = val_and_grad_fn(param)
# outputs
print(f"\tintensity = {value:.4e}")
print(f"\tgrad_norm = {anp.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, param)
param = optax.apply_updates(param, updates)
param = float(param)
# save history
objective_history.append(value)
param_history.append(param)
step = 1 param = -0.5000 size = 2.5012 um intensity = 7.9076e+00 grad_norm = 1.8164e+00 step = 2 param = -0.4500 size = 2.6882 um intensity = 9.8537e+00 grad_norm = 2.1918e+00 step = 3 param = -0.4000 size = 2.8833 um intensity = 1.1362e+01 grad_norm = 1.8857e+00 step = 4 param = -0.3501 size = 3.0855 um intensity = 1.2988e+01 grad_norm = 2.3163e+00 step = 5 param = -0.3000 size = 3.2955 um intensity = 1.6538e+01 grad_norm = 4.7152e+00 step = 6 param = -0.2516 size = 3.5043 um intensity = 1.8165e+01 grad_norm = 2.1677e+00 step = 7 param = -0.2036 size = 3.7162 um intensity = 2.0758e+01 grad_norm = 2.7685e+00 step = 8 param = -0.1552 size = 3.9342 um intensity = 2.3270e+01 grad_norm = 1.3610e+00 step = 9 param = -0.1086 size = 4.1470 um intensity = 2.2610e+01 grad_norm = 3.5083e+00
Analysis¶
Finally we plot our results: optimization progress, field pattern, and box size vs intensity enhancement.
[14]:
# objective function vs iteration number
plt.plot(objective_history)
plt.xlabel("iteration number")
plt.ylabel("intensity enhancement (unitless)")
plt.title("intensity enhancement during optimization")
plt.show()
[15]:
# construct simulation with final parameters
sim_final = make_sim(param=param_history[-1])
# add a field monitor for plotting
fld_mnt = td.FieldMonitor(
center=(+L / 2 - buffer, 0, 0),
size=(td.inf, td.inf, 0),
freqs=[freq0],
name="fields",
)
sim_final = sim_final.updated_copy(monitors=[monitor, fld_mnt])
# run simulation
data_final = run(sim_final, task_name="quickstart_final", verbose=False)
[17]:
# record final intensity
intensity_final = measure_intensity(data_final)
intensity_final_normalized = intensity_final / intensity_norm
objective_history.append(intensity_final_normalized)
[18]:
# plot intensity distribution
ax = data_final.plot_field(
field_monitor_name="fields", field_name="E", val="abs^2", vmax=intensity_final
)
ax.plot(source.center[0], 0, marker="o", mfc="limegreen", mec="black", ms=10)
ax.plot(monitor.center[0], 0, marker="o", mfc="orange", mec="black", ms=10)
plt.show()
[19]:
# scatter the intensity enhancement vs the box size
sizes = [get_size(p) for p in param_history]
objective_history = objective_history
_ = plt.scatter(sizes, objective_history)
ax = plt.gca()
ax.set_xlabel("box size (um)")
ax.set_ylabel("intensity enhancement (unitless)")
plt.title("intensity enhancement vs. box size")
plt.show()
[ ]:
