Modernize 2025 07

sumpy/expansion/__init__.py

@@ -39,6 +39,8 @@
 
     import loopy as lp
 
+    from sumpy.expansion.local import LocalExpansionBase
+    from sumpy.expansion.multipole import MultipoleExpansionBase
     from sumpy.kernel import Kernel
 
 
@@ -915,13 +917,17 @@ class ExpansionFactoryBase(ABC):
     """
 
     @abstractmethod
-    def get_local_expansion_class(self, base_kernel):
+    def get_local_expansion_class(self,
+                base_kernel: Kernel
+            ) -> type[LocalExpansionBase]:
         """
         :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*.
         """
 
     @abstractmethod
-    def get_multipole_expansion_class(self, base_kernel):
+    def get_multipole_expansion_class(self,
+                base_kernel: Kernel
+            ) -> type[MultipoleExpansionBase]:
         """
         :returns: a subclass of :class:`ExpansionBase` suitable for *base_kernel*.
         """
@@ -983,4 +989,22 @@ def get_multipole_expansion_class(self, base_kernel):
 # }}}
 
 
+__all__ = [
+    "BiharmonicConformingVolumeTaylorExpansion",
+    "CSEMatVecOperator",
+    "DefaultExpansionFactory",
+    "ExpansionBase",
+    "ExpansionFactoryBase",
+    "ExpansionTermsWrangler",
+    "FullExpansionTermsWrangler",
+    "HelmholtzConformingVolumeTaylorExpansion",
+    "LaplaceConformingVolumeTaylorExpansion",
+    "LinearPDEBasedExpansionTermsWrangler",
+    "LinearPDEConformingVolumeTaylorExpansion",
+    "VolumeTaylorExpansion",
+    "VolumeTaylorExpansionFactory",
+    "VolumeTaylorExpansionMixin",
+]
+
+
 # vim: fdm=marker

sumpy/expansion/loopy.py

@@ -24,7 +24,7 @@
 """
 
 import logging
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
@@ -109,7 +109,9 @@ def make_e2p_loopy_kernel(
             idx = int(insn.assignee.name[len(result.name):])
             insns[i] = lp.Assignment(
                 assignee=result[idx],
-                expression=result[idx] + insn.expression,
+                expression=(
+                    result[idx]
+                    + cast("pymbolic.ArithmeticExpression", insn.expression)),
                 id=f"result_{idx}",
                 happens_after=insn.happens_after,
             )
@@ -207,7 +209,9 @@ def make_p2e_loopy_kernel(
             idx = int(insn.assignee.name[len(coeffs.name):])
             insns[i] = lp.Assignment(
                 assignee=coeffs[idx],
-                expression=coeffs[idx] + insn.expression,
+                expression=(
+                    coeffs[idx]
+                    + cast("pymbolic.ArithmeticExpression", insn.expression)),
                 id=f"coeff_{idx}",
                 happens_after=insn.happens_after,
             )

sumpy/expansion/multipole.py

@@ -42,6 +42,7 @@
 
 __doc__ = """
 
+.. autoclass:: MultipoleExpansionBase
 .. autoclass:: VolumeTaylorMultipoleExpansion
 .. autoclass:: H2DMultipoleExpansion
 .. autoclass:: Y2DMultipoleExpansion

sumpy/fmm.py

@@ -27,10 +27,21 @@
 
 .. autoclass:: SumpyTreeIndependentDataForWrangler
 .. autoclass:: SumpyExpansionWrangler
-"""
 
+.. autodata:: MultipoleExpansionFactory
+    :noindex:
+.. class:: MultipoleExpansionFactory
+
+    See above.
+.. autodata:: LocalExpansionFactory
+    :noindex:
+.. class:: LocalExpansionFactory
+
+    See above.
+"""
 
-from typing import Protocol, cast
+from collections.abc import Callable, Sequence
+from typing import TYPE_CHECKING, Protocol, TypeAlias, cast
 
 from boxtree.fmm import ExpansionWranglerInterface, TreeIndependentDataForWrangler
 
@@ -52,6 +63,8 @@
     P2EFromSingleBox,
     P2PFromCSR,
 )
+from sumpy.expansion.local import LocalExpansionBase
+from sumpy.expansion.multipole import MultipoleExpansionBase
 from sumpy.tools import (
     AggregateProfilingEvent,
     get_native_event,
@@ -61,6 +74,17 @@
 )
 
 
+if TYPE_CHECKING:
+    from sumpy.kernel import Kernel
+
+
+# parameters here are order, use_rscale
+MultipoleExpansionFactory: TypeAlias = Callable[
+    [int, bool | None], MultipoleExpansionBase]
+LocalExpansionFactory: TypeAlias = Callable[
+    [int, bool | None], LocalExpansionBase]
+
+
 # {{{ tree-independent data for wrangler
 
 class SumpyTreeIndependentDataForWrangler(TreeIndependentDataForWrangler):
@@ -75,11 +99,24 @@ class SumpyTreeIndependentDataForWrangler(TreeIndependentDataForWrangler):
     profiling enabled.
     """
 
-    def __init__(self, cl_context,
-            multipole_expansion_factory,
-            local_expansion_factory,
-            target_kernels, exclude_self=False, use_rscale=None,
-            strength_usage=None, source_kernels=None):
+    cl_context: cl.Context
+    multipole_expansion_factory: MultipoleExpansionFactory
+    local_expansion_factory: LocalExpansionFactory
+    source_kernels: Sequence[Kernel] | None
+    target_kernels: Sequence[Kernel]
+    exclude_self: bool
+    use_rscale: bool | None
+    strength_usage: Sequence[int] | None
+
+    def __init__(self,
+                 cl_context: cl.Context,
+                 multipole_expansion_factory: MultipoleExpansionFactory,
+                 local_expansion_factory: LocalExpansionFactory,
+                 target_kernels: Sequence[Kernel],
+                 exclude_self: bool = False,
+                 use_rscale: bool | None = None,
+                 strength_usage: Sequence[int] | None = None,
+                 source_kernels: Sequence[Kernel] | None = None):
         """
         :arg multipole_expansion_factory: a callable of a single argument (order)
             that returns a multipole expansion.

sumpy/kernel.py

@@ -213,7 +213,7 @@ def replace_base_kernel(self, new_base_kernel: Kernel) -> Kernel:
         """
         return new_base_kernel
 
-    def prepare_loopy_kernel(self, loopy_knl):
+    def prepare_loopy_kernel(self, loopy_knl: lp.TranslationUnit) -> lp.TranslationUnit:
         """Apply some changes (such as registering function
         manglers) to the kernel. Return the new kernel.
         """

sumpy/p2p.py

@@ -252,7 +252,7 @@ def get_kernel(self):
 
         loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
 
-        for knl in self.target_kernels + self.source_kernels:
+        for knl in [*self.target_kernels, *self.source_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
@@ -314,7 +314,7 @@ def get_kernel(self):
 
         loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
 
-        for knl in self.target_kernels + self.source_kernels:
+        for knl in [*self.target_kernels, *self.source_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
@@ -395,7 +395,7 @@ def get_kernel(self):
         loopy_knl = lp.add_dtypes(
                 loopy_knl, {"nsources": np.int32, "ntargets": np.int32})
 
-        for knl in self.target_kernels + self.source_kernels:
+        for knl in [*self.target_kernels, *self.source_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
@@ -657,7 +657,7 @@ def get_kernel(self, max_nsources_in_one_box, max_ntargets_in_one_box,
         loopy_knl = lp.tag_array_axes(loopy_knl, "targets", "sep,C")
         loopy_knl = lp.tag_array_axes(loopy_knl, "sources", "sep,C")
 
-        for knl in self.target_kernels + self.source_kernels:
+        for knl in [*self.target_kernels, *self.source_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl

sumpy/qbx.py

@@ -30,6 +30,7 @@
 import logging
 
 import numpy as np
+from typing_extensions import override
 
 import loopy as lp
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
@@ -145,8 +146,8 @@ def get_loopy_insns_and_result_names(self):
 
         sac.run_global_cse()
 
-        pymbolic_expr_maps = [knl.get_code_transformer() for knl in (
-            self.target_kernels + self.source_kernels)]
+        pymbolic_expr_maps = [knl.get_code_transformer() for knl in [
+            *self.target_kernels, *self.source_kernels]]
 
         from sumpy.codegen import to_loopy_insns
         loopy_insns = to_loopy_insns(
@@ -235,6 +236,7 @@ class LayerPotential(LayerPotentialBase):
     """
 
     @property
+    @override
     def default_name(self):
         return "qbx_apply"
 
@@ -283,7 +285,7 @@ def get_kernel(self):
             lang_version=MOST_RECENT_LANGUAGE_VERSION)
 
         loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
-        for knl in self.target_kernels + self.source_kernels:
+        for knl in [*self.target_kernels, *self.source_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
@@ -360,7 +362,7 @@ def get_kernel(self):
             lang_version=MOST_RECENT_LANGUAGE_VERSION)
 
         loopy_knl = lp.tag_inames(loopy_knl, "idim*:unr")
-        for expn in self.source_kernels + self.target_kernels:
+        for expn in [*self.source_kernels, *self.target_kernels]:
             loopy_knl = expn.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
@@ -443,11 +445,12 @@ def get_kernel(self):
         loopy_knl = lp.add_dtypes(
                 loopy_knl, {"nsources": np.int32, "ntargets": np.int32})
 
-        for knl in self.source_kernels + self.target_kernels:
+        for knl in [*self.source_kernels, *self.target_kernels]:
             loopy_knl = knl.prepare_loopy_kernel(loopy_knl)
 
         return loopy_knl
 
+    @override
     def get_optimized_kernel(self,
             targets_is_obj_array, sources_is_obj_array, centers_is_obj_array):
         loopy_knl = self.get_kernel()

sumpy/tools.py

@@ -49,6 +49,7 @@
 
 if TYPE_CHECKING:
     import numpy
+    from numpy.typing import DTypeLike
 
     import pyopencl
     import pyopencl as cl
@@ -113,6 +114,13 @@
 .. autoclass:: ProfileGetter
 .. autoclass:: AggregateProfilingEvent
 .. autoclass:: MarkerBasedProfilingEvent
+
+References
+----------
+
+.. class:: DTypeLike
+
+    See :data:`numpy.typing.DTypeLike`.
 """
 
 
@@ -291,11 +299,11 @@ class KernelComputation(ABC):
     .. automethod:: get_kernel
     """
 
-    def __init__(self, ctx: Any,
-            target_kernels: list[Kernel],
-            source_kernels: list[Kernel],
-            strength_usage: list[int] | None = None,
-            value_dtypes: list[numpy.dtype[Any]] | None = None,
+    def __init__(self, ctx: cl.Context,
+            target_kernels: Sequence[Kernel],
+            source_kernels: Sequence[Kernel],
+            strength_usage: Sequence[int] | None = None,
+            value_dtypes: Sequence[numpy.dtype[Any]] | numpy.dtype[Any] | None = None,
             name: str | None = None,
             device: Any | None = None) -> None:
         """
@@ -321,7 +329,7 @@ def __init__(self, ctx: Any,
                 else:
                     value_dtypes.append(np.dtype(np.float64))
 
-        if not isinstance(value_dtypes, list | tuple):
+        if not isinstance(value_dtypes, Sequence):
             value_dtypes = [np.dtype(value_dtypes)] * len(target_kernels)
         value_dtypes = [np.dtype(vd) for vd in value_dtypes]
 
@@ -341,14 +349,14 @@ def __init__(self, ctx: Any,
         if device is None:
             device = ctx.devices[0]
 
-        self.context = ctx
-        self.device = device
+        self.context: cl.Context = ctx
+        self.device: cl.Device = device
 
-        self.source_kernels = tuple(source_kernels)
-        self.target_kernels = tuple(target_kernels)
-        self.value_dtypes = value_dtypes
-        self.strength_usage = strength_usage
-        self.strength_count = strength_count
+        self.source_kernels: Sequence[Kernel] = tuple(source_kernels)
+        self.target_kernels: Sequence[Kernel] = tuple(target_kernels)
+        self.value_dtypes: Sequence[np.dtype[Any]] = value_dtypes
+        self.strength_usage: Sequence[int] = strength_usage
+        self.strength_count: int = strength_count
 
         self.name = name or self.default_name
 
@@ -507,7 +515,9 @@ def get_cached_kernel_executor(self, **kwargs) -> lp.ExecutorBase:
         return knl.executor(self.context)
 
     @staticmethod
-    def _allow_redundant_execution_of_knl_scaling(knl):
+    def _allow_redundant_execution_of_knl_scaling(
+                knl: lp.TranslationUnit
+            ) -> lp.TranslationUnit:
         from loopy.match import ObjTagged
         return lp.add_inames_for_unused_hw_axes(
                 knl, within=ObjTagged(ScalingAssignmentTag()))
@@ -734,13 +744,20 @@ def wait(self):
         return self.native_event.wait()
 
 
-def loopy_fft(shape, inverse, complex_dtype, index_dtype=None,
-        name=None):
+def loopy_fft(
+            shape: tuple[int, ...],
+            inverse: bool,
+            complex_dtype: DTypeLike,
+            index_dtype: DTypeLike | None = None,
+            name: str | None = None
+        ):
     from math import pi
 
     from pymbolic import var
     from pymbolic.algorithm import find_factors
 
+    complex_dtype = np.dtype(complex_dtype)
+
     sign = 1 if not inverse else -1
     n = shape[-1]
 
@@ -766,7 +783,7 @@ def loopy_fft(shape, inverse, complex_dtype, index_dtype=None,
     i2 = var("i2")
     i3 = var("i3")
 
-    fixed_parameters = {"const": complex_dtype(sign*(-2j)*pi/n), "n": n}
+    fixed_parameters = {"const": complex_dtype.type(sign*(-2j)*pi/n), "n": int(n)}
 
     index = (*broadcast_dims, i2)
     insns = [