Inverse design quickstart - level 2

This notebook builds upon the previous quickstart notebook by demonstrating a slightly more advanced inverse design optimization using Tidy3D. In this level 2 notebook, we explore inverse design optimization of a multimode interference (MMI) power splitter. Our inverse design goal is to minimize loss by optimizing the width and length of the MMI.

Throughout this notebook, we will set up the simulation, define the objective function clearly, run an optimization loop, and visualize the optimized field distribution, providing insights into the power of inverse design in integrated photonics.

For more advanced inverse design examples and tutorial notebooks, check out

• Tidy3D Autograd Tutorial.

• Topology Optimization.

• Shape Optimization.

• Grating Coupler Inverse Design.

To see all Tidy3D examples and tutorials, as well as other learning materials, please visit our Learning Center.

import autograd as ag
import autograd.numpy as anp
import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.autograd import adam, apply_updates, optimize

Define a function to create a Simulation based on MMI width and length parameters. We set up a ModeSource to excite a guided mode and a ModeMonitor to measure transmission. The width and length will be combined into an array called param, where the first element is the width and the second element is the length. Here we only have two design parameters but in principle, this array can have an arbitrary number of parameters.

The MMI geometry is created by a PolySlab, which is a polygon defined by the coordinates of the vertices. The vertices of the geometry are denoted by the red dots in the schematic below. We manually define the vertices for the top half of the MMI and add a symmetric copy for the vertices on the bottom half.

def make_sim(params: anp.ndarray) -> td.Simulation:
    """Create a Tidy3D simulation given MMI width and length parameters."""

    width_mmi = params[0]
    length_mmi = params[1]

    # add a mode source located at the simulation entrance to excite the guided mode
    source = td.ModeSource(
        center=[-1, 0, 0],
        size=[0, 1.5, 1.5],
        source_time=td.GaussianPulse(freq0=td.C_0 / 1.55, fwidth=0.1 * td.C_0 / 1.55),
        mode_spec=td.ModeSpec(),
        direction="+",
    )

    # add a mode monitor located at the simulation exit to measure transmitted mode amplitude
    mode_monitor = td.ModeMonitor(
        name="mode",
        center=[15, 1, 0],
        size=[0, 1.5, 1.5],
        freqs=[td.C_0 / 1.55],
        mode_spec=td.ModeSpec(),
    )

    # add a field monitor to measure the field distribution
    field_monitor = td.FieldMonitor(
        name="field",
        size=[td.inf, td.inf, 0],
        freqs=[td.C_0 / 1.55],
    )

    # define the vertices for the MMI structure geometry
    top_vertices = [
        (-10, 0.25),
        (0, 0.25),
        (0, width_mmi / 2),
        (length_mmi, width_mmi / 2),
        (length_mmi, 1.25),
        (length_mmi + 10, 1.25),
        (length_mmi + 10, 0.75),
        (length_mmi, 0.75),
    ]

    # create the MMI structure using the defined vertices
    mmi = td.Structure(
        geometry=td.PolySlab(
            vertices=top_vertices + [(x, -y) for x, y in reversed(top_vertices)],
            axis=2,
            slab_bounds=(-0.11, 0.11),
        ),
        medium=td.Medium(permittivity=12),
    )

    # create the simulation
    sim = td.Simulation(
        center=[7, 0, 0],
        size=[17, 8, 3],
        symmetry=[0, -1, 1],
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=15),
        run_time=8e-12,
        medium=td.Medium(permittivity=2),
        sources=[source],
        monitors=[mode_monitor, field_monitor],
        structures=[mmi],
    )

    return sim

Define Objective Function

Define the objective function that calculates transmission based on the MMI structure parameters.

def objective_fn(params: anp.ndarray) -> float:
    """Compute the objective function, the squared magnitude of the mode amplitude (transmission power) in the forward direction."""

    # run the simulation
    sim_data = web.run(simulation=make_sim(params), task_name="adjoint", verbose=False)

    return (
        anp.abs(sim_data["mode"].amps.sel(direction="+").values.item()) ** 2
    )  # calculate the transmission power

Optimization Loop

Next, we use autograd to construct a function that returns the gradient of our objective function.

At the lowest level, we can run Adam explicitly by evaluating the gradient, calling optimizer.update(...), and applying the updates ourselves.

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

# initial design parameters are 4 μm width and 13 μm length
params = anp.array([4.0, 13.0])
optimizer = adam(learning_rate=0.01)
opt_state = optimizer.init(params)

for i in range(7):
    value, gradient = val_grad(params)

    # Adam minimizes by default, so maximize the transmission by stepping along -gradient.
    updates, opt_state_low_level = optimizer.update(-gradient, opt_state, params)
    params[:] = apply_updates(params, updates)

    print(f"step = {i + 1}")
    print(f"\ttransmission = {1e2 * value:.1f}%")

step = 1
        transmission = 28.5%
step = 2
        transmission = 37.8%
step = 3
        transmission = 41.8%
step = 4
        transmission = 43.1%
step = 5
        transmission = 43.8%
step = 6
        transmission = 43.8%
step = 7
        transmission = 43.8%

Tidy3D also provides the helper tidy3d.plugins.autograd.optimize() which enables a convenience usage which is equivalent to the manual loop above as follows:

params0 = anp.array([4.0, 13.0])


def report_step(params, grad, state, step_index, objective_val):
    print(f"step = {step_index + 1}")
    print(f"\ttransmission = {1e2 * objective_val:.1f}%")


params, opt_state, history = optimize(
    objective_fn,
    params0=params0,
    optimizer=adam(learning_rate=0.01),
    num_steps=7,
    callback=report_step,
    direction="max",
)

step = 1
        transmission = 28.5%
step = 2
        transmission = 37.8%
step = 3
        transmission = 41.7%
step = 4
        transmission = 43.0%
step = 5
        transmission = 43.6%
step = 6
        transmission = 43.8%
step = 7
        transmission = 43.8%

Analysis

During the optimization, we can see that the transmission is significantly increased. The final transmission is close to 44%, over 15% higher than the initial design. Finally, we run a simulation using the optimized parameters and visualize the resulting electric field distribution, which shows nice power splitting.

data_final = web.run(simulation=make_sim(params), task_name="optimized_mmi", verbose=False)

data_final.plot_field("field", "E", "abs", vmin=0, vmax=50);