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.
127 lines
3.6 KiB
Python
127 lines
3.6 KiB
Python
"""
|
|
Solvers and solver interface for FDFD problems.
|
|
"""
|
|
|
|
from typing import Callable, Dict, Any
|
|
import logging
|
|
|
|
import numpy # type: ignore
|
|
from numpy.linalg import norm # type: ignore
|
|
import scipy.sparse.linalg # type: ignore
|
|
|
|
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
|
|
'''
|
|
ii = 0
|
|
|
|
def log_residual(xk):
|
|
nonlocal ii
|
|
ii += 1
|
|
if ii % 100 == 0:
|
|
logger.info('Solver residual at iteration {} : {}'.format(ii, 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
|