Use fdfd_tools.solvers.generic for csr solve

This commit is contained in:
jan 2016-08-04 22:28:19 -07:00
parent 80592dff79
commit 848f86f6ee

View File

@ -14,7 +14,7 @@ satisfy the constraints for the 'conjugate gradient' algorithm
(positive definite, symmetric) and some that don't. (positive definite, symmetric) and some that don't.
""" """
from typing import List, Dict, Any from typing import Dict, Any
import time import time
import numpy import numpy
@ -22,7 +22,7 @@ from numpy.linalg import norm
import pyopencl import pyopencl
import pyopencl.array import pyopencl.array
import fdfd_tools.operators import fdfd_tools.solvers
from . import ops from . import ops
@ -43,7 +43,7 @@ class CSRMatrix(object):
self.data = pyopencl.array.to_device(queue, m.data.astype(numpy.complex128)) self.data = pyopencl.array.to_device(queue, m.data.astype(numpy.complex128))
def cg(a: 'scipy.sparse.csr_matrix', def cg(A: 'scipy.sparse.csr_matrix',
b: numpy.ndarray, b: numpy.ndarray,
max_iters: int = 10000, max_iters: int = 10000,
err_threshold: float = 1e-6, err_threshold: float = 1e-6,
@ -54,7 +54,7 @@ def cg(a: 'scipy.sparse.csr_matrix',
""" """
General conjugate-gradient solver for sparse matrices, where A @ x = b. General conjugate-gradient solver for sparse matrices, where A @ x = b.
:param a: Matrix to solve (CSR format) :param A: Matrix to solve (CSR format)
:param b: Right-hand side vector (dense ndarray) :param b: Right-hand side vector (dense ndarray)
:param max_iters: Maximum number of iterations :param max_iters: Maximum number of iterations
:param err_threshold: Error threshold for successful solve, relative to norm(b) :param err_threshold: Error threshold for successful solve, relative to norm(b)
@ -84,7 +84,7 @@ def cg(a: 'scipy.sparse.csr_matrix',
rho = 1.0 + 0j rho = 1.0 + 0j
errs = [] errs = []
m = CSRMatrix(queue, a) m = CSRMatrix(queue, A)
''' '''
Generate OpenCL kernels Generate OpenCL kernels
@ -102,7 +102,8 @@ def cg(a: 'scipy.sparse.csr_matrix',
_, err2 = rhoerr_step(r, []) _, err2 = rhoerr_step(r, [])
b_norm = numpy.sqrt(err2) b_norm = numpy.sqrt(err2)
print('b_norm check: ', b_norm) if verbose:
print('b_norm check: ', b_norm)
success = False success = False
for k in range(max_iters): for k in range(max_iters):
@ -126,7 +127,7 @@ def cg(a: 'scipy.sparse.csr_matrix',
e = a_step(v, m, p, e) e = a_step(v, m, p, e)
alpha = rho / dot(p, v, e) alpha = rho / dot(p, v, e)
if k % 1000 == 0: if verbose and k % 1000 == 0:
print(k) print(k)
''' '''
@ -136,71 +137,43 @@ def cg(a: 'scipy.sparse.csr_matrix',
x = x.get() x = x.get()
if success: if verbose:
print('Success', end='') if success:
else: print('Success', end='')
print('Failure', end=', ') else:
print(', {} iterations in {} sec: {} iterations/sec \ print('Failure', end=', ')
'.format(k, time_elapsed, k / time_elapsed)) print(', {} iterations in {} sec: {} iterations/sec \
print('final error', errs[-1]) '.format(k, time_elapsed, k / time_elapsed))
print('overhead {} sec'.format(start_time2 - start_time)) print('final error', errs[-1])
print('overhead {} sec'.format(start_time2 - start_time))
print('Final residual:', norm(a @ x - b) / norm(b)) print('Final residual:', norm(A @ x - b) / norm(b))
return x return x
def fdfd_cg_solver(omega: complex, def fdfd_cg_solver(solver_opts: Dict[str, Any] = None,
dxes: List[List[numpy.ndarray]], **fdfd_args
J: numpy.ndarray,
epsilon: numpy.ndarray,
mu: numpy.ndarray = None,
pec: numpy.ndarray = None,
pmc: numpy.ndarray = None,
adjoint: bool = False,
solver_opts: Dict[str, Any] = None,
) -> numpy.ndarray: ) -> numpy.ndarray:
""" """
Conjugate gradient FDFD solver using CSR sparse matrices, mainly for Conjugate gradient FDFD solver using CSR sparse matrices, mainly for
testing and development since it's much slower than the solver in main.py. testing and development since it's much slower than the solver in main.py.
All ndarray arguments should be 1D arrays. To linearize a list of 3 3D ndarrays, Calls fdfd_tools.solvers.generic(**fdfd_args,
either use fdfd_tools.vec() or numpy: matrix_solver=opencl_fdfd.csr.cg,
f_1D = numpy.hstack(tuple((fi.flatten(order='F') for fi in [f_x, f_y, f_z]))) matrix_solver_opts=solver_opts)
:param omega: Complex frequency to solve at. :param solver_opts: Passed as matrix_solver_opts to fdfd_tools.solver.generic(...).
:param dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes) Default {}.
:param J: Electric current distribution (at E-field locations) :param fdfd_args: Passed as **fdfd_args to fdfd_tools.solver.generic(...).
:param epsilon: Dielectric constant distribution (at E-field locations) Should include all of the arguments **except** matrix_solver and matrix_solver_opts
: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 solver_opts: Passed as kwargs to opencl_fdfd.csr.cg(**solver_opts)
:return: E-field which solves the system. :return: E-field which solves the system.
""" """
if solver_opts is None: if solver_opts is None:
solver_opts = dict() solver_opts = dict()
b0 = -1j * omega * J x = fdfd_tools.solvers.generic(matrix_solver=cg,
A0 = fdfd_tools.operators.e_full(omega, dxes, epsilon=epsilon, mu=mu, pec=pec, pmc=pmc) matrix_solver_opts=solver_opts,
**fdfd_args)
Pl, Pr = fdfd_tools.operators.e_full_preconditioners(dxes) return x
if adjoint:
A = (Pl @ A0 @ Pr).H
b = Pr.H @ b0
else:
A = Pl @ A0 @ Pr
b = Pl @ b0
x = cg(A.tocsr(), b, **solver_opts)
if adjoint:
x0 = Pl.H @ x
else:
x0 = Pr @ x
return x0