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.

examples/test.py

@@ -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()
+

fdfd_tools/__init__.py

@@ -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'

fdfd_tools/solvers.py

@@ -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