meanas/examples/waveguide_real.py

"""
Real-valued straight-waveguide FDTD/FDFD comparison.

This example shows the user-facing "compare real FDTD against reconstructed real
FDFD" workflow:

1. build a straight waveguide on a uniform Yee grid,
2. drive it with a real-valued continuous-wave mode source,
3. solve the matching FDFD problem from the analytic source phasor, and
4. compare late real monitor slices against `fdtd.reconstruct_real_e/h(...)`.

Unlike the phasor-based examples, this script does not use extracted phasors as
the main output. It is a stricter diagnostic: the comparison target is the raw
real field itself, with full-plane, mode-weighted, guided-mode, and orthogonal-
residual errors reported. Strong phasor agreement can coexist with visibly
larger raw-snapshot error because the latter includes weak nonguided tails on
the monitor plane.
"""

import numpy

from meanas import fdfd, fdtd
from meanas.fdfd import functional, scpml, waveguide_3d
from meanas.fdmath import vec, unvec


DT = 0.25
PERIOD_STEPS = 64
OMEGA = 2 * numpy.pi / (PERIOD_STEPS * DT)
CPML_THICKNESS = 3
SHAPE = (3, 37, 13, 13)
SOURCE_SLICES = (slice(5, 6), slice(None), slice(None))
MONITOR_SLICES = (slice(30, 31), slice(None), slice(None))
WARMUP_PERIODS = 16
SOURCE_PHASE = 0.4
CORE_SLICES = (slice(None), slice(None), slice(4, 9), slice(4, 9))


def build_uniform_dxes(shape: tuple[int, int, int, int]) -> list[list[numpy.ndarray]]:
    return [[numpy.ones(shape[axis + 1]) for axis in range(3)] for _ in range(2)]


def build_epsilon(shape: tuple[int, int, int, int]) -> numpy.ndarray:
    epsilon = numpy.ones(shape, dtype=float)
    y0 = (shape[2] - 3) // 2
    z0 = (shape[3] - 3) // 2
    epsilon[:, :, y0:y0 + 3, z0:z0 + 3] = 12.0
    return epsilon


def build_stretched_dxes(base_dxes: list[list[numpy.ndarray]]) -> list[list[numpy.ndarray]]:
    stretched_dxes = [[dx.copy() for dx in group] for group in base_dxes]
    for axis in (0, 1, 2):
        for polarity in (-1, 1):
            stretched_dxes = scpml.stretch_with_scpml(
                stretched_dxes,
                axis=axis,
                polarity=polarity,
                omega=OMEGA,
                epsilon_effective=1.0,
                thickness=CPML_THICKNESS,
            )
    return stretched_dxes


def build_cpml_params() -> list[list[dict[str, numpy.ndarray | float]]]:
    return [
        [fdtd.cpml_params(axis=axis, polarity=polarity, dt=DT, thickness=CPML_THICKNESS, epsilon_eff=1.0)
         for polarity in (-1, 1)]
        for axis in range(3)
    ]


def weighted_rel_err(observed: numpy.ndarray, reference: numpy.ndarray, weight: numpy.ndarray) -> float:
    return numpy.linalg.norm((observed - reference) * weight) / numpy.linalg.norm(reference * weight)


def project_onto_mode(observed: numpy.ndarray, mode: numpy.ndarray) -> tuple[complex, numpy.ndarray, numpy.ndarray]:
    coefficient = numpy.vdot(mode, observed) / numpy.vdot(mode, mode)
    guided = coefficient * mode
    residual = observed - guided
    return coefficient, guided, residual


def main() -> None:
    epsilon = build_epsilon(SHAPE)
    base_dxes = build_uniform_dxes(SHAPE)
    stretched_dxes = build_stretched_dxes(base_dxes)

    source_mode = waveguide_3d.solve_mode(
        0,
        omega=OMEGA,
        dxes=base_dxes,
        axis=0,
        polarity=1,
        slices=SOURCE_SLICES,
        epsilon=epsilon,
    )
    j_mode = waveguide_3d.compute_source(
        E=source_mode['E'],
        wavenumber=source_mode['wavenumber'],
        omega=OMEGA,
        dxes=base_dxes,
        axis=0,
        polarity=1,
        slices=SOURCE_SLICES,
        epsilon=epsilon,
    )
    # A small global phase aligns the real-valued source with the late-cycle
    # raw-snapshot diagnostic. The underlying phasor problem is unchanged.
    j_mode *= numpy.exp(1j * SOURCE_PHASE)
    monitor_mode = waveguide_3d.solve_mode(
        0,
        omega=OMEGA,
        dxes=base_dxes,
        axis=0,
        polarity=1,
        slices=MONITOR_SLICES,
        epsilon=epsilon,
    )
    e_weight = numpy.abs(monitor_mode['E'][:, MONITOR_SLICES[0], :, :])
    h_weight = numpy.abs(monitor_mode['H'][:, MONITOR_SLICES[0], :, :])
    e_mode = monitor_mode['E'][:, MONITOR_SLICES[0], :, :]
    h_mode = monitor_mode['H'][:, MONITOR_SLICES[0], :, :]

    e_fdfd = unvec(
        fdfd.solvers.generic(
            J=vec(j_mode),
            omega=OMEGA,
            dxes=stretched_dxes,
            epsilon=vec(epsilon),
            matrix_solver_opts={'atol': 1e-10, 'rtol': 1e-7},
        ),
        SHAPE[1:],
    )
    h_fdfd = functional.e2h(OMEGA, stretched_dxes)(e_fdfd)

    update_e, update_h = fdtd.updates_with_cpml(
        cpml_params=build_cpml_params(),
        dt=DT,
        dxes=base_dxes,
        epsilon=epsilon,
    )

    e_field = numpy.zeros_like(epsilon)
    h_field = numpy.zeros_like(epsilon)
    total_steps = (WARMUP_PERIODS + 1) * PERIOD_STEPS
    e_errors: list[float] = []
    h_errors: list[float] = []
    e_core_errors: list[float] = []
    h_core_errors: list[float] = []
    e_weighted_errors: list[float] = []
    h_weighted_errors: list[float] = []
    e_guided_errors: list[float] = []
    h_guided_errors: list[float] = []
    e_residual_errors: list[float] = []
    h_residual_errors: list[float] = []

    for step in range(total_steps):
        update_e(e_field, h_field, epsilon)

        # Real-valued FDTD uses the real part of the analytic mode source.
        t_half = (step + 0.5) * DT
        j_real = (j_mode.real * numpy.cos(OMEGA * t_half) - j_mode.imag * numpy.sin(OMEGA * t_half)).real
        e_field -= DT * j_real / epsilon

        update_h(e_field, h_field)

        if step >= total_steps - PERIOD_STEPS // 4:
            reconstructed_e = fdtd.reconstruct_real_e(
                e_fdfd[:, MONITOR_SLICES[0], :, :],
                OMEGA,
                DT,
                step + 1,
            )
            reconstructed_h = fdtd.reconstruct_real_h(
                h_fdfd[:, MONITOR_SLICES[0], :, :],
                OMEGA,
                DT,
                step + 1,
            )

            e_monitor = e_field[:, MONITOR_SLICES[0], :, :]
            h_monitor = h_field[:, MONITOR_SLICES[0], :, :]
            e_errors.append(numpy.linalg.norm(e_monitor - reconstructed_e) / numpy.linalg.norm(reconstructed_e))
            h_errors.append(numpy.linalg.norm(h_monitor - reconstructed_h) / numpy.linalg.norm(reconstructed_h))
            e_core_errors.append(
                numpy.linalg.norm(e_monitor[CORE_SLICES] - reconstructed_e[CORE_SLICES])
                / numpy.linalg.norm(reconstructed_e[CORE_SLICES]),
            )
            h_core_errors.append(
                numpy.linalg.norm(h_monitor[CORE_SLICES] - reconstructed_h[CORE_SLICES])
                / numpy.linalg.norm(reconstructed_h[CORE_SLICES]),
            )
            e_weighted_errors.append(weighted_rel_err(e_monitor, reconstructed_e, e_weight))
            h_weighted_errors.append(weighted_rel_err(h_monitor, reconstructed_h, h_weight))
            e_guided_coeff, _, e_residual = project_onto_mode(e_monitor, e_mode)
            e_guided_coeff_ref, _, e_residual_ref = project_onto_mode(reconstructed_e, e_mode)
            h_guided_coeff, _, h_residual = project_onto_mode(h_monitor, h_mode)
            h_guided_coeff_ref, _, h_residual_ref = project_onto_mode(reconstructed_h, h_mode)
            e_guided_errors.append(abs(e_guided_coeff - e_guided_coeff_ref) / abs(e_guided_coeff_ref))
            h_guided_errors.append(abs(h_guided_coeff - h_guided_coeff_ref) / abs(h_guided_coeff_ref))
            e_residual_errors.append(numpy.linalg.norm(e_residual - e_residual_ref) / numpy.linalg.norm(e_residual_ref))
            h_residual_errors.append(numpy.linalg.norm(h_residual - h_residual_ref) / numpy.linalg.norm(h_residual_ref))

    print(f'late-cycle monitor E errors: min={min(e_errors):.4f} max={max(e_errors):.4f}')
    print(f'late-cycle monitor H errors: min={min(h_errors):.4f} max={max(h_errors):.4f}')
    print(f'late-cycle core-window E errors: min={min(e_core_errors):.4f} max={max(e_core_errors):.4f}')
    print(f'late-cycle core-window H errors: min={min(h_core_errors):.4f} max={max(h_core_errors):.4f}')
    print(f'late-cycle mode-weighted E errors: min={min(e_weighted_errors):.4f} max={max(e_weighted_errors):.4f}')
    print(f'late-cycle mode-weighted H errors: min={min(h_weighted_errors):.4f} max={max(h_weighted_errors):.4f}')
    print(f'late-cycle guided-mode E coefficient errors: min={min(e_guided_errors):.4f} max={max(e_guided_errors):.4f}')
    print(f'late-cycle guided-mode H coefficient errors: min={min(h_guided_errors):.4f} max={max(h_guided_errors):.4f}')
    print(f'late-cycle orthogonal-residual E errors: min={min(e_residual_errors):.4f} max={max(e_residual_errors):.4f}')
    print(f'late-cycle orthogonal-residual H errors: min={min(h_residual_errors):.4f} max={max(h_residual_errors):.4f}')


if __name__ == '__main__':
    main()