feat(typing): add types to expansion.loopy

Author: alexfikl

sumpy/expansion/loopy.py

@@ -29,7 +29,7 @@
 import numpy as np
 
 import loopy as lp
-import pymbolic
+import pymbolic.primitives as prim
 
 import sumpy.symbolic as sym
 from sumpy.assignment_collection import SymbolicAssignmentCollection
@@ -39,6 +39,8 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
+    from pymbolic.typing import ArithmeticExpression
+
     from sumpy.expansion import ExpansionBase
     from sumpy.kernel import Kernel
 
@@ -49,27 +51,30 @@
 def make_e2p_loopy_kernel(
         expansion: ExpansionBase, kernels: Sequence[Kernel]) -> lp.TranslationUnit:
     """
-    This is a helper function to create a loopy kernel for multipole/local
-    evaluation. This function uses symbolic expressions given by the expansion class,
-    converts them to pymbolic expressions and generates a loopy
-    kernel. Note that the loopy kernel returned has lots of expressions in it and
-    takes a long time. Therefore, this function should be used only as a fallback
-    when there is no "loop-y" kernel to evaluate the expansion.
+    A helper function that creates a :mod:`loopy` kernel for multipole/local evaluation.
+
+    This function uses symbolic expressions given by the expansion class,
+    converts them to :mod:`pymbolic` expressions and generates a :mod:`loopy`
+    kernel. Note that the :mod:`loopy` kernel returned has lots of expressions
+    in it and (likely) takes a long time. Therefore, this function should be
+    used only as a fallback when there is no "loop-y" kernel to evaluate the
+    expansion.
     """
-    dim = expansion.dim
 
-    bvec = sym.make_sym_vector("b", dim)
+    sac = SymbolicAssignmentCollection()
+
+    dim = expansion.dim
     ncoeffs = len(expansion.get_coefficient_identifiers())
 
+    bvec = sym.make_sym_vector("b", dim)
     rscale = sym.Symbol("rscale")
 
-    sac = SymbolicAssignmentCollection()
-
     domains = [
         "{[idim]: 0<=idim<dim}",
         "{[iknl]: 0<=iknl<nresults}",
     ]
-    insns = []
+
+    insns: list[lp.Assignment | lp.CallInstruction] = []
     insns.append(
         lp.Assignment(
             assignee="b[idim]",
@@ -79,15 +84,17 @@ def make_e2p_loopy_kernel(
     target_args = gather_loopy_arguments((expansion, *tuple(kernels)))
 
     coeff_exprs = sym.make_sym_vector("coeffs", ncoeffs)
-    coeff_names = [
+    coeff_names = {
         sac.add_assignment(f"result{i}",
             expansion.evaluate(knl, coeff_exprs, bvec, rscale, sac=sac))
-        for i, knl in enumerate(kernels)]
+        for i, knl in enumerate(kernels)
+    }
 
     sac.run_global_cse()
 
-    code_transformers = [expansion.get_code_transformer()] \
-        + [kernel.get_code_transformer() for kernel in kernels]
+    code_transformers = (
+            [expansion.get_code_transformer()]
+            + [kernel.get_code_transformer() for kernel in kernels])
 
     from sumpy.codegen import to_loopy_insns
     insns += to_loopy_insns(
@@ -98,37 +105,37 @@ def make_e2p_loopy_kernel(
             complex_dtype=np.complex128  # FIXME
             )
 
-    result = pymbolic.var("result")
+    result = prim.Variable("result")
 
     # change result{i} = expr into result[i] += expr
     for i in range(len(insns)):
         insn = insns[i]
-        if isinstance(insn, lp.Assignment) and \
-                isinstance(insn.assignee, pymbolic.var) and \
-                insn.assignee.name.startswith(result.name):
+        if (isinstance(insn, lp.Assignment)
+                and isinstance(insn.assignee, prim.Variable)
+                and insn.assignee.name.startswith(result.name)):
             idx = int(insn.assignee.name[len(result.name):])
             insns[i] = lp.Assignment(
                 assignee=result[idx],
                 expression=(
                     result[idx]
-                    + cast("pymbolic.ArithmeticExpression", insn.expression)),
+                    + cast("ArithmeticExpression", insn.expression)),
                 id=f"result_{idx}",
                 happens_after=insn.happens_after,
             )
 
     loopy_knl = lp.make_function(domains, insns,
             kernel_data=[
-                lp.GlobalArg("result", shape=(len(kernels),), is_input=True,
-                    is_output=True),
-                lp.GlobalArg("coeffs",
-                    shape=(ncoeffs,), is_input=True, is_output=False),
-                lp.GlobalArg("center, target",
-                    shape=(dim,), is_input=True, is_output=False),
+                lp.GlobalArg("result", shape=(len(kernels),),
+                             is_input=True, is_output=True),
+                lp.GlobalArg("coeffs", shape=(ncoeffs,),
+                             is_input=True, is_output=False),
+                lp.GlobalArg("center, target", shape=(dim,),
+                             is_input=True, is_output=False),
                 lp.ValueArg("rscale", is_input=True),
                 lp.ValueArg("itgt", is_input=True),
                 lp.ValueArg("ntargets", is_input=True),
-                lp.GlobalArg("targets",
-                    shape=(dim, "ntargets"), is_input=True, is_output=False),
+                lp.GlobalArg("targets", shape=(dim, "ntargets"),
+                             is_input=True, is_output=False),
                 *target_args,
                 ...],
             name="e2p",
@@ -144,29 +151,33 @@ def make_e2p_loopy_kernel(
 
 
 def make_p2e_loopy_kernel(
-        expansion: ExpansionBase, kernels: Sequence[Kernel],
-        strength_usage: Sequence[int], nstrengths: int) -> lp.TranslationUnit:
+        expansion: ExpansionBase,
+        kernels: Sequence[Kernel],
+        strength_usage: Sequence[int],
+        nstrengths: int) -> lp.TranslationUnit:
     """
-    This is a helper function to create a loopy kernel for multipole/local
-    expression. This function uses symbolic expressions given by the expansion
-    class, converts them to pymbolic expressions and generates a loopy
-    kernel. Note that the loopy kernel returned has lots of expressions in it and
-    takes a long time. Therefore, this function should be used only as a fallback
-    when there is no "loop-y" kernel to evaluate the expansion.
+    A helper function that creates a :mod:`loopy` kernel for multipole/local evaluation.
+
+    This function uses symbolic expressions given by the expansion class,
+    converts them to :mod:`pymbolic` expressions and generates a :mod:`loopy`
+    kernel. Note that the :mod:`loopy` kernel returned has lots of expressions
+    in it and (likely) takes a long time. Therefore, this function should be
+    used only as a fallback when there is no "loop-y" kernel to evaluate the
+    expansion.
     """
-    dim = expansion.dim
+    sac = SymbolicAssignmentCollection()
 
-    avec = sym.make_sym_vector("a", dim)
+    dim = expansion.dim
     ncoeffs = len(expansion.get_coefficient_identifiers())
 
+    avec = sym.make_sym_vector("a", dim)
     rscale = sym.Symbol("rscale")
 
-    sac = SymbolicAssignmentCollection()
-
     domains = [
         "{[idim]: 0<=idim<dim}",
     ]
-    insns = []
+
+    insns: list[lp.Assignment | lp.CallInstruction] = []
     insns.append(
         lp.Assignment(
             assignee="a[idim]",
@@ -179,15 +190,15 @@ def make_p2e_loopy_kernel(
     strengths = [all_strengths[i] for i in strength_usage]
     coeffs = expansion.coefficients_from_source_vec(kernels,
         avec, None, rscale, strengths, sac=sac)
-
-    coeff_names = [
+    coeff_names = {
         sac.add_assignment(f"coeffs{i}", coeff) for i, coeff in enumerate(coeffs)
-    ]
+    }
 
     sac.run_global_cse()
 
-    code_transformers = [expansion.get_code_transformer()] \
-        + [kernel.get_code_transformer() for kernel in kernels]
+    code_transformers = (
+        [expansion.get_code_transformer()]
+        + [kernel.get_code_transformer() for kernel in kernels])
 
     from sumpy.codegen import to_loopy_insns
     insns += to_loopy_insns(
@@ -198,37 +209,38 @@ def make_p2e_loopy_kernel(
             complex_dtype=np.complex128  # FIXME
             )
 
-    coeffs = pymbolic.var("coeffs")
+    coeffs = prim.Variable("coeffs")
 
     # change coeff{i} = expr into coeff[i] += expr
     for i in range(len(insns)):
         insn = insns[i]
-        if isinstance(insn, lp.Assignment) and \
-                isinstance(insn.assignee, pymbolic.var) and \
-                insn.assignee.name.startswith(coeffs.name):
+        if (isinstance(insn, lp.Assignment)
+                and isinstance(insn.assignee, prim.Variable)
+                and insn.assignee.name.startswith(coeffs.name)):
             idx = int(insn.assignee.name[len(coeffs.name):])
+
             insns[i] = lp.Assignment(
                 assignee=coeffs[idx],
                 expression=(
                     coeffs[idx]
-                    + cast("pymbolic.ArithmeticExpression", insn.expression)),
+                    + cast("ArithmeticExpression", insn.expression)),
                 id=f"coeff_{idx}",
                 happens_after=insn.happens_after,
             )
 
     loopy_knl = lp.make_function(domains, insns,
             kernel_data=[
-                lp.GlobalArg("coeffs",
-                    shape=(ncoeffs,), is_input=True, is_output=True),
-                lp.GlobalArg("center, source",
-                    shape=(dim,), is_input=True, is_output=False),
-                lp.GlobalArg("strength",
-                    shape=(nstrengths,), is_input=True, is_output=False),
+                lp.GlobalArg("coeffs", shape=(ncoeffs,),
+                             is_input=True, is_output=True),
+                lp.GlobalArg("center, source", shape=(dim,),
+                             is_input=True, is_output=False),
+                lp.GlobalArg("strength", shape=(nstrengths,),
+                             is_input=True, is_output=False),
                 lp.ValueArg("rscale", is_input=True),
                 lp.ValueArg("isrc", is_input=True),
                 lp.ValueArg("nsources", is_input=True),
-                lp.GlobalArg("sources",
-                    shape=(dim, "nsources"), is_input=True, is_output=False),
+                lp.GlobalArg("sources", shape=(dim, "nsources"),
+                             is_input=True, is_output=False),
                 *source_args,
                 ...],
             name="p2e",