Inverse design optimization of a mode converter#

Note: native autograd support is an experimental feature in tidy3d 2.7. To see the original implementation of this notebook using jax and the adjoint plugin, refer to this notebook.

In this notebook, we will use inverse design and Tidy3D to create an integrated photonics component to convert a fundamental waveguide mode to a higher order mode.

[1]:
from typing import List
import numpy as np
import matplotlib.pylab as plt

import autograd.numpy as anp
from autograd import value_and_grad

# import regular tidy3d
import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.mode import ModeSolver

# set random seed to get same results
np.random.seed(2)

Setup#

We wish to recreate a device like the diagram below:

Schematic of the mode converter

A mode source is injected into a waveguide on the left-hand side. The light propagates through a rectangular region with pixellated permittivity with the value of each pixel independently tunable between 1 (vacuum) and some maximum permittivity. Finally, we measure the transmission of the light into a waveguide on the right-hand side.

The goal of the inverse design exercise is to find the best distribution of permittivities (\(\epsilon_{ij}\)) in the coupling region to maximize the power conversion between the input mode and the output mode.

We also apply our built-in smoothening and binarization filters to ensure that the final device has smooth features, and permitivitty values that are all either 1, or the maximum permittivity of the waveguide material.

Parameters#

First we will define some parameters.

[2]:
# wavelength and frequency
wavelength = 1.0
freq0 = td.C_0 / wavelength
k0 = 2 * np.pi * freq0 / td.C_0

# resolution control
min_steps_per_wvl = 16
# in the design region, we set uniform grid resolution,
# and define the design parameters on the same grid
dl_design_region = 0.01

# space between boxes and PML
buffer = 1.0 * wavelength

# optimize region size
lz = td.inf
lx = 5.0
ly = 3.0

# position of source and monitor (constant for all)
source_x = -lx / 2 - buffer * 0.8
meas_x = lx / 2 + buffer * 0.8

# total size
Lx = lx + 2 * buffer
Ly = ly + 2 * buffer
Lz = 0.0

# permittivity and width of the input/output waveguide
eps_wg = 2.75
wg_width = 0.7

# random starting parameters between 0 and 1
nx = int(lx / dl_design_region)
ny = int(ly / dl_design_region)
params0 = np.random.random((nx, ny))

# frequency width and run time
freqw = freq0 / 10
run_time = 50 / freqw

Static Components#

Next, we will set up the static parts of the geometry, the input source, and the output monitor using these parameters.

[3]:
waveguide = td.Structure(
    geometry=td.Box(size=(2 * Lx, wg_width, lz)), medium=td.Medium(permittivity=eps_wg)
)

mode_size = (0, wg_width * 3, lz)

source_plane = td.Box(
    center=[source_x, 0, 0],
    size=mode_size,
)

measure_plane = td.Box(
    center=[meas_x, 0, 0],
    size=mode_size,
)

Input Structures#

Next, we write a function to return the pixellated array given our flattened tuple of permittivity values \(\epsilon_{ij}\) using the tidy3d.plugins.autograd plugin.

We start with an array of parameters between 0 and 1, apply a conic filter and tanh projection to compute smooth, well-binarized features.

[4]:
from tidy3d.plugins.autograd import rescale, make_filter_and_project

# radius of the circular filter (um) and the threshold strength
radius = 0.120
beta = 50

filter_project = make_filter_and_project(radius, lx / nx)


def get_eps(params, beta):
    """Get the permittivity values (1, eps_wg) array as a function of the parameters (0, 1)"""
    processed_params = filter_project(params, beta)
    eps = rescale(processed_params, 1, eps_wg)
    return eps


def make_input_structures(params, beta) -> List[td.Structure]:

    box = td.Box(center=(0, 0, 0), size=(lx, ly, lz))
    eps_data = get_eps(params, beta=beta).reshape((nx, ny, 1))
    custom_structure = td.Structure.from_permittivity_array(geometry=box, eps_data=eps_data)

    return [custom_structure]

Making the Simulation#

Next, we write a function to return a basic td.Simulation as a function of our parameter values.

We make sure to add the pixellated td.Structure list to input_structures but leave out the sources and monitors for now as we’ll want to add those after the mode solver is run so we can inspect them.

[5]:
def make_sim_base(params: np.ndarray, beta: float) -> td.Simulation:
    """Create the simulation given the parameters and some beta value."""

    structures = make_input_structures(params, beta=beta)
    design_region_mesh = td.MeshOverrideStructure(
        geometry=td.Box(size=(lx, ly, lz)),
        dl=[dl_design_region] * 3,
        enforce=True,
    )
    grid_spec = td.GridSpec.auto(
        wavelength=wavelength,
        min_steps_per_wvl=16,
        override_structures=[design_region_mesh],
    )

    return td.Simulation(
        size=[Lx, Ly, Lz],
        grid_spec=grid_spec,
        structures=[waveguide] + structures,
        sources=[],
        monitors=[],
        run_time=run_time,
        boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=False),
    )

Visualize#

Let’s visualize the simulation to see how it looks

[6]:
sim_start = make_sim_base(params0, beta=1.0)

ax = sim_start.plot_eps(z=0)

plt.show()
../_images/notebooks_Autograd3InverseDesign_11_0.png

Select Input and Output Modes#

Next, let’s visualize the first 4 mode profiles so we can select which mode indices we want to inject and transmit.

[7]:
from tidy3d.plugins.mode.web import run as run_mode_solver

num_modes = 4
mode_spec = td.ModeSpec(num_modes=num_modes)

mode_solver = ModeSolver(
    simulation=sim_start,
    plane=source_plane,
    mode_spec=td.ModeSpec(num_modes=num_modes),
    freqs=[freq0],
)
modes = run_mode_solver(mode_solver, reduce_simulation=True)
09:15:09 EDT Mode solver created with
             task_id='fdve-b5b8b772-96a2-4431-a34d-099f9fd97244',
             solver_id='mo-5e781fc3-596a-44f7-b642-180064bb3c0c'.
09:15:18 EDT Mode solver status: queued
09:15:20 EDT Mode solver status: running
09:15:26 EDT Mode solver status: success

Let’s visualize the modes next.

[8]:
fig, axs = plt.subplots(num_modes, 3, figsize=(12, 12), tight_layout=True)
for mode_index in range(num_modes):
    vmax = 1.1 * max(
        abs(modes.field_components[n].sel(mode_index=mode_index)).max() for n in ("Ex", "Ey", "Ez")
    )
    for field_name, ax in zip(("Ex", "Ey", "Ez"), axs[mode_index]):
        field = modes.field_components[field_name].sel(mode_index=mode_index)
        field.real.plot(label="Real", ax=ax)
        field.imag.plot(ls="--", label="Imag", ax=ax)
        ax.set_title(f"index={mode_index}, {field_name}")
        ax.set_ylim(-vmax, vmax)

axs[0, 0].legend()

print("Effective index of computed modes: ", np.array(modes.n_eff))
Effective index of computed modes:  [[1.5720801 1.535463  1.3032305 1.1848158]]
../_images/notebooks_Autograd3InverseDesign_15_1.png

We want to inject the fundamental, Ez-polarized input into the 1st order Ez-polarized input.

From the plots, we see that these modes correspond to the first and third rows, or mode_index=0 and mode_index=2, respectively.

So we make sure that the mode_index_in and mode_index_out variables are set appropriately and we set a ModeSpec with 3 modes to be able to capture the mode_index_out in our output data.

[9]:
mode_index_in = 0
mode_index_out = 2

num_modes = max(mode_index_in, mode_index_out) + 1

mode_spec = td.ModeSpec(num_modes=num_modes)

Then it is straightforward to generate our source and monitor.

[10]:
# source seeding the simulation
forward_source = td.ModeSource(
    source_time=td.GaussianPulse(freq0=freq0, fwidth=freqw),
    center=[source_x, 0, 0],
    size=mode_size,
    mode_index=mode_index_in,
    mode_spec=mode_spec,
    direction="+",
)

# we'll refer to the measurement monitor by this name often
measurement_monitor_name = "measurement"

# monitor where we compute the objective function from
measurement_monitor = td.ModeMonitor(
    center=[meas_x, 0, 0],
    size=mode_size,
    freqs=[freq0],
    mode_spec=mode_spec,
    name=measurement_monitor_name,
)

Finally, we create a new function that calls our make_sim_base() function and adds the source and monitor to the result. This is the function we will use in our objective function to generate our td.Simulation given the input parameters.

[11]:
def make_sim(params, beta):
    sim = make_sim_base(params, beta=beta)
    return sim.updated_copy(sources=[forward_source], monitors=[measurement_monitor])
[12]:
sim_start = make_sim_base(params0, beta=1.0)

ax = sim_start.plot_eps(z=0)

plt.show()
../_images/notebooks_Autograd3InverseDesign_22_0.png

Post Processing#

Next, we will define a function to tell us how we want to postprocess the output td.SimulationData object to give the conversion power that we are interested in maximizing.

[13]:
def measure_power(sim_data: td.SimulationData) -> float:
    """Return the power in the output_data amplitude at the mode index of interest."""
    output_amps = sim_data["measurement"].amps
    amp = output_amps.sel(direction="+", f=freq0, mode_index=mode_index_out).values
    return anp.sum(anp.abs(amp) ** 2)

Then, we add a penalty to produce structures that are invariant under erosion and dilation, which is a useful approach to implementing minimum length scale features.

[14]:
from tidy3d.plugins.autograd import make_erosion_dilation_penalty

penalty = make_erosion_dilation_penalty(radius, lx / nx)

Define Objective Function#

Finally, we need to define the objective function that we want to maximize as a function of our input parameters (permittivity of each pixel) that returns the conversion power. This is the function we will differentiate later.

[15]:
def J(params: np.ndarray, beta: float, step_num: int = None, verbose: bool = False) -> float:
    sim = make_sim(params, beta=beta)
    task_name = "inv_des"
    if step_num:
        task_name += f"_step_{step_num}"
    sim_data = web.run(sim, task_name=task_name, verbose=verbose)
    penalty_weight = np.minimum(1, beta / 25)
    density = filter_project(params, beta)
    return measure_power(sim_data) - penalty_weight * penalty(density)

Inverse Design#

Now we are ready to perform the optimization.

We use the autograd.value_and_grad function to get the gradient of J with respect to the permittivity of each Box, while also returning the converted power associated with the current iteration, so we can record this value for later.

Let’s try running this function once to make sure it works.

[16]:
dJ_fn = value_and_grad(J)
[17]:
val, grad = dJ_fn(params0, beta=1, verbose=True)
print(grad.shape)
09:15:29 EDT Created task 'inv_des' with task_id
             'fdve-8547cb95-4c0f-4a6f-b3f4-85a160359a0e' and task_type 'FDTD'.
09:15:36 EDT status = success
09:15:38 EDT loading simulation from simulation_data.hdf5
09:15:39 EDT Created task 'inv_des_adjoint' with task_id
             'fdve-657f373b-85fd-48fa-9c43-f3f2bf9378b6' and task_type 'FDTD'.
09:15:41 EDT status = success
09:15:44 EDT loading simulation from simulation_data.hdf5
(500, 300)
[18]:
print(grad)
[[ 1.03973553e-08  2.05572992e-08  1.99225490e-08 ... -3.62808121e-09
  -3.67993793e-09 -1.85471698e-09]
 [ 2.10062527e-08  4.15411755e-08  4.02477491e-08 ... -7.27992507e-09
  -7.41227921e-09 -3.73756529e-09]
 [ 2.17053066e-08  4.29135451e-08  4.15405709e-08 ... -7.42669234e-09
  -7.64574453e-09 -3.86806290e-09]
 ...
 [ 3.51087567e-07  7.04907888e-07  7.13502420e-07 ... -6.70982789e-07
  -6.60610367e-07 -3.28616196e-07]
 [ 3.49886789e-07  7.02537145e-07  7.11220496e-07 ... -6.65168641e-07
  -6.54716039e-07 -3.25649517e-07]
 [ 1.74803716e-07  3.50961722e-07  3.55322763e-07 ... -3.31663157e-07
  -3.26416085e-07 -1.62359310e-07]]

Optimization#

We will use β€œAdam” optimization strategy to perform sequential updates of each of the permittivity values in the td.CustomMedium.

For more information on what we use to implement this method, see this article.

We will run 10 steps and measure both the permittivities and powers at each iteration.

We capture this process in an optimize function, which accepts various parameters that we can tweak.

[19]:
import optax

# hyperparameters
num_steps = 20
learning_rate = 1.0

# initialize adam optimizer with starting parameters
params = np.array(params0)
optimizer = optax.adam(learning_rate=learning_rate)
opt_state = optimizer.init(params)

# store history
Js = []
params_history = [params]
beta_history = []

# gradually increase the binarization strength
beta0 = 1
beta_increment = 1

for i in range(num_steps):
    # compute gradient and current objective funciton value

    density = filter_project(params, beta)
    plt.subplots(figsize=(2,2))
    plt.imshow(np.flipud(1 - density.T), cmap='gray')
    plt.axis('off')
    plt.show()

    beta = beta0 + i * beta_increment
    value, gradient = dJ_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 = anp.clip(params, 0.0, 1.0)

    # save history
    Js.append(value)
    params_history.append(params)
    beta_history.append(beta)

power = J(params_history[-1], beta=beta)
Js.append(power)
../_images/notebooks_Autograd3InverseDesign_34_0.png
step = 1
        beta = 1.0000e+00
        J = -3.9945e-02
        grad_norm = 4.6129e-04
../_images/notebooks_Autograd3InverseDesign_34_2.png
step = 2
        beta = 2.0000e+00
        J = 1.9324e-02
        grad_norm = 1.0155e-02
../_images/notebooks_Autograd3InverseDesign_34_4.png
step = 3
        beta = 3.0000e+00
        J = -7.2511e-02
        grad_norm = 6.9634e-03
../_images/notebooks_Autograd3InverseDesign_34_6.png
step = 4
        beta = 4.0000e+00
        J = -8.3889e-02
        grad_norm = 4.1913e-03
../_images/notebooks_Autograd3InverseDesign_34_8.png
step = 5
        beta = 5.0000e+00
        J = 1.2138e-01
        grad_norm = 1.2703e-02
../_images/notebooks_Autograd3InverseDesign_34_10.png
step = 6
        beta = 6.0000e+00
        J = 3.0753e-01
        grad_norm = 1.9187e-02
../_images/notebooks_Autograd3InverseDesign_34_12.png
step = 7
        beta = 7.0000e+00
        J = 4.2388e-01
        grad_norm = 1.9821e-02
../_images/notebooks_Autograd3InverseDesign_34_14.png
step = 8
        beta = 8.0000e+00
        J = 6.4370e-01
        grad_norm = 1.2999e-02
../_images/notebooks_Autograd3InverseDesign_34_16.png
step = 9
        beta = 9.0000e+00
        J = 7.3787e-01
        grad_norm = 1.4156e-02
../_images/notebooks_Autograd3InverseDesign_34_18.png
step = 10
        beta = 1.0000e+01
        J = 7.4208e-01
        grad_norm = 1.4119e-02
../_images/notebooks_Autograd3InverseDesign_34_20.png
step = 11
        beta = 1.1000e+01
        J = 8.1151e-01
        grad_norm = 7.2219e-03
../_images/notebooks_Autograd3InverseDesign_34_22.png
step = 12
        beta = 1.2000e+01
        J = 8.2154e-01
        grad_norm = 8.7102e-03
../_images/notebooks_Autograd3InverseDesign_34_24.png
step = 13
        beta = 1.3000e+01
        J = 8.4146e-01
        grad_norm = 4.7962e-03
../_images/notebooks_Autograd3InverseDesign_34_26.png
step = 14
        beta = 1.4000e+01
        J = 8.4425e-01
        grad_norm = 5.3352e-03
../_images/notebooks_Autograd3InverseDesign_34_28.png
step = 15
        beta = 1.5000e+01
        J = 8.5054e-01
        grad_norm = 3.7282e-03
../_images/notebooks_Autograd3InverseDesign_34_30.png
step = 16
        beta = 1.6000e+01
        J = 8.5656e-01
        grad_norm = 3.9814e-03
../_images/notebooks_Autograd3InverseDesign_34_32.png
step = 17
        beta = 1.7000e+01
        J = 8.6297e-01
        grad_norm = 3.1513e-03
../_images/notebooks_Autograd3InverseDesign_34_34.png
step = 18
        beta = 1.8000e+01
        J = 8.7055e-01
        grad_norm = 3.4939e-03
../_images/notebooks_Autograd3InverseDesign_34_36.png
step = 19
        beta = 1.9000e+01
        J = 8.7985e-01
        grad_norm = 3.0865e-03
../_images/notebooks_Autograd3InverseDesign_34_38.png
step = 20
        beta = 2.0000e+01
        J = 8.8632e-01
        grad_norm = 3.9383e-03
[20]:
params_final = params_history[-1]

Let’s run the optimize function.

and then record the final power value (including the last iteration’s parameter updates).

Results#

First, we plot the objective function (power converted to 1st order mode) as a function of step and notice that it converges nicely!

[21]:
plt.plot(Js)
plt.xlabel("iterations")
plt.ylabel("objective function")
plt.show()
../_images/notebooks_Autograd3InverseDesign_39_0.png

We then will visualize the final structure, so we convert it to a regular Simulation using the final permittivity values and plot it.

[22]:
sim_final = make_sim(params_final, beta=beta)
[23]:
sim_final.plot_eps(z=0)
plt.show()
../_images/notebooks_Autograd3InverseDesign_42_0.png

Finally, we want to inspect the fields, so we add a field monitor to the Simulation and perform one more run to record the field values for plotting.

[24]:
field_mnt = td.FieldMonitor(
    size=(td.inf, td.inf, 0),
    freqs=[freq0],
    name="field_mnt",
)

sim_final = sim_final.copy(update=dict(monitors=(field_mnt, measurement_monitor)))
[25]:
sim_data_final = web.run(sim_final, task_name="inv_des_final")
09:42:00 EDT Created task 'inv_des_final' with task_id
             'fdve-d4446fc5-bb33-46e8-a241-24cfd2992c6e' and task_type 'FDTD'.
09:42:04 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.
09:42:08 EDT status = preprocess
09:42:09 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
09:42:10 EDT running solver
09:42:14 EDT early shutoff detected at 12%, exiting.
             status = postprocess
09:42:17 EDT status = success
09:42:22 EDT loading simulation from simulation_data.hdf5

We notice that the behavior is as expected and the device performs exactly how we intended!

[26]:
f, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(12, 3), tight_layout=True)
sim_final.plot_eps(z=0.01, ax=ax0)
ax1 = sim_data_final.plot_field("field_mnt", "Ez", z=0, ax=ax1)
ax2 = sim_data_final.plot_field("field_mnt", "E", "abs^2", z=0, ax=ax2)
../_images/notebooks_Autograd3InverseDesign_47_0.png

The final device converts more than 90% of the input power to the 1st mode, up from < 1% when we started.

[27]:
final_power = (
    sim_data_final["measurement"].amps.sel(direction="+", f=freq0, mode_index=mode_index_out).abs
    ** 2
)
print(f"Final power conversion = {final_power*100:.2f}%")
Final power conversion = 94.64%

Export to GDS#

The Simulation object has the .to_gds_file convenience function to export the final design to a GDS file. In addition to a file name, it is necessary to set a cross-sectional plane (z = 0 in this case) on which to evaluate the geometry, a frequency to evaluate the permittivity, and a permittivity_threshold to define the shape boundaries in custom mediums. See the GDS export notebook for a detailed example on using .to_gds_file and other GDS related functions.

[28]:
sim_final.to_gds_file(
    fname="./misc/inv_des_mode_conv.gds", z=0, frequency=freq0, permittivity_threshold=2.6
)
[ ]: