Building an AI surrogate model for a directional coupler¶

# Standard python and external package imports
import gdstk
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import numpy as np
from pathlib import Path
import pickle
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# tidy3D imports
import tidy3d as td
import tidy3d.plugins.design as tdd
from tidy3d import web

# Configure notebook options for re-using simulation data if it exists
use_previous_simulation_data = False

output_dir = Path("./misc/dc_surrogate/")
mc_data_dir = output_dir / "data"
mc_data_dir.mkdir(parents=True, exist_ok=True)

# Wavelength / frequency
lambda0 = 1.550  # all length scales in microns
freq0 = td.constants.C_0 / lambda0
fwidth = freq0 / 10
ldas = np.linspace(1.5, 1.6, 101)  # wavelength range
freqs = td.C_0 / ldas  # frequency range

# Permittivity of waveguide and substrate
wg_n = 3.48
sub_n = 1.45
mat_wg = td.Medium(permittivity=wg_n**2)
mat_sub = td.Medium(permittivity=sub_n**2)

# Waveguide dimensions

# Waveguide height
wg_height = 0.22
# Waveguide width
wg_width = 0.45
# Waveguide separation in the beginning/end
wg_spacing_in = 8
# Reference plane where the cross section of the device is defined
reference_plane = "bottom"
# Angle of the sidewall deviating from the vertical ones, positive values for the base larger than the top
sidewall_angle = np.pi / 20.0
# Total device length along propagation direction
device_length = 100
# Length of the bend region
bend_length = 16
# space between waveguide and PML
pml_spacing = 1
# resolution control: minimum number of grid cells per wavelength in each material
grid_cells_per_wvl = 16

# Coupler optimization parameter ranges to consider
wg_spacing_coup_range = (0.05, 0.25)
coup_length_range = (2, 16)
top_waveguide_transmission_goal = 0.25
bottom_waveguide_transmission_goal = 0.75

# Optimization parameters
num_monte_carlo_sims = 100

filename_mc_data = mc_data_dir / f"result_{num_monte_carlo_sims}.pkl"
if not filename_mc_data.exists() and use_previous_simulation_data:
    print(
        "Warning, MC simulation data does not exist to be loaded. Reconfiguring use_previous_simulation_data flag to False and recomputing."
    )
    use_previous_simulation_data = False

def tanh_interp(max_arg):
    """Interpolator for tanh with adjustable extension"""
    scale = 1 / np.tanh(max_arg)
    return lambda u: 0.5 * (1 + scale * np.tanh(max_arg * (u * 2 - 1)))


def make_coupler(
    length,
    wg_spacing_in,
    wg_width,
    wg_spacing_coup,
    coup_length,
    bend_length,
):
    """Make an integrated coupler using the gdstk RobustPath object."""
    # bend interpolator
    interp = tanh_interp(3)
    delta = wg_width + wg_spacing_coup - wg_spacing_in
    offset = lambda u: wg_spacing_in + interp(u) * delta

    coup = gdstk.RobustPath(
        (-0.5 * length, 0),
        (wg_width, wg_width),
        wg_spacing_in,
        simple_path=True,
        layer=1,
        datatype=[0, 1],
    )
    coup.segment((-0.5 * coup_length - bend_length, 0))
    coup.segment(
        (-0.5 * coup_length, 0),
        offset=[lambda u: -0.5 * offset(u), lambda u: 0.5 * offset(u)],
    )
    coup.segment((0.5 * coup_length, 0))
    coup.segment(
        (0.5 * coup_length + bend_length, 0),
        offset=[lambda u: -0.5 * offset(1 - u), lambda u: 0.5 * offset(1 - u)],
    )
    coup.segment((0.5 * length, 0))
    return coup

def make_sim(coup_length, wg_spacing_coup, add_field_monitor=False):
    """Make a simulation with a given length of the coupling region and
    distance between the waveguides in that region. If ``add_field_monitor``
    is True, a 2D in-plane field monitor will be added.
    """

    # Geometry must be placed in GDS cells to import into Tidy3D
    coup_cell = gdstk.Cell("Coupler")

    substrate = gdstk.rectangle(
        (-device_length / 2, -wg_spacing_in / 2 - 10),
        (device_length / 2, wg_spacing_in / 2 + 10),
        layer=0,
    )
    coup_cell.add(substrate)

    # Add the coupler to a gdstk cell
    gds_coup = make_coupler(
        device_length,
        wg_spacing_in,
        wg_width,
        wg_spacing_coup,
        coup_length,
        bend_length,
    )
    coup_cell.add(gds_coup)

    # Substrate
    (oxide_geo,) = td.PolySlab.from_gds(
        gds_cell=coup_cell,
        gds_layer=0,
        gds_dtype=0,
        slab_bounds=(-10, 0),
        reference_plane=reference_plane,
        axis=2,
    )

    oxide = td.Structure(geometry=oxide_geo, medium=mat_sub)

    # Waveguides (import all datatypes if gds_dtype not specified)
    coupler1_geo, coupler2_geo = td.PolySlab.from_gds(
        gds_cell=coup_cell,
        gds_layer=1,
        slab_bounds=(0, wg_height),
        sidewall_angle=sidewall_angle,
        reference_plane=reference_plane,
        axis=2,
    )

    coupler1 = td.Structure(geometry=coupler1_geo, medium=mat_wg)
    coupler2 = td.Structure(geometry=coupler2_geo, medium=mat_wg)

    # Simulation size along propagation direction
    sim_length = 2 + 2 * bend_length + coup_length

    # Spacing between waveguides and PML
    sim_size = [
        sim_length,
        wg_spacing_in + wg_width + 2 * pml_spacing,
        wg_height + 2 * pml_spacing,
    ]

    # Source
    src_pos = -sim_length / 2 + 0.5
    msource = td.ModeSource(
        center=[src_pos, wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
        direction="+",
        mode_spec=td.ModeSpec(num_modes=1),
        mode_index=0,
    )

    mon_in = td.ModeMonitor(
        center=[(src_pos + 0.5), wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="in",
    )
    mon_ref_bot = td.ModeMonitor(
        center=[(src_pos + 0.5), -wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="reflect_bottom",
    )
    mon_top = td.ModeMonitor(
        center=[-(src_pos + 0.5), wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="top",
    )
    mon_bot = td.ModeMonitor(
        center=[-(src_pos + 0.5), -wg_spacing_in / 2, wg_height / 2],
        size=[0, 3, 2],
        freqs=freqs,
        mode_spec=td.ModeSpec(num_modes=1),
        name="bottom",
    )
    monitors = [mon_in, mon_ref_bot, mon_top, mon_bot]

    if add_field_monitor:
        field_monitor = td.FieldMonitor(
            center=[0, 0, wg_height / 2],
            size=[td.inf, td.inf, 0],
            freqs=freqs,
            name="field",
        )
        monitors.append(field_monitor)

    # Initialize the simulation
    sim = td.Simulation(
        size=sim_size,
        grid_spec=td.GridSpec.auto(min_steps_per_wvl=grid_cells_per_wvl),
        structures=[oxide, coupler1, coupler2],
        sources=[msource],
        monitors=monitors,
        run_time=50 / fwidth,
        boundary_spec=td.BoundarySpec.all_sides(boundary=td.PML()),
    )

    return sim

# Set up the parameters we want to sweep in the Monte Carlo search
param_spc = tdd.ParameterFloat(name="wg_spacing_coup", span=wg_spacing_coup_range)
param_len = tdd.ParameterFloat(name="coup_length", span=coup_length_range)
parameters = [param_spc, param_len]

monte_carlo = tdd.MethodMonteCarlo(
    num_points=num_monte_carlo_sims,
    seed=2,
)
design_space = tdd.DesignSpace(
    method=monte_carlo,
    parameters=parameters,
    task_name="DC_MC_Surrogate",
    path_dir=str(mc_data_dir),
)


def measure_coupling_amplitude(sim_data):
    """Constructs a "row" of the scattering matrix when sourced from top left port"""

    input_amp = sim_data["in"].amps.sel(direction="+")

    amps_dict = {}
    directions = ("-", "-", "+", "+")
    for monitor, direction in zip(sim_data.simulation.monitors[:4], directions):
        amp = sim_data[monitor.name].amps.sel(direction=direction)
        amp_normalized = amp / input_amp
        amps_dict[monitor.name] = amp_normalized

    return amps_dict


def fn_post(*args, **kwargs):
    """Post processing function, but label the outputs using a dictionary."""
    amps_dict = measure_coupling_amplitude(*args, **kwargs)
    return amps_dict

# Neural Network parameters
number_samples = len(expanded_results)
train_samples = int(np.floor(0.6 * number_samples))
validation_samples = int(np.floor(0.2 * number_samples))
test_samples = number_samples - train_samples - validation_samples
batch_size = 16
epochs = 120
final_training_epochs = 1000

training_dataset, validation_dataset, test_dataset = torch.utils.data.random_split(
    expanded_results, [train_samples, validation_samples, test_samples]
)

training_dataloader = DataLoader(training_dataset, batch_size=batch_size, shuffle=True)
validation_dataloader = DataLoader(
    validation_dataset, batch_size=batch_size, shuffle=False
)
testing_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# Assuming the features are coming in as (wg_spacing_coup, coup_length, freq)
# and we can normalize by the min/max possible values to center each feature in [-1, 1]
min_freq = np.min(freqs)
max_freq = np.max(freqs)

subtract_min = torch.tensor(
    [wg_spacing_coup_range[0], coup_length_range[0], min_freq], dtype=torch.float
)
normalize_max = torch.tensor(
    [
        wg_spacing_coup_range[1] - wg_spacing_coup_range[0],
        coup_length_range[1] - coup_length_range[0],
        max_freq - min_freq,
    ],
    dtype=torch.float,
)

network_input_size = 3


# Define a fully connected model in PyTorch with ReLU activation
# where we can set a variable number of hidden layers through
# neurons list. Allow option to specify dropout percentage.
class VariableANN(nn.Module):
    def __init__(self, input_size, neurons, dropout) -> None:
        super().__init__()
        self.layers = nn.ModuleList()
        self.layers.append(nn.Linear(input_size, neurons[0]))

        for i in range(1, len(neurons)):
            self.layers.append(nn.Linear(neurons[i - 1], neurons[i]))

        self.layers.append(nn.Linear(neurons[-1], 2))

        self.drop = nn.Dropout(dropout)
        self.relu = nn.ReLU()

    def forward(self, x):
        subtract_min_unsqueeze = torch.unsqueeze(subtract_min, dim=0)
        normalize_max_unsqueeze = torch.unsqueeze(normalize_max, dim=0)

        # normalize inputs to be between [-1, 1]
        x = x - subtract_min_unsqueeze
        x = x / normalize_max_unsqueeze
        x = 2.0 * (x - 0.5)

        for layer in self.layers[:-1]:
            x = self.drop(self.relu(layer(x)))

        x = self.layers[-1](x)

        return x

loss_fn = nn.MSELoss()


def evaluate_loss(data, network, loss_fn_user=None):
    """Extract inputs/labels from the data and allow for loss evaluation on user defined loss
    function (otherwise use the default loss_fn which is set to MSE)."""

    inputs, labels = data

    if loss_fn_user is None:
        loss_fn_user = loss_fn

    loss = loss_fn_user(
        network(inputs.type(torch.FloatTensor)), labels.type(torch.FloatTensor)
    )

    return loss


def evaluate_loss_from_loader(data_loader, network, loss_fn_user=None):
    """Given a data_loader, run all the batches through the network and evaluate loss."""
    loss_by_iteration = []
    for data_idx, data in enumerate(data_loader):
        loss = evaluate_loss(data, network, loss_fn_user)
        loss_by_iteration.append(loss.item())

    avg_loss = np.mean(loss_by_iteration)
    return avg_loss


def train_network(hyperparameters, epochs):
    """Train a network according to hyperparameters specification for a given number of epochs."""

    # build network and set to training mode
    network = VariableANN(
        network_input_size, hyperparameters["hidden"], hyperparameters["dropout"]
    )
    network.train()
    opt = optim.Adam(
        network.parameters(),
        lr=hyperparameters["learning_rate"],
        weight_decay=hyperparameters["weight_decay"],
    )

    training_loss_by_epoch = np.zeros(1 + epochs)
    validation_loss_by_epoch = np.zeros(1 + epochs)

    avg_training_loss = evaluate_loss_from_loader(training_dataloader, network)
    avg_validation_loss = evaluate_loss_from_loader(validation_dataloader, network)

    training_loss_by_epoch[0] = avg_training_loss
    validation_loss_by_epoch[0] = avg_validation_loss

    final_training_loss = np.inf
    final_validation_loss = np.inf
    for epoch_idx in range(0, epochs):
        training_loss_by_iteration = []

        for data_idx, data in enumerate(training_dataloader):
            opt.zero_grad()
            loss = evaluate_loss(data, network)

            training_loss_by_iteration.append(loss.item())

            loss.backward()
            opt.step()

        # Remove training freatures like dropout for the evaluation section and then turn it back on for training
        network.eval()
        avg_validation_loss = evaluate_loss_from_loader(validation_dataloader, network)
        network.train()

        training_loss_by_epoch[1 + epoch_idx] = np.mean(training_loss_by_iteration)
        validation_loss_by_epoch[1 + epoch_idx] = avg_validation_loss

    return training_loss_by_epoch, validation_loss_by_epoch, network

allowed_hidden_layer_sizes = np.arange(16, 257, 16)
allowed_dropout = np.arange(0, 0.55, 0.05)
num_hidden_layers = 2
allowed_learning_rate_log_range = (-5, -3)
allowed_weight_decay_log_range = (-5, -4)

num_trial_networks = 75


def hyperparemeter_fn(
    dropout, learning_rate_log, weight_decay_log, **hidden_layer_sizes
):
    hidden_layer_sizes = [
        hidden_layer_sizes[f"hidden_layer_size_{idx}"]
        for idx in range(0, len(hidden_layer_sizes))
    ]

    hyperparam = {
        "hidden": hidden_layer_sizes,
        "dropout": dropout,
        "learning_rate": 10**learning_rate_log,
        "weight_decay": 10**weight_decay_log,
    }

    training_loss, validation_loss, network = train_network(hyperparam, epochs)

    return -validation_loss[-1]


param_dropout = tdd.ParameterAny(name="dropout", allowed_values=tuple(allowed_dropout))
param_lr_log = tdd.ParameterFloat(
    name="learning_rate_log", span=allowed_learning_rate_log_range
)
param_wd_log = tdd.ParameterFloat(
    name="weight_decay_log", span=allowed_weight_decay_log_range
)
params_hyperparameter = [
    tdd.ParameterAny(
        name=f"hidden_layer_size_{idx}",
        allowed_values=tuple(allowed_hidden_layer_sizes),
    )
    for idx in range(0, num_hidden_layers)
]
hyperparameters = [param_dropout, param_lr_log, param_wd_log] + params_hyperparameter

hyperparam_optimizer = tdd.MethodMonteCarlo(
    num_points=num_trial_networks,
    seed=2,
)
design_space = tdd.DesignSpace(
    method=hyperparam_optimizer,
    parameters=hyperparameters,
    task_name="DC_Network_MC",
    path_dir=str(mc_data_dir),
)
network_hyperparameter_results = design_space.run(hyperparemeter_fn, verbose=True)