meanas/meanas/fdfd/waveguide_2d.py

r"""
Operators and helper functions for waveguides with unchanging cross-section.

The propagation direction is chosen to be along the z axis, and all fields
are given an implicit z-dependence of the form `exp(-1 * wavenumber * z)`.

As the z-dependence is known, all the functions in this file assume a 2D grid
 (i.e. `dxes = [[[dx_e[0], dx_e[1], ...], [dy_e[0], ...]], [[dx_h[0], ...], [dy_h[0], ...]]]`).


===============

Consider Maxwell's equations in continuous space, in the frequency domain. Assuming
a structure with some (x, y) cross-section extending uniformly into the z dimension,
with a diagonal $\epsilon$ tensor, we have

$$
\begin{aligned}
\nabla \times \vec{E}(x, y, z) &= -\imath \omega \mu \vec{H} \\
\nabla \times \vec{H}(x, y, z) &=  \imath \omega \epsilon \vec{E} \\
\vec{E}(x,y,z) &= (\vec{E}_t(x, y) + E_z(x, y)\vec{z}) e^{-\imath \beta z} \\
\vec{H}(x,y,z) &= (\vec{H}_t(x, y) + H_z(x, y)\vec{z}) e^{-\imath \beta z} \\
\end{aligned}
$$

Expanding the first two equations into vector components, we get

$$
\begin{aligned}
-\imath \omega \mu_{xx} H_x &= \partial_y E_z - \partial_z E_y \\
-\imath \omega \mu_{yy} H_y &= \partial_z E_x - \partial_x E_z \\
-\imath \omega \mu_{zz} H_z &= \partial_x E_y - \partial_y E_x \\
\imath \omega \epsilon_{xx} E_x &= \partial_y H_z - \partial_z H_y \\
\imath \omega \epsilon_{yy} E_y &= \partial_z H_x - \partial_x H_z \\
\imath \omega \epsilon_{zz} E_z &= \partial_x H_y - \partial_y H_x \\
\end{aligned}
$$

Substituting in our expressions for $\vec{E}$, $\vec{H}$ and discretizing:

$$
\begin{aligned}
-\imath \omega \mu_{xx} H_x &= \tilde{\partial}_y E_z + \imath \beta E_y \\
-\imath \omega \mu_{yy} H_y &= -\imath \beta E_x - \tilde{\partial}_x E_z \\
-\imath \omega \mu_{zz} H_z &= \tilde{\partial}_x E_y - \tilde{\partial}_y E_x \\
\imath \omega \epsilon_{xx} E_x &= \hat{\partial}_y H_z + \imath \beta H_y \\
\imath \omega \epsilon_{yy} E_y &= -\imath \beta H_x - \hat{\partial}_x H_z \\
\imath \omega \epsilon_{zz} E_z &= \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
\end{aligned}
$$

Rewrite the last three equations as

$$
\begin{aligned}
\imath \beta H_y &=  \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z \\
\imath \beta H_x &= -\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z \\
\imath \omega E_z &= \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y - \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
\end{aligned}
$$

Now apply $\imath \beta \tilde{\partial}_x$ to the last equation,
then substitute in for $\imath \beta H_x$ and $\imath \beta H_y$:

$$
\begin{aligned}
\imath \beta \tilde{\partial}_x \imath \omega E_z &= \imath \beta \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x H_y
                                                   - \imath \beta \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y H_x \\
        &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x ( \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z)
         - \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (-\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z)  \\
        &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x ( \imath \omega \epsilon_{xx} E_x)
         - \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (-\imath \omega \epsilon_{yy} E_y)  \\
\imath \beta \tilde{\partial}_x E_z &= \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                     + \tilde{\partial}_x \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
\end{aligned}
$$

With a similar approach (but using $\imath \beta \tilde{\partial}_y$ instead), we can get

$$
\begin{aligned}
\imath \beta \tilde{\partial}_y E_z &= \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                     + \tilde{\partial}_y \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y) \\
\end{aligned}
$$

We can combine this equation for $\imath \beta \tilde{\partial}_y E_z$ with
the unused $\imath \omega \mu_{xx} H_x$ and $\imath \omega \mu_{yy} H_y$ equations to get

$$
\begin{aligned}
-\imath \omega \mu_{xx} \imath \beta H_x &=  -\beta^2 E_y + \imath \beta \tilde{\partial}_y E_z \\
-\imath \omega \mu_{xx} \imath \beta H_x &=  -\beta^2 E_y + \tilde{\partial}_y (
                                      \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                    + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
                                    )\\
\end{aligned}
$$

and

$$
\begin{aligned}
-\imath \omega \mu_{yy} \imath \beta H_y &= \beta^2 E_x - \imath \beta \tilde{\partial}_x E_z \\
-\imath \omega \mu_{yy} \imath \beta H_y &= \beta^2 E_x - \tilde{\partial}_x (
                                      \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
                                    + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
                                    )\\
\end{aligned}
$$

However, based on our rewritten equation for $\imath \beta H_x$ and the so-far unused
equation for $\imath \omega \mu_{zz} H_z$ we can also write

$$
\begin{aligned}
-\imath \omega \mu_{xx} (\imath \beta H_x) &= -\imath \omega \mu_{xx} (-\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z) \\
                 &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y + \imath \omega \mu_{xx} \hat{\partial}_x (
                         \frac{1}{-\imath \omega \mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x)) \\
                 &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
                         -\mu_{xx} \hat{\partial}_x \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
\end{aligned}
$$

and, similarly,

$$
\begin{aligned}
-\imath \omega \mu_{yy} (\imath \beta H_y) &= \omega^2 \mu_{yy} \epsilon_{xx} E_x
                                           +\mu_{yy} \hat{\partial}_y \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
\end{aligned}
$$

By combining both pairs of expressions, we get

$$
\begin{aligned}
\beta^2 E_x - \tilde{\partial}_x (
    \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
  + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
    ) &= \omega^2 \mu_{yy} \epsilon_{xx} E_x
        +\mu_{yy} \hat{\partial}_y \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
-\beta^2 E_y + \tilde{\partial}_y (
    \frac{1}{\epsilon_{zz}} \hat{\partial}_x (\epsilon_{xx} E_x)
  + \frac{1}{\epsilon_{zz}} \hat{\partial}_y (\epsilon_{yy} E_y)
    ) &= -\omega^2 \mu_{xx} \epsilon_{yy} E_y
         -\mu_{xx} \hat{\partial}_x \frac{1}{\mu_{zz}} (\tilde{\partial}_x E_y - \tilde{\partial}_y E_x) \\
\end{aligned}
$$

Using these, we can construct the eigenvalue problem

$$
\beta^2 \begin{bmatrix} E_x \\
                        E_y \end{bmatrix} =
    (\omega^2 \begin{bmatrix} \mu_{yy} \epsilon_{xx} & 0 \\
                                                   0 & \mu_{xx} \epsilon_{yy} \end{bmatrix} +
              \begin{bmatrix} -\mu_{yy} \hat{\partial}_y \\
                               \mu_{xx} \hat{\partial}_x \end{bmatrix} \mu_{zz}^{-1}
              \begin{bmatrix} -\tilde{\partial}_y & \tilde{\partial}_x \end{bmatrix} +
      \begin{bmatrix} \tilde{\partial}_x \\
                      \tilde{\partial}_y \end{bmatrix} \epsilon_{zz}^{-1}
                 \begin{bmatrix} \hat{\partial}_x \epsilon_{xx} & \hat{\partial}_y \epsilon_{yy} \end{bmatrix})
    \begin{bmatrix} E_x \\
                    E_y \end{bmatrix}
$$

In the literature, $\beta$ is usually used to denote the lossless/real part of the propagation constant,
but in `meanas` it is allowed to be complex.

An equivalent eigenvalue problem can be formed using the $H_x$ and $H_y$ fields, if those are more convenient.

Note that $E_z$ was never discretized, so $\beta$ will need adjustment to account for numerical dispersion
if the result is introduced into a space with a discretized z-axis.


"""
from typing import Any, Literal
from collections.abc import Sequence
import dataclasses
import numpy
from numpy.typing import NDArray
from numpy.linalg import norm
from scipy import sparse
from scipy import linalg
from scipy.optimize import linear_sum_assignment
from scipy.sparse import linalg as splinalg

from ..fdmath.operators import deriv_forward, deriv_back, cross
from ..fdmath import vec, unvec, dx_lists2_t, vcfdfield2_t, vcfdslice_t, vcfdfield2, vfdslice, vcfdslice
from ..eigensolvers import (
    TargetedEigensolution,
    TargetedEigensolver,
    rayleigh_quotient_iteration,
    scipy_eigs_near,
    signed_eigensolve,
    )


__author__ = 'Jan Petykiewicz'


def operator_e(
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        ) -> sparse.sparray:
    r"""
    Waveguide operator of the form

        omega**2 * mu * epsilon +
        mu * [[-Dy], [Dx]] / mu * [-Dy, Dx] +
        [[Dx], [Dy]] / epsilon * [Dx, Dy] * epsilon

    for use with a field vector of the form `cat([E_x, E_y])`.

    More precisely, the operator is

    $$
    \omega^2 \begin{bmatrix} \mu_{yy} \epsilon_{xx} & 0 \\
                                                      0 & \mu_{xx} \epsilon_{yy} \end{bmatrix} +
                 \begin{bmatrix} -\mu_{yy} \hat{\partial}_y \\
                                   \mu_{xx} \hat{\partial}_x \end{bmatrix} \mu_{zz}^{-1}
                 \begin{bmatrix} -\tilde{\partial}_y & \tilde{\partial}_x \end{bmatrix} +
      \begin{bmatrix} \tilde{\partial}_x \\
                       \tilde{\partial}_y \end{bmatrix} \epsilon_{zz}^{-1}
                 \begin{bmatrix} \hat{\partial}_x \epsilon_{xx} & \hat{\partial}_y \epsilon_{yy} \end{bmatrix}
    $$

    $\tilde{\partial}_x$ and $\hat{\partial}_x$ are the forward and backward derivatives along x,
    and each $\epsilon_{xx}$, $\mu_{yy}$, etc. is a diagonal matrix containing the vectorized material
    property distribution.

    This operator can be used to form an eigenvalue problem of the form
    `operator_e(...) @ [E_x, E_y] = wavenumber**2 * [E_x, E_y]`

    which can then be solved for the eigenmodes of the system (an `exp(-i * wavenumber * z)`
    z-dependence is assumed for the fields).

    Args:
        omega: The angular frequency of the system.
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representation of the operator.
    """
    if mu is None:
        mu = numpy.ones_like(epsilon)

    Dfx, Dfy = deriv_forward(dxes[0])
    Dbx, Dby = deriv_back(dxes[1])

    eps_parts = numpy.split(epsilon, 3)
    eps_xy = sparse.diags_array(numpy.hstack((eps_parts[0], eps_parts[1])))
    eps_z_inv = sparse.diags_array(1 / eps_parts[2])

    mu_parts = numpy.split(mu, 3)
    mu_yx = sparse.diags_array(numpy.hstack((mu_parts[1], mu_parts[0])))
    mu_z_inv = sparse.diags_array(1 / mu_parts[2])

    op = (
        omega * omega * mu_yx @ eps_xy
        + mu_yx @ sparse.vstack((-Dby, Dbx)) @ mu_z_inv @ sparse.hstack((-Dfy, Dfx))
        + sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby)) @ eps_xy
        )
    return op


def operator_h(
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        ) -> sparse.sparray:
    r"""
    Waveguide operator of the form

        omega**2 * epsilon * mu +
        epsilon * [[-Dy], [Dx]] / epsilon * [-Dy, Dx] +
        [[Dx], [Dy]] / mu * [Dx, Dy] * mu

    for use with a field vector of the form `cat([H_x, H_y])`.

    More precisely, the operator is

    $$
    \omega^2 \begin{bmatrix} \epsilon_{yy} \mu_{xx} & 0 \\
                                                  0 & \epsilon_{xx} \mu_{yy} \end{bmatrix} +
                 \begin{bmatrix} -\epsilon_{yy} \tilde{\partial}_y \\
                                  \epsilon_{xx} \tilde{\partial}_x \end{bmatrix} \epsilon_{zz}^{-1}
                 \begin{bmatrix} -\hat{\partial}_y & \hat{\partial}_x \end{bmatrix} +
      \begin{bmatrix} \hat{\partial}_x \\
                      \hat{\partial}_y \end{bmatrix} \mu_{zz}^{-1}
                 \begin{bmatrix} \tilde{\partial}_x \mu_{xx} & \tilde{\partial}_y \mu_{yy} \end{bmatrix}
    $$

    $\tilde{\partial}_x$ and $\hat{\partial}_x$ are the forward and backward derivatives along x,
    and each $\epsilon_{xx}$, $\mu_{yy}$, etc. is a diagonal matrix containing the vectorized material
    property distribution.

    This operator can be used to form an eigenvalue problem of the form
    `operator_h(...) @ [H_x, H_y] = wavenumber**2 * [H_x, H_y]`

    which can then be solved for the eigenmodes of the system (an `exp(-i * wavenumber * z)`
    z-dependence is assumed for the fields).

    Args:
        omega: The angular frequency of the system.
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representation of the operator.
    """
    if mu is None:
        mu = numpy.ones_like(epsilon)

    Dfx, Dfy = deriv_forward(dxes[0])
    Dbx, Dby = deriv_back(dxes[1])

    eps_parts = numpy.split(epsilon, 3)
    eps_yx = sparse.diags_array(numpy.hstack((eps_parts[1], eps_parts[0])))
    eps_z_inv = sparse.diags_array(1 / eps_parts[2])

    mu_parts = numpy.split(mu, 3)
    mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1])))
    mu_z_inv = sparse.diags_array(1 / mu_parts[2])

    op = (
        omega * omega * eps_yx @ mu_xy
        + eps_yx @ sparse.vstack((-Dfy, Dfx)) @ eps_z_inv @ sparse.hstack((-Dby, Dbx))
        + sparse.vstack((Dbx, Dby)) @ mu_z_inv @ sparse.hstack((Dfx, Dfy)) @ mu_xy
        )
    return op


def normalized_fields_e(
        e_xy: vcfdfield2,
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        prop_phase: float = 0,
        *,
        normalization: Literal['power', 'lorentz'] = 'power',
        ) -> tuple[vcfdslice_t, vcfdslice_t]:
    r"""
    Given a vector `e_xy` containing the vectorized E_x and E_y fields,
     returns normalized, vectorized E and H fields for the system.

    Args:
        e_xy: Vector containing E_x and E_y fields
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)
        prop_phase: Phase shift `(dz * corrected_wavenumber)` over 1 cell in propagation direction.
                    Default 0 (continuous propagation direction, i.e. dz->0).
        normalization: Unit physical forward power, or unit unconjugated
            stretched-coordinate self-reaction (``lorentz``). Self-orthogonal
            open modes require paired primal/dual normalization instead.

    Returns:
        `(e, h)`, where each field is vectorized, normalized,
        and contains all three vector components.

    Notes:
        `e_xy` is only the transverse electric eigenvector. This helper first
        reconstructs the full three-component `E` and `H` fields with `exy2e(...)`
        and `exy2h(...)`, then applies the selected normalization using
        `_normalized_fields(...)`.

        The default power normalization target is

        $$
        \Re\left[\mathrm{inner\_product}(e, h, \mathrm{conj\_h}=True)\right] = 1,
        $$

        so the returned fields represent a unit-power forward mode under the
        discrete Yee-grid Poynting inner product.
    """
    e = exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon) @ e_xy
    h = exy2h(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ e_xy
    e_norm, h_norm = _normalized_fields(
        e=e, h=h, dxes=dxes, epsilon=epsilon, prop_phase=prop_phase, normalization=normalization,
        )
    return e_norm, h_norm


def normalized_fields_h(
        h_xy: vcfdfield2,
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        prop_phase: float = 0,
        *,
        normalization: Literal['power', 'lorentz'] = 'power',
        ) -> tuple[vcfdslice_t, vcfdslice_t]:
    r"""
    Given a vector `h_xy` containing the vectorized H_x and H_y fields,
     returns normalized, vectorized E and H fields for the system.

    Args:
        h_xy: Vector containing H_x and H_y fields
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)
        prop_phase: Phase shift `(dz * corrected_wavenumber)` over 1 cell in propagation direction.
                    Default 0 (continuous propagation direction, i.e. dz->0).
        normalization: Unit physical forward power, or unit unconjugated
            stretched-coordinate self-reaction (``lorentz``). Self-orthogonal
            open modes require paired primal/dual normalization instead.

    Returns:
        `(e, h)`, where each field is vectorized, normalized,
        and contains all three vector components.

    Notes:
        This is the `H_x/H_y` analogue of `normalized_fields_e(...)`. The final
        normalized mode should describe the same physical solution, but because
        the overall complex phase and sign are chosen heuristically,
        `normalized_fields_e(...)` and `normalized_fields_h(...)` need not return
        identical representatives for nearly symmetric modes.
    """
    e = hxy2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon, mu=mu) @ h_xy
    h = hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu) @ h_xy
    e_norm, h_norm = _normalized_fields(
        e=e, h=h, dxes=dxes, epsilon=epsilon, prop_phase=prop_phase, normalization=normalization,
    )
    return e_norm, h_norm


def _normalized_fields(
        e: vcfdslice,
        h: vcfdslice,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        prop_phase: float = 0,
        *,
        normalization: Literal['power', 'lorentz'] = 'power',
        ) -> tuple[vcfdslice_t, vcfdslice_t]:
    """Normalize physical forward power or unconjugated complex self-reaction.

    The power convention fixes a dominant-energy phase and uses a quadrant
    sign heuristic. Lorentz normalization applies that phase first, then a
    complex square-root rescaling; only the remaining sign freedom preserves
    B_mm = 1. Self-orthogonal modes need a paired primal/dual construction.
    """
    shape = [s.size for s in dxes[0]]

    if normalization not in ('power', 'lorentz'):
        raise ValueError("normalization must be 'power' or 'lorentz'")
    if not numpy.isfinite(e).all() or not numpy.isfinite(h).all() or not numpy.any(e):
        raise ValueError('normalization requires finite, nonzero fields')

    energy = numpy.abs(epsilon * e.conj() * e)
    norm_angle = -numpy.angle(e[energy.argmax()])
    weighted = unvec(e * energy * numpy.exp(1j * norm_angle), shape)
    quadrant = weighted[:, :max(shape[0] // 2, 1), :max(shape[1] // 2, 1)].real.sum()
    sign = -1 if quadrant < 0 else 1
    phase = sign * numpy.exp(1j * norm_angle)

    if normalization == 'power':
        power = inner_product(e, h, dxes, prop_phase=prop_phase, conj_h=True).real
        if not numpy.isfinite(power) or power <= 0:
            raise ValueError(f'normalization requires positive forward power, got {power}')
        factor = phase / numpy.sqrt(power)
    else:
        reaction = inner_product(e, h, dxes, prop_phase=prop_phase, measure='stretched')
        # Fix the arbitrary phase before reaction normalization. Afterwards only
        # a sign change preserves B_mm = 1.
        scale = numpy.linalg.norm(e) * numpy.linalg.norm(h)
        if not numpy.isfinite(reaction) or abs(reaction) <= numpy.finfo(float).eps * scale:
            raise ValueError('zero self-reaction; use a paired primal/dual normalization')
        factor = phase / numpy.sqrt(complex(reaction * phase ** 2))
    e = e * factor
    h = h * factor

    return vcfdslice_t(e), vcfdslice_t(h)


def exy2h(
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None
        ) -> sparse.sparray:
    """
    Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
     into a vectorized H containing all three H components

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representing the operator.
    """
    e2hop = e2h(wavenumber=wavenumber, omega=omega, dxes=dxes, mu=mu)
    return e2hop @ exy2e(wavenumber=wavenumber, dxes=dxes, epsilon=epsilon)


def hxy2e(
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None
        ) -> sparse.sparray:
    """
    Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
     into a vectorized E containing all three E components

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representing the operator.
    """
    h2eop = h2e(wavenumber=wavenumber, omega=omega, dxes=dxes, epsilon=epsilon)
    return h2eop @ hxy2h(wavenumber=wavenumber, dxes=dxes, mu=mu)


def hxy2h(
        wavenumber: complex,
        dxes: dx_lists2_t,
        mu: vfdslice | None = None
        ) -> sparse.sparray:
    """
    Operator which transforms the vector `h_xy` containing the vectorized H_x and H_y fields,
     into a vectorized H containing all three H components

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`.
                    It should satisfy `operator_h() @ h_xy == wavenumber**2 * h_xy`
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representing the operator.
    """
    Dfx, Dfy = deriv_forward(dxes[0])
    hxy2hz = sparse.hstack((Dfx, Dfy)) / (1j * wavenumber)

    if mu is not None:
        mu_parts = numpy.split(mu, 3)
        mu_xy = sparse.diags_array(numpy.hstack((mu_parts[0], mu_parts[1])))
        mu_z_inv = sparse.diags_array(1 / mu_parts[2])

        hxy2hz = mu_z_inv @ hxy2hz @ mu_xy

    n_pts = dxes[1][0].size * dxes[1][1].size
    op = sparse.vstack((sparse.eye_array(2 * n_pts),
                        hxy2hz))
    return op


def exy2e(
        wavenumber: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        ) -> sparse.sparray:
    r"""
    Operator which transforms the vector `e_xy` containing the vectorized E_x and E_y fields,
     into a vectorized E containing all three E components

    From the operator derivation (see module docs), we have

    $$
    \imath \omega \epsilon_{zz} E_z = \hat{\partial}_x H_y - \hat{\partial}_y H_x \\
    $$

    as well as the intermediate equations

    $$
    \begin{aligned}
    \imath \beta H_y &=  \imath \omega \epsilon_{xx} E_x - \hat{\partial}_y H_z \\
    \imath \beta H_x &= -\imath \omega \epsilon_{yy} E_y - \hat{\partial}_x H_z \\
    \end{aligned}
    $$

    Combining these, we get

    $$
    \begin{aligned}
    E_z &= \frac{1}{- \omega \beta \epsilon_{zz}} ((
             \hat{\partial}_y \hat{\partial}_x H_z
            -\hat{\partial}_x \hat{\partial}_y H_z)
          + \imath \omega (\hat{\partial}_x \epsilon_{xx} E_x + \hat{\partial}_y \epsilon{yy} E_y))
        &= \frac{1}{\imath \beta \epsilon_{zz}} (\hat{\partial}_x \epsilon_{xx} E_x + \hat{\partial}_y \epsilon{yy} E_y)
    \end{aligned}
    $$

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
                    It should satisfy `operator_e() @ e_xy == wavenumber**2 * e_xy`
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid

    Returns:
        Sparse matrix representing the operator.
    """
    Dbx, Dby = deriv_back(dxes[1])
    exy2ez = sparse.hstack((Dbx, Dby)) / (1j * wavenumber)

    if epsilon is not None:
        epsilon_parts = numpy.split(epsilon, 3)
        epsilon_xy = sparse.diags_array(numpy.hstack((epsilon_parts[0], epsilon_parts[1])))
        epsilon_z_inv = sparse.diags_array(1 / epsilon_parts[2])

        exy2ez = epsilon_z_inv @ exy2ez @ epsilon_xy

    n_pts = dxes[0][0].size * dxes[0][1].size
    op = sparse.vstack((sparse.eye_array(2 * n_pts),
                        exy2ez))
    return op


def e2h(
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        mu: vfdslice | None = None
        ) -> sparse.sparray:
    """
    Returns an operator which, when applied to a vectorized E eigenfield, produces
     the vectorized H eigenfield slice.

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representation of the operator.
    """
    op = curl_e(wavenumber, dxes) / (-1j * omega)
    if mu is not None:
        op = sparse.diags_array(1 / mu) @ op
    return op


def h2e(
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        ) -> sparse.sparray:
    """
    Returns an operator which, when applied to a vectorized H eigenfield, produces
     the vectorized E eigenfield slice.

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid

    Returns:
        Sparse matrix representation of the operator.
    """
    op = sparse.diags_array(1 / (1j * omega * epsilon)) @ curl_h(wavenumber, dxes)
    return op


def curl_e(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray:
    """
    Discretized curl operator for use with the waveguide E field slice.

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)

    Returns:
        Sparse matrix representation of the operator.
    """
    nn = 1
    for dd in dxes[0]:
        nn *= len(dd)

    Bz = -1j * wavenumber * sparse.eye_array(nn)
    Dfx, Dfy = deriv_forward(dxes[0])
    return cross([Dfx, Dfy, Bz])


def curl_h(wavenumber: complex, dxes: dx_lists2_t) -> sparse.sparray:
    """
    Discretized curl operator for use with the waveguide H field slice.

    Args:
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)

    Returns:
        Sparse matrix representation of the operator.
    """
    nn = 1
    for dd in dxes[1]:
        nn *= len(dd)

    Bz = -1j * wavenumber * sparse.eye_array(nn)
    Dbx, Dby = deriv_back(dxes[1])
    return cross([Dbx, Dby, Bz])


def h_err(
        h: vcfdslice,
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None
        ) -> float:
    """
    Calculates the relative error in the H field

    Args:
        h: Vectorized H field
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Relative error `norm(A_h @ h) / norm(h)`.
    """
    ce = curl_e(wavenumber, dxes)
    ch = curl_h(wavenumber, dxes)

    eps_inv = sparse.diags_array(1 / epsilon)

    if mu is None:
        op = ce @ eps_inv @ ch @ h - omega ** 2 * h
    else:
        op = ce @ eps_inv @ ch @ h - omega ** 2 * (mu * h)

    return float(norm(op) / norm(h))


def e_err(
        e: vcfdslice,
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        ) -> float:
    """
    Calculates the relative error in the E field

    Args:
        e: Vectorized E field
        wavenumber: Wavenumber assuming fields have z-dependence of `exp(-i * wavenumber * z)`
        omega: The angular frequency of the system
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Relative error `norm(A_e @ e) / norm(e)`.
    """
    ce = curl_e(wavenumber, dxes)
    ch = curl_h(wavenumber, dxes)

    if mu is None:
        op = ch @ ce @ e - omega ** 2 * (epsilon * e)
    else:
        mu_inv = sparse.diags_array(1 / mu)
        op = ch @ mu_inv @ ce @ e - omega ** 2 * (epsilon * e)

    return float(norm(op) / norm(e))


def sensitivity(
        e_norm: vcfdslice,
        h_norm: vcfdslice,
        wavenumber: complex,
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        ) -> vcfdslice_t:
    r"""
      Given a waveguide structure (`dxes`, `epsilon`, `mu`) and mode fields
    (`e_norm`, `h_norm`, `wavenumber`, `omega`), calculates the sensitivity of the wavenumber
    $\beta$ to changes in the dielectric structure $\epsilon$.

    The output is a vector of the same size as `vec(epsilon)`, with each element specifying the
    sensitivity of `wavenumber` to changes in the corresponding element in `vec(epsilon)`, i.e.

    $$sens_{i} = \frac{\partial\beta}{\partial\epsilon_i}$$

    An adjoint approach is used to calculate the sensitivity; the derivation is provided here:

    Starting with the eigenvalue equation

    $$\beta^2 E_{xy} = A_E E_{xy}$$

    where $A_E$ is the waveguide operator from `operator_e()`, and $E_{xy} = \begin{bmatrix} E_x \\
                                                                                             E_y \end{bmatrix}$,
    we can differentiate with respect to one of the $\epsilon$ elements (i.e. at one Yee grid point), $\epsilon_i$:

    $$
    (2 \beta) \partial_{\epsilon_i}(\beta) E_{xy} + \beta^2 \partial_{\epsilon_i} E_{xy}
        = \partial_{\epsilon_i}(A_E) E_{xy} + A_E \partial_{\epsilon_i} E_{xy}
    $$

    We then multiply by $H_{yx}^\star = \begin{bmatrix}H_y^\star \\ -H_x^\star \end{bmatrix}$ from the left:

    $$
    (2 \beta) \partial_{\epsilon_i}(\beta) H_{yx}^\star E_{xy} + \beta^2 H_{yx}^\star \partial_{\epsilon_i} E_{xy}
        = H_{yx}^\star \partial_{\epsilon_i}(A_E) E_{xy} + H_{yx}^\star A_E \partial_{\epsilon_i} E_{xy}
    $$

    However, $H_{yx}^\star$ is actually a left-eigenvector of $A_E$. This can be verified by inspecting
    the form of `operator_h` ($A_H$) and comparing its conjugate transpose to `operator_e` ($A_E$). Also, note
    $H_{yx}^\star \cdot E_{xy} = H^\star \times E$ recalls the mode orthogonality relation. See doi:10.5194/ars-9-85-201
    for a similar approach. Therefore,

    $$
    H_{yx}^\star A_E \partial_{\epsilon_i} E_{xy} = \beta^2 H_{yx}^\star \partial_{\epsilon_i} E_{xy}
    $$

    and we can simplify to

    $$
    \partial_{\epsilon_i}(\beta)
        = \frac{1}{2 \beta} \frac{H_{yx}^\star \partial_{\epsilon_i}(A_E) E_{xy} }{H_{yx}^\star E_{xy}}
    $$

    This expression can be quickly calculated for all $i$ by writing out the various terms of
    $\partial_{\epsilon_i} A_E$ and recognizing that the vector-matrix-vector products (i.e. scalars)
    $sens_i = \vec{v}_{left} \partial_{\epsilon_i} (\epsilon_{xyz}) \vec{v}_{right}$, indexed by $i$, can be expressed as
    elementwise multiplications $\vec{sens} = \vec{v}_{left} \star \vec{v}_{right}$


    Args:
        e_norm: Normalized, vectorized E_xyz field for the mode. E.g. as returned by `normalized_fields_e`.
        h_norm: Normalized, vectorized H_xyz field for the mode. E.g. as returned by `normalized_fields_e`.
        wavenumber: Propagation constant for the mode. The z-axis is assumed to be continuous (i.e. without numerical dispersion).
        omega: The angular frequency of the system.
        dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types` (2D)
        epsilon: Vectorized dielectric constant grid
        mu: Vectorized magnetic permeability grid (default 1 everywhere)

    Returns:
        Sparse matrix representation of the operator.
    """
    if mu is None:
        mu = numpy.ones_like(epsilon)

    Dfx, Dfy = deriv_forward(dxes[0])
    Dbx, Dby = deriv_back(dxes[1])

    eps_x, eps_y, eps_z = numpy.split(epsilon, 3)
    eps_xy = sparse.diags_array(numpy.hstack((eps_x, eps_y)))
    eps_z_inv = sparse.diags_array(1 / eps_z)

    mu_x, mu_y, _mu_z = numpy.split(mu, 3)
    mu_yx = sparse.diags_array(numpy.hstack((mu_y, mu_x)))

    da_exxhyy = vec(dxes[1][0][:, None] * dxes[0][1][None, :])
    da_eyyhxx = vec(dxes[1][1][None, :] * dxes[0][0][:, None])
    ev_xy = numpy.concatenate(numpy.split(e_norm, 3)[:2]) * numpy.concatenate([da_exxhyy, da_eyyhxx])
    hx, hy, hz = numpy.split(h_norm, 3)
    hv_yx_conj = numpy.conj(numpy.concatenate([hy, -hx]))

    sens_xy1 = (hv_yx_conj @ (omega * omega * mu_yx)) * ev_xy
    sens_xy2 = (hv_yx_conj @ sparse.vstack((Dfx, Dfy)) @ eps_z_inv @ sparse.hstack((Dbx, Dby))) * ev_xy
    sens_z   = (hv_yx_conj @ sparse.vstack((Dfx, Dfy)) @ (-eps_z_inv * eps_z_inv)) * (sparse.hstack((Dbx, Dby)) @ eps_xy @ ev_xy)
    norm = hv_yx_conj @ ev_xy

    sens_tot = numpy.concatenate([sens_xy1 + sens_xy2, sens_z]) / (2 * wavenumber * norm)
    return vcfdslice_t(sens_tot)


def solve_modes(
        mode_numbers: Sequence[int],
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        mode_margin: int = 2,
        *,
        rng: numpy.random.Generator | int | None = None,
        v0: NDArray | None = None,
        ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128]]:
    """
    Given a 2D region, attempts to solve for the eigenmode with the specified mode numbers.

    Args:
       mode_numbers: List of 0-indexed mode numbers to solve for
       omega: Angular frequency of the simulation
       dxes: Grid parameters `[dx_e, dx_h]` as described in `meanas.fdmath.types`
       epsilon: Dielectric constant
       mu: Magnetic permeability (default 1 everywhere)
       mode_margin: The eigensolver will actually solve for `(max(mode_number) + mode_margin)`
            modes, but only return the target mode. Increasing this value can improve the solver's
            ability to find the correct mode. Default 2.
       rng: Optional random generator or seed for reproducible initialization.
       v0: Optional starting vector; takes precedence over ``rng``.

    Returns:
        e_xys: NDArray of vfdfield_t specifying fields. First dimension is mode number.
        wavenumbers: list of wavenumbers
    """

    #
    # Solve for the largest-magnitude eigenvalue of the real operator
    #
    dxes_real = [[numpy.real(dx) for dx in dxi] for dxi in dxes]
    mu_real = None if mu is None else numpy.real(mu)
    A_r = operator_e(numpy.real(omega), dxes_real, numpy.real(epsilon), mu_real)

    eigvals, eigvecs = signed_eigensolve(A_r, max(mode_numbers) + mode_margin, rng=rng, v0=v0)
    keep_inds = -(numpy.array(mode_numbers) + 1)
    e_xys = eigvecs[:, keep_inds].T
    eigvals = eigvals[keep_inds]

    #
    # Now solve for the eigenvector of the full operator, using the real operator's
    #  eigenvector as an initial guess for Rayleigh quotient iteration.
    #
    A = operator_e(omega, dxes, epsilon, mu)
    for nn in range(len(mode_numbers)):
        eigvals[nn], e_xys[nn, :] = rayleigh_quotient_iteration(A, e_xys[nn, :])

    # Calculate the wave-vector (force the real part to be positive)
    wavenumbers = numpy.sqrt(eigvals)
    wavenumbers *= numpy.sign(numpy.real(wavenumbers))

    order = wavenumbers.argsort()[::-1]
    e_xys = e_xys[order]
    wavenumbers = wavenumbers[order]

    return e_xys, wavenumbers


def solve_mode(
        mode_number: int,
        *args: Any,
        **kwargs: Any,
        ) -> tuple[vcfdfield2_t, complex]:
    """
    Wrapper around `solve_modes()` that solves for a single mode.

    Args:
       mode_number: 0-indexed mode number to solve for
       *args: passed to `solve_modes()`
       **kwargs: passed to `solve_modes()`

    Returns:
        (e_xy, wavenumber)
    """
    kwargs['mode_numbers'] = [mode_number]
    e_xys, wavenumbers = solve_modes(*args, **kwargs)
    return vcfdfield2_t(e_xys[0]), wavenumbers[0]


def inner_product(
        e1: vcfdfield2,
        h2: vcfdfield2,
        dxes: dx_lists2_t,
        prop_phase: float = 0,
        conj_h: bool = False,
        trapezoid: bool = False,
        *,
        measure: Literal['physical', 'stretched'] = 'physical',
        ) -> complex:
    r"""
    Compute the discrete waveguide overlap / Poynting inner product.

    This is the 2D transverse integral corresponding to the time-averaged
    longitudinal Poynting flux,

    $$
    \frac{1}{2}\int (E_x H_y - E_y H_x) \, dx \, dy
    $$

    with the Yee-grid staggering and optional propagation-phase adjustment used
    by the waveguide helpers in this module.

    Args:
        e1: Vectorized electric field, typically from `exy2e(...)` or
            `normalized_fields_e(...)`.
        h2: Vectorized magnetic field, typically from `hxy2h(...)`,
            `exy2h(...)`, or one of the normalization helpers.
        dxes: Two-dimensional Yee-grid spacing lists `[dx_e, dx_h]`.
        prop_phase: Phase advance over one propagation cell. This is used to
            shift the H field into the same longitudinal reference plane as the
            E field.
        conj_h: Whether to conjugate `h2` before forming the overlap. Use
            `True` for the usual time-averaged power normalization.
        trapezoid: Whether to use trapezoidal quadrature instead of the default
            rectangular Yee-cell sum.
        measure: ``physical`` uses real cell widths for Poynting power;
            ``stretched`` retains complex widths for coordinate reaction.

    Returns:
        Complex overlap / longitudinal power integral.
    """

    shape = [s.size for s in dxes[0]]

    # H phase is adjusted by a half-cell forward shift for Yee cell, and 1-cell reverse shift for Poynting
    phase = numpy.exp(-1j * -prop_phase / 2)

    E1 = unvec(e1, shape)
    H2 = unvec(h2, shape) * phase

    if conj_h:
        H2 = numpy.conj(H2)

    # Find time-averaged Sz and normalize to it
    if measure not in ('physical', 'stretched'):
        raise ValueError("measure must be 'physical' or 'stretched'")
    widths = [[numpy.real(dxyz) if measure == 'physical' else numpy.asarray(dxyz)
               for dxyz in dxeh] for dxeh in dxes]
    if trapezoid:
        Sz_a = numpy.trapezoid(numpy.trapezoid(E1[0] * H2[1], numpy.cumsum(widths[1][1])), numpy.cumsum(widths[0][0]))
        Sz_b = numpy.trapezoid(numpy.trapezoid(E1[1] * H2[0], numpy.cumsum(widths[0][1])), numpy.cumsum(widths[1][0]))
    else:
        # Ex * Hy occupies an electric-x by magnetic-y area; Ey * Hx uses
        # the complementary magnetic-x by electric-y area.
        Sz_a = E1[0] * H2[1] * widths[0][0][:, None] * widths[1][1][None, :]
        Sz_b = E1[1] * H2[0] * widths[1][0][:, None] * widths[0][1][None, :]
    Sz = 0.5 * (Sz_a.sum() - Sz_b.sum())
    return Sz

@dataclasses.dataclass(frozen=True, eq=False)
class ModeDiagnostics:
    """Per-mode numerical and physical diagnostics.

    ``biorthogonality_condition`` is the condition number of the right/left
    Gram matrix after independently equilibrating its vector columns.

    Attributes:
        right_residuals: Relative residual of each primal mode or basis column.
        left_residuals: Relative residual of each transpose-left mode or column.
        powers: Conjugated Poynting power of each primal modal field.
        pml_participation: Fraction of transverse state weight inside the PML.
        target_indices: Index of the nearest requested target for each spectral
            value.
        target_distances: Absolute distance from each spectral value to that
            target.
        biorthogonality_error: Maximum error in the Lorentz dual Gram matrix.
        biorthogonality_condition: Equilibrated right/left Gram condition
            number.
    """

    right_residuals: NDArray[numpy.float64]
    left_residuals: NDArray[numpy.float64]
    powers: NDArray[numpy.complex128]
    pml_participation: NDArray[numpy.float64]
    target_indices: NDArray[numpy.int_]
    target_distances: NDArray[numpy.float64]
    biorthogonality_error: float
    biorthogonality_condition: float


@dataclasses.dataclass(frozen=True, eq=False)
class ModeBasis:
    """Modal fields and propagation data returned by targeted 2D solvers.

    ``wavenumbers`` records the eigenspectrum. ``wavenumber_matrix`` is the
    actual beta operator in the current field basis; it becomes non-diagonal
    when tracking rotates a nearly degenerate subspace. Field arrays have shape
    ``(mode, primal_or_dual, field)``; index 0 of the second axis is the primal
    field and index 1 is its Lorentz dual.

    Attributes:
        wavenumbers: Eigenvalue spectrum associated with the basis.
        wavenumber_matrix: Representation of the beta operator in the current
            possibly mixed basis.
        electric_fields: Primal and Lorentz-dual vectorized electric fields.
        magnetic_fields: Primal and Lorentz-dual vectorized magnetic fields.
    """

    wavenumbers: NDArray[numpy.complex128]
    wavenumber_matrix: NDArray[numpy.complex128]
    electric_fields: NDArray[numpy.complex128] = dataclasses.field(repr=False)
    magnetic_fields: NDArray[numpy.complex128] = dataclasses.field(repr=False)

    def __len__(self) -> int:
        """Return the number of modal basis vectors."""
        return self.wavenumbers.size


def operator_beta(
        omega: complex,
        dxes: dx_lists2_t,
        epsilon: vfdslice,
        mu: vfdslice | None = None,
        ) -> sparse.sparray:
    r"""Build the first-order transverse Maxwell operator ``M psi = beta psi``.

    The state ordering is ``[E_x, E_y, H_x, H_y]``.  Longitudinal fields are
    eliminated, but both transverse field families are retained so the sign of
    ``beta`` remains explicit.  Algebraically, ``M @ M`` has the existing
    electric and magnetic ``beta**2`` waveguide operators on its diagonal.

    Args:
        omega: Nonzero angular frequency.
        dxes: Two-dimensional Yee-grid spacings, which may include complex PML
            stretching.
        epsilon: Vectorized three-component relative permittivity.
        mu: Optional vectorized three-component relative permeability. Unit
            permeability is used when omitted.

    Returns:
        Sparse first-order signed-beta operator on transverse E/H states.

    Raises:
        ValueError: If the material vectors, grid shape, or frequency are
            incompatible.
    """
    epsilon = numpy.asarray(epsilon)
    if mu is None:
        mu = numpy.ones_like(epsilon)
    mu = numpy.asarray(mu)
    if epsilon.ndim != 1 or mu.ndim != 1 or epsilon.shape != mu.shape or epsilon.size % 3:
        raise ValueError('epsilon and mu must be equal-length vectorized three-component fields')
    if omega == 0:
        raise ValueError('omega must be nonzero')

    n_points = epsilon.size // 3
    if n_points != dxes[0][0].size * dxes[0][1].size:
        raise ValueError('material and grid shapes do not match')

    dfx, dfy = deriv_forward(dxes[0])
    dbx, dby = deriv_back(dxes[1])
    zero = sparse.csr_array((n_points, n_points))
    identity = sparse.eye_array(n_points)
    rotation = sparse.block_array([[zero, -identity], [identity, zero]], format='csr')

    eps_x, eps_y, eps_z = numpy.split(epsilon, 3)
    mu_x, mu_y, mu_z = numpy.split(mu, 3)
    eps_xy = sparse.diags_array(numpy.concatenate((eps_x, eps_y)))
    mu_xy = sparse.diags_array(numpy.concatenate((mu_x, mu_y)))
    eps_z_inv = sparse.diags_array(1 / eps_z)
    mu_z_inv = sparse.diags_array(1 / mu_z)

    grad_forward = sparse.vstack((dfx, dfy), format='csr')
    grad_backward = sparse.vstack((dbx, dby), format='csr')
    rotated_div_backward = sparse.hstack((-dby, dbx), format='csr')
    rotated_div_forward = sparse.hstack((-dfy, dfx), format='csr')

    beta_e_from_h = (
        -omega * rotation @ mu_xy
        + grad_forward @ eps_z_inv @ rotated_div_backward / omega
        )
    beta_h_from_e = (
        omega * rotation @ eps_xy
        - grad_backward @ mu_z_inv @ rotated_div_forward / omega
        )
    empty = sparse.csr_array((2 * n_points, 2 * n_points))
    return sparse.block_array(
        [[empty, beta_e_from_h], [beta_h_from_e, empty]],
        format = 'csr',
        )


def lorentz_metric(
        dxes: dx_lists2_t,
        *,
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> sparse.sparray:
    """Return the state metric matching the EME symmetric Lorentz overlap.

    The staggered transverse areas are ``dx_e * dy_h`` for ``Ex * Hy``
    and ``dx_h * dy_e`` for ``Ey * Hx``. These assignments make the
    reciprocal operator self-transpose under the stretched-coordinate metric.

    Args:
        dxes: Two-dimensional Yee-grid spacings.
        measure: Retain complex widths (``stretched``, default), or use their
            real parts (``physical``). Must match the dual-basis measure.

    Returns:
        Sparse bilinear metric on states ordered as ``[Ex, Ey, Hx, Hy]``.
    """
    nx = dxes[0][0].size
    ny = dxes[0][1].size
    if measure not in ('physical', 'stretched'):
        raise ValueError("measure must be 'physical' or 'stretched'")
    widths = [[numpy.real(values) if measure == 'physical' else numpy.asarray(values)
                  for values in family] for family in dxes]
    weight_ex_hy = numpy.outer(widths[0][0], widths[1][1]).reshape(nx * ny)
    weight_ey_hx = numpy.outer(widths[1][0], widths[0][1]).reshape(nx * ny)
    diagonal_a = sparse.diags_array(weight_ex_hy)
    diagonal_b = sparse.diags_array(weight_ey_hx)
    empty_n = sparse.csr_array((nx * ny, nx * ny))
    # inner_product(E, H) = E.T @ q @ H; it already includes a factor 1/2.
    q_matrix = 0.5 * sparse.block_array(
        [[empty_n, diagonal_a], [-diagonal_b, empty_n]],
        format = 'csr',
        )
    empty_2n = sparse.csr_array(q_matrix.shape)
    # lorentz_overlap adds the two reaction products and multiplies by 1/2.
    return 0.5 * sparse.block_array(
        [[empty_2n, q_matrix], [q_matrix.T, empty_2n]],
        format = 'csr',
        )


def infer_physical_mask(dxes: dx_lists2_t) -> NDArray[numpy.bool_]:
    """Infer the non-PML aperture from imaginary grid stretching.

    Args:
        dxes: Two-dimensional Yee-grid spacings used by the mode solve.

    Returns:
        A transverse boolean mask that excludes every row or column containing
        complex stretching in either electric or magnetic spacing family.
    """
    pml_x = numpy.zeros(dxes[0][0].size, dtype=bool)
    pml_y = numpy.zeros(dxes[0][1].size, dtype=bool)
    for family in dxes:
        pml_x |= numpy.abs(numpy.imag(family[0])) > 0
        pml_y |= numpy.abs(numpy.imag(family[1])) > 0
    return ~(pml_x[:, None] | pml_y[None, :])


def _state_weights(dxes: dx_lists2_t, physical_mask: NDArray[numpy.bool_]) -> NDArray[numpy.float64]:
    """Build cell-area weights for all four transverse state components."""
    dx = 0.5 * (numpy.real(dxes[0][0]) + numpy.real(dxes[1][0]))
    dy = 0.5 * (numpy.real(dxes[0][1]) + numpy.real(dxes[1][1]))
    weights = numpy.outer(dx, dy) * physical_mask
    return numpy.tile(weights.ravel(), 4)


def _pml_participation(
        states: NDArray[numpy.complex128],
        dxes: dx_lists2_t,
        physical_mask: NDArray[numpy.bool_],
        ) -> NDArray[numpy.float64]:
    """Return the fraction of each state supported outside the physical mask."""
    grid_shape = physical_mask.shape
    full_weights = _state_weights(dxes, numpy.ones(grid_shape, dtype=bool))
    pml_weights = full_weights * numpy.tile(~physical_mask.ravel(), 4)
    participation = numpy.sum(pml_weights[:, None] * numpy.abs(states) ** 2, axis=0) / numpy.sum(
        full_weights[:, None] * numpy.abs(states) ** 2,
        axis = 0,
        )
    return numpy.asarray(participation, dtype=float)


def _normalize_columns(
        matrix: NDArray[numpy.complex128],
        ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.float64]]:
    """Normalize matrix columns and return their original Euclidean norms."""
    norms = numpy.linalg.norm(matrix, axis=0)
    normalized = numpy.divide(matrix, norms, out=numpy.zeros_like(matrix), where=norms != 0)
    return normalized, numpy.asarray(norms, dtype=float)


def _select_unique_modes(
        solutions: Sequence[TargetedEigensolution],
        targets: NDArray[numpy.complex128],
        modes_per_target: int,
        *,
        eigenvalue_rtol: float,
        subspace_tol: float,
        state_weights: NDArray[numpy.float64],
        ) -> tuple[NDArray[numpy.complex128], NDArray[numpy.complex128], NDArray[numpy.complex128], NDArray[numpy.int_]]:
    """Merge targeted candidates into independent, cross-pairable eigenspaces."""
    sqrt_weights = numpy.sqrt(state_weights)

    def rank_and_pivots(vectors: list[NDArray[numpy.complex128]]) -> tuple[int, NDArray[numpy.int_]]:
        """Estimate weighted subspace rank and QR pivot order."""
        if not vectors:
            return 0, numpy.empty(0, dtype=int)
        matrix = sqrt_weights[:, None] * numpy.column_stack(vectors)
        matrix, _norms = _normalize_columns(matrix)
        _q, triangular, pivots = linalg.qr(matrix, mode='economic', pivoting=True)
        diagonal = numpy.abs(numpy.diag(triangular))
        if diagonal.size == 0 or diagonal[0] == 0:
            return 0, numpy.asarray(pivots, dtype=int)
        rank = int(numpy.count_nonzero(diagonal > subspace_tol * diagonal[0]))
        return rank, numpy.asarray(pivots, dtype=int)

    def normalized_cross_gram(
            right_vectors: list[NDArray[numpy.complex128]],
            left_vectors: list[NDArray[numpy.complex128]],
            ) -> NDArray[numpy.complex128]:
        """Form the bilinear cross Gram matrix after column equilibration."""
        right = numpy.column_stack(right_vectors)
        left = numpy.column_stack(left_vectors)
        right, _right_norms = _normalize_columns(right)
        left, _left_norms = _normalize_columns(left)
        return left.T @ right

    def matrix_rank(matrix: NDArray[numpy.complex128]) -> int:
        """Estimate matrix rank using the configured subspace tolerance."""
        singular_values = linalg.svdvals(matrix)
        if singular_values.size == 0:
            return 0
        cutoff = subspace_tol * max(1.0, float(singular_values[0]))
        return int(numpy.count_nonzero(singular_values > cutoff))

    initial_records: list[tuple[complex, NDArray[numpy.complex128], NDArray[numpy.complex128]]] = []
    remaining_records: list[tuple[complex, NDArray[numpy.complex128], NDArray[numpy.complex128]]] = []
    for target_index, (target, solution) in enumerate(zip(targets, solutions, strict=True)):
        local_right: list[NDArray[numpy.complex128]] = []
        local_left: list[NDArray[numpy.complex128]] = []
        order = numpy.argsort(numpy.abs(solution.eigenvalues - target))
        stop_position: int | None = None
        for position, candidate_index in enumerate(order):
            value = complex(solution.eigenvalues[candidate_index])
            right_vector = solution.right_eigenvectors[:, candidate_index]
            left_vector = solution.left_eigenvectors[:, candidate_index]
            initial_records.append((value, right_vector, left_vector))
            local_right.append(right_vector)
            local_left.append(left_vector)
            right_rank, _right_pivots = rank_and_pivots(local_right)
            left_rank, _left_pivots = rank_and_pivots(local_left)
            if right_rank >= modes_per_target and left_rank >= modes_per_target:
                stop_position = position
                break
        if stop_position is None:
            raise RuntimeError(f'target {target_index} did not produce enough independent right and left modes')
        for candidate_index in order[stop_position + 1:]:
            remaining_records.append((
                complex(solution.eigenvalues[candidate_index]),
                solution.right_eigenvectors[:, candidate_index],
                solution.left_eigenvectors[:, candidate_index],
                ))

    clusters: list[list[int]] = []
    for candidate_index, (value, _right, _left) in enumerate(initial_records):
        for cluster in clusters:
            anchor = initial_records[cluster[0]][0]
            if abs(value - anchor) <= eigenvalue_rtol * max(1.0, abs(value), abs(anchor)):
                cluster.append(candidate_index)
                break
        else:
            clusters.append([candidate_index])

    chosen_values: list[complex] = []
    chosen_right: list[NDArray[numpy.complex128]] = []
    chosen_left: list[NDArray[numpy.complex128]] = []
    for cluster in clusters:
        anchor = initial_records[cluster[0]][0]
        right_values = [initial_records[index][0] for index in cluster]
        left_values = [initial_records[index][0] for index in cluster]
        right_vectors = [initial_records[index][1] for index in cluster]
        left_vectors = [initial_records[index][2] for index in cluster]
        right_rank, _right_pivots = rank_and_pivots(right_vectors)
        left_rank, _left_pivots = rank_and_pivots(left_vectors)
        required_rank = max(right_rank, left_rank)
        cross_gram = normalized_cross_gram(right_vectors, left_vectors)
        cross_rank = matrix_rank(cross_gram)
        for value, right_vector, left_vector in remaining_records:
            if cross_rank >= required_rank:
                break
            if abs(value - anchor) > eigenvalue_rtol * max(1.0, abs(value), abs(anchor)):
                continue
            right_vectors.append(right_vector)
            left_vectors.append(left_vector)
            right_values.append(value)
            left_values.append(value)
            cross_gram = normalized_cross_gram(right_vectors, left_vectors)
            cross_rank = matrix_rank(cross_gram)
        if cross_rank < required_rank:
            raise RuntimeError(
                f'eigenvalue cluster near {anchor!r} has incompatible right/transpose-left '
                f'subspaces (required cross rank {required_rank}, available {cross_rank})',
                )

        _q, _triangular, right_pivots = linalg.qr(cross_gram, mode='economic', pivoting=True)
        right_selection = numpy.sort(numpy.asarray(right_pivots[:required_rank], dtype=int))
        _q, _triangular, left_pivots = linalg.qr(
            cross_gram[:, right_selection].T,
            mode='economic',
            pivoting=True,
            )
        left_selection = numpy.asarray(left_pivots[:required_rank], dtype=int)
        selected_cross_gram = cross_gram[numpy.ix_(left_selection, right_selection)]
        selected_cross_rank = matrix_rank(selected_cross_gram)
        if selected_cross_rank < required_rank:
            raise RuntimeError(
                f'eigenvalue cluster near {anchor!r} has incompatible right/transpose-left '
                f'subspaces (required cross rank {required_rank}, selected {selected_cross_rank})',
                )
        selected_right_values = numpy.asarray(right_values)[right_selection]
        selected_left_values = numpy.asarray(left_values)[left_selection]
        rows, columns = linear_sum_assignment(numpy.abs(selected_right_values[:, None] - selected_left_values[None, :]))
        paired_left = numpy.asarray(left_selection)[columns[numpy.argsort(rows)]]
        chosen_values.extend(numpy.asarray(right_values)[right_selection])
        chosen_right.extend(right_vectors[int(index)] for index in right_selection)
        chosen_left.extend(left_vectors[int(index)] for index in paired_left)

    values = numpy.asarray(chosen_values, dtype=complex)
    target_indices = numpy.argmin(numpy.abs(values[:, None] - targets[None, :]), axis=1)
    return (
        values,
        numpy.column_stack(chosen_right),
        numpy.column_stack(chosen_left),
        numpy.asarray(target_indices, dtype=int),
        )


def solve_modes_near_beta(
        omega: complex,
        solve_dxes: dx_lists2_t,
        epsilon: vfdslice,
        *,
        target_wavenumbers: Sequence[complex],
        modes_per_target: int,
        mu: vfdslice | None = None,
        overlap_dxes: dx_lists2_t | None = None,
        physical_mask: NDArray[numpy.bool_] | None = None,
        candidate_margin: int = 4,
        eigensolver: TargetedEigensolver = scipy_eigs_near,
        eigensolver_opts: dict[str, Any] | None = None,
        residual_tol: float = 1e-8,
        duplicate_eigenvalue_rtol: float = 1e-8,
        duplicate_subspace_tol: float = 1e-7,
        ) -> tuple[ModeBasis, ModeDiagnostics]:
    """Solve targeted modes through the first-order signed-beta operator.

    Each target requests an independent candidate solve. Overlapping results
    are merged by eigenvalue cluster and weighted subspace rank, after which
    transpose-left eigenvectors are converted into a Lorentz-biorthogonal dual
    basis. Returned fields use the full three-component vectorization, while
    the underlying eigenproblem is ordered as ``[Ex, Ey, Hx, Hy]``.

    Args:
        omega: Nonzero angular frequency.
        solve_dxes: Yee-grid spacings used to build the possibly PML-stretched
            eigenproblem.
        epsilon: Vectorized three-component relative permittivity.
        target_wavenumbers: Finite signed-beta targets for shift-invert solves.
        modes_per_target: Independent candidates requested around each target.
        mu: Optional vectorized three-component relative permeability.
        overlap_dxes: Grid used for physical overlaps and normalization. By
            default, the real parts of ``solve_dxes`` are used.
        physical_mask: Optional transverse mask selecting the non-PML aperture.
            By default, cells with imaginary grid stretching are excluded.
        candidate_margin: Extra eigenpairs requested from each targeted solve
            to support deduplication and degenerate-subspace repair.
        eigensolver: Backend returning paired right and transpose-left vectors.
        eigensolver_opts: Keyword arguments forwarded to ``eigensolver``.
        residual_tol: Maximum accepted right, left, and dual-basis residual.
        duplicate_eigenvalue_rtol: Relative tolerance for eigenvalue clustering.
        duplicate_subspace_tol: Relative numerical-rank tolerance used while
            merging repeated or degenerate candidates.

    Returns:
        Modal basis and its numerical and physical diagnostics.

    Raises:
        ValueError: If inputs or backend result shapes contain invalid values.
        RuntimeError: If too few independent modes are found, right/left spaces
            cannot be paired, or requested residual tolerances are not met.
    """
    if modes_per_target < 1:
        raise ValueError('modes_per_target must be positive')
    if candidate_margin < 0:
        raise ValueError('candidate_margin must be nonnegative')
    if not numpy.isfinite(residual_tol) or residual_tol <= 0:
        raise ValueError('residual_tol must be finite and positive')
    if not numpy.isfinite(duplicate_eigenvalue_rtol) or duplicate_eigenvalue_rtol <= 0:
        raise ValueError('duplicate_eigenvalue_rtol must be finite and positive')
    if not numpy.isfinite(duplicate_subspace_tol) or duplicate_subspace_tol <= 0:
        raise ValueError('duplicate_subspace_tol must be finite and positive')
    targets = numpy.asarray(target_wavenumbers, dtype=complex)
    if targets.ndim != 1 or targets.size == 0 or not numpy.isfinite(targets).all():
        raise ValueError('target_wavenumbers must be a finite nonempty one-dimensional sequence')
    if overlap_dxes is None:
        overlap_dxes = [[numpy.real(values) for values in family] for family in solve_dxes]
    if physical_mask is None:
        physical_mask = infer_physical_mask(solve_dxes)
    physical_mask = numpy.asarray(physical_mask, dtype=bool)
    grid_shape = (solve_dxes[0][0].size, solve_dxes[0][1].size)
    if physical_mask.shape != grid_shape:
        raise ValueError('physical_mask must match the transverse grid shape')

    operator = operator_beta(omega, solve_dxes, epsilon, mu)
    options = {} if eigensolver_opts is None else dict(eigensolver_opts)
    solutions = [
        eigensolver(operator, complex(target), modes_per_target + candidate_margin, **options)
        for target in targets
        ]
    for target_index, solution in enumerate(solutions):
        count = numpy.asarray(solution.eigenvalues).size
        if numpy.asarray(solution.eigenvalues).shape != (count,):
            raise ValueError(f'eigensolver result {target_index} eigenvalues must be one-dimensional')
        expected_shape = (operator.shape[0], count)
        if numpy.asarray(solution.right_eigenvectors).shape != expected_shape or numpy.asarray(solution.left_eigenvectors).shape != expected_shape:
            raise ValueError(f'eigensolver result {target_index} eigenvector shapes are invalid')
        if count < modes_per_target:
            raise RuntimeError(f'eigensolver result {target_index} returned too few candidates')
        if not (
                numpy.isfinite(solution.eigenvalues).all()
                and numpy.isfinite(solution.right_eigenvectors).all()
                and numpy.isfinite(solution.left_eigenvectors).all()
                ):
            raise ValueError(f'eigensolver result {target_index} contains non-finite values')
    merge_weights = _state_weights(overlap_dxes, numpy.ones(grid_shape, dtype=bool))
    values, right, left, _target_indices = _select_unique_modes(
        solutions,
        targets,
        modes_per_target,
        eigenvalue_rtol = duplicate_eigenvalue_rtol,
        subspace_tol = duplicate_subspace_tol,
        state_weights = merge_weights,
        )

    right, right_norms = _normalize_columns(right)
    left, left_norms = _normalize_columns(left)
    if numpy.any(right_norms == 0) or numpy.any(left_norms == 0):
        raise RuntimeError('selected right/transpose-left eigenbasis contains a zero vector')
    gram = left.T @ right
    condition = float(numpy.linalg.cond(gram))
    if not numpy.isfinite(condition) or condition > 1e12:
        raise RuntimeError(f'left/right eigenbasis is singular or ill-conditioned ({condition=:.3e})')
    left = linalg.solve(gram, left.T).T

    weights = _state_weights(overlap_dxes, physical_mask)
    for index in range(values.size):
        vector = right[:, index]
        norm_value = numpy.sqrt(numpy.sum(weights * numpy.abs(vector) ** 2))
        if norm_value == 0:
            raise RuntimeError(f'mode {index} has no support in the physical aperture')
        weighted_vector = numpy.where(weights > 0, vector, 0)
        phase_sample = weighted_vector[numpy.argmax(numpy.abs(weighted_vector))]
        scale = numpy.exp(-1j * numpy.angle(phase_sample)) / norm_value
        right[:, index] *= scale
        left[:, index] /= scale

    metric = lorentz_metric(overlap_dxes).tocsc()
    dual_states = numpy.asarray(splinalg.spsolve(metric.T, left))
    if dual_states.ndim == 1:
        dual_states = dual_states[:, None]
    biorthogonal_gram = dual_states.T @ metric @ right
    biorthogonality_error = float(numpy.max(numpy.abs(biorthogonal_gram - numpy.eye(values.size))))
    if biorthogonality_error > residual_tol:
        raise RuntimeError('failed to construct a biorthogonal Lorentz dual basis')

    n_points = grid_shape[0] * grid_shape[1]
    primal_electric = numpy.empty((values.size, 3 * n_points), dtype=complex)
    primal_magnetic = numpy.empty_like(primal_electric)
    dual_electric = numpy.zeros_like(primal_electric)
    dual_magnetic = numpy.zeros_like(primal_electric)
    epsilon_array = numpy.asarray(epsilon)
    mu_array = None if mu is None else numpy.asarray(mu)
    for index, beta in enumerate(values):
        e_xy = right[:2 * n_points, index]
        h_xy = right[2 * n_points:, index]
        primal_electric[index] = exy2e(beta, solve_dxes, epsilon_array) @ e_xy
        primal_magnetic[index] = hxy2h(beta, solve_dxes, mu_array) @ h_xy
        dual_electric[index, :2 * n_points] = dual_states[:2 * n_points, index]
        dual_magnetic[index, :2 * n_points] = dual_states[2 * n_points:, index]

    basis = ModeBasis(
        wavenumbers = values,
        wavenumber_matrix = numpy.diag(values),
        electric_fields = numpy.stack((primal_electric, dual_electric), axis=1),
        magnetic_fields = numpy.stack((primal_magnetic, dual_magnetic), axis=1),
        )
    diagnostics = _compute_mode_diagnostics(
        basis,
        operator,
        targets,
        overlap_dxes,
        physical_mask,
        )
    if (
            numpy.max(diagnostics.right_residuals) > residual_tol
            or numpy.max(diagnostics.left_residuals) > residual_tol
            ):
        raise RuntimeError('targeted eigenmodes did not meet the requested residual tolerance')
    if diagnostics.biorthogonality_error > residual_tol:
        raise RuntimeError('failed to construct a biorthogonal Lorentz dual basis')
    return basis, diagnostics


def solve_modes_near(
        omega: complex,
        solve_dxes: dx_lists2_t,
        epsilon: vfdslice,
        *,
        target_wavenumbers: Sequence[complex],
        modes_per_target: int,
        mu: vfdslice | None = None,
        overlap_dxes: dx_lists2_t | None = None,
        physical_mask: NDArray[numpy.bool_] | None = None,
        candidate_margin: int = 4,
        eigensolver: TargetedEigensolver = scipy_eigs_near,
        eigensolver_opts: dict[str, Any] | None = None,
        residual_tol: float = 1e-8,
        duplicate_eigenvalue_rtol: float = 1e-8,
        duplicate_subspace_tol: float = 1e-7,
        ) -> tuple[ModeBasis, ModeDiagnostics]:
    """Solve targeted complex modes through the electric beta-squared operator.

    The half-sized transverse-electric problem is solved near each
    ``target_wavenumber**2``. Each eigenpair is then lifted to the signed
    first-order state ``[Ex, Ey, Hx, Hy]`` using the square-root branch nearest
    the requested signed target. Opposite signed targets share a beta-squared
    eigensolve.

    Args:
        omega: Nonzero angular frequency.
        solve_dxes: Yee-grid spacings used by the possibly PML-stretched solve.
        epsilon: Vectorized three-component relative permittivity.
        target_wavenumbers: Finite signed-beta targets. Their squares are passed
            to the targeted eigensolver.
        modes_per_target: Independent candidates requested around each target.
        mu: Optional vectorized three-component relative permeability.
        overlap_dxes: Grid used for physical overlaps and normalization.
        physical_mask: Optional transverse mask selecting the non-PML aperture.
        candidate_margin: Extra eigenpairs requested from each targeted solve.
        eigensolver: Backend receiving the beta-squared operator and target.
        eigensolver_opts: Keyword arguments forwarded to ``eigensolver``.
        residual_tol: Maximum accepted full-state and dual-basis residual.
        duplicate_eigenvalue_rtol: Relative tolerance for signed-beta clusters.
        duplicate_subspace_tol: Numerical-rank tolerance for candidate merging.

    Returns:
        Modal basis compatible with tracking and EME helpers, together with
        its numerical and physical diagnostics.

    Raises:
        ValueError: If the backend result is invalid or inputs are incompatible.
        RuntimeError: If a zero-beta candidate cannot be lifted or residual
            checks fail.
    """
    squared_operator = operator_e(omega, solve_dxes, epsilon, mu)
    squared_cache: dict[tuple[complex, int], TargetedEigensolution] = {}

    def lift_squared_solution(
            first_order_operator: sparse.sparray,
            target: complex,
            how_many: int,
            **options: Any,
            ) -> TargetedEigensolution:
        """Solve one beta-squared target and lift it to a signed state basis."""
        squared_target = complex(target) ** 2
        cache_key = (squared_target, how_many)
        if cache_key not in squared_cache:
            squared_cache[cache_key] = eigensolver(
                squared_operator,
                squared_target,
                how_many,
                **options,
                )
        solution = squared_cache[cache_key]
        eigenvalues = numpy.asarray(solution.eigenvalues, dtype=complex)
        right_e = numpy.asarray(solution.right_eigenvectors, dtype=complex)
        left_e = numpy.asarray(solution.left_eigenvectors, dtype=complex)
        expected_shape = (squared_operator.shape[0], eigenvalues.size)
        if eigenvalues.shape != (eigenvalues.size,):
            raise ValueError('beta-squared eigensolver eigenvalues must be one-dimensional')
        if right_e.shape != expected_shape or left_e.shape != expected_shape:
            raise ValueError('beta-squared eigensolver eigenvector shapes are invalid')
        if eigenvalues.size < modes_per_target:
            raise RuntimeError('beta-squared eigensolver returned too few candidates')
        if not (
                numpy.isfinite(eigenvalues).all()
                and numpy.isfinite(right_e).all()
                and numpy.isfinite(left_e).all()
                ):
            raise ValueError('beta-squared eigensolver result contains non-finite values')

        principal = numpy.sqrt(eigenvalues)
        use_negative = numpy.abs(-principal - target) < numpy.abs(principal - target)
        wavenumbers = numpy.where(use_negative, -principal, principal)
        if numpy.any(wavenumbers == 0):
            raise RuntimeError('zero-beta modes cannot be lifted from the beta-squared operator')

        transverse_size = squared_operator.shape[0]
        if first_order_operator.shape != (2 * transverse_size, 2 * transverse_size):
            raise ValueError('first-order and beta-squared operator dimensions are incompatible')
        e_from_h = first_order_operator[:transverse_size, transverse_size:]
        h_from_e = first_order_operator[transverse_size:, :transverse_size]
        right_h = (h_from_e @ right_e) / wavenumbers[None, :]
        left_h = (e_from_h.T @ left_e) / wavenumbers[None, :]
        return TargetedEigensolution(
            eigenvalues = numpy.asarray(wavenumbers, dtype=complex),
            right_eigenvectors = numpy.vstack((right_e, right_h)),
            left_eigenvectors = numpy.vstack((left_e, left_h)),
            )

    return solve_modes_near_beta(
        omega,
        solve_dxes,
        epsilon,
        target_wavenumbers = target_wavenumbers,
        modes_per_target = modes_per_target,
        mu = mu,
        overlap_dxes = overlap_dxes,
        physical_mask = physical_mask,
        candidate_margin = candidate_margin,
        eigensolver = lift_squared_solution,
        eigensolver_opts = eigensolver_opts,
        residual_tol = residual_tol,
        duplicate_eigenvalue_rtol = duplicate_eigenvalue_rtol,
        duplicate_subspace_tol = duplicate_subspace_tol,
        )


def _block_relative_residuals(
        operator: sparse.sparray,
        matrix: NDArray[numpy.complex128],
        vectors: NDArray[numpy.complex128],
        ) -> NDArray[numpy.float64]:
    """Compute column residuals for an invariant subspace representation."""
    image = operator @ vectors
    expected = vectors @ matrix
    numerator = numpy.linalg.norm(image - expected, axis=0)
    denominator = numpy.maximum(numpy.linalg.norm(image, axis=0), numpy.linalg.norm(expected, axis=0))
    return numpy.asarray(numerator / numpy.maximum(denominator, numpy.finfo(float).tiny), dtype=float)


def _compute_mode_diagnostics(
        basis: ModeBasis,
        operator: sparse.sparray,
        targets: NDArray[numpy.complex128],
        overlap_dxes: dx_lists2_t,
        physical_mask: NDArray[numpy.bool_],
        ) -> ModeDiagnostics:
    """Compute diagnostics from an already-built operator and validated context."""
    n_points = basis.electric_fields.shape[2] // 3
    primal_electric = basis.electric_fields[:, 0]
    dual_electric = basis.electric_fields[:, 1]
    primal_magnetic = basis.magnetic_fields[:, 0]
    dual_magnetic = basis.magnetic_fields[:, 1]
    right = numpy.vstack((primal_electric[:, :2 * n_points].T, primal_magnetic[:, :2 * n_points].T))
    dual = numpy.vstack((dual_electric[:, :2 * n_points].T, dual_magnetic[:, :2 * n_points].T))
    metric = lorentz_metric(overlap_dxes)
    left = metric.T @ dual
    gram = dual.T @ metric @ right
    normalized_right, right_norms = _normalize_columns(right)
    normalized_left, left_norms = _normalize_columns(left)
    if numpy.any(right_norms == 0) or numpy.any(left_norms == 0):
        biorthogonality_condition = numpy.inf
    else:
        normalized_gram = normalized_left.T @ normalized_right
        biorthogonality_condition = float(numpy.linalg.cond(normalized_gram))
    powers = numpy.asarray([
        inner_product(e_field, h_field, overlap_dxes, conj_h=True)
        for e_field, h_field in zip(primal_electric, primal_magnetic, strict=True)
        ])
    pml_participation = _pml_participation(right, overlap_dxes, physical_mask)
    target_indices = numpy.argmin(numpy.abs(basis.wavenumbers[:, None] - targets[None, :]), axis=1)
    target_distances = numpy.abs(basis.wavenumbers - targets[target_indices])
    return ModeDiagnostics(
        right_residuals = _block_relative_residuals(operator, basis.wavenumber_matrix, right),
        left_residuals = _block_relative_residuals(operator.T, basis.wavenumber_matrix.T, left),
        powers = powers,
        pml_participation = numpy.asarray(pml_participation, dtype=float),
        target_indices = numpy.asarray(target_indices, dtype=int),
        target_distances = numpy.asarray(target_distances, dtype=float),
        biorthogonality_error = float(numpy.max(numpy.abs(gram - numpy.eye(len(basis))))),
        biorthogonality_condition = biorthogonality_condition,
        )


def compute_mode_diagnostics(
        basis: ModeBasis,
        omega: complex,
        solve_dxes: dx_lists2_t,
        epsilon: vfdslice,
        *,
        target_wavenumbers: Sequence[complex],
        mu: vfdslice | None = None,
        overlap_dxes: dx_lists2_t | None = None,
        physical_mask: NDArray[numpy.bool_] | None = None,
        ) -> ModeDiagnostics:
    """Compute numerical and physical diagnostics for a modal basis.

    Args:
        basis: Modal basis to diagnose.
        omega: Nonzero angular frequency used by the mode solve.
        solve_dxes: Yee-grid spacings used to build the eigenproblem.
        epsilon: Vectorized three-component relative permittivity.
        target_wavenumbers: Signed-beta targets used to label the spectrum.
        mu: Optional vectorized three-component relative permeability.
        overlap_dxes: Grid used for physical overlaps and normalization. The
            real parts of ``solve_dxes`` are used by default.
        physical_mask: Optional transverse non-PML mask. It is inferred from
            complex stretching in ``solve_dxes`` by default.

    Returns:
        Per-mode residual, power, PML, target, and biorthogonality diagnostics.

    Raises:
        ValueError: If the basis or diagnostic context has invalid dimensions.
    """
    targets = numpy.asarray(target_wavenumbers, dtype=complex)
    if targets.ndim != 1 or targets.size == 0 or not numpy.isfinite(targets).all():
        raise ValueError('target_wavenumbers must be a finite nonempty one-dimensional sequence')
    if overlap_dxes is None:
        overlap_dxes = [[numpy.real(values) for values in family] for family in solve_dxes]
    grid_shape = (solve_dxes[0][0].size, solve_dxes[0][1].size)
    if physical_mask is None:
        physical_mask = infer_physical_mask(solve_dxes)
    physical_mask = numpy.asarray(physical_mask, dtype=bool)
    if physical_mask.shape != grid_shape:
        raise ValueError('physical_mask must match the transverse grid shape')
    field_shape = (len(basis), 2, 3 * grid_shape[0] * grid_shape[1])
    if basis.electric_fields.shape != field_shape or basis.magnetic_fields.shape != field_shape:
        raise ValueError('basis field arrays must have shape (mode, primal_or_dual, field)')
    if basis.wavenumber_matrix.shape != (len(basis), len(basis)):
        raise ValueError('basis wavenumber_matrix must match the mode count')
    operator = operator_beta(omega, solve_dxes, epsilon, mu)
    return _compute_mode_diagnostics(basis, operator, targets, overlap_dxes, physical_mask)


def transform_mode_basis(
        basis: ModeBasis,
        transform: NDArray[numpy.complex128],
        *,
        spectrum_order: NDArray[numpy.int_] | None = None,
        ) -> ModeBasis:
    """Apply a modal basis transform and its contragredient dual transform.

    Args:
        basis: Modal basis to transform.
        transform: Finite invertible square matrix acting on primal basis
            columns. Dual fields receive the contragredient transformation.
        spectrum_order: Optional permutation assigning spectral labels to the
            transformed basis order.

    Returns:
        A transformed modal basis.

    Raises:
        ValueError: If ``transform`` has an invalid shape or ``spectrum_order``
            is not a complete integer permutation.
        numpy.linalg.LinAlgError: If ``transform`` is singular.
    """
    transform = numpy.asarray(transform, dtype=complex)
    if transform.shape != (len(basis), len(basis)) or not numpy.isfinite(transform).all():
        raise ValueError('basis transform must be a finite square mode matrix')
    inverse = numpy.linalg.inv(transform)
    wavenumbers = basis.wavenumbers
    if spectrum_order is not None:
        raw_order = numpy.asarray(spectrum_order)
        if (
                raw_order.shape != (len(basis),)
                or not numpy.issubdtype(raw_order.dtype, numpy.integer)
                or not numpy.array_equal(numpy.sort(raw_order), numpy.arange(len(basis)))
                ):
            raise ValueError('spectrum_order must be a complete integer permutation')
        spectrum_order = numpy.asarray(raw_order, dtype=int)
        wavenumbers = wavenumbers[spectrum_order]
    primal_electric = transform.T @ basis.electric_fields[:, 0]
    dual_electric = inverse @ basis.electric_fields[:, 1]
    primal_magnetic = transform.T @ basis.magnetic_fields[:, 0]
    dual_magnetic = inverse @ basis.magnetic_fields[:, 1]
    return ModeBasis(
        wavenumbers = wavenumbers,
        wavenumber_matrix = inverse @ basis.wavenumber_matrix @ transform,
        electric_fields = numpy.stack((primal_electric, dual_electric), axis=1),
        magnetic_fields = numpy.stack((primal_magnetic, dual_magnetic), axis=1),
        )


def power_normalize_modes(
        basis: ModeBasis,
        indices: Sequence[int],
        *,
        overlap_dxes: dx_lists2_t,
        power_atol: float = 1e-12,
        power_rtol: float = 1e-6,
        ) -> ModeBasis:
    """Normalize selected positive-forward-power modes to unit power.

    Primal fields are scaled by the inverse square root of their real power and
    Lorentz duals are scaled reciprocally, preserving biorthogonality and the
    represented beta operator.

    Args:
        basis: Modal basis to copy and rescale.
        indices: Basis-vector indices to normalize.
        overlap_dxes: Grid used to evaluate modal power.
        power_atol: Absolute threshold for positive power and negligible
            imaginary power.
        power_rtol: Relative threshold for negligible imaginary power.

    Returns:
        A power-normalized modal basis.

    Raises:
        ValueError: If tolerances are invalid or a selected mode does not carry
            positive, nearly real forward power.
    """
    if not numpy.isfinite(power_atol) or power_atol <= 0:
        raise ValueError('power_atol must be finite and positive')
    if not numpy.isfinite(power_rtol) or power_rtol <= 0:
        raise ValueError('power_rtol must be finite and positive')
    powers = numpy.asarray([
        inner_product(e_field, h_field, overlap_dxes, conj_h=True)
        for e_field, h_field in zip(basis.electric_fields[:, 0], basis.magnetic_fields[:, 0], strict=True)
        ])
    scales = numpy.ones(len(basis), dtype=complex)
    for index in indices:
        power = powers[index]
        if power.real <= power_atol or abs(power.imag) > max(power_atol, power_rtol * abs(power.real)):
            raise ValueError(f'mode {index} does not have positive nearly-real forward power')
        scales[index] = 1 / numpy.sqrt(power.real)
    return transform_mode_basis(basis, numpy.diag(scales))


def lorentz_overlap(
        dual_mode: tuple[vcfdfield2, vcfdfield2],
        mode: tuple[vcfdfield2, vcfdfield2],
        dxes: dx_lists2_t,
        *,
        measure: Literal['physical', 'stretched'] = 'stretched',
        ) -> complex:
    """Evaluate a symmetric unconjugated reaction using the selected measure.

    Use the same grid and measure that defined the supplied Lorentz duals.
    Complex coordinate widths are retained by default.
    """
    return 0.5 * (
        inner_product(dual_mode[0], mode[1], dxes, conj_h=False, measure=measure)
        + inner_product(mode[0], dual_mode[1], dxes, conj_h=False, measure=measure)
        )


def _degenerate_clusters(
        previous: NDArray[numpy.complex128],
        current: NDArray[numpy.complex128],
        rtol: float,
        ) -> list[list[int]]:
    """Find connected mode groups degenerate in two adjacent spectra."""
    count = previous.size
    remaining = set(range(count))
    clusters: list[list[int]] = []
    while remaining:
        seed = remaining.pop()
        cluster = {seed}
        changed = True
        while changed:
            changed = False
            for candidate in tuple(remaining):
                if any(
                        abs(previous[candidate] - previous[item]) <= rtol * max(1.0, abs(previous[candidate]), abs(previous[item]))
                        and abs(current[candidate] - current[item]) <= rtol * max(1.0, abs(current[candidate]), abs(current[item]))
                        for item in cluster
                        ):
                    remaining.remove(candidate)
                    cluster.add(candidate)
                    changed = True
        clusters.append(sorted(cluster))
    return clusters


def track_mode_path(
        bases: Sequence[ModeBasis],
        *,
        overlap_dxes: dx_lists2_t,
        degeneracy_rtol: float = 1e-6,
        minimum_overlap: float = 1e-3,
        ) -> tuple[ModeBasis, ...]:
    """Assign, align, and parallel-transport modal bases along a path.

    Adjacent modes are assigned by Lorentz overlap, with beta distance used as
    a deterministic tie-breaker. Nearly degenerate groups are aligned by a
    unitary Procrustes rotation, so their propagation operator may become
    non-diagonal in the transported basis.

    Args:
        bases: Equal-sized modal bases sampled in path order.
        overlap_dxes: Common grid used for adjacent Lorentz overlaps.
        degeneracy_rtol: Relative spectral tolerance used to group modes for
            subspace rotation.
        minimum_overlap: Minimum accepted adjacent diagonal overlap after
            assignment and subspace alignment.

    Returns:
        Basis sequence in tracked order with aligned phases, contragredient
        duals, and transformed beta matrices.

    Raises:
        ValueError: If the sequence, tolerances, or field shapes are incompatible.
        RuntimeError: If an assigned mode has insufficient adjacent overlap.
    """
    if not bases:
        raise ValueError('bases must not be empty')
    if not numpy.isfinite(degeneracy_rtol) or degeneracy_rtol <= 0:
        raise ValueError('degeneracy_rtol must be finite and positive')
    if not numpy.isfinite(minimum_overlap) or minimum_overlap <= 0:
        raise ValueError('minimum_overlap must be finite and positive')
    count = len(bases[0])
    if count == 0 or any(len(basis) != count for basis in bases):
        raise ValueError('all bases must contain the same nonzero mode count')
    reference = bases[0]
    for index, basis in enumerate(bases[1:], start=1):
        if basis.electric_fields.shape != reference.electric_fields.shape or basis.magnetic_fields.shape != reference.magnetic_fields.shape:
            raise ValueError(f'bases[{index}] modal field shapes do not match bases[0]')

    tracked = [bases[0]]
    for current_unordered in bases[1:]:
        previous = tracked[-1]
        previous_dual_e = previous.electric_fields[:, 1]
        previous_dual_h = previous.magnetic_fields[:, 1]
        current_primal_e = current_unordered.electric_fields[:, 0]
        current_primal_h = current_unordered.magnetic_fields[:, 0]
        overlap = numpy.asarray([
            [
                lorentz_overlap((dual_e, dual_h), (mode_e, mode_h), overlap_dxes)
                for mode_e, mode_h in zip(current_primal_e, current_primal_h, strict=True)
                ]
            for dual_e, dual_h in zip(previous_dual_e, previous_dual_h, strict=True)
            ])
        beta_scale = numpy.maximum(1.0, numpy.maximum(
            numpy.abs(previous.wavenumbers[:, None]),
            numpy.abs(current_unordered.wavenumbers[None, :]),
            ))
        beta_distance = numpy.abs(previous.wavenumbers[:, None] - current_unordered.wavenumbers[None, :]) / beta_scale
        cost = -numpy.abs(overlap) + 1e-9 * beta_distance
        rows, columns = linear_sum_assignment(cost)
        order = columns[numpy.argsort(rows)]
        permutation = numpy.eye(count, dtype=complex)[:, order]
        reordered_e = current_primal_e[order]
        reordered_h = current_primal_h[order]
        reordered_wavenumbers = current_unordered.wavenumbers[order]

        cluster_transform = numpy.eye(count, dtype=complex)
        for cluster in _degenerate_clusters(previous.wavenumbers, reordered_wavenumbers, degeneracy_rtol):
            if len(cluster) == 1:
                continue
            indices = numpy.asarray(cluster)
            cross = numpy.asarray([
                [
                    lorentz_overlap(
                        (previous_dual_e[ii], previous_dual_h[ii]),
                        (reordered_e[jj], reordered_h[jj]),
                        overlap_dxes,
                        )
                    for jj in indices
                    ]
                for ii in indices
                ])
            left_u, _singular_values, right_vh = numpy.linalg.svd(cross)
            transform = right_vh.conj().T @ left_u.conj().T
            cluster_transform[numpy.ix_(indices, indices)] = transform

        aligned_e = cluster_transform.T @ reordered_e
        aligned_h = cluster_transform.T @ reordered_h
        diagonal_overlap = numpy.asarray([
            lorentz_overlap(
                (previous_dual_e[index], previous_dual_h[index]),
                (aligned_e[index], aligned_h[index]),
                overlap_dxes,
                )
            for index in range(count)
            ])
        if numpy.any(numpy.abs(diagonal_overlap) < minimum_overlap):
            bad = numpy.flatnonzero(numpy.abs(diagonal_overlap) < minimum_overlap)
            raise RuntimeError(f'could not track modes {bad.tolist()}: adjacent overlap is too small')
        phase_scales = numpy.exp(-1j * numpy.angle(diagonal_overlap))
        transform = permutation @ cluster_transform @ numpy.diag(phase_scales)
        tracked.append(transform_mode_basis(current_unordered, transform, spectrum_order=order))

    return tuple(tracked)