from pathlib import Path

import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import pandas as pd
from PIL import Image
from IPython.display import display, clear_output

# Use Autograd's NumPy wrapper so operations on design variables remain differentiable.
import autograd.numpy as np
from autograd import grad_and_aux
from autograd.scipy.special import erf

import tidy3d as td
import tidy3d.web as web
from tidy3d.plugins.autograd import adam, apply_updates
from tidy3d.plugins.autograd.invdes import ConicFilter

td.config.logging.level = "ERROR"

# All lengths are in µm. Tidy3D uses frequency, so convert wavelengths with f = c / λ.
wl_blue = 0.450
wl_green = 0.550
wl_red = 0.650
wls = np.array([wl_blue, wl_green, wl_red])
freqs = td.C_0 / wls
freq_min, freq_max = min(freqs), max(freqs)
# Broadband Gaussian pulse covering all three target colors.
freq0 = 0.5 * (freq_min + freq_max)
fwidth = 0.65 * (freq_max - freq_min)
freq_sweep = td.C_0 / np.linspace(
    0.7, 0.4, 301
)  # broadband sweep for validation, 400-700 nm

# Nondispersive background media; dispersive SiN is loaded below.
n_glass, n_air = 1.46, 1.0
eps_glass = n_glass**2
eps_air = n_air**2

# The periodic domain is one 2 × 2 Bayer supercell with side length 2P.
P = 0.6  # pixel pitch (µm)
period = 2 * P  # design domain side = 1.2 µm (one 2×2 Bayer cell)
h = 0.6  # metasurface (post) height (µm)
f_focal = 0.6  # focal length (µm), measured from post top
target_side = 0.4  # A_lambda square side (µm)
z_src = -0.1  # plane-wave source (in glass)
z_focal = h + f_focal  # 1.2
# Leave one longest-wavelength margin above the focal plane for the PML.
sim_size = (period, period, f_focal + h + (td.C_0 / freq_min))

# Design and optimization parameters
PIXEL_SIZE = 0.01  # design grid (µm) → 120x120 design pixels
R_BLUR = 0.08  # conic filter radius ≈ min feature size (~80 nm)
ADAM_B1 = 0.9
ADAM_B2 = 0.99
ADAM_EPS = 1e-20
BETA_START = 2.0
LOG_FLOOR = 1e-20  # avoids log(0) in the figure of merit
DETA = 0.05  # eta = +/- DETA for eroded/dilated patterns
ETA_BIAS = -0.05  # Initial bias toward air

RANDOM_SEED = 0  # seeds the demo pattern below and the optimizer's initial guess (params_init); change for a different random start
N_STEP = 70  # total optimization iterations; set to a quick ~70-step demo here
LR = 0.025  # Adam learning rate
N_KEEP_BETA = 40  # steps held at BETA_START before ramping beta; scale up alongside N_STEP for a full paper-scale run
BETA_INCREMENT = 2.5  # binarization ramp

FOLDER_NAME = "ColorRouterInverseDesign"  # Tidy3D cloud task folder
SAVE_FIG = False  # save intermediate density plot
FIG_DIR = Path(f"./data_{RANDOM_SEED}/figs")  # per-step design plots
SIM_DIR = Path(f"./data_{RANDOM_SEED}/sims")  # simulation files + optimization history

# Use nondispersive background media and a measured dispersive model for SiN.
glass = td.Medium(permittivity=eps_glass)
SiN = td.material_library["SiN"]["Horiba"]

# Semi-infinite glass substrate below the metasurface
substrate = td.Structure(
    geometry=td.Box(center=(0, 0, -1.5), size=(td.inf, td.inf, 3.0)),
    medium=glass,
    name="glass substrate",
)

# Normally incident plane wave launched upward from the glass.
plane_wave = td.PlaneWave(
    source_time=td.GaussianPulse(freq0=freq0, fwidth=fwidth),
    center=(0, 0, z_src),
    size=(td.inf, td.inf, 0),
    direction="+",
    pol_angle=0,
)

# Broadband fields used for validation spectra and focal-plane maps.
focal_monitor = td.FieldMonitor(
    center=(0, 0, z_focal),
    size=(td.inf, td.inf, 0),
    freqs=freq_sweep,
    name="focal",
)

# Narrowband target monitors reduce the cost of each optimization simulation.
flux_r = td.FieldMonitor(
    center=(-P / 2, -P / 2, z_focal),
    size=(target_side, target_side, 0),
    freqs=freqs[2],
    name="flux_red",
)
flux_b = td.FieldMonitor(
    center=(P / 2, P / 2, z_focal),
    size=(target_side, target_side, 0),
    freqs=freqs[0],
    name="flux_blue",
)
flux_g1 = td.FieldMonitor(
    center=(+P / 2, -P / 2, z_focal),
    size=(target_side, target_side, 0),
    freqs=freqs[1],
    name="flux_green1",
)
flux_g2 = td.FieldMonitor(
    center=(-P / 2, +P / 2, z_focal),
    size=(target_side, target_side, 0),
    freqs=freqs[1],
    name="flux_green2",
)

# Periodic x/y boundaries model the tiled array; z-directed PMLs absorb outgoing waves.
boundary_spec = td.BoundarySpec(
    x=td.Boundary.periodic(),
    y=td.Boundary.periodic(),
    z=td.Boundary.pml(),
)

# Match the x/y FDTD grid to the design pixels and let Tidy3D refine z automatically.
grid_spec = td.GridSpec(
    grid_x=td.UniformGrid(dl=PIXEL_SIZE),
    grid_y=td.UniformGrid(dl=PIXEL_SIZE),
    grid_z=td.AutoGrid(min_steps_per_wvl=20),
)

# Shared template; make_sim later inserts a design and selects the required monitors.
sim_base = td.Simulation(
    size=sim_size,
    center=(0, 0, h),
    grid_spec=grid_spec,
    structures=[substrate],
    sources=[plane_wave],
    monitors=[focal_monitor, flux_r, flux_g1, flux_g2, flux_b],
    run_time=td.RunTimeSpec(quality_factor=200),
    boundary_spec=boundary_spec,
)

# Convert the blur radius to an odd pixel width with a unique center.
kernel_size = (
    int(R_BLUR / PIXEL_SIZE) * 2 + 1
)  # conic filter footprint, in design pixels
N_pixel = int(period / PIXEL_SIZE)  # design grid resolution (N_pixel x N_pixel)

# wrap padding makes the filter continuous across periodic cell boundaries.
conic_filter = ConicFilter(
    kernel_size=kernel_size,
    normalize=True,
    padding="wrap",
)

# Filter a unit impulse to visualize the conic kernel.
delta = np.zeros((N_pixel, N_pixel))
delta[60, 60] = 1.0
kernel = conic_filter(delta)

x_grid = np.arange(-P + PIXEL_SIZE / 2, P, PIXEL_SIZE)
y_grid = np.arange(-P + PIXEL_SIZE / 2, P, PIXEL_SIZE)
x_grid, y_grid = np.meshgrid(x_grid, y_grid, indexing="ij")

fig, ax = plt.subplots()
ax.pcolormesh(x_grid, y_grid, kernel)
ax.set(
    aspect="equal",
    xlabel=r"$x$ ($\mu$m)",
    ylabel=r"$y$ ($\mu$m)",
    xlim=(-0.2, 0.2),
    ylim=(-0.2, 0.2),
    title=r"Conic filter kernel ($r=" + f"{R_BLUR * 1e3:.0f}$ nm)",
)


# Differentiable projection from continuous parameters toward air/SiN values.
def erf_proj(arr, beta=BETA_START, eta=0.0):
    # Smoothed step function: beta sharpens toward a binary (air/SiN) pattern,
    # eta shifts the threshold to synthesize eroded/dilated variants.
    return erf(beta * (arr - eta))


arr = np.linspace(-2.5, 2.5, 501)
fig, ax = plt.subplots()
[
    ax.plot(arr, erf_proj(arr, beta=beta), label=r"$\beta = " + f"{beta}$")
    for beta in [1, 2, 10, 100]
]
ax.set(
    xlabel=r"$x$",
    ylabel=r"erf($\beta x$)",
    title="Binary projection",
    xlim=(-2.5, 2.5),
    ylim=(-1, 1),
)
ax.legend(frameon=False)


# Differentiable preprocessing pipeline; return the final projection first.
def preprocess(arr, beta=BETA_START, eta=0.0, symmetrize=True):
    result = []
    if symmetrize:
        arr_sym = (arr + arr.T) / 2  # diagonal mirror symmetry, matches the Bayer cell
        result.append(arr_sym)
    # Normalization keeps beta and eta comparable across iterations.
    arr_norm = arr_sym / np.std(arr_sym) if symmetrize else arr / np.std(arr)
    arr_blurred = conic_filter(arr_norm)  # enforces the minimum feature size
    arr_proj = erf_proj(arr_blurred, beta=beta, eta=eta)  # final binarized density
    result.append(arr_norm)
    result.append(arr_blurred)
    result.append(arr_proj)
    return result[::-1]  # most- to least-processed: [proj, blurred, norm, sym]

# Fixed seed for a reproducible preprocessing example.
rng = np.random.default_rng(RANDOM_SEED)
params_raw = rng.normal(loc=0, scale=1, size=(N_pixel, N_pixel))
# beta/eta here are arbitrarily set to make each pipeline stage visually distinct
# (actual optimization starts at beta=BETA_START and eta=0/±DETA)
params_nominal, params_blurred, params_norm, params_sym = preprocess(
    params_raw, beta=100, symmetrize=True
)
params_eroded, _, _, _ = preprocess(params_raw, beta=100, eta=0.01, symmetrize=True)
params_dilated, _, _, _ = preprocess(params_raw, beta=100, eta=-0.01, symmetrize=True)

# Use a diverging map for parameters and grayscale for air/SiN.
cmap = plt.cm.RdBu
cmap2 = plt.cm.Grays
fig, ax = plt.subplots(
    3, 5, figsize=(12, 6), sharex=True, sharey=True, tight_layout=True
)
ax[1, 0].pcolormesh(x_grid, y_grid, params_raw, vmin=-1.5, vmax=1.5, cmap=cmap)
ax[1, 0].set_title("Raw (Gaussian)")
ax[1, 1].pcolormesh(x_grid, y_grid, params_sym, vmin=-1.5, vmax=1.5, cmap=cmap)
ax[1, 1].set_title("Symmetrized")
ax[1, 1].plot([-P, P], [-P, P], "k--")
ax[1, 2].pcolormesh(x_grid, y_grid, params_norm, vmin=-1.5, vmax=1.5, cmap=cmap)
ax[1, 2].set_title("Normalized")
pc1 = ax[1, 3].pcolormesh(
    x_grid, y_grid, params_blurred, vmin=-1.5, vmax=1.5, cmap=cmap
)
ax[1, 3].set_title("Blurred")
ax[0, 4].pcolormesh(x_grid, y_grid, params_eroded, vmin=-1, vmax=1, cmap=cmap2)
ax[0, 4].set_title("Projection (eroded)")
ax[1, 4].pcolormesh(x_grid, y_grid, params_nominal, vmin=-1, vmax=1, cmap=cmap2)
ax[1, 4].set_title("Projection (nominal)")
pc2 = ax[2, 4].pcolormesh(x_grid, y_grid, params_dilated, vmin=-1, vmax=1, cmap=cmap2)
ax[2, 4].set_title("Projection (dilated)")
[ax[ii, jj].set_aspect("equal") for ii in range(3) for jj in range(5)]
[ax[ii, jj].remove() for ii in (0, 2) for jj in range(4)]

ax[1, 3].set(xticks=[-0.5, 0, 0.5], yticks=[-0.5, 0, 0.5])

cax = ax[1, 3].inset_axes([1.05, 0, 0.05, 1])
cbar = plt.colorbar(pc1, cax=cax, extend="both")
cbar.ax.set_yticks([-1, 0, 1])

cax = ax[1, 4].inset_axes([1.05, 0, 0.05, 1])
cbar = plt.colorbar(pc2, cax=cax)
cbar.ax.set_yticks([-1, 1])
cbar.ax.set_yticklabels(["air", "SiN"]);

# Tidy3D material coordinates: one sample per x/y pixel and one extruded z slice.
coords = dict(
    x=np.arange(-P + PIXEL_SIZE / 2, P, PIXEL_SIZE),
    y=np.arange(-P + PIXEL_SIZE / 2, P, PIXEL_SIZE),
    z=[h / 2],  # design layer is a single z-slice at mid-height
)


# Convert latent parameters into a differentiable Tidy3D material structure.
def get_design_structure(params, beta=BETA_START, eta=0):
    # Density in [0, 1] linearly interpolates eps_inf and each residue of the pole-residue model
    # between air and bulk SiN, so gray (partially blurred) pixels map to an effective medium.
    density = (1 + preprocess(params, beta=beta, eta=eta)[0]) / 2
    density = density.reshape(N_pixel, N_pixel, 1)
    eps_new = td.SpatialDataArray(eps_air + (SiN.eps_inf - 1) * density, coords=coords)
    # Scale SiN pole residues by density: 0 gives air and 1 gives bulk SiN.
    poles_new = []
    for a_SiN, c_SiN in SiN.poles:
        a_new = td.SpatialDataArray(a_SiN * np.ones_like(density), coords=coords)
        c_new = td.SpatialDataArray(density * c_SiN, coords=coords)
        poles_new.append((a_new, c_new))
    return td.Structure(
        geometry=td.Box(center=(0, 0, h / 2), size=(td.inf, td.inf, h)),
        medium=td.CustomPoleResidue(eps_inf=eps_new, poles=poles_new),
    )


# Select narrowband optimization monitors or the broadband validation monitor.
def make_sim(params, beta=BETA_START, eta=0, validation=False):
    design = get_design_structure(params, beta=beta, eta=eta)
    return sim_base.updated_copy(
        structures=[substrate, design],
        # Narrowband flux monitors for optimization; broadband focal monitor for validation
        monitors=[focal_monitor] if validation else [flux_r, flux_g1, flux_g2, flux_b],
    )


# Build a random candidate to check the geometry and monitor locations.
sim = make_sim(params_raw, beta=100)

# Side view of the simulation geometry.
fig, ax = plt.subplots()
sim.plot_eps(x=P / 2, freq=freqs[1], ax=ax)
ax.set_title("Side view")


# Top view of the RGGB target-monitor locations.
fig, ax = plt.subplots()
sim.plot(z=z_focal, ax=ax)

ax.text(
    -P / 2,
    -P / 2,
    "Red",
    color="r",
    horizontalalignment="center",
    verticalalignment="center",
)
ax.text(
    P / 2,
    -P / 2,
    "Green",
    color="g",
    horizontalalignment="center",
    verticalalignment="center",
)
ax.text(
    -P / 2,
    P / 2,
    "Green",
    color="g",
    horizontalalignment="center",
    verticalalignment="center",
)
ax.text(
    P / 2,
    P / 2,
    "Blue",
    color="b",
    horizontalalignment="center",
    verticalalignment="center",
)
ax.set_title("Focal plane");

# Combine Tidy3D monitor fluxes into one differentiable objective.
def figure_of_merit(sim_data):
    RE_red = sim_data["flux_red"].flux.values.mean()
    RE_blue = sim_data["flux_blue"].flux.values.mean()
    # Add the two green subpixels in the RGGB cell.
    RE_green = (
        sim_data["flux_green1"].flux.values.mean()
        + sim_data["flux_green2"].flux.values.mean()
    )  # two green subpixels per Bayer cell
    # Sum of log10(efficiency): balances channels so one strong channel can't mask two weak ones
    FOM = (
        np.log10(RE_red + LOG_FLOOR)
        + np.log10(RE_green + LOG_FLOOR)
        + np.log10(RE_blue + LOG_FLOOR)
    )
    return FOM, (RE_red, RE_green, RE_blue)


# Return a differentiable FOM plus auxiliary metrics for logging.
def param2fom(params, beta=BETA_START, step=0, verbose=False, robust=False):
    if robust:
        # Fabrication-robust evaluation: eroded/nominal/dilated variants run as one async batch
        sims = {
            f"seed{RANDOM_SEED}_step{step} eroded": make_sim(
                params, beta=beta, eta=DETA
            ),
            f"seed{RANDOM_SEED}_step{step} nominal": make_sim(params, beta=beta, eta=0),
            f"seed{RANDOM_SEED}_step{step} dilated": make_sim(
                params, beta=beta, eta=-DETA
            ),
        }
        batch_data = web.run_async(
            simulations=sims,
            path_dir=SIM_DIR / f"step_{step}",
            folder_name=FOLDER_NAME,
            verbose=verbose,
        )
        # Log channel efficiencies for the nominal device only.
        FOMn, (RE_red, RE_green, RE_blue) = figure_of_merit(
            batch_data[f"seed{RANDOM_SEED}_step{step} nominal"]
        )
        FOMe, _ = figure_of_merit(batch_data[f"seed{RANDOM_SEED}_step{step} eroded"])
        FOMd, _ = figure_of_merit(batch_data[f"seed{RANDOM_SEED}_step{step} dilated"])
        FOM_tot = 0.5 * FOMn + 0.25 * (FOMe + FOMd)  # weight toward robustness
        # Extract plain values from Autograd-traced results for storage.
        return FOM_tot, [
            FOM_tot._value,
            FOMn._value,
            FOMe._value,
            FOMd._value,
            RE_red._value,
            RE_green._value,
            RE_blue._value,
        ]
    else:
        # Nominal-only evaluation (used before the robustness ramp begins)
        sim = make_sim(params, beta=beta, eta=0)
        sim_data = web.run(
            simulation=sim,
            path=SIM_DIR / f"step_{step}.hdf5",
            folder_name=FOLDER_NAME,
            verbose=verbose,
            task_name=f"seed{RANDOM_SEED}_step{step}",
        )
        FOM_tot, (RE_red, RE_green, RE_blue) = figure_of_merit(sim_data)
        return FOM_tot, [
            FOM_tot._value,
            FOM_tot._value,
            np.nan,
            np.nan,
            RE_red._value,
            RE_green._value,
            RE_blue._value,
        ]


# Adjoint gradient of the FOM w.r.t. params, plus the auxiliary FOM/efficiency outputs
grad_aux_p2f = grad_and_aux(param2fom, argnum=(0,))

job = web.Job(
    simulation=sim, task_name="cost estimation", folder_name=FOLDER_NAME, verbose=False
)
cost = web.estimate_cost(job.task_id, verbose=False)

print(
    f"Estimated FlexCredits per nominal-only iteration (1 fwd. + 3 adj.): up to {4 * cost:.2f}.",
    f"Estimated total FlexCredits: up to {4 * cost * (N_KEEP_BETA + (N_STEP - N_KEEP_BETA) * 3):.2f}.",
    sep="\n",
)

Estimated FlexCredits per nominal-only iteration (1 fwd. + 3 adj.): up to 0.26.
Estimated total FlexCredits: up to 33.73.

# Rows used by the live optimization table.
optimization_display_rows = []


def print_optimization_table_row(
    step,
    beta,
    FOM_tot,
    FOMn,
    FOMe,
    FOMd,
    RE_red,
    RE_green,
    RE_blue,
    grad_max_abs,
):
    # Convert metrics to floats and redraw the current table.
    optimization_display_rows.append(
        {
            "step": int(step),
            "beta": float(beta),
            "FOM total": float(FOM_tot),
            "nominal": float(FOMn),
            "eroded": float(FOMe),
            "dilated": float(FOMd),
            "RE red (%)": float(RE_red * 100),
            "green (%)": float(RE_green * 100),
            "blue (%)": float(RE_blue * 100),
            "max grad": float(grad_max_abs),
        }
    )
    df = pd.DataFrame(optimization_display_rows).set_index("step")
    clear_output(wait=True)
    display(
        df.style.format(
            {
                "beta": "{:.1f}",
                "FOM total": "{:.4f}",
                "nominal": "{:.3f}",
                "eroded": "{:.3f}",
                "dilated": "{:.3f}",
                "RE red (%)": "{:.3f}",
                "green (%)": "{:.3f}",
                "blue (%)": "{:.3f}",
                "max grad": "{:.2e}",
            }
        )
    )


# Save the eroded, nominal, and dilated patterns for this iteration.
def plot_density(params, beta=BETA_START, step=0):
    # Renders and saves the eroded/nominal/dilated pattern for this step
    params_nominal, _, _, _ = preprocess(params, beta=beta, symmetrize=True)
    params_eroded, _, _, _ = preprocess(params, beta=beta, eta=DETA, symmetrize=True)
    params_dilated, _, _, _ = preprocess(params, beta=beta, eta=-DETA, symmetrize=True)
    fig, ax = plt.subplots(
        1, 3, figsize=(9, 3), sharex=True, sharey=True, tight_layout=True
    )
    ax[0].pcolormesh(x_grid, y_grid, params_eroded, vmin=-1, vmax=1, cmap=cmap2)
    ax[1].pcolormesh(x_grid, y_grid, params_nominal, vmin=-1, vmax=1, cmap=cmap2)
    pc = ax[2].pcolormesh(x_grid, y_grid, params_dilated, vmin=-1, vmax=1, cmap=cmap2)
    cax = ax[2].inset_axes([1.05, 0, 0.05, 1])
    cbar = plt.colorbar(pc, cax=cax)
    cbar.ax.set_yticks([-1, 1])
    cbar.ax.set_yticklabels(["air", "SiN"])
    ax[0].set_title(f"step {step} eroded")
    ax[1].set_title(f"step {step} nominal")
    ax[2].set_title(f"step {step} dilated")
    [a.set_aspect("equal") for a in ax]
    ax[0].set(
        xticks=[],
        yticks=[],
    )
    FIG_DIR.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIG_DIR / f"step{step}.png")
    plt.close(fig)


# Reproducible initial field with a small bias toward air.
rng = np.random.default_rng(RANDOM_SEED)
params_init = rng.normal(
    loc=ETA_BIAS, scale=1, size=(N_pixel, N_pixel)
)  # biased toward air
optimizer = adam(learning_rate=LR, beta1=ADAM_B1, beta2=ADAM_B2, eps=ADAM_EPS)

# Track convergence and retain the best pre-update design.
params = params_init.copy()
opt_state = optimizer.init(params)
params_history = []
grad_history = []
FOM_history = []
FOMned_history = []
RE_red_history = []
RE_green_history = []
RE_blue_history = []
beta_history = []

# Each iteration runs differentiable FDTD, evaluates the adjoint gradient, and updates all pixels.
for step in range(N_STEP):
    params_history.append(params.copy())
    # Binarization schedule: hold beta flat for N_KEEP_BETA steps, then ramp it up
    beta = (
        BETA_START
        if step < N_KEEP_BETA
        else BETA_START + BETA_INCREMENT * (step - N_KEEP_BETA + 1)
    )
    if SAVE_FIG:
        # Save the geometry evaluated at this step before updating it.
        plot_density(params, beta=beta, step=step)

    # Switch on the eroded/nominal/dilated robustness ensemble once binarization begins
    robust = False if step < N_KEEP_BETA else True
    (grad,), (FOM_tot, FOMn, FOMe, FOMd, RE_red, RE_green, RE_blue) = grad_aux_p2f(
        params, beta=beta, step=step, robust=robust
    )

    grad_max_abs = np.abs(grad).max()
    # Adam is written for gradient descent; negate grad to ascend the FOM
    updates, opt_state = optimizer.update(-grad, opt_state, params)
    params = apply_updates(params, updates)

    # Store history.
    grad_history.append(grad.copy())
    FOM_history.append(float(FOM_tot))
    FOMned_history.append([FOMn, FOMe, FOMd])
    RE_red_history.append(float(RE_red))
    RE_green_history.append(float(RE_green))
    RE_blue_history.append(float(RE_blue))
    beta_history.append(float(beta))

    # Print progress table.
    print_optimization_table_row(
        step=step,
        beta=beta,
        FOM_tot=FOM_tot,
        FOMn=FOMn,
        FOMe=FOMe,
        FOMd=FOMd,
        RE_red=RE_red,
        RE_green=RE_green,
        RE_blue=RE_blue,
        grad_max_abs=grad_max_abs,
    )

print(f"Best FOM_tot: {float(np.max(FOM_history)):.6f}")

# Save the history for analysis without rerunning cloud simulations.
np.savez(
    SIM_DIR / "opt_history.npz",
    params=np.array(params_history),
    grad=np.array(grad_history),
    FOM=np.array(FOM_history),
    FOMned=np.array(FOMned_history),
    RE_red=np.array(RE_red_history),
    RE_green=np.array(RE_green_history),
    RE_blue=np.array(RE_blue_history),
    beta=np.array(beta_history),
)

Best FOM_tot: -1.457629

fig, ax = plt.subplots(2, 1, sharex=True, tight_layout=True)
N_step_performed = len(FOM_history)
steps = np.arange(N_step_performed)
ax[0].plot(steps, np.array(FOM_history), "k")
ax[1].plot(steps, np.array(RE_red_history) * 100, "r")
ax[1].plot(steps, np.array(RE_green_history) * 100, "green")
ax[1].plot(steps, np.array(RE_blue_history) * 100, "b")

# Shade phase 1 (fixed beta and nominal-only evaluation).
ymin, ymax = np.array(FOM_history).min(), np.array(FOM_history).max() + 0.01
ax[0].fill_between(
    [0, N_KEEP_BETA - 0.5], [ymin] * 2, [ymax] * 2, color="gray", alpha=0.2
)
ax[0].text(0, ymax + 0.1, r"Phase 1: fixed $\beta$, nominal-only")
ax[0].text(N_KEEP_BETA, ymax + 0.1, r"Phase 2: $\beta$ ramp, robust design")
ax[0].set(ylabel=r"FoM", ylim=(ymin, ymax))

ymax = np.array([RE_red_history, RE_green_history, RE_blue_history]).max() * 100 + 1
ax[1].fill_between([0, N_KEEP_BETA - 0.5], [0] * 2, [ymax] * 2, color="gray", alpha=0.2)

ax[1].set(
    xlim=(0, N_step_performed - 1),
    xlabel=r"Iteration",
    ylabel=r"Routing Efficiency (%)",
    ylim=(0, ymax),
)

baseline = target_side**2 / period**2 * 100
ax[1].axhline(y=baseline * 2, color="g", ls="--", label=r"Baseline for green")
ax[1].axhline(y=baseline, color="m", ls="--", label=r"Baseline for red/blue")
ax[1].legend(frameon=False);

IMAGE_PATH = Path(
    "img/ColorRouter_ref.png"
)  # cropped image of the design shown in the reference paper
BLACK_THRESHOLD = 128  # pixels darker than this are treated as SiN
FILL_FRACTION_THRESHOLD = (
    0.5  # output cell becomes SiN once this fraction of its source pixels is black
)
BORDER_IGNORE_PIXELS = (
    4  # ignore this many pixels along each outer edge (thin black frame in the image)
)

# Convert the reference image to a grayscale SiN/air mask.
image_array = np.asarray(Image.open(IMAGE_PATH).convert("L"))  # black = 0, white = 255
height, width = image_array.shape

# Remove the figure's thin black border before downsampling.
black_mask = image_array < BLACK_THRESHOLD
black_mask[:BORDER_IGNORE_PIXELS, :] = False
black_mask[-BORDER_IGNORE_PIXELS:, :] = False
black_mask[:, :BORDER_IGNORE_PIXELS] = False
black_mask[:, -BORDER_IGNORE_PIXELS:] = False

# Downsample onto the design grid using the black fraction in each cell.
y_edges = np.linspace(0, height, N_pixel + 1)
x_edges = np.linspace(0, width, N_pixel + 1)

black_fraction = np.zeros((N_pixel, N_pixel), dtype=float)
for iy in range(N_pixel):
    y0, y1 = int(np.floor(y_edges[iy])), int(np.ceil(y_edges[iy + 1]))
    for ix in range(N_pixel):
        x0, x1 = int(np.floor(x_edges[ix])), int(np.ceil(x_edges[ix + 1]))
        cell = black_mask[y0:y1, x0:x1]
        if cell.size > 0:
            black_fraction[iy, ix] = cell.mean()

# Create a binary mask and flip the image to match simulation coordinates.
params_ref = (black_fraction >= FILL_FRACTION_THRESHOLD).astype(np.float32)[
    ::-1
]  # 1 = SiN, 0 = air
params_ref = (
    params_ref + params_ref.T
) / 2  # match the symmetrization used for our own design

# Build the Tidy3D material directly because the reference is already binary.
# Use the same pole-residue interpolation as get_design_structure.
density = params_ref.reshape(N_pixel, N_pixel, 1)
eps_new = td.SpatialDataArray(1 + (SiN.eps_inf - 1) * density, coords=coords)
poles_new = []
for a_SiN, c_SiN in SiN.poles:
    a_new = td.SpatialDataArray(a_SiN * np.ones_like(density), coords=coords)
    c_new = td.SpatialDataArray(density * c_SiN, coords=coords)
    poles_new.append((a_new, c_new))
design_ref = td.Structure(
    geometry=td.Box(center=(0, 0, h / 2), size=(td.inf, td.inf, h)),
    medium=td.CustomPoleResidue(eps_inf=eps_new, poles=poles_new),
)

# Use the same broadband Tidy3D setup as the optimized design.
sim_ref = sim_base.updated_copy(
    structures=[substrate, design_ref], monitors=[focal_monitor]
)

# Select the highest recorded FOM rather than assuming the final step is best.
idx_best = np.argmax(FOM_history)  # step with the highest recorded figure of merit

params_best = params_history[idx_best]
beta_best = beta_history[idx_best]

# Rebuild the nominal design with the broadband validation monitor.
sim_best = make_sim(params_best, beta=beta_best, eta=0, validation=True)
sim_best.plot_eps(z=h / 2, freq=freqs[1]);

# Run both validation simulations concurrently with Tidy3D web.run_async.
batch_data = web.run_async(
    simulations={"best": sim_best, "ref": sim_ref},
    folder_name=FOLDER_NAME,
    path_dir=Path(SIM_DIR),
    verbose=False,
)

sim_data_best, sim_data_ref = batch_data["best"], batch_data["ref"]

# Integrate Tidy3D focal-plane fields over each Bayer target.
def RoutingEfficiency(sim_data):
    # squeeze preserves the labeled x, y, and frequency coordinates.
    Ex = sim_data["focal"].Ex.squeeze()
    Ey = sim_data["focal"].Ey.squeeze()
    Hx = sim_data["focal"].Hx.squeeze()
    Hy = sim_data["focal"].Hy.squeeze()
    # Time-averaged z-directed Poynting flux from complex phasor fields.
    Sz = 0.5 * np.real(
        Ex * Hy.conj() - Ey * Hx.conj()
    )  # time-averaged Poynting flux, z-component
    x, y, f = Sz.x, Sz.y, Sz.f
    # Select each square detector aperture and integrate over x and y.
    Pz_red = Sz.sel(
        x=x[np.abs(x + P / 2) < target_side / 2],
        y=y[np.abs(y + P / 2) < target_side / 2],
    ).integrate(coord=["x", "y"])
    Pz_green1 = Sz.sel(
        x=x[np.abs(x - P / 2) < target_side / 2],
        y=y[np.abs(y + P / 2) < target_side / 2],
    ).integrate(coord=["x", "y"])
    Pz_green2 = Sz.sel(
        x=x[np.abs(x + P / 2) < target_side / 2],
        y=y[np.abs(y - P / 2) < target_side / 2],
    ).integrate(coord=["x", "y"])
    Pz_blue = Sz.sel(
        x=x[np.abs(x - P / 2) < target_side / 2],
        y=y[np.abs(y - P / 2) < target_side / 2],
    ).integrate(coord=["x", "y"])
    # Combine both green targets and convert frequency back to wavelength.
    Pz_green = Pz_green1 + Pz_green2
    return td.C_0 / f.values, Pz_red.values, Pz_green.values, Pz_blue.values


wl_sweep, RE_red, RE_green, RE_blue = RoutingEfficiency(sim_data_best)
wl_sweep_ref, RE_red_ref, RE_green_ref, RE_blue_ref = RoutingEfficiency(sim_data_ref)
fig, ax = plt.subplots()
ax.plot(wl_sweep, RE_red * 100, "r", label="R")
ax.plot(wl_sweep, RE_green * 100, "g", label="G1+G2")
ax.plot(wl_sweep, RE_blue * 100, "b", label="B")

ax.fill_between(
    wl_sweep_ref, RE_red_ref * 100, alpha=0.1, color="r", label="R (ref)", lw=0
)
ax.fill_between(
    wl_sweep_ref,
    RE_green_ref * 100,
    alpha=0.1,
    color="g",
    label="G1+G2 (ref)",
    lw=0,
)
ax.fill_between(
    wl_sweep_ref, RE_blue_ref * 100, alpha=0.1, color="b", label="B (ref)", lw=0
)

ymax = (
    max(
        np.max(RE_red),
        np.max(RE_green),
        np.max(RE_blue),
        np.max(RE_red_ref),
        np.max(RE_green_ref),
        np.max(RE_blue_ref),
    )
    * 100
)
# Mark the three optimization wavelengths.
[ax.axvline(x=wl, color="gray", ls="--") for wl in wls]
ax.set(
    xlim=(0.4, 0.7),
    ylim=(0, ymax),
    xlabel=r"Wavelength ($\mu$m)",
    ylabel=r"Routing Efficiency (%)",
)
ax.legend(ncol=2);

# Compute the z-directed, time-averaged Poynting flux.
def power_flux(sim_data):
    Ex = sim_data["focal"].Ex.squeeze()
    Ey = sim_data["focal"].Ey.squeeze()
    Hx = sim_data["focal"].Hx.squeeze()
    Hy = sim_data["focal"].Hy.squeeze()
    Sz = 0.5 * np.real(
        Ex * Hy.conj() - Ey * Hx.conj()
    )  # time-averaged Poynting flux, z-component
    return Sz


# Use one color scale for a direct optimized/reference comparison.
Sz = power_flux(sim_data_best)
Sz_ref = power_flux(sim_data_ref)
vmax = max(Sz.max(), Sz_ref.max())
x, y = np.meshgrid(Sz.x, Sz.y, indexing="ij")

# Rows: optimized/reference; columns: blue/green/red.
fig, ax = plt.subplots(2, 3, sharex=True, sharey=True, tight_layout=True)
x_det = target_side * np.array([1, 1, -1, -1, 1]) / 2
y_det = target_side * np.array([1, -1, -1, 1, 1]) / 2

for jj in range(2):
    for ii, (f, c) in enumerate(zip(freqs, "bgr")):
        # Single-channel (b/g/r) colormap so each wavelength's flux overlays as its own color
        cmap = ListedColormap(
            np.linspace(0, 1, 256).reshape(-1, 1)
            * (np.arange(3) == (2 - ii)).reshape(1, -1)
        )
        # Select the nearest sampled Tidy3D frequency coordinate.
        focal_image = (
            Sz.sel(f=f, method="nearest")
            if jj == 0
            else Sz_ref.sel(f=f, method="nearest")
        )
        ax[jj, ii].pcolormesh(x, y, focal_image, vmin=0, vmax=vmax, cmap=cmap)
        ax[jj, ii].set_aspect("equal")
        ax[jj, ii].set(xlim=(-P, P), ylim=(-P, P), xticks=[], yticks=[])
        ax[jj, ii].set_title(
            f"λ = {td.C_0 / f * 1000:.0f} nm" + (" (Opt.)" if jj == 0 else " (Ref.)")
        )

    # Outline each Bayer subpixel's target square in white
    ax[jj, 0].plot(P / 2 + x_det, P / 2 + y_det, color="w", ls="--", lw=0.75)
    ax[jj, 1].plot(-P / 2 + x_det, P / 2 + y_det, color="w", ls="--", lw=0.75)
    ax[jj, 1].plot(P / 2 + x_det, -P / 2 + y_det, color="w", ls="--", lw=0.75)
    ax[jj, 2].plot(-P / 2 + x_det, -P / 2 + y_det, color="w", ls="--", lw=0.75)

TIDY3D
LEARNING CENTER

Metasurface color router design for image sensors¶

Simulation Setup

Base Simulation

Design Parameters

From Design Parameters to a Simulation

Figure of Merit and Adjoint Gradient

Cost Estimation

Optimization

Results and Analysis

Reference Design

Validating the Best Design

Optimized Routing Efficiencies

Concluding Remarks

More Examples

TIDY3D LEARNING CENTER

Metasurface color router design for image sensors¶

Simulation Setup

Base Simulation

Design Parameters

From Design Parameters to a Simulation

Figure of Merit and Adjoint Gradient

Cost Estimation

Optimization

Results and Analysis

Reference Design

Validating the Best Design

Optimized Routing Efficiencies

Concluding Remarks

More Examples

Cookie preferences

Privacy choices

TIDY3D
LEARNING CENTER