Validates the adjoint gradient of the diffraction efficiency of a 1D grating.

binary_grating_adjoint_gradient.py

"""Validates the adjoint gradient of the diffraction efficiency of a 1D grating.                                  
                                                                                                                  
The 2D test involves computing the gradient of the diffraction efficiency of                                      
the first transmitted order in air with respect to the height of a 1D binary
grating. A linearly polarized planewave is from incident from the substrate on
which the grating is placed. The adjoint gradient is validated using the brute-force
finite difference via the directional derivative.                                                                                                       
"""

rom enum import Enum
from typing import Tuple

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


RESOLUTION_UM = 100
PML_UM = 1.0
SUBSTRATE_UM = 1.1
AIR_UM = 1.2
WAVELENGTH_UM = 0.5
N_GLASS = 1.5
GRATING_PERIOD_UM = 1.427
GRATING_HEIGHT_UM = 0.518
GRATING_DUTY_CYCLE = 0.665
DESIGN_REGION_SIZE = mp.Vector3(
    GRATING_HEIGHT_UM,
    GRATING_DUTY_CYCLE * GRATING_PERIOD_UM,
    0
)
DESIGN_REGION_RESOLUTION = int(2 * RESOLUTION_UM)
NX = int(DESIGN_REGION_SIZE.x * DESIGN_REGION_RESOLUTION) + 1
NY = int(DESIGN_REGION_SIZE.y * DESIGN_REGION_RESOLUTION) + 1


Polarization = Enum("Polarization", "S P")


def binary_grating_1d(pol: Polarization) -> mpa.OptimizationProblem:
    """Sets up the adjoint problem for a 1D binary grating."""

    frequency = 1 / WAVELENGTH_UM

    pml_layers = [mp.PML(direction=mp.X, thickness=PML_UM)]

    glass = mp.Medium(index=N_GLASS)

    size_x_um = PML_UM + SUBSTRATE_UM + GRATING_HEIGHT_UM + AIR_UM + PML_UM
    size_y_um = GRATING_PERIOD_UM
    cell_size = mp.Vector3(size_x_um, size_y_um, 0)

    k_point = mp.Vector3()

    if pol.name == "S":
         eig_parity = mp.ODD_Z
         src_cmpt = mp.Ez
    else:
         eig_parity = mp.EVEN_Z
         src_cmpt = mp.Hz

    src_pt = mp.Vector3(-0.5 * size_x_um + PML_UM, 0, 0)
    sources = [
        mp.Source(
            mp.GaussianSource(frequency, fwidth=0.1 * frequency),
            component=src_cmpt,
            center=src_pt,
            size=mp.Vector3(0, size_y_um, 0),
        )
    ]

    matgrid = mp.MaterialGrid(
        mp.Vector3(NX, NY),
        mp.air,
        glass,
        weights=np.ones((NX, NY)),
    )

    matgrid_region = mpa.DesignRegion(
        matgrid,
        volume=mp.Volume(
            center=mp.Vector3(
                (-0.5 * size_x_um + PML_UM + SUBSTRATE_UM +
                 0.5 * GRATING_HEIGHT_UM),
                0,
                0
            ),
            size=DESIGN_REGION_SIZE,
        ),
    )
    
    geometry = [
        mp.Block(
            material=glass,
            size=mp.Vector3(PML_UM + SUBSTRATE_UM, mp.inf, mp.inf),
            center=mp.Vector3(
                -0.5 * size_x_um + 0.5 * (PML_UM + SUBSTRATE_UM), 0, 0
            ),
        ),
        mp.Block(
            material=matgrid,
            size=matgrid_region.size,
            center=matgrid_region.center,
        ),
    ]

    sim = mp.Simulation(
        resolution=RESOLUTION_UM,
        cell_size=cell_size,
        boundary_layers=pml_layers,
        k_point=k_point,
        sources=sources,
        geometry=geometry,
    )

    tran_pt = mp.Vector3(0.5 * size_x_um - PML_UM, 0, 0)

    order_y = 1
    kdiff = mp.Vector3(
        (frequency**2 - (order_y / GRATING_PERIOD_UM)**2)**0.5,
        order_y / GRATING_PERIOD_UM,
        0
    )

    obj_list = [
        mpa.EigenmodeCoefficient(
            sim,
            mp.Volume(
                center=tran_pt,
                size=mp.Vector3(0, size_y_um, 0),
            ),
            mode=1,
            kpoint_func=lambda *not_used: kdiff,
            eig_parity=eig_parity,
            eig_vol=mp.Volume(
                center=tran_pt,
                size=mp.Vector3(0, 1 / RESOLUTION_UM, 0)
            ),
        ),
    ]

    def obj_func(mode_coeff):
        return npa.abs(mode_coeff)**2

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

    return opt


if __name__ == "__main__":
    opt = binary_grating_1d(Polarization.P)

    weights = 0.9 * np.ones((NX * NY))
    weights_perturb_mag = 1e-5
    weights_perturb = np.zeros((NX, NY))
    weights_perturb[:, -1] += weights_perturb_mag
    weights_perturb[:, -2] += weights_perturb_mag
    weights_perturb = weights_perturb.flatten()

    unperturbed_val, unperturbed_grad = opt([weights], need_gradient=True)

    fig, ax = plt.subplots()
    opt.plot2D(init_opt=False, ax=ax)
    fig.savefig(
        'binary_grating_1D_plot2D.png',
        dpi=150,
        bbox_inches='tight'
    )

    perturbed_val, _ = opt([weights + weights_perturb], need_gradient=False)

    if unperturbed_grad.ndim < 2:
        unperturbed_grad = np.expand_dims(unperturbed_grad, axis=1)

    adj_directional_deriv = (weights_perturb[None, :] @ unperturbed_grad)[0][0]
    fnd_directional_deriv = perturbed_val[0] - unperturbed_val[0]
    err = (abs(fnd_directional_deriv - adj_directional_deriv) /
           abs(fnd_directional_deriv))

    print(f"dir-deriv:, {fnd_directional_deriv} (finite difference), "
          f"{adj_directional_deriv} (adjoint), {err} (error)")