diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index 2af736b3c..bc042efe5 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -16,6 +16,14 @@ "endColumn": 14, "lineCount": 1 } + }, + { + "code": "reportCallIssue", + "range": { + "startColumn": 13, + "endColumn": 78, + "lineCount": 1 + } } ], "./examples/expansion-toys.py": [ @@ -43111,22 +43119,6 @@ "lineCount": 1 } }, - { - "code": "reportMissingTypeStubs", - "range": { - "startColumn": 17, - "endColumn": 30, - "lineCount": 1 - } - }, - { - "code": "reportUnknownMemberType", - "range": { - "startColumn": 39, - "endColumn": 64, - "lineCount": 1 - } - }, { "code": "reportUnknownMemberType", "range": { @@ -43135,14 +43127,6 @@ "lineCount": 1 } }, - { - "code": "reportAttributeAccessIssue", - "range": { - "startColumn": 57, - "endColumn": 64, - "lineCount": 1 - } - }, { "code": "reportUnannotatedClassAttribute", "range": { @@ -43288,10 +43272,18 @@ } }, { - "code": "reportMissingTypeStubs", + "code": "reportUnknownParameterType", "range": { - "startColumn": 13, - "endColumn": 26, + "startColumn": 8, + "endColumn": 13, + "lineCount": 1 + } + }, + { + "code": "reportUnknownParameterType", + "range": { + "startColumn": 12, + "endColumn": 22, "lineCount": 1 } }, @@ -43328,10 +43320,18 @@ } }, { - "code": "reportAny", + "code": "reportUnknownArgumentType", "range": { - "startColumn": 16, - "endColumn": 22, + "startColumn": 38, + "endColumn": 44, + "lineCount": 1 + } + }, + { + "code": "reportUnknownArgumentType", + "range": { + "startColumn": 46, + "endColumn": 51, "lineCount": 1 } }, @@ -52586,8 +52586,8 @@ { "code": "reportUnknownArgumentType", "range": { - "startColumn": 59, - "endColumn": 66, + "startColumn": 61, + "endColumn": 68, "lineCount": 1 } }, diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ba38b8344..6db9e167b 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -42,7 +42,7 @@ jobs: curl -L -O https://tiker.net/ci-support-v0 . ./ci-support-v0 build_py_project_in_conda_env - cipip install pytest pyfmmlib scipy matplotlib pyvisfile + cipip install pytest pyfmmlib scipy scipy-stubs matplotlib pyvisfile cipip install basedpyright basedpyright diff --git a/examples/curve-pot.py b/examples/curve-pot.py index c1c113662..e51380d3f 100644 --- a/examples/curve-pot.py +++ b/examples/curve-pot.py @@ -1,5 +1,6 @@ import numpy as np import numpy.linalg as la +from numpy.typing import NDArray import pyopencl as cl @@ -120,7 +121,7 @@ def draw_pot_figure(aspect_ratio, a = 1 b = 1/aspect_ratio - def map_to_curve(t): + def map_to_curve(t: NDArray[np.floating]): t = t*(2*np.pi) x = a*np.cos(t) @@ -185,8 +186,8 @@ def map_to_curve(t): from sumpy.tools import build_matrix - def apply_lpot(x): - xovsmp = np.dot(fim, x) + def apply_lpot(x: NDArray[np.inexact]) -> NDArray[np.inexact]: + xovsmp = fim @ x _evt, (y,) = lpot(actx.queue, sources, ovsmp_sources, @@ -197,7 +198,7 @@ def apply_lpot(x): return actx.to_numpy(y) - op = LinearOperator((nsrc, nsrc), apply_lpot) + op = LinearOperator((nsrc, nsrc), np.dtype(np.complex128), apply_lpot) mat = build_matrix(op, dtype=np.complex128) w, _v = la.eig(mat) plt.plot(w.real, "o-") diff --git a/examples/curve.py b/examples/curve.py index 760db5f77..2a7e53b88 100644 --- a/examples/curve.py +++ b/examples/curve.py @@ -1,8 +1,11 @@ +from typing import final + import numpy as np import scipy as sp import scipy.fftpack +@final class CurveGrid: def __init__(self, x, y): self.pos = np.vstack([x, y]).copy() diff --git a/examples/fourier.py b/examples/fourier.py index f637c2a4f..9e51c8449 100644 --- a/examples/fourier.py +++ b/examples/fourier.py @@ -1,13 +1,14 @@ import numpy as np +from numpy.typing import NDArray -def make_fourier_vdm(n, inverse): +def make_fourier_vdm(n: int, inverse: bool) -> NDArray[np.complex128]: i = np.arange(n, dtype=np.float64) imat = i[:, np.newaxis]*i/n result = np.exp((2j*np.pi)*imat) if inverse: - result = result.T.conj()/n + result = np.conj(result.T)/n return result @@ -30,9 +31,7 @@ def make_fourier_mode_extender(m, n, dtype): return result -def make_fourier_interp_matrix(m, n): - return np.dot( - np.dot( - make_fourier_vdm(m, inverse=False), - make_fourier_mode_extender(m, n, np.float64)), - make_fourier_vdm(n, inverse=True)) +def make_fourier_interp_matrix(m: int, n: int): + return (make_fourier_vdm(m, inverse=False) + @ make_fourier_mode_extender(m, n, np.float64) + @ make_fourier_vdm(n, inverse=True)) diff --git a/sumpy/distributed.py b/sumpy/distributed.py index 563de10a8..df66f308f 100644 --- a/sumpy/distributed.py +++ b/sumpy/distributed.py @@ -26,6 +26,7 @@ from boxtree.distributed.calculation import DistributedExpansionWrangler import pyopencl.array as cl_array +import pytools.obj_array as obj_array from sumpy.fmm import SumpyExpansionWrangler @@ -66,8 +67,7 @@ def gather_potential_results(self, potentials, tgt_idx_all_ranks): super().gather_potential_results(potentials_host, tgt_idx_all_ranks)) if mpi_rank == 0: - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ cl_array.to_device(potentials_dev.queue, gathered_potentials_host) for gathered_potentials_host, potentials_dev in zip(gathered_potentials_host_vec, potentials, strict=True)]) @@ -85,7 +85,6 @@ def reorder_potentials(self, potentials): if self.comm.Get_rank() == 0: import numpy as np - from pytools.obj_array import obj_array_vectorize assert ( isinstance(potentials, np.ndarray) and potentials.dtype.char == "O") @@ -93,7 +92,7 @@ def reorder_potentials(self, potentials): def reorder(x): return x[self.global_traversal.tree.sorted_target_ids] - return obj_array_vectorize(reorder, potentials) + return obj_array.vectorize(reorder, potentials) else: return None diff --git a/sumpy/fmm.py b/sumpy/fmm.py index d56bfbc12..7c9587e77 100644 --- a/sumpy/fmm.py +++ b/sumpy/fmm.py @@ -47,6 +47,7 @@ import pyopencl as cl import pyopencl.array as cl_array +import pytools.obj_array as obj_array from pytools import memoize_method from sumpy import ( @@ -546,8 +547,7 @@ def output_zeros(self, template_ary): (e.g. :class:`pyopencl.CommandQueue`) the returned array should reuse. """ - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ cl_array.zeros( template_ary.queue, self.tree.ntargets, @@ -560,7 +560,6 @@ def reorder_sources(self, source_array): def reorder_potentials(self, potentials): import numpy as np - from pytools.obj_array import obj_array_vectorize assert ( isinstance(potentials, np.ndarray) and potentials.dtype.char == "O") @@ -568,7 +567,7 @@ def reorder_potentials(self, potentials): def reorder(x): return x[self.tree.sorted_target_ids] - return obj_array_vectorize(reorder, potentials) + return obj_array.vectorize(reorder, potentials) @property @memoize_method diff --git a/sumpy/point_calculus.py b/sumpy/point_calculus.py index 6e59e7f92..f2c3d1d29 100644 --- a/sumpy/point_calculus.py +++ b/sumpy/point_calculus.py @@ -26,6 +26,7 @@ import numpy as np import numpy.linalg as la +import pytools.obj_array as obj_array from pytools import memoize_method @@ -232,8 +233,7 @@ def curl(self, arg): :class:`numpy.ndarray`\ s with shape ``(npoints_total,)``. """ from pytools import levi_civita - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ sum( levi_civita((k, m, n)) * self.diff(m, arg[n]) for m in range(3) for n in range(3)) diff --git a/sumpy/qbx.py b/sumpy/qbx.py index 3e44c0c07..8d6092ebd 100644 --- a/sumpy/qbx.py +++ b/sumpy/qbx.py @@ -33,6 +33,7 @@ from typing_extensions import override import loopy as lp +import pytools.obj_array as obj_array from loopy.version import MOST_RECENT_LANGUAGE_VERSION from pymbolic import parse, var from pytools import memoize_method @@ -646,8 +647,7 @@ def normal(self): self.arguments["normal"] = ( lp.GlobalArg("normal", self.geometry_dtype, shape=("ntargets", self.dim), order="C")) - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ parse(f"normal[itgt, {i}]") for i in range(self.dim)]) @@ -657,8 +657,7 @@ def tangent(self): self.arguments["tangent"] = ( lp.GlobalArg("tangent", self.geometry_dtype, shape=("ntargets", self.dim), order="C")) - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ parse(f"tangent[itgt, {i}]") for i in range(self.dim)]) @@ -678,8 +677,7 @@ def src_derivative_dir(self): lp.GlobalArg("src_derivative_dir", self.geometry_dtype, shape=("ntargets", self.dim), order="C")) - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ parse(f"src_derivative_dir[itgt, {i}]") for i in range(self.dim)]) @@ -690,8 +688,7 @@ def tgt_derivative_dir(self): lp.GlobalArg("tgt_derivative_dir", self.geometry_dtype, shape=("ntargets", self.dim), order="C")) - from pytools.obj_array import make_obj_array - return make_obj_array([ + return obj_array.new_1d([ parse(f"tgt_derivative_dir[itgt, {i}]") for i in range(self.dim)]) diff --git a/sumpy/tools.py b/sumpy/tools.py index c28214bd2..06f3072b7 100644 --- a/sumpy/tools.py +++ b/sumpy/tools.py @@ -1,7 +1,5 @@ from __future__ import annotations -from pyopencl import MemoryObjectHolder - __copyright__ = """ Copyright (C) 2012 Andreas Kloeckner @@ -29,6 +27,7 @@ THE SOFTWARE. """ + import enum import logging import warnings @@ -40,7 +39,9 @@ import numpy as np import loopy as lp +import pytools.obj_array as obj_array from pymbolic.mapper import WalkMapper +from pyopencl import MemoryObjectHolder from pytools import memoize_method from pytools.tag import Tag, tag_dataclass @@ -226,17 +227,14 @@ def build_matrix(op, dtype=None, shape=None): def vector_to_device(queue, vec): from pyopencl.array import to_device - from pytools.obj_array import obj_array_vectorize def to_dev(ary): return to_device(queue, ary) - return obj_array_vectorize(to_dev, vec) + return obj_array.vectorize(to_dev, vec) def vector_from_device(queue, vec): - from pytools.obj_array import obj_array_vectorize - def from_dev(ary): from numbers import Number if isinstance(ary, np.number | Number): @@ -245,7 +243,7 @@ def from_dev(ary): return ary.get(queue=queue) - return obj_array_vectorize(from_dev, vec) + return obj_array.vectorize(from_dev, vec) def _merge_kernel_arguments(dictionary, arg): diff --git a/sumpy/toys.py b/sumpy/toys.py index dd20c17ce..ad7ee5517 100644 --- a/sumpy/toys.py +++ b/sumpy/toys.py @@ -30,6 +30,7 @@ from numbers import Number from typing import TYPE_CHECKING +import pytools.obj_array as obj_array from pytools import memoize_method from sumpy.kernel import TargetTransformationRemover @@ -309,8 +310,6 @@ def _e2p(psource, targets, e2p): queue, np.array(psource.center, dtype=np.float64).reshape(toy_ctx.kernel.dim, 1)) - from pytools.obj_array import make_obj_array - from sumpy.tools import vector_to_device coeffs = cl_array.to_device(queue, np.array([psource.coeffs])) @@ -323,7 +322,7 @@ def _e2p(psource, targets, e2p): box_target_counts_nonchild=box_target_counts_nonchild, centers=centers, rscale=psource.rscale, - targets=vector_to_device(queue, make_obj_array(targets)), + targets=vector_to_device(queue, obj_array.new_1d(targets)), **toy_ctx.extra_kernel_kwargs) return pot.get(queue) diff --git a/test/test_fmm.py b/test/test_fmm.py index d85eafe60..b79b800b7 100644 --- a/test/test_fmm.py +++ b/test/test_fmm.py @@ -32,6 +32,7 @@ import numpy.linalg as la import pytest +import pytools.obj_array as obj_array from arraycontext import pytest_generate_tests_for_array_contexts from sumpy.array_context import PytestPyOpenCLArrayContextFactory, _acf # noqa: F401 @@ -142,8 +143,7 @@ def _test_sumpy_fmm(actx_factory, knl, local_expn_class, mpole_expn_class, from sumpy.visualization import FieldPlotter fp = FieldPlotter(np.array([0.5, 0]), extent=3, npoints=200) - from pytools.obj_array import make_obj_array - targets = make_obj_array([fp.points[i] for i in range(knl.dim)]) + targets = obj_array.new_1d([fp.points[i] for i in range(knl.dim)]) from boxtree import TreeBuilder tb = TreeBuilder(actx.context) diff --git a/test/test_kernels.py b/test/test_kernels.py index 41d9da0b5..9c7b23cc0 100644 --- a/test/test_kernels.py +++ b/test/test_kernels.py @@ -31,9 +31,9 @@ import numpy.linalg as la import pytest +import pytools.obj_array as obj_array from arraycontext import pytest_generate_tests_for_array_contexts from pytools.convergence import PConvergenceVerifier -from pytools.obj_array import make_obj_array import sumpy.symbolic as sym import sumpy.toys as t @@ -350,7 +350,7 @@ def test_p2e2p(actx_factory, base_knl, expn_class, order, with_source_derivative + center[:, np.newaxis]) centers = actx.from_numpy(centers) - targets = actx.from_numpy(make_obj_array(fp.points)) + targets = actx.from_numpy(obj_array.new_1d(fp.points)) rscale = 0.5 # pick something non-1 diff --git a/test/test_matrixgen.py b/test/test_matrixgen.py index 1dc0dd854..87a085af0 100644 --- a/test/test_matrixgen.py +++ b/test/test_matrixgen.py @@ -30,6 +30,7 @@ import numpy.linalg as la import pytest +import pytools.obj_array as obj_array from arraycontext import pytest_generate_tests_for_array_contexts from sumpy.array_context import PytestPyOpenCLArrayContextFactory, _acf # noqa: F401 @@ -141,9 +142,8 @@ def test_qbx_direct(actx_factory, factor, lpot_id, visualize=False): extra_kwargs = {} if lpot_id == 2: - from pytools.obj_array import make_obj_array extra_kwargs["dsource_vec"] = ( - actx.from_numpy(make_obj_array(np.ones((ndim, n)))) + actx.from_numpy(obj_array.new_1d(np.ones((ndim, n)))) ) _, (result_lpot,) = lpot(actx.queue, @@ -229,9 +229,8 @@ def test_p2p_direct(actx_factory, exclude_self, factor, lpot_id, visualize=False actx.from_numpy(np.arange(n, dtype=np.int32)) ) if lpot_id == 2: - from pytools.obj_array import make_obj_array extra_kwargs["dsource_vec"] = ( - actx.from_numpy(make_obj_array(np.ones((ndim, n))))) + actx.from_numpy(obj_array.new_1d(np.ones((ndim, n))))) _, (result_lpot,) = lpot(actx.queue, targets=targets,