Adjoint optimization of a photonic crystal¶

import numpy as np
import matplotlib.pyplot as plt

import tidy3d as td
import tidy3d.web as web

# we'll use autograd for automatic differentiation, so derivative-traced numpy operations will use autograd.numpy
import autograd.numpy as anp
import autograd

nm = 1e-3

wvl0 = 1550 * nm
a = 400 * nm

num_freqs = 151

wvl_min, wvl_max = (1400 * nm, 1700 * nm)
freq_min = td.C_0 / wvl_max
freq_max = td.C_0 / wvl_min

freqs = np.linspace(freq_min, freq_max, num_freqs)
freq0 = td.C_0 / wvl0

fwidth = freq0 / 10
run_time = 300 / fwidth

sqrt3_div2 = np.sqrt(3) / 2.0

# silicon material (for slab)
n_si = 3.48
si = td.Medium(permittivity=n_si**2)
air = td.Medium()

# radius of holes
r0 = 90 * nm

# when we optimize, the radius will vary between these two values
# for now we just constrain all radius to r0, but change these values to add more degrees of freedom
r_range = rmin, rmax = (r0, r0)
rmid = r0

# how much centers can move from the center
drmax = a/4

# number of holes in x and y
N_rows = 15
N_cols = 19
N_cols_static = 5

# total length of the PhC region
Lx_phc = N_cols * a + a/2
Ly_phc = N_rows * sqrt3_div2 * a + a/2

# buffer on each side of the design region
buffer = 1.0 * wvl0

# thickness of slab
t_slab = 220 * nm

# waveguide width
w_wg = 1.3 * a

# size of simulation
Lx = Lx_phc + 2 * buffer
Ly = Ly_phc + 2 * buffer
Lz = t_slab + 2 * buffer

# define grid resolution
steps_per_unit_cell = 14
dx = a / steps_per_unit_cell
dy = a * sqrt3_div2 / steps_per_unit_cell

grid_spec = td.GridSpec(
    grid_x = td.UniformGrid(dl=dx),
    grid_y = td.UniformGrid(dl=dy),
    grid_z = td.AutoGrid(min_steps_per_wvl=steps_per_unit_cell)
)

def make_holes(params) -> td.Structure:
    """Convenience function to make the phc holes given the design parameters."""

    hole_spacing_x = a
    hole_spacing_y = a * sqrt3_div2
    
    x_slab_length, y_slab_length = hole_spacing_x * (N_cols + 0.5), hole_spacing_y * N_rows
    start_x, start_y = (
        - x_slab_length / 2 + hole_spacing_x / 2,
        - y_slab_length / 2 + hole_spacing_y / 2,
    )

    cylinders = []

    for i in range(0, N_cols):
        for j in range(0, N_rows):

            # depending on distance from central column, hole is either static
            i_dist = abs(i - N_cols // 2)
            if i_dist < ((N_cols_static) / 2):
                radius = rmid
                dx = dy = 0

            # or optimizable
            else:
                radius = params[0,i,j]
                dx = params[1,i,j]
                dy = params[2,i,j]

            x0 = dx + start_x + (i + (j % 2) * 0.5) * hole_spacing_x
            y0 = dy + start_y + j * hole_spacing_y
            if j != N_rows // 2:
                c = td.Cylinder(
                    axis=2,
                    radius=radius,
                    center=(x0, y0, 0),
                    length=td.inf,
                )
                cylinders.append(c)

    # use GeometryGroup since all same medium, for performance
    structure = td.Structure(
        geometry=td.GeometryGroup(geometries=cylinders),
        medium=air,
        background_medium=si, # note: we need this for correct gradients when embedded in slab
    )

    return structure

source = td.ModeSource(
    center=(-Lx/2 + wvl0/10, 0, 0),
    size=(0, Ly_phc/2.0, td.inf),
    source_time=td.GaussianPulse(
        freq0=freq0,
        fwidth=fwidth
    ),
    direction="+",
)

# these monitors are just for plotting the broadband flux response, not optimizing
# FluxMonitor is not supported in autograd

slab = td.Structure(
    geometry=td.Box(
        center=(0,0,0),
        size=(Lx_phc, Ly_phc, t_slab),
    ),
    medium=si,
)

waveguide = td.Structure(
    geometry=td.Box(
        center=(0,0,0),
        size=(td.inf, w_wg, t_slab),
    ),
    medium=si,
)


mnt_flux = td.FluxMonitor(
    center=(-source.center[0], 0, 0),
    size=(0, td.inf, td.inf),
    freqs=freqs,
    name='flux',
)

# this monitor is for visualizing field patterns only
mnt_field = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, td.inf, 0),
    freqs=[freq0],
    name='field',
)

# These monitors are used for the optimization
# FieldMonitor.flux is used because FluxMonitor not supported in autograd

mnt_field_flux = td.FieldMonitor(
    center=(+Lx_phc/2 - a, 0, 0),
    size=(0, td.inf, td.inf),
    freqs=[freq0],
    name='flux',
)

def make_sim(params, optimization_mode: bool = False) -> td.Simulation:
    """Function to generate a simulation with different monitors depending on whether optimizing."""
    
    # create the holes
    holes = make_holes(params)

    # decide which monitors to include
    if not optimization_mode:
        monitors = [mnt_flux, mnt_field]
    else:
        monitors = [mnt_field_flux]

    return td.Simulation(
        center=[0, 0, 0],
        size=[Lx, Ly, Lz],
        grid_spec=grid_spec,
        structures=[slab, waveguide, holes],
        sources=[source],
        monitors=monitors,
        run_time=run_time,
        boundary_spec=td.BoundarySpec.pml(x=True, y=True, z=True),
        symmetry=(0,-1,1),
    )

params0 = rmid * np.ones((3, N_cols, N_rows))
params0[1:] = 0.0

sim0 = make_sim(params0, optimization_mode=False)

sim_wg = sim0.updated_copy(structures=[waveguide])

def objective(params: anp.ndarray) -> float:
    """Maximize flux in -y, minimize flux in +x."""
    
    sim = make_sim(params, optimization_mode=True)
    sim_data = web.run(sim, task_name='phc_adjoint', verbose=False)
    
    flux_measure_freq0 = anp.sum(anp.abs(sim_data['flux'].flux.data))

    return flux_measure_freq0 / flux_wg.interp(f=freq0).item()

val_grad = autograd.value_and_grad(objective)

import optax
from autograd.tracer import getval

# hyperparameters
num_steps = 10
learning_rate = a/40
# note: the step size needs to be quite low because of the direct modification of geometric parameter

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

# store history
objective_history = []
param_history = [params]
data_history = []