meanas/meanas/fdfd/eme.py

"""
Low-level mode-matching helpers for waveguide / EME workflows.

These helpers operate on already-solved and already-normalized port fields.
They do not build geometries or solve modes themselves; downstream users are
expected to supply compatible `(E, H)` modal field pairs from
`waveguide_2d`, `waveguide_3d`, or `waveguide_cyl`.

The returned matrices follow the usual port ordering:

- `get_tr(...)` returns `(T, R)` for left-incident modes.
- `get_abcd(...)` returns the 2-port block transfer matrix built from the
  interface continuity solve.
- `get_s(...)` returns the full block scattering matrix
  `[[R_L, T_R_to_L], [T_L_to_R, R_R]]`.
- `get_taper_abcd(...)` and `get_taper_s(...)` build the same transfer /
  scattering objects for sampled tapers using midpoint interfaces and stable
  Redheffer cascading.

This module is intentionally a thin library layer rather than an integrated
simulation suite. It provides the overlap algebra that downstream users can
compose into larger workflows.
"""

from collections.abc import Sequence
from typing import Literal
import dataclasses
import operator
import numpy
from numpy.typing import NDArray
from scipy import linalg
from scipy import sparse

from ..fdmath import dx_lists2_t, vcfdfield2
from .waveguide_2d import ModeBasis, inner_product, lorentz_overlap

type wavenumber_seq = Sequence[complex] | NDArray[numpy.complexfloating] | NDArray[numpy.floating]
type mode_seq = Sequence[Sequence[vcfdfield2]]


@dataclasses.dataclass(frozen=True, eq=False)
class ModeSection:
    """
    Local modal basis at one longitudinal sample of a continuously varying guide.

    Args:
        z: Longitudinal coordinate of this section.
        modes: Forward modes as `(E, H)` field pairs.
        wavenumbers: Forward propagation constants for `modes`.
        backward_modes: Optional explicit backward modes. If omitted, backward
            modes are synthesized as `(E, -H)`.
        backward_wavenumbers: Optional propagation constants for
            `backward_modes`. If omitted, they are synthesized as `-wavenumbers`.
        dual_modes: Optional forward dual / adjoint projection modes. If
            omitted, `modes` are used as their own projection basis.
        dual_backward_modes: Optional backward dual / adjoint projection modes.
            If omitted, they are synthesized from `dual_modes` when available,
            otherwise from `backward_modes`.
        wavenumber_matrix: Optional dense beta operator in the supplied forward
            basis. If omitted, ``diag(wavenumbers)`` is used.
        backward_wavenumber_matrix: Optional dense beta operator for explicit
            backward modes. Synthesized reciprocal branches use ``-wavenumber_matrix``.
    """

    z: float
    modes: mode_seq = dataclasses.field(repr=False)
    wavenumbers: wavenumber_seq
    backward_modes: mode_seq | None = dataclasses.field(default=None, repr=False)
    backward_wavenumbers: wavenumber_seq | None = None
    dual_modes: mode_seq | None = dataclasses.field(default=None, repr=False)
    dual_backward_modes: mode_seq | None = dataclasses.field(default=None, repr=False)
    wavenumber_matrix: NDArray[numpy.complex128] | None = None
    backward_wavenumber_matrix: NDArray[numpy.complex128] | None = None


def _validate_port_modes(
        name: str,
        ehs: mode_seq,
        wavenumbers: wavenumber_seq,
        ) -> tuple[tuple[int, ...], tuple[int, ...]]:
    """Validate a nonempty modal list and return its common E/H shapes."""
    if len(ehs) != len(wavenumbers):
        raise ValueError(f'{name} mode list and wavenumber list must have the same length')
    if not ehs:
        raise ValueError(f'{name} must contain at least one mode')

    e_shape: tuple[int, ...] | None = None
    h_shape: tuple[int, ...] | None = None
    for index, mode in enumerate(ehs):
        if len(mode) != 2:
            raise ValueError(f'{name}[{index}] must be a 2-tuple of (E, H) modal fields')
        e_field, h_field = mode
        mode_e_shape = numpy.shape(e_field)
        mode_h_shape = numpy.shape(h_field)
        if mode_e_shape != mode_h_shape:
            raise ValueError(f'{name}[{index}] has mismatched E/H field shapes')
        if e_shape is None:
            e_shape = mode_e_shape
            h_shape = mode_h_shape
        elif mode_e_shape != e_shape or mode_h_shape != h_shape:
            raise ValueError(f'{name} modal fields must all share the same shape')

    assert e_shape is not None
    assert h_shape is not None
    return e_shape, h_shape


def _validate_dual_modes(
        name: str,
        dual_ehs: mode_seq | None,
        reference_shape: tuple[int, ...],
        wavenumbers: wavenumber_seq,
        ) -> mode_seq | None:
    """Validate optional dual modes against a primal modal field shape."""
    if dual_ehs is None:
        return None

    dual_e_shape, dual_h_shape = _validate_port_modes(name, dual_ehs, wavenumbers)
    if dual_e_shape != reference_shape or dual_h_shape != reference_shape:
        raise ValueError(f'{name} modal fields must share the same E/H shapes as the corresponding modes')
    return dual_ehs


def _as_wavenumber_array(
        name: str,
        wavenumbers: wavenumber_seq,
        ) -> NDArray[numpy.complex128]:
    """Convert finite propagation constants to a one-dimensional array."""
    array = numpy.asarray(wavenumbers, dtype=complex)
    if array.ndim != 1:
        raise ValueError(f'{name} must be a one-dimensional sequence')
    if not numpy.isfinite(array).all():
        raise ValueError(f'{name} must contain only finite values')
    return array


def _as_wavenumber_matrix(
        name: str,
        matrix: NDArray[numpy.complex128] | None,
        wavenumbers: NDArray[numpy.complex128],
        ) -> NDArray[numpy.complex128]:
    """Return a validated propagation matrix or the diagonal spectrum."""
    if matrix is None:
        return numpy.diag(wavenumbers)
    array = numpy.asarray(matrix, dtype=complex)
    if array.shape != (wavenumbers.size, wavenumbers.size):
        raise ValueError(f'{name} must have one row and column per mode')
    if not numpy.isfinite(array).all():
        raise ValueError(f'{name} must contain only finite values')
    return array


def _as_mode_arrays(
        ehs: mode_seq,
        ) -> list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]]:
    """Copy a modal sequence into complex NumPy array views."""
    return [
        (numpy.array(e_field, dtype=complex, copy=True), numpy.array(h_field, dtype=complex, copy=True))
        for e_field, h_field in ehs
        ]


def _phase_align_modes(
        previous: Sequence[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
        current: Sequence[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
        dxes: dx_lists2_t,
        previous_dual: Sequence[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]] | None = None,
        *,
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> tuple[list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]], NDArray[numpy.complex128]]:
    """Align current modal phases and return their diagonal basis transform."""
    aligned = []
    scales = []
    test_modes = previous if previous_dual is None else previous_dual
    for index, (previous_mode, current_mode, test_mode) in enumerate(zip(previous, current, test_modes, strict=True)):
        overlap = lorentz_overlap(test_mode, current_mode, dxes, measure=measure)
        self_overlap = lorentz_overlap(test_mode, previous_mode, dxes, measure=measure)
        if overlap == 0:
            raise ValueError(f'cannot phase-track mode {index}: adjacent section overlap is zero')
        if self_overlap == 0:
            raise ValueError(f'cannot phase-track mode {index}: mode dual-overlap is zero')
        phase = (overlap / abs(overlap)) / (self_overlap / abs(self_overlap))
        e_field, h_field = current_mode
        scale = 1 / phase
        scales.append(scale)
        aligned.append((e_field * scale, h_field * scale))
    return aligned, numpy.diag(scales)


def _section_branches(
        section: ModeSection,
        index: int,
        expected_count: int | None,
        expected_shape: tuple[int, ...] | None,
        ) -> tuple[
            float,
            list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
            NDArray[numpy.complex128],
            list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
            NDArray[numpy.complex128],
            list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
            list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]],
            NDArray[numpy.complex128],
            NDArray[numpy.complex128],
            tuple[int, ...],
            ]:
    """Validate one section and construct its forward and backward branches."""
    z_coord = float(section.z)
    if not numpy.isfinite(z_coord):
        raise ValueError(f'sections[{index}].z must be finite')

    shape, _h_shape = _validate_port_modes(f'sections[{index}].modes', section.modes, section.wavenumbers)
    wavenumbers = _as_wavenumber_array(f'sections[{index}].wavenumbers', section.wavenumbers)
    wavenumber_matrix = _as_wavenumber_matrix(f'sections[{index}].wavenumber_matrix', section.wavenumber_matrix, wavenumbers)
    if expected_count is not None and len(wavenumbers) != expected_count:
        raise ValueError('all taper sections must contain the same number of modes')
    if expected_shape is not None and shape != expected_shape:
        raise ValueError('all taper section modal fields must share the same E/H shapes')

    if (section.backward_modes is None) != (section.backward_wavenumbers is None):
        raise ValueError('backward_modes and backward_wavenumbers must be supplied together')

    forward_modes = _as_mode_arrays(section.modes)
    if section.backward_modes is None:
        backward_modes = [(e_field.copy(), -h_field.copy()) for e_field, h_field in forward_modes]
        backward_wavenumbers = -wavenumbers
        if section.backward_wavenumber_matrix is not None:
            raise ValueError('backward_wavenumber_matrix requires explicit backward modes and wavenumbers')
        backward_wavenumber_matrix = -wavenumber_matrix
    else:
        backward_shape, _backward_h_shape = _validate_port_modes(
            f'sections[{index}].backward_modes',
            section.backward_modes,
            section.backward_wavenumbers,
            )
        if backward_shape != shape:
            raise ValueError('forward and backward modal fields must share the same E/H shapes')
        backward_wavenumbers = _as_wavenumber_array(
            f'sections[{index}].backward_wavenumbers',
            section.backward_wavenumbers,
            )
        backward_modes = _as_mode_arrays(section.backward_modes)
        backward_wavenumber_matrix = _as_wavenumber_matrix(
            f'sections[{index}].backward_wavenumber_matrix',
            section.backward_wavenumber_matrix,
            backward_wavenumbers,
        )

    if len(backward_wavenumbers) != len(wavenumbers):
        raise ValueError('forward and backward mode counts must match')

    if section.dual_modes is None:
        dual_modes = forward_modes
    else:
        dual_shape, _dual_h_shape = _validate_port_modes(
            f'sections[{index}].dual_modes',
            section.dual_modes,
            section.wavenumbers,
            )
        if dual_shape != shape:
            raise ValueError('modal fields and dual modal fields must share the same E/H shapes')
        dual_modes = _as_mode_arrays(section.dual_modes)

    if section.dual_backward_modes is None:
        if section.dual_modes is None and section.backward_modes is not None:
            dual_backward_modes = backward_modes
        else:
            dual_backward_modes = [(e_field.copy(), -h_field.copy()) for e_field, h_field in dual_modes]
    else:
        dual_backward_shape, _dual_backward_h_shape = _validate_port_modes(
            f'sections[{index}].dual_backward_modes',
            section.dual_backward_modes,
            section.backward_wavenumbers if section.backward_wavenumbers is not None else backward_wavenumbers,
            )
        if dual_backward_shape != shape:
            raise ValueError('backward modal fields and dual backward modal fields must share the same E/H shapes')
        dual_backward_modes = _as_mode_arrays(section.dual_backward_modes)

    if len(dual_modes) != len(forward_modes) or len(dual_backward_modes) != len(backward_modes):
        raise ValueError('dual mode counts must match the corresponding modal basis counts')

    return (
        z_coord, forward_modes, wavenumbers, backward_modes, backward_wavenumbers,
        dual_modes, dual_backward_modes, wavenumber_matrix, backward_wavenumber_matrix, shape,
    )


def _validate_taper_sections(
        sections: Sequence[ModeSection],
        dxes: dx_lists2_t,
        *,
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> tuple[
            NDArray[numpy.float64],
            list[list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]]],
            list[list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]]],
            list[NDArray[numpy.complex128]],
            int,
            list[NDArray[numpy.complex128]],
            ]:
    """Validate, phase-align, and assemble all taper branch data."""
    if len(sections) < 2:
        raise ValueError('at least two taper sections are required')

    z_coords: list[float] = []
    branch_modes: list[list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]]] = []
    branch_dual_modes: list[list[tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]]] = []
    branch_wavenumbers: list[NDArray[numpy.complex128]] = []
    expected_count: int | None = None
    expected_shape: tuple[int, ...] | None = None

    for index, section in enumerate(sections):
        z_coord, forward_modes, forward_wavenumbers, backward_modes, backward_wavenumbers, dual_modes, dual_backward_modes, forward_matrix, backward_matrix, shape = _section_branches(
            section,
            index,
            expected_count,
            expected_shape,
            )
        if expected_count is None:
            expected_count = len(forward_wavenumbers)
            expected_shape = shape
        z_coords.append(z_coord)
        branch_modes.append([*forward_modes, *backward_modes])
        branch_dual_modes.append([*dual_modes, *dual_backward_modes])
        branch_wavenumbers.append(linalg.block_diag(forward_matrix, backward_matrix))

    z_array = numpy.asarray(z_coords, dtype=float)
    if not (numpy.diff(z_array) > 0).all():
        raise ValueError('taper section z coordinates must be strictly increasing')

    transforms = [numpy.eye(2 * expected_count, dtype=complex) for _ in sections]
    for index in range(1, len(branch_modes)):
        branch_modes[index], transform = _phase_align_modes(
            branch_modes[index - 1], branch_modes[index], dxes, branch_dual_modes[index - 1], measure=measure,
        )
        transforms[index] = transform
        inverse = numpy.linalg.inv(transform)
        branch_dual_modes[index] = [
            (e_field * inverse[row, row], h_field * inverse[row, row])
            for row, (e_field, h_field) in enumerate(branch_dual_modes[index])
        ]
        branch_wavenumbers[index] = inverse @ branch_wavenumbers[index] @ transform

    assert expected_count is not None
    return z_array, branch_modes, branch_dual_modes, branch_wavenumbers, expected_count, transforms


def _interface_s(
        left: mode_seq,
        right: mode_seq,
        left_dual: mode_seq,
        right_dual: mode_seq,
        dxes: dx_lists2_t,
        *,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> NDArray[numpy.complex128]:
    """Solve projected tangential continuity for both incidences together.

    Branches are ordered [forward, backward]. Mixed testing uses left electric
    duals on H continuity and right magnetic duals on E continuity. In contrast,
    Galerkin tests both equations against the left dual basis.
    """
    nl, nr = len(left) // 2, len(right) // 2
    if formulation not in ('mixed', 'galerkin'):
        raise ValueError("formulation must be 'mixed' or 'galerkin'")
    if formulation == 'galerkin' and nl != nr:
        raise ValueError('Galerkin interfaces require equal port mode counts')
    electric_tests = left_dual[:nl]
    magnetic_tests = right_dual[:nr] if formulation == 'mixed' else left_dual[:nl]

    def h_projection(modes: mode_seq) -> NDArray[numpy.complex128]:
        return numpy.asarray([
            [inner_product(test[0], mode[1], dxes, measure=measure) for mode in modes]
            for test in electric_tests
        ])

    def e_projection(modes: mode_seq) -> NDArray[numpy.complex128]:
        return numpy.asarray([
            [inner_product(mode[0], test[1], dxes, measure=measure) for mode in modes]
            for test in magnetic_tests
        ])

    hl, hr = h_projection(left), h_projection(right)
    el, er = e_projection(left), e_projection(right)
    outgoing = numpy.block([[hl[:, nl:], -hr[:, :nr]],
                            [el[:, nl:], -er[:, :nr]]])
    incoming = numpy.block([[-hl[:, :nl], hr[:, nr:]],
                            [-el[:, :nl], er[:, nr:]]])
    return _solve(outgoing, incoming, 'projected interface continuity')


def _solve(matrix: NDArray, rhs: NDArray, context: str) -> NDArray[numpy.complex128]:
    """Solve without rank truncation or hidden regularization."""
    if not numpy.isfinite(matrix).all() or not numpy.isfinite(rhs).all():
        raise ValueError(f'{context} requires finite matrices')
    try:
        result = numpy.linalg.solve(matrix, rhs)
    except numpy.linalg.LinAlgError as error:
        raise numpy.linalg.LinAlgError(f'{context}: singular system') from error
    if not numpy.isfinite(result).all():
        raise numpy.linalg.LinAlgError(f'{context}: result is not representable')
    return result


def _two_port_matrix(matrix: NDArray | sparse.sparray, n_modes: int | None) -> tuple[NDArray, int]:
    array = numpy.asarray(matrix.toarray() if sparse.issparse(matrix) else matrix, dtype=complex)
    if array.ndim != 2 or array.shape[0] != array.shape[1] or array.shape[0] % 2:
        raise ValueError('transfer conversion requires a square matrix with equal port mode counts')
    n = array.shape[0] // 2
    if n == 0 or (n_modes is not None and n_modes != n):
        raise ValueError('n_modes must match the nonempty equal-sized matrix blocks')
    if not numpy.isfinite(array).all():
        raise ValueError('matrix must contain only finite values')
    return array, n


def abcd_to_s(abcd: NDArray | sparse.sparray, n_modes: int | None = None) -> NDArray[numpy.complex128]:
    """Convert [f_R, b_R] = ABCD @ [f_L, b_L] to standard port scattering.

    Requires equal port counts and an invertible D block. Singular systems
    raise ``LinAlgError`` rather than silently using a pseudoinverse.
    """
    matrix, n = _two_port_matrix(abcd, n_modes)
    a, b, c, d = matrix[:n, :n], matrix[:n, n:], matrix[n:, :n], matrix[n:, n:]
    solved = _solve(d, numpy.concatenate((c, numpy.eye(n)), axis=1), 'ABCD to S conversion')
    dc, di = solved[:, :n], solved[:, n:]
    result = numpy.block([[-dc, di], [a - b @ dc, b @ di]])
    if not numpy.isfinite(result).all():
        raise numpy.linalg.LinAlgError('ABCD to S conversion: result is not representable')
    return result


def s_to_abcd(ss: NDArray, n_modes: int | None = None) -> NDArray[numpy.complex128]:
    """Convert standard port scattering to [f_R, b_R] = ABCD @ [f_L, b_L].

    Requires equal port counts and invertible right-to-left transmission.
    Strongly attenuating devices can have finite S but unrepresentable ABCD.
    """
    matrix, n = _two_port_matrix(ss, n_modes)
    rl, trl, tlr, rr = matrix[:n, :n], matrix[:n, n:], matrix[n:, :n], matrix[n:, n:]
    solved = _solve(trl, numpy.concatenate((rl, numpy.eye(n)), axis=1), 'S to ABCD conversion')
    tr, ti = solved[:, :n], solved[:, n:]
    result = numpy.block([[tlr - rr @ tr, rr @ ti], [-tr, ti]])
    if not numpy.isfinite(result).all():
        raise numpy.linalg.LinAlgError('S to ABCD conversion: result is not representable')
    return result


def cascade_s(left: NDArray, right: NDArray, *, internal_modes: int) -> NDArray[numpy.complex128]:
    """Cascade spatially ordered two-port S matrices using the Redheffer product.

    ``internal_modes`` is the common right port of ``left`` and left port of
    ``right``. External port counts may differ. Matrices use outgoing rows and
    incoming columns in [left, right] order. No transfer matrix is constructed.
    """
    try:
        k = operator.index(internal_modes)
    except TypeError as error:
        raise ValueError('internal_modes must be a positive integer') from error
    if isinstance(internal_modes, bool | numpy.bool_) or k <= 0:
        raise ValueError('internal_modes must be a positive integer')
    left, right = numpy.asarray(left, dtype=complex), numpy.asarray(right, dtype=complex)
    for matrix in (left, right):
        if matrix.ndim != 2 or matrix.shape[0] != matrix.shape[1] or matrix.shape[0] <= k:
            raise ValueError('each S matrix must be square with a nonempty external port')
        if not numpy.isfinite(matrix).all():
            raise ValueError('S matrices must contain only finite values')
    n = left.shape[0] - k
    ral, tarl, talr, rar = left[:n, :n], left[:n, n:], left[n:, :n], left[n:, n:]
    rbl, tbrl, tblr, rbr = right[:k, :k], right[:k, k:], right[k:, :k], right[k:, k:]
    rhs = numpy.concatenate((talr, rar @ tbrl), axis=1)
    solved = _solve(numpy.eye(k) - rar @ rbl, rhs, 'Redheffer internal feedback')
    x, y = solved[:, :n], solved[:, n:]
    result = numpy.block([[ral + tarl @ rbl @ x, tarl @ (tbrl + rbl @ y)],
                          [tblr @ x, rbr + tblr @ y]])
    if not numpy.isfinite(result).all():
        raise numpy.linalg.LinAlgError('Redheffer composition: result is not representable')
    return result


def _postprocess_s(ss: NDArray, force_nogain: bool, force_reciprocal: bool) -> NDArray[numpy.complex128]:
    if force_nogain:
        u, singular, vh = numpy.linalg.svd(ss)
        ss = (u * numpy.minimum(singular, 1)) @ vh
    if force_reciprocal:
        ss = 0.5 * (ss + ss.T)
    return ss


def get_s(
        ehLs: mode_seq,
        wavenumbers_L: wavenumber_seq,
        ehRs: mode_seq,
        wavenumbers_R: wavenumber_seq,
        dxes: dx_lists2_t,
        *,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        force_nogain: bool = False,
        force_reciprocal: bool = False,
        dual_ehLs: mode_seq | None = None,
        dual_ehRs: mode_seq | None = None,
        ) -> NDArray[numpy.complex128]:
    """Solve an interface in standard [left, right] scattering order.

    ``S = [[R_L, T_R_to_L], [T_L_to_R, R_R]]`` maps incoming to outgoing
    amplitudes. Mixed testing is the default and supports unequal port counts.
    Galerkin testing requires equal counts. Both use full projection matrices.

    ``dxes`` and ``measure`` must match the supplied modal duals; omitted duals
    use the primal fields. Wavenumbers label the modes and are not used in the
    interface continuity equations. Forward modes imply backward (E, -H).

    Mixed testing assigns different test families to the two ports. Incomplete
    modal subspaces can therefore give different approximations after a port
    exchange, even when each result is reciprocal and unitary. Establish basis
    convergence, especially within coupled or degenerate blocks; do not infer
    physical accuracy from these matrix constraints alone.

    ``force_nogain`` caps Euclidean singular values; ``force_reciprocal``
    symmetrizes S. These optional numerical projections are not general power
    or reciprocity tests for open, evanescent, or non-power-normalized bases.
    No such projection is applied by default.
    """
    shape, _ = _validate_port_modes('ehLs', ehLs, wavenumbers_L)
    right_shape, _ = _validate_port_modes('ehRs', ehRs, wavenumbers_R)
    if shape != right_shape:
        raise ValueError('left and right modal fields must share the same E/H shapes')
    _as_wavenumber_array('wavenumbers_L', wavenumbers_L)
    _as_wavenumber_array('wavenumbers_R', wavenumbers_R)
    dl = _validate_dual_modes('dual_ehLs', dual_ehLs, shape, wavenumbers_L)
    dr = _validate_dual_modes('dual_ehRs', dual_ehRs, shape, wavenumbers_R)
    left, right = _as_mode_arrays(ehLs), _as_mode_arrays(ehRs)
    left_dual = left if dl is None else _as_mode_arrays(dl)
    right_dual = right if dr is None else _as_mode_arrays(dr)
    ss = _interface_s(
        [*left, *((e, -h) for e, h in left)],
        [*right, *((e, -h) for e, h in right)],
        left_dual, right_dual, dxes, formulation=formulation, measure=measure,
    )
    return _postprocess_s(ss, force_nogain, force_reciprocal)


def get_tr(
        ehLs: mode_seq,
        wavenumbers_L: wavenumber_seq,
        ehRs: mode_seq,
        wavenumbers_R: wavenumber_seq,
        dxes: dx_lists2_t,
        *,
        dual_ehLs: mode_seq | None = None,
        dual_ehRs: mode_seq | None = None,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
    """Return (T_L_to_R, R_L) from the shared two-sided continuity solve.

    Unlike the old implementation, transmission singular values are never
    modified. Galerkin truncation can therefore expose apparent gain; mixed
    testing is the default. Dual modes and the overlap measure follow get_s.
    """
    ss = get_s(ehLs, wavenumbers_L, ehRs, wavenumbers_R, dxes,
               dual_ehLs=dual_ehLs, dual_ehRs=dual_ehRs,
               formulation=formulation, measure=measure)
    n = len(ehLs)
    return ss[n:, :n], ss[:n, :n]


def get_abcd(
        ehLs: mode_seq,
        wavenumbers_L: wavenumber_seq,
        ehRs: mode_seq,
        wavenumbers_R: wavenumber_seq,
        dxes: dx_lists2_t,
        *,
        dual_ehLs: mode_seq | None = None,
        dual_ehRs: mode_seq | None = None,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> sparse.sparray:
    """Return [f_R, b_R] = ABCD @ [f_L, b_L] for an equal-count interface.

    Prefer get_s for unequal ports or strongly attenuating devices.
    """
    if len(ehLs) != len(ehRs):
        raise ValueError('transfer conversion requires equal port mode counts')
    ss = get_s(ehLs, wavenumbers_L, ehRs, wavenumbers_R, dxes,
               dual_ehLs=dual_ehLs, dual_ehRs=dual_ehRs,
               formulation=formulation, measure=measure)
    return sparse.csr_array(s_to_abcd(ss))


@dataclasses.dataclass(frozen=True, eq=False)
class PreparedTaper:
    """Read-only field-dependent data for repeated frozen-basis evaluation.

    Arrays are copied and made read-only, including when constructed directly.
    Beta matrices are stored in aligned bases. Evaluation overrides are in the
    original section bases and receive the recorded similarity transforms.
    Changing modal fields, grids, duals, or formulation requires preparation
    again. Reusing these fields across wavelengths is an approximation.
    """

    z_coords: NDArray[numpy.float64]
    interfaces: NDArray[numpy.complex128] = dataclasses.field(repr=False)
    wavenumber_matrices: NDArray[numpy.complex128] = dataclasses.field(repr=False)
    backward_wavenumber_matrices: NDArray[numpy.complex128] = dataclasses.field(repr=False)
    basis_transforms: NDArray[numpy.complex128] = dataclasses.field(repr=False)
    reciprocal_branches: tuple[bool, ...]

    def __post_init__(self) -> None:
        for field in dataclasses.fields(self):
            if field.name == 'reciprocal_branches':
                object.__setattr__(self, field.name, tuple(self.reciprocal_branches))
                continue
            value = numpy.array(getattr(self, field.name), dtype=float if field.name == 'z_coords' else complex, copy=True)
            if not numpy.isfinite(value).all():
                raise ValueError(f'{field.name} must contain only finite values')
            value.flags.writeable = False
            object.__setattr__(self, field.name, value)
        z, beta = self.z_coords, self.wavenumber_matrices
        if z.ndim != 1 or z.size < 2 or not (numpy.diff(z) > 0).all():
            raise ValueError('z_coords must contain at least two strictly increasing finite coordinates')
        if beta.ndim != 3 or beta.shape[0] != z.size or beta.shape[1] != beta.shape[2] or beta.shape[1] == 0:
            raise ValueError('wavenumber_matrices must have shape (section, mode, mode)')
        n = beta.shape[1]
        if self.backward_wavenumber_matrices.shape != beta.shape:
            raise ValueError('backward matrix shape must match forward matrices')
        if self.interfaces.shape != (z.size - 1, 2 * n, 2 * n):
            raise ValueError('interface shape must match section and mode counts')
        if self.basis_transforms.shape != (z.size, 2 * n, 2 * n):
            raise ValueError('basis transform shape must match section and mode counts')
        if len(self.reciprocal_branches) != z.size:
            raise ValueError('reciprocal branch flags must match section count')

    @property
    def n_modes(self) -> int:
        """Number of modes in either external port."""
        return self.wavenumber_matrices.shape[1]


def prepare_taper(
        sections: Sequence[ModeSection],
        dxes: dx_lists2_t,
        *,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> PreparedTaper:
    """Validate, align, and cache midpoint interfaces for sampled sections.

    Section mode counts must agree. Fields and duals must already use a
    consistent normalization and measure. No eigensolver is invoked.
    """
    z, modes, duals, beta, n, transforms = _validate_taper_sections(sections, dxes, measure=measure)
    interfaces = [
        _interface_s(modes[i], modes[i + 1], duals[i], duals[i + 1], dxes,
                     formulation=formulation, measure=measure)
        for i in range(len(sections) - 1)
    ]
    return PreparedTaper(
        z, numpy.asarray(interfaces), numpy.asarray([b[:n, :n] for b in beta]),
        numpy.asarray([b[n:, n:] for b in beta]), numpy.asarray(transforms),
        tuple(section.backward_modes is None for section in sections),
    )


def _propagation_s(forward: NDArray, backward: NDArray, length: float) -> NDArray[numpy.complex128]:
    n = forward.shape[0]
    zero = numpy.zeros((n, n), dtype=complex)
    # Backward propagation is evaluated toward decreasing z, avoiding the
    # exponentially growing transfer branch of passive leaky modes.
    result = numpy.block([[zero, linalg.expm(1j * backward * length)],
                          [linalg.expm(-1j * forward * length), zero]])
    if not numpy.isfinite(result).all():
        raise numpy.linalg.LinAlgError('propagation S matrix is not representable')
    return result


def evaluate_taper_s(
        prepared: PreparedTaper,
        *,
        z_coords: Sequence[float] | None = None,
        wavenumber_matrices: NDArray | None = None,
        backward_wavenumber_matrices: NDArray | None = None,
        force_nogain: bool = False,
        force_reciprocal: bool = False,
        ) -> NDArray[numpy.complex128]:
    """Evaluate cached interfaces with optional lengths and beta overrides.

    Overrides have shape (section, mode, mode) in the original section bases.
    Forward overrides synthesize negative backward beta only for reciprocal
    branches; explicit branches require backward overrides as well. This call
    performs no field overlaps or mode solves. It is suitable for frozen-basis
    wavelength approximations and exact length sweeps within the sampled model.
    """
    z = prepared.z_coords if z_coords is None else numpy.asarray(z_coords, dtype=float)
    if z.shape != prepared.z_coords.shape or not numpy.isfinite(z).all() or not (numpy.diff(z) > 0).all():
        raise ValueError('z_coords must match section count and be finite and strictly increasing')
    n = prepared.n_modes

    def override(value: NDArray | None, default: NDArray, start: int) -> NDArray:
        if value is None:
            return default
        value = numpy.asarray(value, dtype=complex)
        if value.shape != default.shape or not numpy.isfinite(value).all():
            raise ValueError('beta overrides must be finite with shape (section, mode, mode)')
        transforms = prepared.basis_transforms[:, start:start + n, start:start + n]
        return numpy.asarray([_solve(t, b @ t, 'beta basis transform') for t, b in zip(transforms, value, strict=True)])

    forward = override(wavenumber_matrices, prepared.wavenumber_matrices, 0)
    if wavenumber_matrices is not None and backward_wavenumber_matrices is None:
        if not all(prepared.reciprocal_branches):
            raise ValueError('forward overrides require backward overrides for explicit backward branches')
        # Forward and backward phase tracking may choose different gauges.
        backward_wavenumber_matrices = -numpy.asarray(wavenumber_matrices)
    backward = override(backward_wavenumber_matrices, prepared.backward_wavenumber_matrices, n)
    total = _propagation_s(forward[0], backward[0], 0)
    for index, length in enumerate(numpy.diff(z)):
        for ss in (
                _propagation_s(forward[index], backward[index], length / 2),
                prepared.interfaces[index],
                _propagation_s(forward[index + 1], backward[index + 1], length / 2),
                ):
            total = cascade_s(total, ss, internal_modes=n)
    return _postprocess_s(total, force_nogain, force_reciprocal)


def get_taper_s(
        sections: Sequence[ModeSection],
        dxes: dx_lists2_t,
        *,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        force_nogain: bool = False,
        force_reciprocal: bool = False,
        ) -> NDArray[numpy.complex128]:
    """Sampled-taper S matrix using midpoint interfaces and Redheffer cascading.

    Each interval propagates half its length in either adjacent basis. Accuracy
    is controlled by section sampling and basis/PML convergence, not ODE
    tolerances. Standard outgoing/incoming [left, right] ordering is used.
    Longitudinal refinement alone cannot repair omitted coupled/degenerate
    modes or remove the orientation bias of incomplete interface projections.
    """
    prepared = prepare_taper(sections, dxes, formulation=formulation, measure=measure)
    return evaluate_taper_s(prepared, force_nogain=force_nogain, force_reciprocal=force_reciprocal)


def get_taper_abcd(
        sections: Sequence[ModeSection],
        dxes: dx_lists2_t,
        *,
        formulation: Literal['mixed', 'galerkin'] = 'mixed',
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> sparse.sparray:
    """Convert sampled-taper S to ABCD, where representable.

    Prefer get_taper_s for long lossy devices. Zero transmission or an
    unrepresentable transfer matrix raises LinAlgError. The removed rtol, atol,
    and max_step arguments never refined the former sampled CMT model.
    """
    return sparse.csr_array(s_to_abcd(get_taper_s(sections, dxes, formulation=formulation, measure=measure)))


def mode_sections_from_bases(
        z_coords: Sequence[float],
        bases: Sequence[ModeBasis],
        *,
        mode_indices: Sequence[int] | None = None,
        backward_bases: Sequence[ModeBasis] | None = None,
        coupling_rtol: float = 1e-10,
        ) -> tuple[ModeSection, ...]:
    """Convert tracked complex-mode bases into EME taper sections.

    The caller is responsible for tracking the bases first and for explicitly
    power-normalizing any port modes whose scattering amplitudes will be
    interpreted as power waves.  When ``backward_bases`` is omitted, the normal
    reciprocal ``(E, -H, -beta)`` branches are synthesized by ``ModeSection``.
    A selected mode subset must be invariant under each basis's wavenumber
    matrix; ``coupling_rtol`` controls the relative numerical coupling threshold.

    Args:
        z_coords: Strictly increasing longitudinal coordinates for the bases.
        bases: Equal-sized, already tracked forward modal bases.
        mode_indices: Optional ordered subset of modal basis indices. Omit to
            retain every mode.
        backward_bases: Optional explicitly solved backward bases. When
            omitted, reciprocal backward branches are synthesized later.
        coupling_rtol: Relative threshold below which propagation-matrix
            coupling across the selected subset is treated as numerical zero.

    Returns:
        Immutable sequence of EME mode sections carrying primal, dual, and
        full propagation-matrix data.

    Raises:
        ValueError: If sequence lengths, mode indices, tolerances, or selected
            propagation blocks are incompatible.
    """
    if len(z_coords) != len(bases):
        raise ValueError('z_coords and bases must have the same length')
    if len(bases) < 2:
        raise ValueError('at least two bases are required')
    if backward_bases is not None and len(backward_bases) != len(bases):
        raise ValueError('backward_bases must have the same length as bases')
    if not numpy.isfinite(coupling_rtol) or coupling_rtol <= 0:
        raise ValueError('coupling_rtol must be finite and positive')

    if mode_indices is None:
        indices = numpy.arange(len(bases[0]), dtype=int)
    else:
        raw_indices = numpy.asarray(mode_indices)
        if raw_indices.ndim != 1 or raw_indices.size == 0:
            raise ValueError('mode_indices must be a nonempty one-dimensional sequence')
        if any(isinstance(index, bool | numpy.bool_) for index in mode_indices):
            raise ValueError('mode_indices entries must be integers')
        try:
            indices = numpy.asarray([operator.index(index) for index in mode_indices], dtype=int)
        except TypeError as error:
            raise ValueError('mode_indices entries must be integers') from error
    if indices.ndim != 1 or indices.size == 0:
        raise ValueError('mode_indices must be a nonempty one-dimensional sequence')
    if numpy.any(indices < 0):
        raise ValueError('mode_indices must be nonnegative')
    if numpy.unique(indices).size != indices.size:
        raise ValueError('mode_indices must not contain duplicates')

    def validate_invariant_subset(matrix: NDArray[numpy.complex128], size: int, section_name: str) -> None:
        """Reject a subset coupled to excluded coordinates by its beta matrix."""
        excluded = numpy.setdiff1d(numpy.arange(size), indices, assume_unique=False)
        if excluded.size == 0:
            return
        scale = max(1.0, float(numpy.linalg.norm(matrix)))
        coupling = max(
            float(numpy.linalg.norm(matrix[numpy.ix_(indices, excluded)])),
            float(numpy.linalg.norm(matrix[numpy.ix_(excluded, indices)])),
        )
        if coupling > coupling_rtol * scale:
            raise ValueError(
                f'mode_indices cut through a coupled wavenumber block in {section_name}; '
                'select the complete coupled block',
            )

    sections: list[ModeSection] = []
    for section_index, (z_coord, basis) in enumerate(zip(z_coords, bases, strict=True)):
        if numpy.any(indices >= len(basis)):
            raise ValueError(f'mode_indices are out of range for bases[{section_index}]')
        validate_invariant_subset(basis.wavenumber_matrix, len(basis), f'bases[{section_index}]')
        modes = [
            (basis.electric_fields[index, 0], basis.magnetic_fields[index, 0])
            for index in indices
            ]
        dual_modes = [
            (basis.electric_fields[index, 1], basis.magnetic_fields[index, 1])
            for index in indices
            ]
        wavenumbers = basis.wavenumbers[indices]
        wavenumber_matrix = basis.wavenumber_matrix[numpy.ix_(indices, indices)]
        if backward_bases is None:
            sections.append(ModeSection(
                z=float(z_coord),
                modes=modes,
                wavenumbers=wavenumbers,
                dual_modes=dual_modes,
                wavenumber_matrix=wavenumber_matrix,
            ))
            continue

        backward_basis = backward_bases[section_index]
        if numpy.any(indices >= len(backward_basis)):
            raise ValueError(f'mode_indices are out of range for backward_bases[{section_index}]')
        validate_invariant_subset(
            backward_basis.wavenumber_matrix,
            len(backward_basis),
            f'backward_bases[{section_index}]',
        )
        sections.append(ModeSection(
            z=float(z_coord),
            modes=modes,
            wavenumbers=wavenumbers,
            backward_modes=[
                (backward_basis.electric_fields[index, 0], backward_basis.magnetic_fields[index, 0])
                for index in indices
                ],
            backward_wavenumbers=backward_basis.wavenumbers[indices],
            dual_modes=dual_modes,
            dual_backward_modes=[
                (backward_basis.electric_fields[index, 1], backward_basis.magnetic_fields[index, 1])
                for index in indices
                ],
            wavenumber_matrix=wavenumber_matrix,
            backward_wavenumber_matrix=backward_basis.wavenumber_matrix[numpy.ix_(indices, indices)],
        ))
    return tuple(sections)