Inverse design optimization of a plasmonic nanoantenna metasurface¶

import matplotlib.pylab as plt

import autograd
import autograd.numpy as np

import tidy3d as td
import tidy3d.web as web

# whether to run checks for setting up sim and normalization enhancement
run_pre_sims = True

um = 1e0
nm = 1e-3

wavelength = 0.910

freq = td.C_0 / wavelength
fwidth = freq / 10
run_time = 20 / fwidth

dl = 3 * nm

# min_steps_per_wvl = 30

medium_SiO2 = td.Medium(permittivity=1.44**2)
medium_gold = td.material_library["Au"]["JohnsonChristy1972"]
eps_background = 1.0

size_design_x = 500 * nm
size_design_y = 500 * nm
thick_metal = 50 * nm
thick_sub = 0.8 * wavelength
thick_air = 0.8 * wavelength
buffer_xy = 0.0 * wavelength

hole_radius = 12 * nm

size_x = size_design_x + 2 * buffer_xy
size_y = size_design_y + 2 * buffer_xy
size_z = thick_sub + thick_metal + thick_air

top_z = size_z / 2.0

center_air_z = top_z - thick_air / 2.0
center_metal_z = top_z - thick_air - thick_metal / 2.0
center_sub_z = top_z - thick_air - thick_metal - thick_sub / 2.0

design_region_geometry = td.Box(
    center=(0, 0, center_metal_z),
    size=(td.inf, td.inf, thick_metal),
)

substrate = td.Structure(
    geometry=td.Box(
        center=(0, 0, -size_z + thick_sub),
        size=(td.inf, td.inf, size_z),
    ),
    medium=medium_SiO2,
)

plane_wave = td.PlaneWave(
    center=(0, 0, center_air_z),
    size=(td.inf, td.inf, 0),
    source_time=td.GaussianPulse(freq0=freq, fwidth=fwidth),
    pol_angle=0,
    direction="-",
)

point_monitor = td.FieldMonitor(
    size=(0, 0, 0),
    center=(0, 0, center_metal_z),
    freqs=[freq],
    name="point",
)

field_monitor_xy = td.FieldMonitor(
    center=(0, 0, center_metal_z),
    size=(td.inf, td.inf, 0),
    freqs=[freq],
    name="field_xy",
)

eps_monitor_xy = td.PermittivityMonitor(
    center=(0, 0, center_metal_z),
    size=(td.inf, td.inf, 0),
    freqs=[freq],
    name="eps_xy",
)

field_monitor_xz = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, 0, td.inf),
    freqs=[freq],
    name="field_xz",
)

sim_no_antenna = td.Simulation(
    size=(size_x, size_y, size_z),
    run_time=run_time,
    structures=[substrate],
    sources=[plane_wave],
    monitors=[point_monitor, field_monitor_xy, field_monitor_xz, eps_monitor_xy],
    grid_spec=td.GridSpec.uniform(dl=dl),
    boundary_spec=td.BoundarySpec.pml(x=False, y=False, z=True),
    # grid_spec=td.GridSpec.auto(wavelength=wavelength, min_steps_per_wvl=min_steps_per_wvl),
)

# resolution of the design region pixels
# pixel_size = 10 * nm
pixel_size = dl

# radius of the circular filter (um) (higher = larger features)
radius = 24 * nm

# projection strengths (higher = more binarized)
beta_project = 10
beta_penalty = 10

from tidy3d.plugins.autograd.invdes import make_filter_and_project, get_kernel_size_px

filter_size = get_kernel_size_px(radius, pixel_size)


def get_density(params: np.ndarray, beta: float = beta_project) -> np.ndarray:
    """Generic function to get the etch density as a function of design parameters, using filter and projection."""
    filter_size = get_kernel_size_px(radius, pixel_size)
    filter_project = make_filter_and_project(filter_size, beta=beta)
    return filter_project(params)

def make_antenna_from_density(density: np.ndarray) -> td.Structure:
    """Make a `td.Structure` containing a `td.CustomPoleResidue` corresponding to a density array."""

    rmin, rmax = sim_no_antenna.bounds
    coords = {}
    for key, pt_min, pt_max, num_pts in zip("xyz", rmin, rmax, density.shape):
        coord_edges = np.linspace(pt_min, pt_max, num_pts + 1)
        coord_centers = (coord_edges[1:] + coord_edges[:-1]) / 2.0
        coords[key] = coord_centers

    mask = td.ScalarFieldDataArray(
        np.ones_like(density),
        coords=coords,
    )
    is_in_hole = mask.x**2 + mask.y**2 >= hole_radius**2
    mask = mask.where(is_in_hole, 0)

    density_masked = density * mask.values

    eps_inf_scaled_array = eps_background + density_masked * (medium_gold.eps_inf - eps_background)

    poles_arrays_scaled = [
        (a * np.ones_like(density_masked), c * density_masked) for (a, c) in medium_gold.poles
    ]

    eps_inf_scaled = td.ScalarFieldDataArray(
        eps_inf_scaled_array,
        coords=coords,
    )

    poles_scaled = []
    for a_array, c_array in poles_arrays_scaled:
        a_scaled = td.ScalarFieldDataArray(a_array, coords=coords)
        c_scaled = td.ScalarFieldDataArray(c_array, coords=coords)
        poles_scaled.append((a_scaled, c_scaled))

    medium_gold_scaled = td.CustomPoleResidue(
        eps_inf=eps_inf_scaled, poles=poles_scaled, interp_method="linear"
    )

    return td.Structure(geometry=design_region_geometry, medium=medium_gold_scaled)

def make_antenna(params: np.ndarray, beta: float = beta_project) -> td.Structure:
    """Make the antenna structure."""
    density = get_density(params, beta=beta)
    return make_antenna_from_density(density)

def make_sim_with_antenna(
    params: np.ndarray, include_field_mnts: bool = True, beta: float = beta_project
) -> td.Simulation:
    """Make a simulation as a function of the density parameters for the antenna regions."""

    antenna = make_antenna(params, beta=beta)

    # add uniform mesh override structures to simulation (if desired)
    design_region_mesh = td.MeshOverrideStructure(
        geometry=antenna.geometry.updated_copy(size=(td.inf, td.inf, thick_metal)),
        dl=[dl] * 3,
        enforce=True,
    )

    sim = sim_no_antenna.updated_copy(
        structures=list(sim_no_antenna.structures) + [antenna],
    )

    if not include_field_mnts:
        sim = sim.updated_copy(monitors=[point_monitor])

    return sim

# some variables we might need later

nx = int(size_design_x // pixel_size)
ny = int(size_design_y // pixel_size)


# some intial parameters to test with
def make_symmetric_x(arr: np.ndarray) -> np.ndarray:
    """make an array symmetric in x."""
    return (arr + np.flipud(arr)) / 2.0


params0 = make_symmetric_x(0.5 * np.ones((nx, ny, 1)))

sim_antenna_params0 = make_sim_with_antenna(params0)

from tidy3d.plugins.autograd.invdes import make_erosion_dilation_penalty

penalty_fn = make_erosion_dilation_penalty(filter_size, beta=beta_penalty)


def penalty(params: np.ndarray, beta: float = None) -> float:
    """Define the erosion dilation invariance penalty for Si density parameters."""
    beta_kwargs = {}
    if beta is not None:
        beta_kwargs["beta"] = beta
    density = get_density(params, **beta_kwargs)
    pen_val = penalty_fn(density)
    return pen_val

# let's us grab and print autograd values while they're in the objective function
from autograd.tracer import getval


def intensity_enhancement(
    params: np.ndarray, task_name: str = "antenna_intensity", verbose=False
) -> float:
    intensity_with_params = measure_intensity(params, task_name=task_name, verbose=verbose)
    return intensity_with_params / intensity0


def objective(
    params, beta: float = beta_project, verbose: bool = False, penalty_only: bool = False
) -> float:
    if penalty_only:
        enhancement_factor = 0.0
    else:
        enhancement_factor = intensity_enhancement(params, verbose=verbose, task_name="antenna")
        print(f"\tenhancement = {getval(enhancement_factor):.2e}")

    # penalty_value = 0
    penalty_value = penalty(params, beta=beta)
    print(f"\tpenalty val = {getval(penalty_value):.2e}")
    objective_value = enhancement_factor * (1.0 - penalty_value)
    print(f"\tobjective = {getval(objective_value):.2e}")
    return objective_value

val_grad_fn = autograd.value_and_grad(objective)