meanas/meanas/fdfd/solvers.py

"""
Solvers and solver interface for FDFD problems.
"""

from typing import Any
from collections.abc import Callable
import logging

import numpy
from numpy.typing import ArrayLike, NDArray
from numpy.linalg import norm
import scipy.sparse.linalg

from ..fdmath import dx_lists_t, vfdfield, vcfdfield, vcfdfield_t
from . import operators


logger = logging.getLogger(__name__)


def _scipy_qmr(
        A: scipy.sparse.csr_array,
        b: ArrayLike,
        **kwargs: Any,
        ) -> NDArray[numpy.float64]:
    """
    Wrapper for scipy.sparse.linalg.qmr

    Args:
        A: Sparse matrix
        b: Right-hand-side vector
        kwargs: Passed as **kwargs to the wrapped function

    Returns:
        Converged solution. Raises RuntimeError on breakdown or nonconvergence.
    """

    #
    #Report on our progress
    #
    ii = 0

    def log_residual(xk: ArrayLike) -> None:
        nonlocal ii
        ii += 1
        if ii % 100 == 0:
            cur_norm = norm(A @ xk - b) / norm(b)
            logger.info(f'Solver residual at iteration {ii} : {cur_norm}')

    if 'callback' in kwargs:
        callback = kwargs['callback']

        def augmented_callback(xk: ArrayLike) -> None:
            log_residual(xk)
            callback(xk)

        kwargs['callback'] = augmented_callback
    else:
        kwargs['callback'] = log_residual

    #
    # Run the actual solve
    #
    x, info = scipy.sparse.linalg.qmr(A, b, **kwargs)
    if info != 0:
        residual = norm(A @ x - b)
        raise RuntimeError(f'QMR failed to converge ({info=}, residual={residual:.6g})')
    return x


def generic(
        omega: complex,
        dxes: dx_lists_t,
        J: vcfdfield,
        epsilon: vfdfield,
        mu: vfdfield | None = None,
        *,
        pec: vfdfield | None = None,
        pmc: vfdfield | None = None,
        adjoint: bool = False,
        matrix_solver: Callable[..., ArrayLike] = _scipy_qmr,
        matrix_solver_opts: dict[str, Any] | None = None,
        E_guess: vcfdfield | None = None,
        ) -> vcfdfield_t:
    """
    Preconditioned FDFD solver using CSR sparse matrices.

    All ndarray arguments should be 1D arrays, as returned by `meanas.fdmath.vectorization.vec()`.

    Args:
        omega: Complex frequency to solve at.
        dxes: `[[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]]` (complex cell sizes) as
            discussed in `meanas.fdmath.types`
        J: Electric current distribution (at E-field locations)
        epsilon: Dielectric constant distribution (at E-field locations)
        mu: Magnetic permeability distribution (at H-field locations)
        pec: Perfect electric conductor distribution
             (at E-field locations; non-zero value indicates PEC is present).
             Sources and initial guesses are zeroed at PEC locations.
        pmc: Perfect magnetic conductor distribution
             (at H-field locations; non-zero value indicates PMC is present)
        adjoint: If true, solves the adjoint problem.
        matrix_solver: Called as `matrix_solver(A, b, **matrix_solver_opts) -> x`,
                where `A`: `scipy.sparse.csr_array`;
                      `b`: `ArrayLike`;
                      `x`: `ArrayLike`;
                Default is a wrapped version of `scipy.sparse.linalg.qmr()`
                 which raises on nonconvergence and logs the residual every
                 100 iterations.
        matrix_solver_opts: Passed as kwargs to `matrix_solver(...)`
        E_guess: Guess at the solution E-field. `matrix_solver` must accept an
                `x0` argument with the same purpose.

    Returns:
        E-field which solves the system.
    """

    matrix_solver_opts = {} if matrix_solver_opts is None else dict(matrix_solver_opts)

    b0 = -1j * omega * J
    if pec is not None:
        b0 = numpy.where(pec, 0, b0)
    A0 = operators.e_full(omega, dxes, epsilon=epsilon, mu=mu, pec=pec, pmc=pmc)

    Pl, Pr = operators.e_full_preconditioners(dxes)

    if adjoint:
        A = (Pl @ A0 @ Pr).T.conjugate()
        b = Pr.T.conjugate() @ b0
    else:
        A = Pl @ A0 @ Pr
        b = Pl @ b0

    if E_guess is not None:
        if pec is not None:
            E_guess = numpy.where(pec, 0, E_guess)
        if adjoint:
            x0 = Pr.T.conjugate() @ E_guess
        else:
            x0 = Pl @ E_guess
        matrix_solver_opts['x0'] = x0

    x = numpy.asarray(matrix_solver(A.tocsr(), b, **matrix_solver_opts))
    if x.shape != b.shape or not numpy.isfinite(x).all():
        raise RuntimeError('matrix_solver must return a finite solution vector matching the source shape')

    if adjoint:
        x0 = Pl.T.conjugate() @ x
    else:
        x0 = Pr @ x

    if pec is not None:
        x0 = numpy.where(pec, 0, x0)
    return vcfdfield_t(x0)