dispersive_gradient.py

import meep as mp
try:
    import meep.adjoint as mpa
except:
    import adjoint as mpa
import numpy as np
from autograd import numpy as npa
from autograd import tensor_jacobian_product
from enum import Enum
from matplotlib import pyplot as plt
from meep.materials import LiNbO3 as LN

resolution = 30
silicon = mp.Medium(epsilon=12)

sxy = 5.0
cell_size = mp.Vector3(sxy,sxy,0)

dpml = 1.0
boundary_layers = [mp.PML(thickness=dpml)]

eig_parity = mp.EVEN_Y + mp.ODD_Z

dx = 1.5
dy = 2.0
design_region_size = mp.Vector3(dx,dy)
design_region_resolution = int(2*resolution)
Nx = int(design_region_resolution*design_region_size.x) + 1
Ny = int(design_region_resolution*design_region_size.y) + 1

## ensure reproducible results
np.random.seed(314159)

## random design region
p = np.random.rand(Nx*Ny)

deps = 1e-5
dp = deps*np.random.rand(Nx*Ny)

w = 1.0
waveguide_geometry = [mp.Block(material=silicon,
                               center=mp.Vector3(),
                               size=mp.Vector3(mp.inf,w,mp.inf))]

fcen = 1/1.55
df = 0.23*fcen
mode = 1
sources = [mp.EigenModeSource(src=mp.GaussianSource(fcen,fwidth=df),
                              center=mp.Vector3(-0.5*sxy+dpml+0.1,0),
                              size=mp.Vector3(0,sxy-2*dpml),
                              eig_band=mode,
                              eig_parity=eig_parity)]

matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
                            mp.air,
                            LN,
                            grid_type='U_MEAN',
                            damping=0,
                            weights=np.ones((Nx*Ny,)),
                            do_averaging=False)

matgrid_region = mpa.DesignRegion(matgrid,
                                    volume=mp.Volume(center=mp.Vector3(),
                                                    size=mp.Vector3(design_region_size.x,design_region_size.y,0)))

matgrid_geometry =  [mp.Block(center=matgrid_region.center,
                                size=matgrid_region.size,
                                material=matgrid)]
matgrid_geometry += [mp.Block(center=matgrid_region.center,
                                size=matgrid_region.size,
                                material=matgrid,e2=mp.Vector3(y=-1))]

geometry = waveguide_geometry + matgrid_geometry

sim = mp.Simulation(resolution=resolution,
                    cell_size=cell_size,
                    boundary_layers=boundary_layers,
                    sources=sources,
                    extra_materials=[LN],
                    geometry=geometry)

frequencies = [fcen]
obj_list = [mpa.EigenmodeCoefficient(sim,
                                    mp.Volume(center=mp.Vector3(0.5*sxy-dpml-0.1),
                                    size=mp.Vector3(0,sxy-2*dpml,0)),
                                    mode,
                                    eig_parity=eig_parity)]

def J(mode_mon):
    return npa.power(npa.abs(mode_mon),2)

opt = mpa.OptimizationProblem(
    simulation=sim,
    objective_functions=J,
    objective_arguments=obj_list,
    design_regions=[matgrid_region],
    frequencies=frequencies)

def mapping(x,filter_radius):
    filtered_field = mpa.conic_filter(x,
                                      filter_radius,
                                      design_region_size.x,
                                      design_region_size.y,
                                      design_region_resolution)

    return filtered_field.flatten()
filter_radius = 0.1
beta = 0
opt.update_design([mapping(p,filter_radius)])

f, adjsol_grad = opt()
if np.asarray(frequencies).size > 1:
    bp_adjsol_grad = np.zeros(adjsol_grad.shape)
    for i in range(len(frequencies)):
        bp_adjsol_grad[:,i] = tensor_jacobian_product(mapping,0)(p,filter_radius,adjsol_grad[:,i])
else:
    bp_adjsol_grad = tensor_jacobian_product(mapping,0)(p,filter_radius,adjsol_grad)
if bp_adjsol_grad.ndim < 2:
    bp_adjsol_grad = np.expand_dims(bp_adjsol_grad,axis=1)

bp_adjsol_grad_m = (dp[None,:]@bp_adjsol_grad).flatten()

opt.update_design([mapping(p+dp/2,filter_radius)])
f_pp, _ = opt(need_gradient=False)
opt.update_design([mapping(p-dp/2,filter_radius)])
f_pm, _ = opt(need_gradient=False)
fd_grad = f_pp-f_pm
rel_error = np.abs(bp_adjsol_grad_m-fd_grad) / np.abs(fd_grad)

print("Directional derivative -- adjoint solver: {}, FD: {}|| rel_error = {}".format(bp_adjsol_grad_m,fd_grad,rel_error))