You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
meanas/meanas/fdfd/solvers.py

127 lines
3.5 KiB
Python

"""
Solvers and solver interface for FDFD problems.
"""
from typing import List, Callable, Dict, Any
import logging
import numpy
from numpy.linalg import norm
import scipy.sparse.linalg
from ..fdmath import dx_lists_t, vfdfield_t
from . import operators
logger = logging.getLogger(__name__)
def _scipy_qmr(A: scipy.sparse.csr_matrix,
b: numpy.ndarray,
**kwargs
) -> numpy.ndarray:
"""
Wrapper for scipy.sparse.linalg.qmr
Args:
A: Sparse matrix
b: Right-hand-side vector
kwargs: Passed as **kwargs to the wrapped function
Returns:
Guess for solution (returned even if didn't converge)
"""
'''
Report on our progress
'''
iter = 0
def log_residual(xk):
nonlocal iter
iter += 1
if iter % 100 == 0:
logger.info('Solver residual at iteration {} : {}'.format(iter, norm(A @ xk - b)))
if 'callback' in kwargs:
def augmented_callback(xk):
log_residual(xk)
kwargs['callback'](xk)
kwargs['callback'] = augmented_callback
else:
kwargs['callback'] = log_residual
'''
Run the actual solve
'''
x, _ = scipy.sparse.linalg.qmr(A, b, **kwargs)
return x
def generic(omega: complex,
dxes: dx_lists_t,
J: vfdfield_t,
epsilon: vfdfield_t,
mu: vfdfield_t = None,
pec: vfdfield_t = None,
pmc: vfdfield_t = None,
adjoint: bool = False,
matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
matrix_solver_opts: Dict[str, Any] = None,
) -> vfdfield_t:
"""
Conjugate gradient FDFD solver using CSR sparse matrices.
All ndarray arguments should be 1D arrays, as returned by `meanas.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)
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_matrix`;
`b`: `numpy.ndarray`;
`x`: `numpy.ndarray`;
Default is a wrapped version of `scipy.sparse.linalg.qmr()`
which doesn't return convergence info and logs the residual
every 100 iterations.
matrix_solver_opts: Passed as kwargs to `matrix_solver(...)`
Returns:
E-field which solves the system.
"""
if matrix_solver_opts is None:
matrix_solver_opts = dict()
b0 = -1j * omega * J
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).H
b = Pr.H @ b0
else:
A = Pl @ A0 @ Pr
b = Pl @ b0
x = matrix_solver(A.tocsr(), b, **matrix_solver_opts)
if adjoint:
x0 = Pl.H @ x
else:
x0 = Pr @ x
return x0