Diffractive metasurface inverse design with topology optimization¶

import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import tidy3d as td
import tidy3d.plugins.adjoint as tda

# wavelength and source properties
wavelength = 1.0
freq0 = td.C_0 / wavelength
fwidth = freq0 / 10
run_time = 100 / fwidth

# permittivity of the mask and substrate
permittivity = 2.0

# side length on x and y
length = 20

# thickess of the metalens, enough to apply a relative phase shift of just over pi
k0 = 2 * np.pi / wavelength
delta_n = np.sqrt(permittivity) - 1
thickness = 4 / k0 / delta_n

# distances between PML and source / monitor
buffer = 1.5 * wavelength

# distances between source / monitor and the mask
dist_src = 1.5 * wavelength
dist_mnt = 6.1 * wavelength

# resolution
min_steps_per_wvl = 17

# total z size and the center of the slab
Lz = buffer + dist_src + thickness + dist_mnt + buffer
z_center_slab = -Lz / 2 + buffer + dist_src + thickness / 2.0

# resolution of the design region
dl_design_region = 2 * wavelength / min_steps_per_wvl / np.sqrt(permittivity)

# number of pixel cells in the design region (in x and y)
nx = ny = int(length / dl_design_region)

# substrate of the same permittivity as the mask
substrate = td.Structure(
    geometry=td.Box.from_bounds(
        rmin=(-td.inf, -td.inf, -1000), rmax=(+td.inf, +td.inf, z_center_slab - thickness / 2)
    ),
    medium=td.Medium(permittivity=permittivity),
)

# plane wave
src = td.PlaneWave(
    center=(0, 0, -Lz / 2 + buffer),
    size=(td.inf, td.inf, 0),
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    direction="+",
)

# monitor we use to measure the intensity pattern above the device
mnt_out = td.FieldMonitor(
    center=(0, 0, +Lz / 2 - buffer),
    size=(td.inf, td.inf, 0),
    freqs=[freq0],
    colocate=False,
    name="output",
)

# monitor we use to inspect the field pattern from the side for visualization
mnt_side = td.FieldMonitor(
    center=(0, 0, 0),
    size=(td.inf, 0, td.inf),
    freqs=[freq0],
    name="side",
)

from tidy3d.plugins.adjoint.utils.filter import BinaryProjector, ConicFilter

radius = 0.20
beta = 50

conic_filter = ConicFilter(radius=radius, design_region_dl=dl_design_region)


def filter_project(params: jnp.ndarray, beta: float, eta=0.5) -> jnp.ndarray:
    """Apply conic filter and binarization to the raw params."""
    params_smooth = conic_filter.evaluate(params)
    binary_projector = BinaryProjector(vmin=0, vmax=1, beta=beta, eta=eta)
    params_smooth_binarized = binary_projector.evaluate(params_smooth)
    return params_smooth_binarized


def get_eps(params: jnp.ndarray, beta: float) -> jnp.ndarray:
    """Get the permittivity values (1, permittivity) array as a function of the parameters (0, 1)"""
    mask = filter_project(params, beta)
    eps = 1 + mask * (permittivity - 1)
    return eps.reshape((nx, ny, 1, 1))


def make_slab(params: jnp.ndarray, beta: float) -> tda.JaxStructure:
    """make the phase mask as a function of the parameters for a given `beta` value."""

    # construct the coordinates
    x0_max = +length / 2 - dl_design_region / 2
    y0_max = +length / 2 - dl_design_region / 2
    coords_x = np.linspace(-x0_max, x0_max, nx).tolist()
    coords_y = np.linspace(-y0_max, y0_max, ny).tolist()
    coords = dict(x=coords_x, y=coords_y, z=[z_center_slab], f=[freq0])

    # construct the data array for the permittivity
    eps_values = get_eps(params, beta)
    eps_data_array = tda.JaxDataArray(values=eps_values, coords=coords)

    # construct the permittiviy dataset
    field_components = {f"eps_{dim}{dim}": eps_data_array for dim in "xyz"}
    eps_dataset = tda.JaxPermittivityDataset(**field_components)

    # construct the phase mask slab
    custom_medium = tda.JaxCustomMedium(eps_dataset=eps_dataset)
    box = td.Box(center=(0, 0, z_center_slab), size=(td.inf, td.inf, thickness))
    return tda.JaxStructureStaticGeometry(geometry=box, medium=custom_medium)

def make_sim(params: jnp.ndarray, beta: float, pml_xy: bool = False) -> tda.JaxSimulation:
    """The `JaxSimulation` as a function of the design parameters."""

    slab = make_slab(params, beta)

    # put a mesh override structure to ensure uniform dl across the slab
    design_region_mesh = td.MeshOverrideStructure(
        geometry=slab.geometry,
        dl=[dl_design_region] * 3,
        enforce=True,
    )

    return tda.JaxSimulation(
        size=(length, length, Lz),
        grid_spec=td.GridSpec.auto(
            min_steps_per_wvl=min_steps_per_wvl, override_structures=[design_region_mesh]
        ),
        boundary_spec=td.BoundarySpec.pml(x=pml_xy, y=pml_xy, z=True),
        input_structures=[slab],
        structures=[substrate],
        output_monitors=[mnt_out],
        monitors=[mnt_side],
        sources=[src],
        run_time=run_time,
    )

params0 = np.random.random((nx, ny))
beta0 = 1.0

symmetrize = True

# symmetrize the starting parameters (optional)
if symmetrize:
    params0 += np.fliplr(params0)
    params0 += np.flipud(params0)
    params0 /= 4.0

sim = make_sim(params=params0, beta=beta0)

import cv2
import xarray as xr

logo_fname = "misc/logo.png"


def get_logo() -> np.ndarray:
    """Get the Flexcompute logo from file, load it into a numpy array, rescale it to (0, 1)."""
    im = cv2.imread(logo_fname, cv2.IMREAD_GRAYSCALE).astype(float)
    im -= np.min(im)
    im /= np.max(im)
    return im


def intensity_desired_fn_logo(xs: list, ys: list, rescale: float = 0.5) -> np.ndarray:
    """Return the 'value' of the flexcompute logo as a function of (x,y) with some rescaling."""
    logo_values = get_logo()

    # some rotations to get the logo in the right orientation for the final intensity pattern
    logo_values = np.rot90(np.rot90(np.rot90(logo_values)))

    # re-interpolate the logo data at the supplied x,y points using xarray
    nx, ny = logo_values.shape
    xs_logo = np.linspace(rescale * min(xs), rescale * max(xs), nx)
    ys_logo = np.linspace(rescale * min(ys), rescale * max(ys), ny)
    logo_dataarray = xr.DataArray(logo_values, coords=dict(x=xs_logo, y=ys_logo))
    logo_interp = logo_dataarray.interp(x=xs, y=ys)

    # handle any nans for out of bounds (replace with 0)
    return np.nan_to_num(logo_interp.values, nan=np.min(logo_interp))

xs = ys = np.linspace(-length / 2, length / 2, nx)
intensity_desired = intensity_desired_fn_logo(xs, ys)

def get_intensities(sim_data: tda.JaxSimulationData) -> tuple[jnp.ndarray, np.ndarray]:
    """Convenience function to grab the (unnormalized) intensity patterns from the data."""

    # first, grab the dataset storing the intensity values and coordinates
    intensity_dataset = sim_data.get_intensity(mnt_out.name)
    xs = intensity_dataset.coords["x"]
    ys = intensity_dataset.coords["y"]

    # the "measured" values are just the raw data
    intensity_measured = jnp.squeeze(intensity_dataset.values)

    # the "desired" or "target" values are the logo function evaluated at the data coords
    intensity_desired = intensity_desired_fn_logo(xs, ys)

    return intensity_measured, intensity_desired


# range within which to consider intensity as part of the objective function
# eg. if the measured intensity is above int_max, we just consider it at the target value of 1.0
intensity_range = int_min, int_max = (0.0, 1.0)


def intensity_diff_fn(sim_data: tda.JaxSimulationData) -> float:
    """Returns a measure for the amount of difference between desired and target intensity patterns."""

    intensity_measured, intensity_desired = get_intensities(sim_data)

    # normalize the measured intensity such that there's the same "power" in the signal as expected in the logo
    intensity_measured *= np.mean(intensity_desired) / intensity_norm_mean

    # apply the "capping" within intensity_range (optional)
    int_range_magnitude = abs(int_max - int_min)
    intensity_measured = jnp.minimum(intensity_measured, int_max)
    intensity_measured = jnp.maximum(intensity_measured, int_min)
    intensity_desired = int_range_magnitude * intensity_desired + int_min

    # take the elementwise difference
    difference = intensity_measured - intensity_desired

    # normalized by the 'worst case' (difference if measured was exact inverse of the target)
    difference_denominator = int_range_magnitude * np.ones_like(intensity_desired)

    # return the normalized norm of the difference
    return jnp.linalg.norm(difference) / jnp.linalg.norm(difference_denominator)

from tidy3d.plugins.adjoint.utils.penalty import ErosionDilationPenalty


def penalty_fn(params, beta):
    processed_params = filter_project(params, beta=beta)

    penalty = ErosionDilationPenalty(pixel_size=dl_design_region, length_scale=radius)
    return penalty.evaluate(processed_params)

penalty_weight = 0.1

# let's ignore the penalty for the demo, but set to True to explore how it changes the final device
include_penalty = False


def loss_fn(params: jnp.ndarray, beta: float) -> tuple[float, dict]:
    """Loss function for the design, the difference in intensity + the feature size penalty."""

    # construct and run the simulation
    sim = make_sim(params, beta=beta)
    sim_data = tda.web.run(sim, task_name="phase_mask_example", verbose=False)

    # grab the respective and total losses
    intensity_diff_loss = intensity_diff_fn(sim_data)
    if include_penalty:
        penalty_loss = penalty_weight * penalty_fn(params, beta)
    else:
        penalty_loss = 0.0

    total_loss = intensity_diff_loss + penalty_loss

    # dictionary to stash results if we want to use them in the optimization loop
    aux_data = dict(intensity_diff=intensity_diff_loss, penalty=penalty_loss, sim_data=sim_data)

    return total_loss, aux_data

# construct a function of `params` and `beta` that returns the loss value, gradient, and the aux_data
loss_fn_val_grad = jax.value_and_grad(loss_fn, has_aux=True)

# call this on our initial parameters
(val, aux_data), grad = loss_fn_val_grad(params0, beta=beta0)