Plots the fields and computes the transmittance of the stack configuration obtained from optimization

multilayer_validate.py

from typing import Tuple

import matplotlib.pyplot as plt
import meep as mp
import numpy as np


RESOLUTION_UM = 800
WAVELENGTH_UM = 1.0
AIR_UM = 1.0
PML_UM = 1.0
NUM_LAYERS = 9
MIN_LAYER_UM = 0.1
MAX_LAYER_UM = 0.5
N_LAYER_1 = 1.0 # TODO (oskooi): support arbitrary materials.                                                     
N_LAYER_2 = 1.3
LAYER_PERTURBATION_UM = 1.0 / RESOLUTION_UM
DESIGN_REGION_RESOLUTION_UM = 10 * RESOLUTION_UM
DESIGN_REGION_SIZE = mp.Vector3(0, 0, NUM_LAYERS * MAX_LAYER_UM)
NZ_DESIGN_GRID = int(DESIGN_REGION_SIZE.z * DESIGN_REGION_RESOLUTION_UM) + 1
NZ_SIM_GRID = int(DESIGN_REGION_SIZE.z * RESOLUTION_UM) + 1


def design_region_to_grid(nz: int) -> np.ndarray:
    """Returns the coordinates of the grid for the design region.                                                 
                                                                                                                  
    Args:                                                                                                         
      nz: number of grid points.                                                                                  
                                                                                                                  
    Returns:                                                                                                      
      The 1D coordinates of the grid points.                                                                      
    """

    z_grid = np.linspace(
        -0.5 * DESIGN_REGION_SIZE.z,
        +0.5 * DESIGN_REGION_SIZE.z,
	nz
    )

    return z_grid
    

def levelset_and_smoothing(layer_thickness_um: np.ndarray) -> np.ndarray:
    """Returns the density weights for a multilayer stack as a levelset.                                          
                                                                                                                  
    Args:                                                                                                         
      layer_thickness_um: thickness of each layer in the stack.                                                   
                                                                                                                  
    Returns:                                                                                                      
      The density weights as a flattened (1D) array.                                                              
    """

    air_padding_um = 0.5 * (DESIGN_REGION_SIZE.z - np.sum(layer_thickness_um))

    weights = np.zeros(NZ_DESIGN_GRID)

    # Air padding at left edge                                                                                    
    z_start = 0
    z_end = int(air_padding_um * DESIGN_REGION_RESOLUTION_UM)
    weights[z_start:z_end] = 0

    z_start = z_end
    for j in range(NUM_LAYERS):
        z_end = z_start + int(layer_thickness_um[j] * DESIGN_REGION_RESOLUTION_UM)
        weights[z_start:z_end] = 1 if (j % 2 == 0) else 0
        z_start = z_end

    # Air padding at right edge                                                                                   
    z_end = z_start + int(air_padding_um * DESIGN_REGION_RESOLUTION_UM)
    weights[z_start:z_end] = 0

    # Smooth the design weights by downsampling from the design grid                                              
    # to the simulation grid using bilinear interpolation.                                                        
    z_sim_grid = design_region_to_grid(NZ_SIM_GRID)
    z_design_grid = design_region_to_grid(NZ_DESIGN_GRID)
    smoothed_weights = np.interp(z_sim_grid, z_design_grid, weights)

    return smoothed_weights.flatten()
    
 
def input_flux() -> float:
    """Returns the flux generated by the source in vacuum."""

    frequency = 1 / WAVELENGTH_UM
    pml_layers = [mp.PML(direction=mp.Z, thickness=PML_UM)]

    size_z_um = PML_UM + AIR_UM + DESIGN_REGION_SIZE.z + AIR_UM + PML_UM
    cell_size = mp.Vector3(0, 0, size_z_um)

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

    sim = mp.Simulation(
        resolution=RESOLUTION_UM,
        cell_size=cell_size,
        dimensions=1,
        boundary_layers=pml_layers,
        sources=sources
    )

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

    flux_mon = sim.add_flux(
        frequency,
        0,
        1,
        mp.FluxRegion(center=tran_pt)
    )

    sim.run(
        until_after_sources=mp.stop_when_fields_decayed(
            25.0, src_cmpt, tran_pt, 1e-6
        )
    )

    flux = mp.get_fluxes(flux_mon)[0]

    return flux


def multilayer_stack(
        input_flux: float,
        layer_thickness_um: np.ndarray
) -> Tuple[float, np.ndarray]:
    """Returns the DFT fields of a multilayer stack."""

    frequency = 1 / WAVELENGTH_UM
    pml_layers = [mp.PML(thickness=PML_UM)]

    size_z_um = PML_UM + AIR_UM + DESIGN_REGION_SIZE.z + AIR_UM + PML_UM
    cell_size = mp.Vector3(0, 0, size_z_um)

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

    mat_1 = mp.Medium(index=N_LAYER_1)
    mat_2 = mp.Medium(index=N_LAYER_2)

    weights = levelset_and_smoothing(layer_thickness_um)
    matgrid = mp.MaterialGrid(
        mp.Vector3(0, 0, NZ_SIM_GRID),
        mat_1,
        mat_2,
        weights
    )

    geometry = [
        mp.Block(
            material=matgrid,
            size=mp.Vector3(0, 0, DESIGN_REGION_SIZE.z),
            center=mp.Vector3()
        )
    ]

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

    dft_field_mon = sim.add_dft_fields(
        [mp.Ex],
        frequency,
        0,
        1,
        center=mp.Vector3(),
        size=mp.Vector3(0, 0, DESIGN_REGION_SIZE.z)
    )

    tran_pt = mp.Vector3(0, 0, 0.5 * size_z_um - PML_UM)
    
    flux_mon = sim.add_flux(
        frequency,
        0,
        1,
        mp.FluxRegion(center=tran_pt)
    )

    sim.run(
        until_after_sources=mp.stop_when_fields_decayed(
            25.0, src_cmpt, tran_pt, 1e-8
        )
    )

    ex_dft = sim.get_dft_array(dft_field_mon, mp.Ex, 0)

    tran_flux = mp.get_fluxes(flux_mon)[0]
    transmittance = tran_flux / norm_flux

    return transmittance, ex_dft


if __name__ == "__main__":
    norm_flux = input_flux()

    layer_thickness_um = [0.20503, 0.22066, 0.20547, 0.24602, 0.17276, 0.28224, 0.21971, 0.44517, 0.33954]
    tran, ex_dft = multilayer_stack(norm_flux, layer_thickness_um)
    print(f"tran:, {tran:.6f}")

    layer_thickness_um = [0.50000, 0.33898, 0.50000, 0.30020, 0.19164, 0.24854, 0.19684, 0.24210, 0.19353]
    tran_2, ex_dft_2 = multilayer_stack(norm_flux, layer_thickness_um)
    print(f"tran:, {tran_2:.6f}")

    z_grid = design_region_to_grid(ex_dft.size)
    fig, ax = plt.subplots()
    ax.plot(
        z_grid,
        np.absolute(ex_dft)**2,
        'bo-',
        label="obj. func. = fields in stack"
    )
    ax.plot(
        z_grid,
        np.absolute(ex_dft_2)**2,
        'ro-',
        label="obj. func. = transmittance"
    )
    ax.set_xlabel("z coordinate (μm)")
    ax.set_ylabel("$|E_x|^2$ in multilayer stack")
    ax.legend()
    if mp.am_master():
        fig.savefig(
            "field_decay_in_multilayer_stack.png",
            dpi=150,
            bbox_inches="tight"
        )