Inverse design quickstart

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

• Tidy3D Autograd Tutorial.

• Topology Optimization.

• Shape Optimization.

• Grating Coupler Inverse Design.

# To install other packages needed, uncomment lines below.
# !pip install optax

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.

# 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

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

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

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

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

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

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

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

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.

# use autograd to get function that returns objective function and its gradient
val_and_grad_fn = ag.value_and_grad(objective_fn)

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

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

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

# record final intensity
intensity_final = measure_intensity(data_final)
intensity_final_normalized = intensity_final / intensity_norm

objective_history.append(intensity_final_normalized)

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

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