meanas/examples/eme.py

"""
Mode-matching / EME example for a straight rib-waveguide interface.

This example shows the intended user-facing workflow for `meanas.fdfd.eme` on a
simple straight interface:

1. build two nearby waveguide cross-sections on a Yee grid,
2. solve a small set of guided modes on each side,
3. normalize those modes into E/H port fields,
4. assemble the interface scattering matrix with `meanas.fdfd.eme.get_s(...)`,
5. inspect the dominant modal coupling numerically and visually.
"""

from __future__ import annotations

import importlib

import numpy
from numpy import pi

import gridlock
from gridlock import Extent

from meanas.fdfd import eme, waveguide_2d
from meanas.fdmath import unvec


WL = 1310.0
DX = 40.0
WIDTH = 400.0
THF = 161.0
THP = 77.0
EPS_SI = 3.51 ** 2
EPS_OX = 1.453 ** 2
MODE_NUMBERS = numpy.array([0])


def require_optional(name: str, package_name: str | None = None):
    package_name = package_name or name
    try:
        return importlib.import_module(name)
    except ImportError as exc:  # pragma: no cover - user environment guard
        raise SystemExit(
            f"This example requires the optional package '{package_name}'. "
            "Install example dependencies with `pip install -e './meanas[examples]'`.",
        ) from exc


def build_geometry(
        *,
        dx: float = DX,
        width: float = WIDTH,
        thf: float = THF,
        thp: float = THP,
        eps_si: float = EPS_SI,
        eps_ox: float = EPS_OX,
        ) -> tuple[gridlock.Grid, numpy.ndarray, list[list[numpy.ndarray]], float]:
    x0 = (width / 2) % dx
    omega = 2 * pi / WL

    grid = gridlock.Grid(
        [
            numpy.arange(-800, 800 + dx, dx),
            numpy.arange(-400, 400 + dx, dx),
            numpy.arange(-2 * dx, 2 * dx + dx, dx),
        ],
        periodic=True,
    )
    epsilon = grid.allocate(eps_ox)

    grid.draw_cuboid(
        epsilon,
        foreground=eps_si,
        x=Extent(center=x0, span=width + 1200),
        y=Extent(min=0, max=thf),
        z=Extent(min=-1e6, max=0),
    )
    grid.draw_cuboid(
        epsilon,
        foreground=eps_ox,
        x=Extent(max=-width / 2, span=300),
        y=Extent(min=thp, max=1e6),
        z=Extent(min=-1e6, max=0),
    )
    grid.draw_cuboid(
        epsilon,
        foreground=eps_ox,
        x=Extent(min=width / 2, span=300),
        y=Extent(min=thp, max=1e6),
        z=Extent(min=-1e6, max=0),
    )

    grid.draw_cuboid(
        epsilon,
        foreground=eps_si,
        x=Extent(max=-(width / 2 + 600), span=240),
        y=Extent(min=0, max=thf),
        z=Extent(min=0, max=1e6),
    )
    grid.draw_cuboid(
        epsilon,
        foreground=eps_si,
        x=Extent(max=width / 2 + 600, span=240),
        y=Extent(min=0, max=thf),
        z=Extent(min=0, max=1e6),
    )

    dxes = [grid.dxyz, grid.autoshifted_dxyz()]
    dxes_2d = [[d[0], d[1]] for d in dxes]
    return grid, epsilon, dxes_2d, omega


def solve_cross_section_modes(
        epsilon_slice: numpy.ndarray,
        *,
        omega: float,
        dxes_2d: list[list[numpy.ndarray]],
        mode_numbers: numpy.ndarray = MODE_NUMBERS,
        ) -> tuple[list[tuple[numpy.ndarray, numpy.ndarray]], numpy.ndarray]:
    e_xys, wavenumbers = waveguide_2d.solve_modes(
        epsilon=epsilon_slice.ravel(),
        omega=omega,
        dxes=dxes_2d,
        mode_numbers=mode_numbers,
    )
    eh_fields = [
        waveguide_2d.normalized_fields_e(
            e_xy,
            wavenumber=wavenumber,
            dxes=dxes_2d,
            omega=omega,
            epsilon=epsilon_slice.ravel(),
        )
        for e_xy, wavenumber in zip(e_xys, wavenumbers, strict=True)
    ]
    return eh_fields, wavenumbers


def print_summary(ss: numpy.ndarray, wavenumbers_left: numpy.ndarray, wavenumbers_right: numpy.ndarray, omega: float) -> None:
    n_left = len(wavenumbers_left)
    left_neff = numpy.real(wavenumbers_left / omega)
    right_neff = numpy.real(wavenumbers_right / omega)

    print('left effective indices:', ', '.join(f'{value:.5f}' for value in left_neff[:4]))
    print('right effective indices:', ', '.join(f'{value:.5f}' for value in right_neff[:4]))

    reflection = abs(ss[0, 0]) ** 2
    transmission = abs(ss[n_left, 0]) ** 2
    total_output = numpy.sum(abs(ss[:, 0]) ** 2)
    print(f'fundamental left-incident reflection |S_00|^2      = {reflection:.6f}')
    print(f'fundamental left-to-right transmission |S_{n_left},0|^2 = {transmission:.6f}')
    print(f'fundamental left-incident total output power      = {total_output:.6f}')

    strongest = numpy.argsort(abs(ss[n_left:, 0]) ** 2)[::-1][:3]
    print('dominant transmitted right-side modes for left mode 0:')
    for mode_index in strongest:
        print(f'  mode {mode_index}: |S|^2 = {abs(ss[n_left + mode_index, 0]) ** 2:.6f}')


def plot_results(
        *,
        pyplot,
        ss: numpy.ndarray,
        left_mode: tuple[numpy.ndarray, numpy.ndarray],
        right_mode: tuple[numpy.ndarray, numpy.ndarray],
        shape_2d: tuple[int, int],
        ) -> None:
    fig_s, ax_s = pyplot.subplots()
    image = ax_s.imshow(abs(ss) ** 2, origin='lower', cmap='magma')
    fig_s.colorbar(image, ax=ax_s)
    ax_s.set_title('Interface scattering magnitude |S|^2')
    ax_s.set_xlabel('Incident mode index')
    ax_s.set_ylabel('Outgoing mode index')

    e_left = unvec(left_mode[0], shape_2d)
    e_right = unvec(right_mode[0], shape_2d)
    left_intensity = numpy.sum(abs(e_left) ** 2, axis=0).T
    right_intensity = numpy.sum(abs(e_right) ** 2, axis=0).T

    fig_modes, axes = pyplot.subplots(1, 2, figsize=(10, 4))
    left_plot = axes[0].imshow(left_intensity, origin='lower', cmap='viridis')
    fig_modes.colorbar(left_plot, ax=axes[0])
    axes[0].set_title('Left fundamental mode |E|^2')
    right_plot = axes[1].imshow(right_intensity, origin='lower', cmap='viridis')
    fig_modes.colorbar(right_plot, ax=axes[1])
    axes[1].set_title('Right fundamental mode |E|^2')
    if pyplot.get_backend().lower().endswith('agg'):
        pyplot.close(fig_s)
        pyplot.close(fig_modes)
    else:
        pyplot.show()


def main() -> None:
    pyplot = require_optional('matplotlib.pyplot', package_name='matplotlib')

    grid, epsilon, dxes_2d, omega = build_geometry()
    left_slice = epsilon[:, :, :, 1]
    right_slice = epsilon[:, :, :, -2]

    left_modes, wavenumbers_left = solve_cross_section_modes(left_slice, omega=omega, dxes_2d=dxes_2d)
    right_modes, wavenumbers_right = solve_cross_section_modes(right_slice, omega=omega, dxes_2d=dxes_2d)

    ss = eme.get_s(left_modes, wavenumbers_left, right_modes, wavenumbers_right, dxes=dxes_2d)

    print_summary(ss, wavenumbers_left, wavenumbers_right, omega)
    plot_results(
        pyplot=pyplot,
        ss=ss,
        left_mode=left_modes[0],
        right_mode=right_modes[0],
        shape_2d=grid.shape[:2],
    )


if __name__ == '__main__':
    main()