diff --git a/examples/test.py b/examples/test.py index 09f425a..b71e859 100644 --- a/examples/test.py +++ b/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() + diff --git a/fdfd_tools/__init__.py b/fdfd_tools/__init__.py index 4d75b2c..d19a4b9 100644 --- a/fdfd_tools/__init__.py +++ b/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' \ No newline at end of file +__author__ = 'Jan Petykiewicz' diff --git a/fdfd_tools/solvers.py b/fdfd_tools/solvers.py new file mode 100644 index 0000000..3ee4080 --- /dev/null +++ b/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