Topology optimization of a waveguide crossing

In this example, we will walk you through performing a simple inverse design optimization of a 3D waveguide crossing.

In this example, we’ll be using a pixelated material grid to define the design region. However, one could also use shape parameterization to solve the same problem. You can see how this is done for a waveguide bend here.

If you are unfamiliar with inverse design, we also recommend our intro to inverse design tutorials and our primer on automatic differentiation with tidy3d.

Before using Tidy3D, you must first sign up for a user account. See this link for installation instructions.

In this tutorial we are demonstrating the smoothed projection by Alec Hammond et al. from the paper “Unifying and accelerating level-set and density-based topology optimization by subpixel-smoothed projection”. You can find the article here.

Setup

First we import the packages we need and also define some of the global parameters that define our problem.

Important note: we use autograd.numpy instead of regular numpy, this allows numpy functions to be differentiable. If one forgets this step, the error may be a bit opaque, just a heads up.

import math

import autograd.numpy as np
import matplotlib.pylab as plt
import tidy3d as td
import tidy3d.web as web
from tidy3d import config
from tidy3d.plugins.autograd import adam, optimize
from tidy3d.plugins.autograd.invdes.filters import ConicFilter
from tidy3d.plugins.autograd.invdes.parametrizations import initialize_params_from_simulation
from tidy3d.plugins.autograd.invdes.projections import smoothed_projection
from tidy3d.plugins.autograd.invdes.symmetries import symmetrize_diagonal, symmetrize_mirror

config.simulation.use_local_subpixel = False

# spectral parameters
wvl0 = 1.0
freq0 = td.C_0 / wvl0

# geometric parameters
eps_mat = 4.0
wg_width = 0.5 * wvl0
wg_length = 1.0 * wvl0
design_size = 3 * wvl0
thick = 0.2 * wvl0

# resolution
min_steps_per_wvl = 50
pixel_size = wvl0 / (min_steps_per_wvl * np.sqrt(eps_mat))
print(f"{pixel_size=}")

radius = 0.150
floor_radius = math.floor(radius / pixel_size)
discrete_radius = floor_radius if floor_radius % 2 == 1 else floor_radius + 1
discrete_num_pixels = round(design_size / pixel_size)
print(f"{floor_radius=}")
print(f"{discrete_radius=}")
print(f"{discrete_num_pixels=}")

pixel_size=np.float64(0.01)
floor_radius=15
discrete_radius=15
discrete_num_pixels=300

Next, we will define all the components that make up our “base” simulation. This simulation defines the static portion of our optimization, which doesn’t change over the iterations. Since our simulation scene has a symmetry in the y and z-axis, we will include the corresponding symmetry flag in the Simulation definition. This makes the simulation four times as fast and also reduces cost by a factor of four. For a tutorial on symmetries in simulations check out this notebook here.

For now, we’ll include a definition of the design region geometry, just to have that on hand later, but will not include a design region in the base simulation as we’ll add it later.

waveguide_long_horizontal = td.Structure(
    geometry=td.Box(
        center=(0, 0, 0),
        size=(td.inf, wg_width, thick),
    ),
    medium=td.Medium(permittivity=eps_mat),
)

waveguide_long_vertical = td.Structure(
    geometry=td.Box(
        center=(0, 0, 0),
        size=(wg_width, td.inf, thick),
    ),
    medium=td.Medium(permittivity=eps_mat),
)

design_region_geometry = td.Box(
    center=(0, 0, 0),
    size=(design_size, design_size, thick),
)

mode_source_horizontal = td.ModeSource(
    center=(-design_size / 2.0 - wg_length + wvl0 / 3, 0, 0),
    size=(0, wg_width * 3, td.inf),
    source_time=td.GaussianPulse(
        freq0=freq0,
        fwidth=freq0 / 5,
    ),
    mode_index=0,
    direction="+",
)

mode_monitor = td.ModeMonitor(
    center=(-mode_source_horizontal.center[0] + 0.2, 0, 0),
    size=(0, mode_source_horizontal.size[1], td.inf),
    freqs=[freq0],
    mode_spec=td.ModeSpec(num_modes=1),
    name="mode",
)

mode_monitor_multiple_wl = td.ModeMonitor(
    center=(-mode_source_horizontal.center[0] + 0.2, 0, 0),
    size=(0, mode_source_horizontal.size[1], td.inf),
    freqs=td.C_0 / np.linspace(wvl0 * 1.2, wvl0 / 1.2, 101),
    mode_spec=td.ModeSpec(num_modes=1),
    name="mode_multiple_wl",
)

field_monitor = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, td.inf, 0),
    freqs=[freq0],
    name="field",
)

sim_base = td.Simulation(
    size=(2 * wg_length + design_size, 2 * wg_length + design_size, thick + 2 * wvl0),
    symmetry=(0, -1, 1),
    run_time=100 / mode_source_horizontal.source_time.fwidth,
    structures=[
        waveguide_long_horizontal,
        waveguide_long_vertical,
    ],
    sources=[mode_source_horizontal],
    monitors=[mode_monitor],
    boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True),
    grid_spec=td.GridSpec.auto(
        min_steps_per_wvl=min_steps_per_wvl,
        override_structures=[
            td.MeshOverrideStructure(geometry=design_region_geometry, dl=3 * [pixel_size])
        ],
    ),
)

Let’s visualize the base simulations to verify that they look correct. The shaded region indicates that we are using a symmetry condition to determine the fields in this region, such that this region does not need to be simulated.

fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(10, 7), tight_layout=True)
sim_base.plot(z=0, ax=ax0)
sim_base.plot(y=0, ax=ax1)
plt.show()

Define Parameterization

Next, we will define how our design region is constructed as a function of our optimization parameters.

We will define a structure containing a grid of permittivity values defined by an array. Since we have a symmetric waveguide crossing, we only need to optimize 1/8 of the design volume. Therefore, we use some symmetry helper functions to constrain our device to these symmetries.

We will convolve our optimization parameters with a conic filter to smooth the features over a given radius. Then we will add a smoothed projection function (see the paper here for full details). The advantage of this projection is that we can work with fully binarized structures in our simulation. Since we assume that our topology should not change during optimization, we can use a value of \(\beta = \infty\) during the whole simulation process, which completely binarizes the design except for a small boundary at the interface change. If we want to change the topology during optimization, we could also add a scheduling for the beta parameter, increasing it stepwise towards \(\beta = \infty\). However, for our example this is not necessary.

For more details on the parameterization process, we highly recommend our short tutorial, which explains the process in more detail.

def get_density(params: np.ndarray) -> np.ndarray:
    """Get the density of the material in the design region as a function of optimization parameters."""
    arr = symmetrize_mirror(params, axis=(0, 1))
    arr = symmetrize_diagonal(arr)
    filter = ConicFilter(kernel_size=discrete_radius)
    arr_filtered = filter(arr)

    arr_projected = smoothed_projection(arr_filtered, beta=np.inf, eta=0.5)
    return arr_projected


def get_design_region(params: np.ndarray) -> td.Structure:
    """Get design region structure as a function of optimization parameters."""
    density = np.clip(get_density(params), 0.0, 1.0)
    eps_data = 1 + (eps_mat - 1) * density[:, :, None]
    return td.Structure.from_permittivity_array(eps_data=eps_data, geometry=design_region_geometry)

Next, it is very convenient to wrap this in a function that returns an updated copy of the base simulation with the design region added. We’ll be calling this in our objective function. We’ll also add some logic to exclude field monitors if they aren’t needed, for example during the optimization.

def get_sim(params: np.ndarray, with_fld_mnt: bool = False) -> td.Simulation:
    """Get simulation as a function of optimization parameters."""
    design_region = get_design_region(params)
    sim = sim_base.updated_copy(structures=sim_base.structures + (design_region,))
    if with_fld_mnt:
        sim = sim.updated_copy(
            monitors=sim_base.monitors + (field_monitor, mode_monitor_multiple_wl)
        )
    return sim

To start with the optimization, we need some initial parameters that can be optimized. These initial parameters could be chosen randomly, but the optimization will run faster and need less iterations to converge if we start from a design that is already performing a little better than random. In this case it is very easy to come up with such a design: We just initialize our parameters such that they represent the waveguides of the base simulations, which we defined above. In other words, we initialize our parameters such that they represent a naive waveguide crossing where the two waveguides simply overlap.

However, it can be a bit difficult to work out how the parameters should look to form this structure in the simulation since there is a smoothing and projection involved. Luckily, we have a helper function, which automatically determines how the parameters should look to represent the base simulation. This helper function takes the desired simulation as input and optimizes our initial parameters to reflect this simulation as best as possible.

rng = np.random.default_rng(seed=42)
params0 = rng.random(size=(discrete_num_pixels, discrete_num_pixels))
params0 = initialize_params_from_simulation(
    sim=sim_base,
    param_to_structure=get_design_region,
    params0=params0,
    maxiter=1000,
    bounds=(0.0, 1.0),
)

08:22:12 UTC WARNING: Design coordinates do not include the geometry center     
             (x=0, y=0). Centered features may appear half-cell shifted or      
             kinked when combined with symmetry or projection constraints. If   
             you need a design pixel centered on the geometry center, use an odd
             number of pixels along those axes.                                 

Let’s take a look at the resulting parameters which form the initial design. The raw parameters look noisy, but due to the conic filter and smoothed projection the parameters used in simulations form the initial cross design.

fig, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(10, 7), tight_layout=True)

im0 = ax0.imshow(params0)
ax0.set_title("raw parameters")
fig.colorbar(im0, ax=ax0, shrink=0.3)

density0 = get_density(params0)
im1 = ax1.imshow(density0)
ax1.set_title("material density")
fig.colorbar(im1, ax=ax1, shrink=0.3)


temporary_sim = get_sim(params0)
temporary_sim.plot_eps(z=0, ax=ax2)

plt.show()

We can also run a quick simulation with a field monitor added to verify how poorly the initial device is transmitting. This gives us lots of room to improve things through optimization!

sim0 = get_sim(params0, with_fld_mnt=True)
sim_data_init = web.run(sim0, task_name="initial_waveguide_crossing")

08:22:18 UTC Created task 'initial_waveguide_crossing' with resource_id
             'fdve-c77fb03b-d032-4c6a-a16e-d99f7ef7c99d' and task_type 'FDTD'.

             View task using web UI at
             'https://tidy3d.simulation.cloud/workbench?taskId=fdve-c77fb03b-d03
             2-4c6a-a16e-d99f7ef7c99d'.

             Task folder: 'default'.

08:22:20 UTC Estimated FlexCredit cost: 0.298. This assumes the FDTD solver runs
             for the full simulation time; if early shutoff is reached, the
             billed cost can be lower. Use 'web.real_cost(task_id)' to get the
             billed FlexCredit cost after a simulation run.

             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.

08:22:30 UTC status = preprocess

08:22:35 UTC starting up solver

             running solver

08:22:42 UTC early shutoff detected at 3%, exiting.

             status = postprocess

08:22:51 UTC status = success

08:22:53 UTC View simulation result at
             'https://tidy3d.simulation.cloud/workbench?taskId=fdve-c77fb03b-d03
             2-4c6a-a16e-d99f7ef7c99d'.

08:22:55 UTC Loading results from simulation_data.hdf5

fig, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2, 2, figsize=(10, 7), tight_layout=True)
sim0.plot_eps(z=0.01, ax=ax0)
ax1 = sim_data_init.plot_field("field", "E", "abs^2", z=0, ax=ax1)
ax2 = sim_data_init.plot_field("field", "Ex", z=0, ax=ax2)
ax3 = sim_data_init.plot_field("field", "Ey", z=0, ax=ax3)

Define Objective Function

The next step is to define the metric that we want to optimize, or our “figure of merit”. We will first write a function to compute the transmission (between 0-1) of the output mode given the simulation data.

def get_transmission_from_data(data):
    mode_amps = data["mode"].amps.sel(direction="+").values
    energy = np.sum(np.abs(mode_amps) ** 2)
    return energy

print(
    f"Initial transmission at center wavelength: {get_transmission_from_data(sim_data_init).item()}"
)

transmission_init = abs(sim_data_init["mode_multiple_wl"].amps.sel(direction="+")) ** 2
transmission_init_db = 10 * np.log10(transmission_init)
wavelengths = td.C_0 / transmission_init.f

plt.plot(wavelengths, transmission_init_db, c="black", linewidth=3)
plt.xlabel("wavelength (μm)")
plt.ylabel("transmission (dB)")
plt.grid()
plt.show()

Initial transmission at center wavelength: 0.7746249481935996

Next we can put everything together into a single objective function.

def objective(params: np.ndarray) -> float:
    """Objective function."""
    sim = get_sim(params)
    data = web.run(sim, task_name="crossing", verbose=False)
    transmission = get_transmission_from_data(data)
    return transmission

Optimization

We first define the parameters to optimize and initialize our Tidy3D Adam optimizer.

# hyperparameters
num_steps = 25
learning_rate = 0.1

# initialize Adam optimizer with starting parameters
params = np.copy(params0)
optimizer = adam(learning_rate=learning_rate)

And then we can run the optimization using the optimize helper, which handles the iterative parameter updates, optimizer state, and bookkeeping internally.

At each step, the gradient of the objective is used to update the parameters within the specified bounds. The callback function is used to log progress and visualize the device material density, allowing us to monitor how the design evolves over time.

Note: the following optimization loop will take about half an hour. To run fewer iterations, just change num_steps to something smaller above.

%%time

# store history
param_history = []


def record_step(current_params, grad, state, step_index, objective_val):
    print(f"step = {step_index + 1}")

    sim_i = get_sim(current_params)
    _, ax = plt.subplots(figsize=(2, 2))
    sim_i.plot_eps(z=0, ax=ax, monitor_alpha=0.0, source_alpha=0.0)
    plt.axis("off")
    plt.show()

    print(f"    J = {float(objective_val):.4e}")
    print(f"    grad_norm = {np.linalg.norm(grad):.4e}")

    param_history.append(np.array(current_params).copy())


params, opt_state, history = optimize(
    objective,
    params0=params,
    optimizer=optimizer,
    num_steps=num_steps,
    bounds=(0.0, 1.0),
    callback=record_step,
    direction="max",
)
param_history.append(np.array(params).copy())
objective_history = history["objective_fn_val"]

step = 1

    J = 7.7462e-01
    grad_norm = 2.6051e-03
step = 2

    J = 7.9235e-01
    grad_norm = 5.0813e-03
step = 3

    J = 8.0888e-01
    grad_norm = 3.4315e-03
step = 4

    J = 8.2333e-01
    grad_norm = 4.1143e-03
step = 5

    J = 8.4046e-01
    grad_norm = 4.2516e-03
step = 6

    J = 8.6078e-01
    grad_norm = 6.2112e-03
step = 7

    J = 8.6467e-01
    grad_norm = 1.5697e-02
step = 8

    J = 8.8626e-01
    grad_norm = 4.3691e-03
step = 9

    J = 8.9397e-01
    grad_norm = 4.4066e-03
step = 10

    J = 9.0177e-01
    grad_norm = 5.3548e-03
step = 11

    J = 9.1082e-01
    grad_norm = 4.1091e-03
step = 12

    J = 9.1801e-01
    grad_norm = 4.5727e-03
step = 13

    J = 9.2295e-01
    grad_norm = 4.0186e-03
step = 14

    J = 9.2503e-01
    grad_norm = 3.2287e-03
step = 15

    J = 9.2774e-01
    grad_norm = 3.9265e-03
step = 16

    J = 9.3180e-01
    grad_norm = 2.0518e-03
step = 17

    J = 9.3222e-01
    grad_norm = 5.8458e-03
step = 18

    J = 9.3809e-01
    grad_norm = 1.7284e-03
step = 19

    J = 9.3940e-01
    grad_norm = 5.2914e-03
step = 20

    J = 9.4337e-01
    grad_norm = 1.2844e-03
step = 21

    J = 9.4471e-01
    grad_norm = 2.0487e-03
step = 22

    J = 9.4601e-01
    grad_norm = 2.1767e-03
step = 23

    J = 9.4775e-01
    grad_norm = 1.2717e-03
step = 24

    J = 9.4901e-01
    grad_norm = 2.3540e-03
step = 25

    J = 9.5104e-01
    grad_norm = 1.3481e-03
CPU times: user 1min, sys: 2.28 s, total: 1min 3s
Wall time: 28min 10s

Analysis

Now is the fun part! We get to take a look at our optimization results.

We first plot the objective function values over the course of optimization, which should show a steady increase. Since the curve has not yet leveled off completely, we can conclude that we could run the optimization for longer to achieve even more improvements.

plt.plot(objective_history)
plt.xlabel("iterations")
plt.ylabel("objective function")
plt.show()

Next, we can look at the performance of our optimized device, we first construct it using the final parameter values and then take a look at the design.

params_final = param_history[-1]
sim_final = get_sim(params_final)
sim_final.plot_eps(z=0)
plt.show()

This optimized device seems to have some very intricate structures, which might be difficult to fabricate. This can be recifified by adding an additional loss term penalizing small structures. You can take a look at this notebook on how this may be done.

Let’s add a multi-frequency mode monitor and a field monitor to inspect the performance.

sim0 = get_sim(params_final, with_fld_mnt=True)
sim_data_final = web.run(sim0, task_name="final_waveguide_crossing")

08:51:09 UTC Created task 'final_waveguide_crossing' with resource_id
             'fdve-1d8ee4be-b4fe-44ab-9ef8-de643009552a' and task_type 'FDTD'.

             View task using web UI at
             'https://tidy3d.simulation.cloud/workbench?taskId=fdve-1d8ee4be-b4f
             e-44ab-9ef8-de643009552a'.

             Task folder: 'default'.

08:51:11 UTC Estimated FlexCredit cost: 0.298. This assumes the FDTD solver runs
             for the full simulation time; if early shutoff is reached, the
             billed cost can be lower. Use 'web.real_cost(task_id)' to get the
             billed FlexCredit cost after a simulation run.

08:51:12 UTC 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.

08:51:21 UTC status = preprocess

08:51:26 UTC starting up solver

             running solver

08:51:35 UTC early shutoff detected at 5%, exiting.

             status = postprocess

08:51:42 UTC status = success

08:51:44 UTC View simulation result at
             'https://tidy3d.simulation.cloud/workbench?taskId=fdve-1d8ee4be-b4f
             e-44ab-9ef8-de643009552a'.

08:51:45 UTC Loading results from simulation_data.hdf5

f, ((ax0, ax1), (ax2, ax3)) = plt.subplots(2, 2, figsize=(10, 7), tight_layout=True)
sim_final.plot_eps(z=0.01, ax=ax0)
ax1 = sim_data_final.plot_field("field", "E", "abs^2", z=0, ax=ax1)
ax2 = sim_data_final.plot_field("field", "Ex", z=0, ax=ax2)
ax3 = sim_data_final.plot_field("field", "Ey", z=0, ax=ax3)
plt.show()

transmission_final = abs(sim_data_final["mode_multiple_wl"].amps.sel(direction="+")) ** 2
transmission_final_db = 10 * np.log10(transmission_final)

plt.plot(wavelengths, transmission_init_db, c="black", label="Initial Design", linewidth=3)
plt.plot(wavelengths, transmission_final_db, c="red", label="Optimized Design", linewidth=3)
plt.xlabel("wavelength (μm)")
plt.ylabel("transmission (dB)")
plt.legend()
plt.grid()
plt.show()

It seems to perform quite well! Let’s export the simulation to a GDS file for fabrication!

# uncomment below to export

# sim_final.to_gds_file(
#     fname="./misc/invdes_crossing.gds",
#     z=0,
#     frequency=freq0,
#     permittivity_threshold=(eps_mat + 1) / 2.0
# )

Other Examples

Here are some other selected inverse design examples if you want to explore more!

• Mode converter

• Shape optimized taper

• Metalens optimization

• Wavelength division de-multiplexer