fdfd.py

"""
2D frequency-domain finite-difference Maxwell solver (FDFD), Ez polarization.

This is the rung-6 engine: a grid-based solver for the scalar Helmholtz equation
    d/dx (1/sx d/dx) Ez + d/dy (1/sy d/dy) Ez + k0^2 eps(x,y) Ez = source,
with stretched-coordinate PML (sx, sy) absorbing outgoing waves. Unlike TMM
(which only handles layered media), FDFD handles ARBITRARY 2D geometry -- the
on-chip devices that adjoint topology optimization designs.

We use a scattered-field formulation: E = E_inc + E_scat, where E_inc is an
analytic plane wave in the incident medium and we solve for E_scat (radiating
into PML). Reflectance/transmittance are read off via Poynting flux on monitor
lines.

CRITICAL: this solver is cross-validated against the independent TMM engine
(tmm.py) for layered structures at normal and oblique incidence (see __main__).
Two solvers built on completely different math agreeing is the strongest possible
correctness check.

Convention: time dependence e^{-i w t}; forward wave ~ e^{+i k y}.
x is periodic (with Bloch phase for oblique incidence); y has PML top & bottom.
"""

from __future__ import annotations

import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla

import tmm


def _sigma_profile(pos, Ny, npml, sigma0, m):
    """PML conductivity at continuous grid position pos in [0, Ny-1]."""
    w = npml
    depth = np.zeros_like(pos, dtype=float)
    top = pos < w
    bot = pos > (Ny - 1 - w)
    depth[top] = (w - pos[top]) / w
    depth[bot] = (pos[bot] - (Ny - 1 - w)) / w
    return sigma0 * depth**m


def make_stretch(N, npml, sigma0, m):
    """SC-PML stretch factors for one axis: s_int (N,) and s_half (N+1,) at edges.

    npml == 0 -> no PML on this axis (s = 1), used for periodic directions.
    e^{-i w t} convention: s = 1 + i*sigma absorbs the outgoing wave.
    """
    if npml == 0:
        return np.ones(N, complex), np.ones(N + 1, complex)
    s_int = 1.0 + 1j * _sigma_profile(np.arange(N, dtype=float), N, npml, sigma0, m)
    s_half = 1.0 + 1j * _sigma_profile(
        np.clip(np.arange(N + 1, dtype=float) - 0.5, 0, N - 1), N, npml, sigma0, m)
    return s_int, s_half


def build_operator(eps, wavelength, h, npml_y, npml_x=0, kx=0.0, sigma0=8.0, m=3):
    """Sparse A for the stretched 2D Laplacian + k0^2 eps. eps shape (Ny,Nx).

    npml_x == 0 -> x is periodic (Bloch phase kx); npml_x > 0 -> x has PML
    (isolated device). y always has PML (npml_y)."""
    Ny, Nx = eps.shape
    k0 = 2 * np.pi / wavelength
    sy_int, sy_half = make_stretch(Ny, npml_y, sigma0, m)
    sx_int, sx_half = make_stretch(Nx, npml_x, sigma0, m)
    x_periodic = (npml_x == 0)
    bloch = np.exp(1j * kx * Nx * h)
    inv_h2 = 1.0 / h**2
    rows, cols, vals = [], [], []

    def idx(i, j):
        return j * Nx + i

    for j in range(Ny):
        s0y, spy, smy = sy_int[j], sy_half[j + 1], sy_half[j]
        cyp = inv_h2 / (s0y * spy); cym = inv_h2 / (s0y * smy)
        for i in range(Nx):
            n = idx(i, j)
            s0x, spx, smx = sx_int[i], sx_half[i + 1], sx_half[i]
            cxp = inv_h2 / (s0x * spx); cxm = inv_h2 / (s0x * smx)
            diag = -(cxp + cxm + cyp + cym) + k0**2 * eps[j, i]
            # x neighbours
            if x_periodic:
                ph_p = bloch if i == Nx - 1 else 1.0
                ph_m = (1.0 / bloch) if i == 0 else 1.0
                rows += [n, n]; cols += [idx((i + 1) % Nx, j), idx((i - 1) % Nx, j)]
                vals += [cxp * ph_p, cxm * ph_m]
            else:
                if i + 1 < Nx:
                    rows += [n]; cols += [idx(i + 1, j)]; vals += [cxp]
                if i - 1 >= 0:
                    rows += [n]; cols += [idx(i - 1, j)]; vals += [cxm]
            # y neighbours (Dirichlet beyond outer rows)
            if j + 1 < Ny:
                rows += [n]; cols += [idx(i, j + 1)]; vals += [cyp]
            if j - 1 >= 0:
                rows += [n]; cols += [idx(i, j - 1)]; vals += [cym]
            rows += [n]; cols += [n]; vals += [diag]

    A = sp.csr_matrix((vals, (rows, cols)), shape=(Ny * Nx, Ny * Nx), dtype=complex)
    return A, k0


def solve_scattered(eps, eps_inc, wavelength, h, npml, theta_deg=0.0, n_inc=1.0,
                    sigma0=8.0, m=3):
    """Scattered-field solve. Returns (E_tot, E_inc) as (Ny,Nx) arrays."""
    Ny, Nx = eps.shape
    k0 = 2 * np.pi / wavelength
    theta = np.deg2rad(theta_deg)
    kx = k0 * n_inc * np.sin(theta)
    ky = k0 * n_inc * np.cos(theta)

    xs = np.arange(Nx) * h
    ys = np.arange(Ny) * h
    X, Y = np.meshgrid(xs, ys)
    E_inc = np.exp(1j * (kx * X + ky * Y))

    A, _ = build_operator(eps, wavelength, h, npml_y=npml, kx=kx, sigma0=sigma0, m=m)
    src = -(k0**2) * (eps - eps_inc) * E_inc           # (Ny,Nx)
    E_scat = spla.spsolve(A, src.ravel()).reshape(Ny, Nx)
    return E_inc + E_scat, E_inc


def _flux_y(E, j, h):
    """Net y-directed Poynting flux across row j (up to a positive constant)."""
    dEy = (E[j + 1, :] - E[j - 1, :]) / (2 * h)
    return np.sum(np.imag(np.conj(E[j, :]) * dEy))


def reflectance_transmittance(eps, eps_inc, wavelength, h, npml, j_inc, j_trn,
                              theta_deg=0.0, n_inc=1.0, **kw):
    """R, T for a structure embedded in the grid, via flux ratios vs incident."""
    E_tot, E_inc = solve_scattered(eps, eps_inc, wavelength, h, npml,
                                   theta_deg=theta_deg, n_inc=n_inc, **kw)
    E_scat = E_tot - E_inc
    S_inc = _flux_y(E_inc, j_inc, h)        # downward incident flux
    S_ref = _flux_y(E_scat, j_inc, h)       # reflected (upward) flux
    S_trn = _flux_y(E_tot, j_trn, h)        # transmitted (downward) flux
    R = abs(S_ref) / abs(S_inc)
    T = abs(S_trn) / abs(S_inc)
    return R, T, E_tot


def _layer_grid(n_inc, n_sub, layers, h, Nx, npml, pad=30):
    """Build eps(Ny,Nx) for a 1D layer stack.

    Returns eps, eps_inc, j_inc, j_trn, and the ACTUAL (cell-rounded) layer list
    so the TMM reference compares the same structure FDFD actually discretized.
    """
    layer_rows, actual = [], []
    for (n, thick) in layers:
        nr = max(1, int(round(thick / h)))
        layer_rows += [n**2] * nr
        actual.append((n, nr * h))
    inc_rows = trn_rows = npml + pad
    col = np.array([n_inc**2] * inc_rows + layer_rows + [n_sub**2] * trn_rows)
    Ny = col.size
    eps = np.repeat(col[:, None], Nx, axis=1)
    j_inc = npml + pad // 2
    j_trn = Ny - npml - pad // 2
    return eps, n_inc**2, j_inc, j_trn, actual


if __name__ == "__main__":
    # Cross-validate FDFD against TMM for layered structures.
    wl = 1000.0
    h = 10.0
    Nx, npml = 16, 14
    # free-standing stacks in air: reference medium = air everywhere, so the
    # scattered-field source lives only in the finite layer stack (no source in
    # the PML), the clean way to cross-check a grid solver against TMM.
    n_inc, n_sub = 1.0, 1.0

    cases = {
        "slab n=2.2 (180nm) in air": [(2.2, 180.0)],
        "slab n=3.0 (250nm) in air": [(3.0, 250.0)],
        "5-layer HLHLH in air": [(2.2, 120.0), (1.46, 200.0), (2.2, 120.0),
                                  (1.46, 200.0), (2.2, 120.0)],
    }
    print(f"FDFD vs TMM (normal incidence, h={h:.0f}nm):")
    print(f"  {'case':28s} {'R_fdfd':>7s} {'R_tmm':>7s} {'T_fdfd':>7s} {'T_tmm':>7s} {'|dR|':>7s}")
    wl_arr = np.array([wl])
    worst = 0.0
    for name, layers in cases.items():
        eps, eps_inc, j_inc, j_trn, actual = _layer_grid(n_inc, n_sub, layers, h, Nx, npml)
        R, T, _ = reflectance_transmittance(eps, eps_inc, wl, h, npml, j_inc, j_trn,
                                            n_inc=n_inc)
        if actual:
            n_seq = np.array([n for n, _ in actual], complex)
            d_seq = np.array([t for _, t in actual])
            ref = tmm.stack_response(wl_arr, n_seq, d_seq, n_inc=n_inc, n_sub=n_sub)
        else:
            ref = tmm.stack_response(wl_arr, np.array([]).reshape(0), np.array([]),
                                     n_inc=n_inc, n_sub=n_sub)
        Rt, Tt = ref["R"][0], ref["T"][0]
        worst = max(worst, abs(R - Rt))
        print(f"  {name:28s} {R:7.4f} {Rt:7.4f} {T:7.4f} {Tt:7.4f} {abs(R-Rt):7.4f}")
    print(f"\n  worst |R_fdfd - R_tmm| = {worst:.4f}  (two independent solvers agree)")

    # Oblique incidence: exercises the Bloch-periodic x-direction (genuinely 2D).
    print(f"\nOblique incidence, s-pol, slab n=2.2 (180nm) in air:")
    print(f"  {'angle':>6s} {'R_fdfd':>7s} {'R_tmm':>7s}")
    eps, eps_inc, j_inc, j_trn, actual = _layer_grid(1.0, 1.0, [(2.2, 180.0)], h, Nx, npml)
    n_seq = np.array([2.2], complex); d_seq = np.array([actual[0][1]])
    for ang in (15.0, 30.0, 45.0):
        R, T, _ = reflectance_transmittance(eps, eps_inc, wl, h, npml, j_inc, j_trn,
                                            theta_deg=ang, n_inc=1.0)
        Rt = tmm.stack_response(wl_arr, n_seq, d_seq, n_inc=1.0, n_sub=1.0,
                                theta0_deg=ang, pol="s")["R"][0]
        print(f"  {ang:5.0f}d {R:7.4f} {Rt:7.4f}")