196 lines
5.2 KiB
Python
Raw Permalink Normal View History

2016-08-04 20:14:17 -07:00
"""
Sparse matrix solvers
This file holds the sparse matrix solvers, as well as the
CSRMatrix sparse matrix representation.
The FDFD solver (fdfd_cg_solver()) solves an FDFD problem by
creating a sparse matrix representing the problem (using
2021-07-11 17:07:46 -07:00
meanas) and then passing it to cg(), which performs a
2016-08-04 20:14:17 -07:00
conjugate gradient solve.
cg() is capable of solving arbitrary sparse matrices which
satisfy the constraints for the 'conjugate gradient' algorithm
(positive definite, symmetric) and some that don't.
"""
2024-07-30 22:41:27 -07:00
from typing import Any, TYPE_CHECKING
2016-08-04 17:43:01 -07:00
import time
import logging
2016-08-04 17:43:01 -07:00
2016-07-04 19:30:36 -07:00
import numpy
from numpy.typing import NDArray, ArrayLike
2016-07-04 22:19:21 -07:00
from numpy.linalg import norm
2024-07-30 22:41:27 -07:00
from numpy import complexfloating
2016-07-04 19:30:36 -07:00
import pyopencl
import pyopencl.array
2021-07-11 17:07:46 -07:00
import meanas.fdfd.solvers
2016-07-04 19:30:36 -07:00
from . import ops
2016-07-04 19:30:36 -07:00
2024-07-30 22:41:27 -07:00
if TYPE_CHECKING:
import scipy
2016-07-04 19:30:36 -07:00
logger = logging.getLogger(__name__)
class CSRMatrix:
2016-08-04 17:43:01 -07:00
"""
Matrix stored in Compressed Sparse Row format, in GPU RAM.
"""
row_ptr: pyopencl.array.Array
col_ind: pyopencl.array.Array
data: pyopencl.array.Array
def __init__(
self,
queue: pyopencl.CommandQueue,
m: 'scipy.sparse.csr_matrix',
) -> None:
2016-07-04 19:30:36 -07:00
self.row_ptr = pyopencl.array.to_device(queue, m.indptr)
self.col_ind = pyopencl.array.to_device(queue, m.indices)
self.data = pyopencl.array.to_device(queue, m.data.astype(numpy.complex128))
def cg(
A: 'scipy.sparse.csr_matrix',
b: ArrayLike,
max_iters: int = 10000,
err_threshold: float = 1e-6,
2024-07-30 22:41:27 -07:00
context: pyopencl.Context | None = None,
queue: pyopencl.CommandQueue | None = None,
) -> NDArray[complexfloating]:
2016-08-04 17:43:01 -07:00
"""
General conjugate-gradient solver for sparse matrices, where A @ x = b.
Args:
A: Matrix to solve (CSR format)
b: Right-hand side vector (dense ndarray)
max_iters: Maximum number of iterations
err_threshold: Error threshold for successful solve, relative to norm(b)
context: PyOpenCL context. Will be created if not given.
queue: PyOpenCL command queue. Will be created if not given.
Returns:
Solution vector x; returned even if solve doesn't converge.
2016-08-04 17:43:01 -07:00
"""
2016-07-04 19:30:36 -07:00
start_time = time.perf_counter()
if context is None:
context = pyopencl.create_some_context(False)
if queue is None:
queue = pyopencl.CommandQueue(context)
2024-07-30 22:41:27 -07:00
def load_field(v: NDArray[numpy.complexfloating], dtype: type = numpy.complex128) -> pyopencl.array.Array:
return pyopencl.array.to_device(queue, v.astype(dtype))
r = load_field(numpy.asarray(b))
x = pyopencl.array.zeros_like(r)
v = pyopencl.array.zeros_like(r)
p = pyopencl.array.zeros_like(r)
2016-07-04 19:30:36 -07:00
alpha = 1.0 + 0j
rho = 1.0 + 0j
errs = []
m = CSRMatrix(queue, A)
2016-07-04 19:30:36 -07:00
2024-07-30 22:43:29 -07:00
#
# Generate OpenCL kernels
#
a_step = ops.create_a_csr(context)
xr_step = ops.create_xr_step(context)
rhoerr_step = ops.create_rhoerr_step(context)
p_step = ops.create_p_step(context)
dot = ops.create_dot(context)
2016-07-04 19:30:36 -07:00
2024-07-30 22:43:29 -07:00
#
# Start the solve
#
2016-07-04 19:30:36 -07:00
start_time2 = time.perf_counter()
_, err2 = rhoerr_step(r, [])
b_norm = numpy.sqrt(err2)
2022-11-20 21:57:54 -08:00
logging.debug(f'b_norm check: {b_norm}')
2016-07-04 22:19:21 -07:00
success = False
2016-07-04 19:30:36 -07:00
for k in range(max_iters):
2022-11-20 21:57:54 -08:00
logging.debug(f'[{k:06d}] rho {rho:.4} alpha {alpha:4.4}')
2016-07-04 19:30:36 -07:00
rho_prev = rho
e = xr_step(x, p, r, v, alpha, [])
rho, err2 = rhoerr_step(r, e)
2016-07-04 19:30:36 -07:00
errs += [numpy.sqrt(err2) / b_norm]
2022-11-20 21:57:54 -08:00
logging.debug(f'err {errs[-1]}')
2016-07-04 19:30:36 -07:00
2016-08-04 17:43:01 -07:00
if errs[-1] < err_threshold:
2016-07-04 22:19:21 -07:00
success = True
break
e = p_step(p, r, rho/rho_prev, [])
e = a_step(v, m, p, e)
alpha = rho / dot(p, v, e)
2016-07-04 19:30:36 -07:00
2022-11-20 21:57:54 -08:00
if k % 1000 == 0:
logger.info(f'iteration {k}')
2016-07-04 19:30:36 -07:00
2024-07-30 22:43:29 -07:00
#
# Done solving
#
2016-07-04 22:19:21 -07:00
time_elapsed = time.perf_counter() - start_time
x = x.get()
if success:
logging.info('Solve success')
else:
logging.warning('Solve failure')
2022-11-20 21:57:54 -08:00
logging.info(f'{k} iterations in {time_elapsed} sec: {k / time_elapsed} iterations/sec')
logging.debug(f'final error {errs[-1]}')
logging.debug(f'overhead {start_time2 - start_time} sec')
2022-11-20 21:57:54 -08:00
residual = norm(A @ x - b) / norm(b)
logging.info(f'Final residual: {residual}')
2016-07-04 22:19:21 -07:00
return x
2016-07-04 23:30:06 -07:00
def fdfd_cg_solver(
2024-07-30 22:41:27 -07:00
solver_opts: dict[str, Any] | None = None,
**fdfd_args,
) -> NDArray[complexfloating]:
2016-08-04 17:43:01 -07:00
"""
Conjugate gradient FDFD solver using CSR sparse matrices, mainly for
testing and development since it's much slower than the solver in main.py.
Calls meanas.fdfd.solvers.generic(
**fdfd_args,
matrix_solver=opencl_fdfd.csr.cg,
matrix_solver_opts=solver_opts,
)
Args:
solver_opts: Passed as matrix_solver_opts to fdfd_tools.solver.generic(...).
Default {}.
fdfd_args: Passed as **fdfd_args to fdfd_tools.solver.generic(...).
Should include all of the arguments **except** matrix_solver and matrix_solver_opts
Returns:
E-field which solves the system.
2016-08-04 17:43:01 -07:00
"""
2016-07-04 23:55:43 -07:00
if solver_opts is None:
solver_opts = dict()
x = meanas.fdfd.solvers.generic(
matrix_solver=cg,
matrix_solver_opts=solver_opts,
**fdfd_args,
)
2016-07-04 23:30:06 -07:00
return x