Add solvers submodule and clean up examples.
Solvers submodule includes a generic solver in case you already have a sparse matrix solver, or in case you have no solver at all. Example file now uses alternate solvers if available, and has a nicer way of picking which solver gets used.
This commit is contained in:
parent
85880c859e
commit
ec674fe3f4
@ -1,18 +1,22 @@
|
||||
import importlib
|
||||
import numpy
|
||||
from numpy.linalg import norm
|
||||
|
||||
from fdfd_tools import vec, unvec, waveguide_mode
|
||||
import fdfd_tools, fdfd_tools.functional, fdfd_tools.grid
|
||||
import fdfd_tools
|
||||
import fdfd_tools.functional
|
||||
import fdfd_tools.grid
|
||||
from fdfd_tools.solvers import generic as generic_solver
|
||||
|
||||
import gridlock
|
||||
|
||||
from matplotlib import pyplot
|
||||
|
||||
#import magma_fdfd
|
||||
from opencl_fdfd import cg_solver, csr
|
||||
|
||||
__author__ = 'Jan Petykiewicz'
|
||||
|
||||
|
||||
def test0():
|
||||
def test0(solver=generic_solver):
|
||||
dx = 50 # discretization (nm/cell)
|
||||
pml_thickness = 10 # (number of cells)
|
||||
|
||||
@ -59,21 +63,27 @@ def test0():
|
||||
J = [numpy.zeros_like(grid.grids[0], dtype=complex) for _ in range(3)]
|
||||
J[1][15, grid.shape[1]//2, grid.shape[2]//2] = 1e5
|
||||
|
||||
'''
|
||||
Solve!
|
||||
'''
|
||||
x = solver(J=vec(J), **sim_args)
|
||||
|
||||
A = fdfd_tools.functional.e_full(omega, dxes, vec(grid.grids)).tocsr()
|
||||
b = -1j * omega * vec(J)
|
||||
print('Norm of the residual is ', norm(A @ x - b))
|
||||
|
||||
x = solve_A(A, b)
|
||||
E = unvec(x, grid.shape)
|
||||
|
||||
print('Norm of the residual is {}'.format(numpy.linalg.norm(A.dot(x) - b)/numpy.linalg.norm(b)))
|
||||
|
||||
'''
|
||||
Plot results
|
||||
'''
|
||||
pyplot.figure()
|
||||
pyplot.pcolor(numpy.real(E[1][:, :, grid.shape[2]//2]), cmap='seismic')
|
||||
pyplot.axis('equal')
|
||||
pyplot.show()
|
||||
|
||||
|
||||
def test1():
|
||||
def test1(solver=generic_solver):
|
||||
dx = 40 # discretization (nm/cell)
|
||||
pml_thickness = 10 # (number of cells)
|
||||
|
||||
@ -142,17 +152,14 @@ def test1():
|
||||
'pmc': vec(pmcg.grids),
|
||||
}
|
||||
|
||||
x = solver(J=vec(J), **sim_args)
|
||||
|
||||
b = -1j * omega * vec(J)
|
||||
A = fdfd_tools.operators.e_full(**sim_args).tocsr()
|
||||
# x = magma_fdfd.solve_A(A, b)
|
||||
|
||||
# x = csr.cg_solver(J=vec(J), **sim_args)
|
||||
x = cg_solver(J=vec(J), **sim_args)
|
||||
print('Norm of the residual is ', norm(A @ x - b))
|
||||
|
||||
E = unvec(x, grid.shape)
|
||||
|
||||
print('Norm of the residual is ', numpy.linalg.norm(A @ x - b))
|
||||
|
||||
'''
|
||||
Plot results
|
||||
'''
|
||||
@ -197,6 +204,22 @@ def test1():
|
||||
pyplot.show()
|
||||
print('Average overlap with mode:', sum(q)/len(q))
|
||||
|
||||
|
||||
def module_available(name):
|
||||
return importlib.util.find_spec(name) is not None
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
# test0()
|
||||
test1()
|
||||
|
||||
if module_available('opencl_fdfd'):
|
||||
from opencl_fdfd import cg_solver as opencl_solver
|
||||
test1(opencl_solver)
|
||||
# from opencl_fdfd.csr import fdfd_cg_solver as opencl_csr_solver
|
||||
# test1(opencl_csr_solver)
|
||||
# elif module_available('magma_fdfd'):
|
||||
# from magma_fdfd import solver as magma_solver
|
||||
# test1(magma_solver)
|
||||
else:
|
||||
test1()
|
||||
|
||||
|
@ -22,4 +22,4 @@ Dependencies:
|
||||
from .vectorization import vec, unvec, field_t, vfield_t
|
||||
from .grid import dx_lists_t
|
||||
|
||||
__author__ = 'Jan Petykiewicz'
|
||||
__author__ = 'Jan Petykiewicz'
|
||||
|
114
fdfd_tools/solvers.py
Normal file
114
fdfd_tools/solvers.py
Normal file
@ -0,0 +1,114 @@
|
||||
"""
|
||||
Solvers for FDFD problems.
|
||||
"""
|
||||
|
||||
from typing import List, Callable, Dict, Any
|
||||
|
||||
import numpy
|
||||
from numpy.linalg import norm
|
||||
import scipy.sparse.linalg
|
||||
|
||||
from . import operators
|
||||
|
||||
|
||||
def _scipy_qmr(A: scipy.sparse.csr_matrix,
|
||||
b: numpy.ndarray,
|
||||
**kwargs
|
||||
) -> numpy.ndarray:
|
||||
"""
|
||||
Wrapper for scipy.sparse.linalg.qmr
|
||||
|
||||
:param A: Sparse matrix
|
||||
:param b: Right-hand-side vector
|
||||
:param kwargs: Passed as **kwargs to the wrapped function
|
||||
:return: Guess for solution (returned even if didn't converge)
|
||||
"""
|
||||
|
||||
'''
|
||||
Report on our progress
|
||||
'''
|
||||
iter = 0
|
||||
|
||||
def print_residual(xk):
|
||||
nonlocal iter
|
||||
iter += 1
|
||||
if iter % 100 == 0:
|
||||
print('Solver residual at iteration', iter, ':', norm(A @ xk - b))
|
||||
|
||||
if 'callback' in kwargs:
|
||||
def augmented_callback(xk):
|
||||
print_residual(xk)
|
||||
kwargs['callback'](xk)
|
||||
|
||||
kwargs['callback'] = augmented_callback
|
||||
else:
|
||||
kwargs['callback'] = print_residual
|
||||
|
||||
'''
|
||||
Run the actual solve
|
||||
'''
|
||||
|
||||
x, _ = scipy.sparse.linalg.qmr(A, b, **kwargs)
|
||||
return x
|
||||
|
||||
|
||||
def generic(omega: complex,
|
||||
dxes: List[List[numpy.ndarray]],
|
||||
J: numpy.ndarray,
|
||||
epsilon: numpy.ndarray,
|
||||
mu: numpy.ndarray = None,
|
||||
pec: numpy.ndarray = None,
|
||||
pmc: numpy.ndarray = None,
|
||||
adjoint: bool = False,
|
||||
matrix_solver: Callable[..., numpy.ndarray] = _scipy_qmr,
|
||||
matrix_solver_opts: Dict[str, Any] = None,
|
||||
) -> numpy.ndarray:
|
||||
"""
|
||||
Conjugate gradient FDFD solver using CSR sparse matrices.
|
||||
|
||||
All ndarray arguments should be 1D array, as returned by fdfd_tools.vec().
|
||||
|
||||
:param omega: Complex frequency to solve at.
|
||||
:param dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes)
|
||||
:param J: Electric current distribution (at E-field locations)
|
||||
:param epsilon: Dielectric constant distribution (at E-field locations)
|
||||
:param mu: Magnetic permeability distribution (at H-field locations)
|
||||
:param pec: Perfect electric conductor distribution
|
||||
(at E-field locations; non-zero value indicates PEC is present)
|
||||
:param pmc: Perfect magnetic conductor distribution
|
||||
(at H-field locations; non-zero value indicates PMC is present)
|
||||
:param adjoint: If true, solves the adjoint problem.
|
||||
:param 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 prints the residual
|
||||
every 100 iterations.
|
||||
:param matrix_solver_opts: Passed as kwargs to matrix_solver(...)
|
||||
:return: 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
|
Loading…
Reference in New Issue
Block a user