154 lines
4.7 KiB
Python
154 lines
4.7 KiB
Python
"""
|
|
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)
|