adding support for `load_minus_flux` to `EigenmodeCoefficient` objective function of adjoint solver

[oskooi]

In calculations involving the $S_{11}$ scattering parameter (i.e., the reflectance into a given port from a source located in the same port), it is often useful prior to the start of the run to first subtract out the incident fields using the load_minus_flux feature, particularly if the reflectance is smaller than e.g. 1e-6. While the computation of $S_{11}$ does not typically require require a normalization run because in most telecom applications reflectance values less than 40 dB are considered negligible, adding support for load_minus_flux to the adjoint solver's objective function EigenmodeCoefficient (and also FourierFields and Poynting flux in #2191) would be useful as e.g. a means to validate the adjoint gradient against the brute-force finite-difference result for arbitrary values and problems.

[smartalecH]

Are you seeing a discrepancy in gradients with the typical approach?

[oskooi]

I am trying to set up a test problem involving a 2d waveguide mode converter based on Schubert et al., (2022) (in particular, Figures 6a and 7a). The objective function involves minimizing the reflectance $|S_{11}|^2$ across several frequencies (and also maximizing $|S_{21}|^2$ but the transmittance is not relevant here). The problem, though, is that at a (typically coarse) resolution suitable for fast throughput, there is no guarantee whether small values of the $S_{11}$ computed by the adjoint solver are accurate and not dominated by discretization error. The only way to ensure the accuracy of the $S_{11}$ values at any resolution is to subtract away the incident fields first in the mode monitor.

[smartalecH]

The problem, though, is that at a (typically coarse) resolution suitable for fast throughput, there is no guarantee whether small values of the $S_{11}$ computed by the adjoint solver are accurate and not dominated by discretization error.

But if you're running at a coarse resolution, does it really matter if you are truly 60 dB down?

[oskooi]

The issue here is that at a given resolution, we do not know apriori what the discretization error of the $S_{11}$ calculation is unless we do a convergence analysis similar to the one in the tutorial. If the discretization error is 60 dB or smaller then it will likely not affect the optimization results. However, it can be larger, e.g., the discretization error is 50 dB at a resolution of 50 pixels/μm in the convergence analysis of the directional coupler. A large discretization error could also make the optimization fail to converge.

[oskooi]

Here are the results for a convergence analysis of the relative error in $|S_{11}|^2$ (the objective function) as well as its directional derivative for a 2d mode converter with a random design region. This is a similar setup to the existing unit tests in python/tests/test_adjoint_solver.py. No filtering of the design region is used. The reference result uses the incident fields from a normalization run and load_minus_flux (forward_sim function in the script below)

Results are shown for two wavelengths and three resolution values (50, 100, and 200). At the largest resolution of 200, $|S_{11}|^2$ is accurate to about three digits (which is above the 40 dB telecom limit).

The point in this example is that the discretization errors can be large even at modest resolutions.

import numpy as np
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
from autograd import numpy as npa

import meep as mp
import meep.adjoint as mpa


resolution = 200 # pixels/μm                                                                                                                             

w = 0.4          # waveguide width                                                                                                                       
l = 3.0          # waveguide length (on each side of design region)                                                                                      
dpad = 0.6       # padding length above/below design region                                                                                              
dpml = 1.0       # PML thickness                                                                                                                         
dx = 1.6         # length of design region                                                                                                               
dy = 1.6         # width of design region                                                                                                                

sx = dpml+l+dx+l+dpml
sy = dpml+dpad+dy+dpad+dpml
cell_size = mp.Vector3(sx,sy,0)

pml_layers = [mp.PML(thickness=dpml)]

# wavelengths to monitor                                                                                                                                 
wvls = (1.265, 1.270, 1.275, 1.285, 1.290, 1.295)
frqs = [1/wvl for wvl in wvls]

# pulsed source center frequency and bandwidth                                                                                                           
wvl_min = 1.26
wvl_max = 1.30
frq_min = 1/wvl_max
frq_max = 1/wvl_min
fcen = 0.5*(frq_min+frq_max)
df = frq_max-frq_min

eig_parity = mp.ODD_Z
src_pt = mp.Vector3(-0.5*sx+dpml,0)
refl_pt = mp.Vector3(-0.5*sx+dpml+0.5*l)
tran_pt = mp.Vector3(0.5*sx-dpml-0.5*l)

nSiO2 = 1.5
SiO2 = mp.Medium(index=nSiO2)
nSi = 3.5
Si = mp.Medium(index=nSi)

design_region_size = mp.Vector3(dx,dy,0)
design_region_resolution = int(2*resolution)
Nx = int(design_region_size.x*design_region_resolution)
Ny = int(design_region_size.y*design_region_resolution)

stop_cond = mp.stop_when_fields_decayed(50, mp.Ez, refl_pt, 1e-8)

# ensure reproducible results                                                                                                                            
rng = np.random.RandomState(59387385)

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

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

sources = [mp.EigenModeSource(src=mp.GaussianSource(fcen,fwidth=df),
                              size=mp.Vector3(0,sy,0),
                              center=src_pt,
                              eig_band=1,
                              eig_parity=eig_parity)]


def straight_waveguide():
    geometry = [mp.Block(size=mp.Vector3(mp.inf,w,mp.inf),
                         center=mp.Vector3(),
                         material=Si)]

    sim = mp.Simulation(resolution=resolution,
                        default_material=SiO2,
                        cell_size=cell_size,
                        sources=sources,
                        geometry=geometry,
                        boundary_layers=pml_layers,
                        k_point=mp.Vector3())

    refl_mon = sim.add_mode_monitor(frqs,
                                    mp.ModeRegion(center=refl_pt,
                                                  size=mp.Vector3(0,sy,0)))

    sim.run(until_after_sources=stop_cond)

    res = sim.get_eigenmode_coefficients(refl_mon,
                                         [1],
                                         eig_parity=eig_parity)

    coeffs = res.alpha
    input_flux = np.abs(coeffs[0,:,0])**2
    input_flux_data = sim.get_flux_data(refl_mon)

    return input_flux, input_flux_data

def forward_sim(design_params, input_flux, input_flux_data):
    geometry = [mp.Block(size=mp.Vector3(mp.inf,w,mp.inf),
                         center=mp.Vector3(),
                         material=Si),
                mp.Block(size=mp.Vector3(dx,dy,mp.inf),
                         center=mp.Vector3(),
                         material=mp.MaterialGrid(grid_size=mp.Vector3(Nx,Ny),
                                                  medium1=SiO2,
                                                  medium2=Si,
                                                  weights=design_params))]

    sim = mp.Simulation(resolution=resolution,
                        default_material=SiO2,
                        cell_size=cell_size,
                        sources=sources,
                        geometry=geometry,
                        boundary_layers=pml_layers,
                        k_point=mp.Vector3())

    refl_mon = sim.add_mode_monitor(frqs,
                                    mp.ModeRegion(center=refl_pt,
                                                  size=mp.Vector3(0,sy,0)))

    sim.load_minus_flux_data(refl_mon, input_flux_data)

    tran_mon = sim.add_mode_monitor(frqs,
                                    mp.ModeRegion(center=tran_pt,
                                                  size=mp.Vector3(0,sy,0)))

    sim.run(until_after_sources=stop_cond)

    plt.figure()
    sim.plot2D()
    if mp.am_master():
        plt.savefig('forward_sim.png',dpi=150,bbox_inches='tight')

    res = sim.get_eigenmode_coefficients(refl_mon,
                                         [1],
                                         eig_parity=eig_parity)
    coeffs = res.alpha
    refl = np.abs(coeffs[0,:,1])**2 / input_flux

    res = sim.get_eigenmode_coefficients(tran_mon,
                                         [2],
                                         eig_parity=eig_parity)
    coeffs = res.alpha
    tran = np.abs(coeffs[0,:,0])**2 / input_flux

    return refl, tran

def adjoint_sim(design_params, input_flux, input_flux_data):
    matgrid = mp.MaterialGrid(mp.Vector3(Nx,Ny,0),
                              SiO2,
                              Si,
                              weights=np.ones((Nx,Ny)))

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

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

    geometry = [mp.Block(size=mp.Vector3(mp.inf,w,mp.inf),
                         center=mp.Vector3(),
                         material=Si)]

    geometry += matgrid_geometry

    sim = mp.Simulation(resolution=resolution,
                        default_material=SiO2,
                        cell_size=cell_size,
                        sources=sources,
                        geometry=geometry,
                        boundary_layers=pml_layers,
                        k_point=mp.Vector3())

    obj_list = [
        mpa.EigenmodeCoefficient(
            sim,
            mp.Volume(
                center=refl_pt,
                size=mp.Vector3(0,sy,0),
            ),
            1,
            forward=False,
            eig_parity=eig_parity,
        )
    ]

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

   opt = mpa.OptimizationProblem(
        simulation=sim,
        objective_functions=J,
        objective_arguments=obj_list,
        design_regions=[matgrid_region],
        frequencies=frqs,
        decay_by=1e-12,
    )

    f, dJ_du = opt([design_params])

    return f, dJ_du


if __name__ == '__main__':
    input_flux, input_flux_data = straight_waveguide()

    R_unp, T_unp = forward_sim(p, input_flux, input_flux_data)
    R_per, T_per = forward_sim(p+dp, input_flux, input_flux_data)
    fd_grad = R_per - R_unp

    adjsol_obj, adjsol_grad = adjoint_sim(p, input_flux, input_flux_data)
    dir_deriv = (dp[None,:] @ adjsol_grad).flatten()

    rel_err = np.abs((fd_grad - dir_deriv) / fd_grad)
    print(f"rel_err:, {rel_err}")