Obj array renames

.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
                 }
             },

.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
 

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-")

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()

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))

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,15 +85,14 @@ 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")
 
             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
 

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,15 +560,14 @@ 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")
 
         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

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))

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)])
 

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):