Use fdfd_tools.solvers.generic for csr solve
This commit is contained in:
parent
80592dff79
commit
848f86f6ee
@ -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,6 +102,7 @@ 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)
|
||||||
|
if verbose:
|
||||||
print('b_norm check: ', b_norm)
|
print('b_norm check: ', b_norm)
|
||||||
|
|
||||||
success = False
|
success = False
|
||||||
@ -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,6 +137,7 @@ def cg(a: 'scipy.sparse.csr_matrix',
|
|||||||
|
|
||||||
x = x.get()
|
x = x.get()
|
||||||
|
|
||||||
|
if verbose:
|
||||||
if success:
|
if success:
|
||||||
print('Success', end='')
|
print('Success', end='')
|
||||||
else:
|
else:
|
||||||
@ -145,62 +147,33 @@ def cg(a: 'scipy.sparse.csr_matrix',
|
|||||||
print('final error', errs[-1])
|
print('final error', errs[-1])
|
||||||
print('overhead {} sec'.format(start_time2 - start_time))
|
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
|
|
||||||
|
Loading…
Reference in New Issue
Block a user