From 5062dd180806dbb608156634f467bf432a0b7fa6 Mon Sep 17 00:00:00 2001 From: Alexandru Fikl Date: Thu, 3 Aug 2023 21:30:19 +0300 Subject: [PATCH] fix mypy errors --- pytential/linalg/hmatrix.py | 51 ++++++++++++++++--------------------- pytential/linalg/utils.py | 4 +-- 2 files changed, 24 insertions(+), 31 deletions(-) diff --git a/pytential/linalg/hmatrix.py b/pytential/linalg/hmatrix.py index cd32ad379..19df0a1f9 100644 --- a/pytential/linalg/hmatrix.py +++ b/pytential/linalg/hmatrix.py @@ -21,10 +21,11 @@ """ from dataclasses import dataclass -from typing import Any, Callable, Dict, Iterable, Optional, Union +from typing import Any, Callable, Dict, Iterable, Optional, Sequence, Union import numpy as np import numpy.linalg as la +from scipy.sparse.linalg import LinearOperator from arraycontext import PyOpenCLArrayContext, ArrayOrContainerT, flatten, unflatten from meshmode.dof_array import DOFArray @@ -35,13 +36,6 @@ from pytential.linalg.skeletonization import ( SkeletonizationWrangler, SkeletonizationResult) -try: - from scipy.sparse.linalg import LinearOperator -except ImportError: - # NOTE: scipy should be available (for interp_decomp), but just in case - class LinearOperator: - pass - import logging logger = logging.getLogger(__name__) @@ -124,7 +118,7 @@ def _update_skeleton_diagonal( targets, sources = parent.skel_tgt_src_index # FIXME: nicer way to do this? - mat = np.empty(skeleton.nclusters, dtype=object) + mat: np.ndarray = np.empty(skeleton.nclusters, dtype=object) for k in range(skeleton.nclusters): D = skeleton.D[k].copy() @@ -146,9 +140,9 @@ def _update_skeleton_diagonal( def _update_skeletons_diagonal( wrangler: "ProxyHierarchicalMatrixWrangler", - func: Callable[[SkeletonizationResult], np.ndarray], + func: Callable[[SkeletonizationResult], Optional[np.ndarray]], ) -> np.ndarray: - skeletons = np.empty(wrangler.skeletons.shape, dtype=object) + skeletons: np.ndarray = np.empty(wrangler.skeletons.shape, dtype=object) skeletons[0] = wrangler.skeletons[0] for i in range(1, wrangler.ctree.nlevels): @@ -263,11 +257,14 @@ def _matvec(self, x: ArrayOrContainerT) -> ArrayOrContainerT: else: raise TypeError(f"unsupported input type: {type(x)}") + assert actx is None or isinstance(actx, PyOpenCLArrayContext) result = apply_skeleton_forward_matvec(self, ary) + if isinstance(x, DOFArray): + assert actx is not None result = unflatten(x, actx.from_numpy(result), actx) - return result + return result # type: ignore[return-value] def apply_skeleton_forward_matvec( @@ -276,7 +273,7 @@ def apply_skeleton_forward_matvec( ) -> ArrayOrContainerT: from pytential.linalg.cluster import split_array targets, sources = hmat.skeletons[0].tgt_src_index - x = split_array(ary, sources) + x = split_array(ary, sources) # type: ignore[arg-type] # NOTE: this computes a telescoping product of the form # @@ -297,7 +294,7 @@ def apply_skeleton_forward_matvec( # # which gives back the desired product when we reach the leaf level again. - d_dot_x = np.empty(hmat.nlevels, dtype=object) + d_dot_x: np.ndarray = np.empty(hmat.nlevels, dtype=object) # {{{ recurse down @@ -307,8 +304,8 @@ def apply_skeleton_forward_matvec( assert x.shape == (skeleton.nclusters,) assert skeleton.tgt_src_index.shape[1] == sum([xi.size for xi in x]) - d_dot_x_k = np.empty(skeleton.nclusters, dtype=object) - r_dot_x_k = np.empty(skeleton.nclusters, dtype=object) + d_dot_x_k: np.ndarray = np.empty(skeleton.nclusters, dtype=object) + r_dot_x_k: np.ndarray = np.empty(skeleton.nclusters, dtype=object) for i in range(skeleton.nclusters): r_dot_x_k[i] = skeleton.R[i] @ x[i] @@ -366,23 +363,26 @@ def _matvec(self, x: ArrayOrContainerT) -> ArrayOrContainerT: else: raise TypeError(f"unsupported input type: {type(x)}") + assert actx is None or isinstance(actx, PyOpenCLArrayContext) result = apply_skeleton_backward_matvec(actx, self, ary) + if isinstance(x, DOFArray): + assert actx is not None result = unflatten(x, actx.from_numpy(result), actx) - return result + return result # type: ignore[return-value] def apply_skeleton_backward_matvec( - actx: PyOpenCLArrayContext, + actx: Optional[PyOpenCLArrayContext], hmat: ProxyHierarchicalMatrix, ary: ArrayOrContainerT, ) -> ArrayOrContainerT: from pytential.linalg.cluster import split_array targets, sources = hmat.skeletons[0].tgt_src_index - b = split_array(ary, targets) - r_dot_b = np.empty(hmat.nlevels, dtype=object) + b = split_array(ary, targets) # type: ignore[arg-type] + r_dot_b: np.ndarray = np.empty(hmat.nlevels, dtype=object) # {{{ recurse down @@ -412,7 +412,7 @@ def apply_skeleton_backward_matvec( assert b.shape == (skeleton.nclusters,) assert skeleton.tgt_src_index.shape[0] == sum([bi.size for bi in b]) - dhat_dot_b_k = np.empty(skeleton.nclusters, dtype=object) + dhat_dot_b_k: np.ndarray = np.empty(skeleton.nclusters, dtype=object) for i in range(skeleton.nclusters): dhat_dot_b_k[i] = ( skeleton.Dhat[i] @ (skeleton.R[i] @ (skeleton.invD[i] @ b[i])) @@ -467,7 +467,7 @@ def build_hmatrix_by_proxy( exprs: Union[sym.Expression, Iterable[sym.Expression]], input_exprs: Union[sym.Expression, Iterable[sym.Expression]], *, auto_where: Optional[sym.DOFDescriptorLike] = None, - domains: Optional[Iterable[sym.DOFDescriptorLike]] = None, + domains: Optional[Sequence[sym.DOFDescriptorLike]] = None, context: Optional[Dict[str, Any]] = None, id_eps: float = 1.0e-8, @@ -483,13 +483,6 @@ def build_hmatrix_by_proxy( _approx_nproxy: Optional[int] = None, _proxy_radius_factor: Optional[float] = None, ) -> ProxyHierarchicalMatrixWrangler: - try: - import scipy # noqa: F401 - except ImportError: - raise ImportError( - "The direct solver requires 'scipy' for the interpolative " - "decomposition used in skeletonization") - from pytential.symbolic.matrix import P2PClusterMatrixBuilder from pytential.linalg.skeletonization import make_skeletonization_wrangler diff --git a/pytential/linalg/utils.py b/pytential/linalg/utils.py index 64e43a575..bdfce8f78 100644 --- a/pytential/linalg/utils.py +++ b/pytential/linalg/utils.py @@ -443,8 +443,8 @@ def mnorm(x: np.ndarray, y: np.ndarray) -> "np.floating[Any]": def skeletonization_matrix( mat: np.ndarray, skeleton: "SkeletonizationResult", ) -> Tuple[np.ndarray, np.ndarray]: - D = np.empty(skeleton.nclusters, dtype=object) - S = np.empty((skeleton.nclusters, skeleton.nclusters), dtype=object) + D: np.ndarray = np.empty(skeleton.nclusters, dtype=object) + S: np.ndarray = np.empty((skeleton.nclusters, skeleton.nclusters), dtype=object) from itertools import product for i, j in product(range(skeleton.nclusters), repeat=2):