Naive Inverse Design

Let’s leave the conditional generator out of it for now…

This notebook was adapted from Ceviche’s inverse design introduction to use a JAX-based optimization loop in stead of the default Ceviche optimization loop.

Introduction: multi-mode waveguides

The ceviche package has a built-in method insert_mode() that allows different modes to be inserted as sources.

Below, we demonstrate how this functionality can be used to excite the first and second order modes of a straight waveguide:

# Define simulation parameters (see above)
omega = 2 * np.pi * 200e12
dl = 25e-9

Nx = 200
Ny = 120
Npml = 20

# Define permittivity for a straight waveguide
epsr = np.ones((Nx, Ny))
epsr[:, 50:67] = 12.0

# Source position
src_y = np.arange(20, 100)
src_x = 30 * np.ones(src_y.shape, dtype=int)

# Source for mode 1
source1 = insert_mode(omega, dl, src_x, src_y, epsr, m=1)

# Source for mode 2
source2 = insert_mode(omega, dl, src_x, src_y, epsr, m=2)

# Run the simulation exciting mode 1
simulation = ceviche.fdfd_ez(omega, dl, epsr, [Npml, Npml])
Hx, Hy, Ez = simulation.solve(source1)

# Visualize the electric field
ax = ceviche.viz.real(Ez, outline=epsr, cmap="RdBu_r")
ax.plot(src_x, src_y, "k")

# Run the simulation exciting mode 2
simulation = ceviche.fdfd_ez(omega, dl, epsr, [Npml, Npml])
Hx, Hy, Ez = simulation.solve(source2)

# Visualize the electric field
ax = ceviche.viz.real(Ez, outline=epsr, cmap="RdBu_r")
ax.plot(src_x, src_y, "k")

plt.show()

Simulation and optimization parameters

Our toy optimization problem will be to design a device that converts an input in the first-order mode into an output as the second-order mode. First, we define the parameters of our device and optimization:

# Angular frequency of the source in Hz
omega = 2 * np.pi * 200e12
# Spatial resolution in meters
dl = 40e-9
# Number of pixels in x-direction
Nx = 100
# Number of pixels in y-direction
Ny = 100
# Number of pixels in the PMLs in each direction
Npml = 20
# Initial value of the structure's relative permittivity
epsr_init = 12.0
# Space between the PMLs and the design region (in pixels)
space = 10
# Width of the waveguide (in pixels)
wg_width = 12
# Length in pixels of the source/probe slices on each side of the center point
space_slice = 8
# Number of epochs in the optimization
Nsteps = 100
# Step size for the Adam optimizer
step_size = 1e-2

Utility functions

We now define some utility functions for initialization and optimization:

source

init_domain

 init_domain (Nx=100, Ny=100, Npml=20, space=10, wg_width=12,
              space_slice=8)

Initializes the domain and design region

space : The space between the PML and the structure wg_width : The feed and probe waveguide width space_slice : The added space for the probe and source slices

source

mask_combine_epsr

 mask_combine_epsr (epsr, bg_epsr, design_region)

Utility function for combining the design region epsr and the background epsr

source

viz_sim

 viz_sim (epsr, source, slices=[])

Solve and visualize a simulation with permittivity ‘epsr’

source

mode_overlap

 mode_overlap (E1, E2)

Defines an overlap integral between the simulated field and desired field

Visualizing the starting device

We can visualize what our starting device looks like and how it behaves. Our device is initialized by the init_domain() function which was defined several cells above.

# Initialize the parametrization rho and the design region
epsr, bg_epsr, design_region, input_slice, output_slice = init_domain(
    Nx, Ny, Npml, space=space, wg_width=wg_width, space_slice=space_slice
)

epsr_total = mask_combine_epsr(epsr, bg_epsr, design_region)

# Setup source
source = insert_mode(omega, dl, input_slice.x, input_slice.y, epsr_total, m=1)

# Setup probe
probe = insert_mode(omega, dl, output_slice.x, output_slice.y, epsr_total, m=2)

# Simulate initial device
simulation, ax = viz_sim(epsr_total, source, slices=[input_slice, output_slice])

# get normalization factor (field overlap before optimizing)
_, _, Ez = simulation.solve(source)
E0 = mode_overlap(Ez, probe)

Define objective function

We will now define our objective function. This is a scalar-valued function which our optimizer uses to improve the device’s performance.

Our objective function will consist of maximizing an overlap integral of the field in the output waveguide of the simulated device and the field of the waveguide’s second order mode (minimizing the negative overlap). The function takes in a single argument, epsr and returns the value of the overlap integral. The details of setting the permittivity and solving for the fields happens inside the objective function.

source

loss_fn

 loss_fn (epsr)

Objective function called by optimizer

Takes the epsr distribution as input
Runs the simulation
Returns the overlap integral between the output wg field and the desired mode field

Run optimization

This is where our discussion deviates from the original discussion by the ceviche maintainers. In our case, we would like the optimization to fit in a JAX-based optimization scheme:

# Simulate initial device
simulation, ax = viz_sim(epsr_total, source, slices=[input_slice, output_slice])

grad_fn = jacobian(loss_fn, mode='reverse')

init_fn, update_fn, params_fn = adam(step_size)
state = init_fn(epsr.reshape(1, -1))

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

this is the optimization step:

source

step_fn

 step_fn (step, state)

we can now loop over the optimization:

range_ = trange(30)
for step in range_:
    loss, state = step_fn(step, state)
    range_.set_postfix(loss=float(loss))

loss

-6146280.518657835

epsr_optimum = params_fn(state)
epsr_optimum = epsr_optimum.reshape((Nx, Ny))

# Simulate and show the optimal device
epsr_optimum_total = mask_combine_epsr(epsr_optimum, bg_epsr, design_region)
simulation, ax = viz_sim(epsr_optimum_total, source, slices=[input_slice, output_slice])

At the end of the optimization we can see our final device. From the field pattern, we can easily observe that the device is doing what we intend: the even mode enters from the left and exits as the odd mode on the right.

However, an additional observation is that our device’s permittivity changes continuously. This is not ideal if we wanted to fabricated our device. We’re also not constraining the minimum and maximum values of \(\epsilon_r\). Thus, we need to consider alternative ways of parameterizing our device.

plt.imshow(np.sqrt(epsr_optimum_total.T), cmap="plasma", vmin=1, vmax=4)
plt.ylim(*plt.ylim()[::-1])
plt.colorbar(ticks=[1,2,3,4], label="n")
plt.xlabel("x")
plt.xlabel("y")
plt.grid(True)