adjoint_gradiient_finite_difference.py

import meep as mp
import meep.adjoint as mpa
import numpy as np
from autograd import numpy as npa
from autograd import tensor_jacobian_product

silicon = mp.Medium(epsilon=12)

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

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

eig_parity = mp.EVEN_Y + mp.ODD_Z

design_region_size = mp.Vector3(1.5,1.5)
design_region_resolution = 256
Nx = int(design_region_resolution*design_region_size.x) + 1
Ny = int(design_region_resolution*design_region_size.y) + 1

fcen = 1/1.55
df = 0.23*fcen
source = [mp.Source(src=mp.GaussianSource(fcen,fwidth=df,is_integrated=True),
                    center=mp.Vector3(-0.5*sxy+dpml,0),
                    size=mp.Vector3(0,sxy),
                    component=mp.Ez)]

x = np.linspace(-0.5*design_region_size.x,0.5*design_region_size.x,Nx)
y = np.linspace(-0.5*design_region_size.y,0.5*design_region_size.y,Ny)
xv, yv = np.meshgrid(x,y)

rad = 0.538948295
wdt = 0.194838432
weights = np.where(np.logical_and(np.sqrt(np.square(xv) + np.square(yv)) > rad,
                                  np.sqrt(np.square(xv) + np.square(yv)) < rad+wdt),
                   1.,
                   0.)

filter_radius = 20/design_region_resolution

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

    return filtered_weights.flatten()

def forward_solver(design_params, resolution):
    matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
                              mp.air,
                              silicon,
                              weights=design_params.reshape(Nx,Ny),
                              do_averaging=True,
                              beta=mp.inf)

    matgrid_geometry = [mp.Block(center=mp.Vector3(),
                                 size=mp.Vector3(design_region_size.x,
                                                 design_region_size.y,
                                                 0),
                                 material=matgrid)]

    sim = mp.Simulation(resolution=resolution,
                        cell_size=cell_size,
                        boundary_layers=pml,
                        k_point=mp.Vector3(),
                        sources=source,
                        geometry=matgrid_geometry)

    dft_mon = sim.add_dft_fields([mp.Ez],
                                 [fcen],
                                 center=mp.Vector3(1.25),
                                 size=mp.Vector3(0.25, 1.),
                                 yee_grid=False)

    sim.run(until_after_sources=2000)
    
    Ez_dft = sim.get_dft_array(dft_mon, mp.Ez, 0)
    Ez2 = np.power(np.abs(Ez_dft[4,10]),2)

    sim.reset_meep()

    return Ez2


def adjoint_solver(design_params, resolution):
    matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny),
                              mp.air,
                              silicon,
                              weights=np.ones((Nx,Ny)),
                              do_averaging=True,
                              beta=mp.inf)

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

    geometry = [mp.Block(center=matgrid_region.center,
                         size=matgrid_region.size,
                         material=matgrid)]

    sim = mp.Simulation(resolution=resolution,
                        cell_size=cell_size,
                        boundary_layers=pml,
                        k_point=mp.Vector3(),
                        sources=source,
                        geometry=geometry)
                        
    obj_list = [mpa.FourierFields(sim,
                                  mp.Volume(center=mp.Vector3(1.25),
                                            size=mp.Vector3(0.25, 1.)),
                                  mp.Ez)]

    def J(mode_mon):
        return npa.power(npa.abs(mode_mon[:,4,10]),2)

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

    f, dJ_du = opt([design_params])

    sim.reset_meep()

    return f, dJ_du
    
 # ensure reproducible results                                                                                                                                
rng = np.random.RandomState(9861548)

# random epsilon perturbation for design region                                                                                                              
deps = 1e-5
dp = deps*rng.rand(Nx*Ny)

p = 0.5*weights.flatten()
mapped_p = mapping(p)
print(f"mapped_p:, min={np.amin(mapped_p)}, max={np.amax(mapped_p)}")

for res in [25, 50, 100, 200]:
    f_unperturbed = forward_solver(mapped_p, res)
    f_perturbed = forward_solver(mapping(p+dp), res)

    adjsol_obj, adjsol_grad = adjoint_solver(mapped_p, res)
    bp_adjsol_grad = tensor_jacobian_product(mapping,0)(p,adjsol_grad)

    print(f"obj. val.:, {res}, {f_unperturbed}, {adjsol_obj[0]}")

    if bp_adjsol_grad.ndim < 2:
        bp_adjsol_grad = np.expand_dims(bp_adjsol_grad,axis=1)
    adj_scale = (dp[None,:]@bp_adjsol_grad).flatten()
    fd_grad = f_perturbed-f_unperturbed

    rel_err = abs((fd_grad - adj_scale[0])/fd_grad)
    print(f"dir_deriv:, {res}, {fd_grad}, {adj_scale[0]}, {rel_err}")