From efeb29479b30335adbd44820ff8343b008ca3ea5 Mon Sep 17 00:00:00 2001 From: Jan Petykiewicz Date: Sun, 20 Nov 2022 21:57:43 -0800 Subject: [PATCH] improve type annotations, formatting, comment styles --- opencl_fdfd/csr.py | 88 ++++++++------- opencl_fdfd/main.py | 78 +++++++------- opencl_fdfd/ops.py | 257 ++++++++++++++++++++++++++++---------------- 3 files changed, 261 insertions(+), 162 deletions(-) diff --git a/opencl_fdfd/csr.py b/opencl_fdfd/csr.py index a34f503..df95822 100644 --- a/opencl_fdfd/csr.py +++ b/opencl_fdfd/csr.py @@ -14,14 +14,16 @@ satisfy the constraints for the 'conjugate gradient' algorithm (positive definite, symmetric) and some that don't. """ -from typing import Dict, Any +from typing import Dict, Any, Optional import time import logging import numpy +from numpy.typing import NDArray, ArrayLike from numpy.linalg import norm import pyopencl import pyopencl.array +import scipy import meanas.fdfd.solvers @@ -33,39 +35,45 @@ __author__ = 'Jan Petykiewicz' logger = logging.getLogger(__name__) -class CSRMatrix(object): +class CSRMatrix: """ Matrix stored in Compressed Sparse Row format, in GPU RAM. """ - row_ptr = None # type: pyopencl.array.Array - col_ind = None # type: pyopencl.array.Array - data = None # type: pyopencl.array.Array + 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'): + def __init__( + self, + queue: pyopencl.CommandQueue, + m: 'scipy.sparse.csr_matrix', + ) -> None: 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: numpy.ndarray, - max_iters: int = 10000, - err_threshold: float = 1e-6, - context: pyopencl.Context = None, - queue: pyopencl.CommandQueue = None, - ) -> numpy.ndarray: +def cg( + A: 'scipy.sparse.csr_matrix', + b: ArrayLike, + max_iters: int = 10000, + err_threshold: float = 1e-6, + context: Optional[pyopencl.Context] = None, + queue: Optional[pyopencl.CommandQueue] = None, + ) -> NDArray: """ General conjugate-gradient solver for sparse matrices, where A @ x = b. - :param A: Matrix to solve (CSR format) - :param b: Right-hand side vector (dense ndarray) - :param max_iters: Maximum number of iterations - :param err_threshold: Error threshold for successful solve, relative to norm(b) - :param context: PyOpenCL context. Will be created if not given. - :param queue: PyOpenCL command queue. Will be created if not given. - :return: Solution vector x; returned even if solve doesn't converge. + 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. """ start_time = time.perf_counter() @@ -151,29 +159,37 @@ def cg(A: 'scipy.sparse.csr_matrix', return x -def fdfd_cg_solver(solver_opts: Dict[str, Any] = None, - **fdfd_args - ) -> numpy.ndarray: +def fdfd_cg_solver( + solver_opts: Optional[Dict[str, Any]] = None, + **fdfd_args + ) -> NDArray: """ 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) + Calls meanas.fdfd.solvers.generic( + **fdfd_args, + matrix_solver=opencl_fdfd.csr.cg, + matrix_solver_opts=solver_opts, + ) - :param solver_opts: Passed as matrix_solver_opts to fdfd_tools.solver.generic(...). - Default {}. - :param fdfd_args: Passed as **fdfd_args to fdfd_tools.solver.generic(...). - Should include all of the arguments **except** matrix_solver and matrix_solver_opts - :return: E-field which solves the system. + 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. """ if solver_opts is None: solver_opts = dict() - x = meanas.fdfd.solvers.generic(matrix_solver=cg, - matrix_solver_opts=solver_opts, - **fdfd_args) + x = meanas.fdfd.solvers.generic( + matrix_solver=cg, + matrix_solver_opts=solver_opts, + **fdfd_args, + ) return x diff --git a/opencl_fdfd/main.py b/opencl_fdfd/main.py index 3de469b..d6e0d83 100644 --- a/opencl_fdfd/main.py +++ b/opencl_fdfd/main.py @@ -6,11 +6,12 @@ operator implemented directly as OpenCL arithmetic (rather than as a matrix). """ -from typing import List +from typing import List, Optional, cast import time import logging import numpy +from numpy.typing import NDArray, ArrayLike from numpy.linalg import norm import pyopencl import pyopencl.array @@ -25,18 +26,19 @@ __author__ = 'Jan Petykiewicz' logger = logging.getLogger(__name__) -def cg_solver(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, - max_iters: int = 40000, - err_threshold: float = 1e-6, - context: pyopencl.Context = None, - ) -> numpy.ndarray: +def cg_solver( + omega: complex, + dxes: List[List[NDArray]], + J: ArrayLike, + epsilon: ArrayLike, + mu: Optional[ArrayLike] = None, + pec: Optional[ArrayLike] = None, + pmc: Optional[ArrayLike] = None, + adjoint: bool = False, + max_iters: int = 40000, + err_threshold: float = 1e-6, + context: Optional[pyopencl.Context] = None, + ) -> NDArray: """ OpenCL FDFD solver using the iterative conjugate gradient (cg) method and implementing the diagonalized E-field wave operator directly in @@ -46,28 +48,30 @@ def cg_solver(omega: complex, either use meanas.fdmath.vec() or numpy: f_1D = numpy.hstack(tuple((fi.flatten(order='F') for fi in [f_x, f_y, f_z]))) - :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 max_iters: Maximum number of iterations. Default 40,000. - :param err_threshold: If (r @ r.conj()) / norm(1j * omega * J) < err_threshold, success. - Default 1e-6. - :param context: PyOpenCL context to run in. If not given, construct a new context. - :return: E-field which solves the system. Returned even if we did not converge. - """ + Args: + omega: Complex frequency to solve at. + dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes) + J: Electric current distribution (at E-field locations) + epsilon: Dielectric constant distribution (at E-field locations) + mu: Magnetic permeability distribution (at H-field locations) + pec: Perfect electric conductor distribution + (at E-field locations; non-zero value indicates PEC is present) + pmc: Perfect magnetic conductor distribution + (at H-field locations; non-zero value indicates PMC is present) + adjoint: If true, solves the adjoint problem. + max_iters: Maximum number of iterations. Default 40,000. + err_threshold: If (r @ r.conj()) / norm(1j * omega * J) < err_threshold, success. + Default 1e-6. + context: PyOpenCL context to run in. If not given, construct a new context. + Returns: + E-field which solves the system. Returned even if we did not converge. + """ start_time = time.perf_counter() - b = -1j * omega * J + shape = [dd.size for dd in dxes[0]] - shape = [d.size for d in dxes[0]] + b = -1j * omega * numpy.array(J, copy=False) ''' ** In this comment, I use the following notation: @@ -96,9 +100,10 @@ def cg_solver(omega: complex, We can accomplish all this simply by conjugating everything (except J) and reversing the order of L and R ''' + epsilon = numpy.array(epsilon, copy=False) if adjoint: # Conjugate everything - dxes = [[numpy.conj(d) for d in dd] for dd in dxes] + dxes = [[numpy.conj(dd) for dd in dds] for dds in dxes] omega = numpy.conj(omega) epsilon = numpy.conj(epsilon) if mu is not None: @@ -132,7 +137,7 @@ def cg_solver(omega: complex, rho = 1.0 + 0j errs = [] - inv_dxes = [[load_field(1 / d) for d in dd] for dd in dxes] + inv_dxes = [[load_field(1 / numpy.array(dd, copy=False)) for dd in dds] for dds in dxes] oeps = load_field(-omega ** 2 * epsilon) Pl = load_field(L.diagonal()) Pr = load_field(R.diagonal()) @@ -140,17 +145,18 @@ def cg_solver(omega: complex, if mu is None: invm = load_field(numpy.array([])) else: - invm = load_field(1 / mu) + invm = load_field(1 / numpy.array(mu, copy=False)) + mu = numpy.array(mu, copy=False) if pec is None: gpec = load_field(numpy.array([]), dtype=numpy.int8) else: - gpec = load_field(pec.astype(bool), dtype=numpy.int8) + gpec = load_field(numpy.array(pec, dtype=bool, copy=False), dtype=numpy.int8) if pmc is None: gpmc = load_field(numpy.array([]), dtype=numpy.int8) else: - gpmc = load_field(pmc.astype(bool), dtype=numpy.int8) + gpmc = load_field(numpy.array(pmc, dtype=bool, copy=False), dtype=numpy.int8) ''' Generate OpenCL kernels diff --git a/opencl_fdfd/ops.py b/opencl_fdfd/ops.py index 2e0f160..32f2400 100644 --- a/opencl_fdfd/ops.py +++ b/opencl_fdfd/ops.py @@ -7,10 +7,11 @@ kernels for use by the other solvers. See kernels/ for any of the .cl files loaded in this file. """ -from typing import List, Callable +from typing import List, Callable, Union, Type, Sequence, Optional, Tuple import logging import numpy +from numpy.typing import NDArray, ArrayLike import jinja2 import pyopencl @@ -28,12 +29,17 @@ jinja_env = jinja2.Environment(loader=jinja2.PackageLoader(__name__, 'kernels')) operation = Callable[..., List[pyopencl.Event]] -def type_to_C(float_type: numpy.float32 or numpy.float64) -> str: +def type_to_C( + float_type: Type, + ) -> str: """ Returns a string corresponding to the C equivalent of a numpy type. - :param float_type: numpy type: float32, float64, complex64, complex128 - :return: string containing the corresponding C type (eg. 'double') + Args: + float_type: numpy type: float32, float64, complex64, complex128 + + Returns: + string containing the corresponding C type (eg. 'double') """ types = { numpy.float32: 'float', @@ -68,12 +74,13 @@ def ptrs(*args: str) -> List[str]: return [ctype + ' *' + s for s in args] -def create_a(context: pyopencl.Context, - shape: numpy.ndarray, - mu: bool = False, - pec: bool = False, - pmc: bool = False, - ) -> operation: +def create_a( + context: pyopencl.Context, + shape: ArrayLike, + mu: bool = False, + pec: bool = False, + pmc: bool = False, + ) -> operation: """ Return a function which performs (A @ p), where A is the FDFD wave equation for E-field. @@ -94,12 +101,15 @@ def create_a(context: pyopencl.Context, and returns a list of pyopencl.Event. - :param context: PyOpenCL context - :param shape: Dimensions of the E-field - :param mu: False iff (mu == 1) everywhere - :param pec: False iff no PEC anywhere - :param pmc: False iff no PMC anywhere - :return: Function for computing (A @ p) + Args: + context: PyOpenCL context + shape: Dimensions of the E-field + mu: False iff (mu == 1) everywhere + pec: False iff no PEC anywhere + pmc: False iff no PMC anywhere + + Returns: + Function for computing (A @ p) """ common_source = jinja_env.get_template('common.cl').render(shape=shape) @@ -113,45 +123,67 @@ def create_a(context: pyopencl.Context, Convert p to initial E (ie, apply right preconditioner and PEC) ''' p2e_source = jinja_env.get_template('p2e.cl').render(pec=pec) - P2E_kernel = ElementwiseKernel(context, - name='P2E', - preamble=preamble, - operation=p2e_source, - arguments=', '.join(ptrs('E', 'p', 'Pr') + pec_arg)) + P2E_kernel = ElementwiseKernel( + context, + name='P2E', + preamble=preamble, + operation=p2e_source, + arguments=', '.join(ptrs('E', 'p', 'Pr') + pec_arg), + ) ''' Calculate intermediate H from intermediate E ''' - e2h_source = jinja_env.get_template('e2h.cl').render(mu=mu, - pmc=pmc, - common_cl=common_source) - E2H_kernel = ElementwiseKernel(context, - name='E2H', - preamble=preamble, - operation=e2h_source, - arguments=', '.join(ptrs('E', 'H', 'inv_mu') + pmc_arg + des)) + e2h_source = jinja_env.get_template('e2h.cl').render( + mu=mu, + pmc=pmc, + common_cl=common_source, + ) + E2H_kernel = ElementwiseKernel( + context, + name='E2H', + preamble=preamble, + operation=e2h_source, + arguments=', '.join(ptrs('E', 'H', 'inv_mu') + pmc_arg + des), + ) ''' Calculate final E (including left preconditioner) ''' - h2e_source = jinja_env.get_template('h2e.cl').render(pec=pec, - common_cl=common_source) - H2E_kernel = ElementwiseKernel(context, - name='H2E', - preamble=preamble, - operation=h2e_source, - arguments=', '.join(ptrs('E', 'H', 'oeps', 'Pl') + pec_arg + dhs)) + h2e_source = jinja_env.get_template('h2e.cl').render( + pec=pec, + common_cl=common_source, + ) + H2E_kernel = ElementwiseKernel( + context, + name='H2E', + preamble=preamble, + operation=h2e_source, + arguments=', '.join(ptrs('E', 'H', 'oeps', 'Pl') + pec_arg + dhs), + ) - def spmv(E, H, p, idxes, oeps, inv_mu, pec, pmc, Pl, Pr, e): + def spmv( + E: pyopencl.array.Array, + H: pyopencl.array.Array, + p: pyopencl.array.Array, + idxes: Sequence[Sequence[pyopencl.array.Array]], + oeps: pyopencl.array.Array, + inv_mu: Optional[pyopencl.array.Array], + pec: Optional[pyopencl.array.Array], + pmc: Optional[pyopencl.array.Array], + Pl: pyopencl.array.Array, + Pr: pyopencl.array.Array, + e: List[pyopencl.Event], + ) -> List[pyopencl.Event]: e2 = P2E_kernel(E, p, Pr, pec, wait_for=e) e2 = E2H_kernel(E, H, inv_mu, pmc, *idxes[0], wait_for=[e2]) e2 = H2E_kernel(E, H, oeps, Pl, pec, *idxes[1], wait_for=[e2]) return [e2] - logger.debug('Preamble: \n{}'.format(preamble)) - logger.debug('p2e: \n{}'.format(p2e_source)) - logger.debug('e2h: \n{}'.format(e2h_source)) - logger.debug('h2e: \n{}'.format(h2e_source)) + logger.debug(f'Preamble: \n{preamble}') + logger.debug(f'p2e: \n{p2e_source}') + logger.debug(f'e2h: \n{e2h_source}') + logger.debug(f'h2e: \n{h2e_source}') return spmv @@ -167,8 +199,11 @@ def create_xr_step(context: pyopencl.Context) -> operation: after waiting for all in the list e and returns a list of pyopencl.Event - :param context: PyOpenCL context - :return: Function for performing x and r updates + Args: + context: PyOpenCL context + + Returns: + Function for performing x and r updates """ update_xr_source = ''' x[i] = add(x[i], mul(alpha, p[i])); @@ -177,19 +212,28 @@ def create_xr_step(context: pyopencl.Context) -> operation: xr_args = ', '.join(ptrs('x', 'p', 'r', 'v') + [ctype + ' alpha']) - xr_kernel = ElementwiseKernel(context, - name='XR', - preamble=preamble, - operation=update_xr_source, - arguments=xr_args) + xr_kernel = ElementwiseKernel( + context, + name='XR', + preamble=preamble, + operation=update_xr_source, + arguments=xr_args, + ) - def xr_update(x, p, r, v, alpha, e): + def xr_update( + x: pyopencl.array.Array, + p: pyopencl.array.Array, + r: pyopencl.array.Array, + v: pyopencl.array.Array, + alpha: complex, + e: List[pyopencl.Event], + ) -> List[pyopencl.Event]: return [xr_kernel(x, p, r, v, alpha, wait_for=e)] return xr_update -def create_rhoerr_step(context: pyopencl.Context) -> operation: +def create_rhoerr_step(context: pyopencl.Context) -> Callable[..., Tuple[complex, complex]]: """ Return a function ri_update(r, e) @@ -200,8 +244,11 @@ def create_rhoerr_step(context: pyopencl.Context) -> operation: after waiting for all pyopencl.Event in the list e and returns a list of pyopencl.Event - :param context: PyOpenCL context - :return: Function for performing x and r updates + Args: + context: PyOpenCL context + + Returns: + Function for performing x and r updates """ update_ri_source = ''' @@ -213,16 +260,18 @@ def create_rhoerr_step(context: pyopencl.Context) -> operation: # Use a vector type (double3) to make the reduction simpler ri_dtype = pyopencl.array.vec.double3 - ri_kernel = ReductionKernel(context, - name='RHOERR', - preamble=preamble, - dtype_out=ri_dtype, - neutral='(double3)(0.0, 0.0, 0.0)', - map_expr=update_ri_source, - reduce_expr='a+b', - arguments=ctype + ' *r') + ri_kernel = ReductionKernel( + context, + name='RHOERR', + preamble=preamble, + dtype_out=ri_dtype, + neutral='(double3)(0.0, 0.0, 0.0)', + map_expr=update_ri_source, + reduce_expr='a+b', + arguments=ctype + ' *r', + ) - def ri_update(r, e): + def ri_update(r: pyopencl.array.Array, e: List[pyopencl.Event]) -> Tuple[complex, complex]: g = ri_kernel(r, wait_for=e).astype(ri_dtype).get() rr, ri, ii = [g[q] for q in 'xyz'] rho = rr + 2j * ri - ii @@ -242,48 +291,66 @@ def create_p_step(context: pyopencl.Context) -> operation: after waiting for all pyopencl.Event in the list e and returns a list of pyopencl.Event - :param context: PyOpenCL context - :return: Function for performing the p update + Args: + context: PyOpenCL context + + Returns: + Function for performing the p update """ update_p_source = ''' p[i] = add(r[i], mul(beta, p[i])); ''' p_args = ptrs('p', 'r') + [ctype + ' beta'] - p_kernel = ElementwiseKernel(context, - name='P', - preamble=preamble, - operation=update_p_source, - arguments=', '.join(p_args)) + p_kernel = ElementwiseKernel( + context, + name='P', + preamble=preamble, + operation=update_p_source, + arguments=', '.join(p_args), + ) - def p_update(p, r, beta, e): + def p_update( + p: pyopencl.array.Array, + r: pyopencl.array.Array, + beta: complex, + e: List[pyopencl.Event]) -> List[pyopencl.Event]: return [p_kernel(p, r, beta, wait_for=e)] return p_update -def create_dot(context: pyopencl.Context) -> operation: +def create_dot(context: pyopencl.Context) -> Callable[..., complex]: """ Return a function for performing the dot product p @ v with the signature - dot(p, v, e) -> float + dot(p, v, e) -> complex - :param context: PyOpenCL context - :return: Function for performing the dot product + Args: + context: PyOpenCL context + + Returns: + Function for performing the dot product """ dot_dtype = numpy.complex128 - dot_kernel = ReductionKernel(context, - name='dot', - preamble=preamble, - dtype_out=dot_dtype, - neutral='zero', - map_expr='mul(p[i], v[i])', - reduce_expr='add(a, b)', - arguments=ptrs('p', 'v')) + dot_kernel = ReductionKernel( + context, + name='dot', + preamble=preamble, + dtype_out=dot_dtype, + neutral='zero', + map_expr='mul(p[i], v[i])', + reduce_expr='add(a, b)', + arguments=ptrs('p', 'v'), + ) - def dot(p, v, e): + def dot( + p: pyopencl.array.Array, + v: pyopencl.array.Array, + e: List[pyopencl.Event], + ) -> complex: g = dot_kernel(p, v, wait_for=e) return g.get() @@ -304,8 +371,11 @@ def create_a_csr(context: pyopencl.Context) -> operation: The function waits on all the pyopencl.Event in e before running, and returns a list of pyopencl.Event. - :param context: PyOpenCL context - :return: Function for sparse (M @ v) operation where M is in CSR format + Args: + context: PyOpenCL context + + Returns: + Function for sparse (M @ v) operation where M is in CSR format """ spmv_source = ''' int start = m_row_ptr[i]; @@ -326,13 +396,20 @@ def create_a_csr(context: pyopencl.Context) -> operation: m_args = 'int *m_row_ptr, int *m_col_ind, ' + ctype + ' *m_data' v_in_args = ctype + ' *v_in' - spmv_kernel = ElementwiseKernel(context, - name='csr_spmv', - preamble=preamble, - operation=spmv_source, - arguments=', '.join((v_out_args, m_args, v_in_args))) + spmv_kernel = ElementwiseKernel( + context, + name='csr_spmv', + preamble=preamble, + operation=spmv_source, + arguments=', '.join((v_out_args, m_args, v_in_args)), + ) - def spmv(v_out, m, v_in, e): + def spmv( + v_out, + m, + v_in, + e: List[pyopencl.Event], + ) -> List[pyopencl.Event]: return [spmv_kernel(v_out, m.row_ptr, m.col_ind, m.data, v_in, wait_for=e)] return spmv