Fourier/Hankel/LineTaylor: Use tuples as coeff identifiers

Author: inducer

sumpy/expansion/local.py

@@ -47,6 +47,7 @@
     from collections.abc import Sequence
 
     from sumpy.assignment_collection import SymbolicAssignmentCollection
+    from sumpy.expansion.diff_op import MultiIndex
     from sumpy.expansion.m2l import M2LTranslationBase
     from sumpy.expansion.multipole import (
         MultipoleExpansionBase,
@@ -117,12 +118,14 @@ def m2l_translation(self):
         # FIXME: Um...
         raise NotImplementedError()
 
-    def get_storage_index(self, k):
-        return k
-
+    @override
+    def get_storage_index(self, mi: MultiIndex):
+        ind, = mi
+        return ind
 
+    @override
     def get_coefficient_identifiers(self):
-        return list(range(self.order+1))
+        return [(i,) for i in range(self.order+1)]
 
     @override
     def coefficients_from_source(self,
@@ -172,7 +175,7 @@ def evaluate(self, tgt_kernel, coeffs, bvec, rscale, sac=None):
         # NOTE: We can't meaningfully apply target derivatives here.
         # Instead, this is handled in LayerPotentialBase._evaluate.
         return sym.Add(*(
-                    coeffs[self.get_storage_index(i)] / math.factorial(i)
+                    coeffs[self.get_storage_index(i)] / math.factorial(i[0])
                     for i in self.get_coefficient_identifiers()))
 
     @override
@@ -551,11 +554,14 @@ def m2l_translation(self) -> M2LTranslationBase:
     def get_bessel_arg_scaling(self):
         pass
 
-    def get_storage_index(self, k):
-        return self.order+k
+    @override
+    def get_storage_index(self, mi: MultiIndex):
+        ind, = mi
+        return self.order+ind
 
+    @override
     def get_coefficient_identifiers(self):
-        return list(range(-self.order, self.order+1))
+        return [(i,) for i in range(-self.order, self.order+1)]
 
     @override
     def coefficients_from_source(self,
@@ -579,7 +585,7 @@ def coefficients_from_source(self,
                     Hankel1(c, arg_scale * avec_len, 0)
                     * rscale ** abs(c)
                     * sym.exp(sym.I * c * source_angle_rel_center), avec)
-                    for c in self.get_coefficient_identifiers()]
+                    for c, in self.get_coefficient_identifiers()]
 
     @override
     def evaluate(self,
@@ -598,12 +604,12 @@ def evaluate(self,
 
         arg_scale = self.get_bessel_arg_scaling()
 
-        return sym.sympify(sum(coeffs[self.get_storage_index(c)]
+        return sym.sympify(sum(coeffs[self.get_storage_index((c,))]
                    * kernel.postprocess_at_target(
                        BesselJ(c, arg_scale * bvec_len, 0)
                        / rscale ** abs(c)
                        * sym.exp(sym.I * c * -target_angle_rel_center), bvec)
-                for c in self.get_coefficient_identifiers()))
+                for c, in self.get_coefficient_identifiers()))
 
     @override
     def translate_from(self,
@@ -628,14 +634,14 @@ def translate_from(self,
             new_center_angle_rel_old_center = sym.atan2(dvec[1], dvec[0])
             translated_coeffs = []
 
-            for j in self.get_coefficient_identifiers():
+            for j, in self.get_coefficient_identifiers():
                 translated_coeffs.append(
-                    sum(src_coeff_exprs[src_expansion.get_storage_index(m)]
+                    sum(src_coeff_exprs[src_expansion.get_storage_index((m,))]
                         * BesselJ(m - j, arg_scale * dvec_len, 0)
                         / src_rscale ** abs(m)
                         * tgt_rscale ** abs(j)
                         * sym.exp(sym.I * (m - j) * -new_center_angle_rel_old_center)
-                    for m in src_expansion.get_coefficient_identifiers()))
+                    for m, in src_expansion.get_coefficient_identifiers()))
             return translated_coeffs
 
         if isinstance(src_expansion, self.mpole_expn_class):

sumpy/expansion/m2l.py

@@ -850,9 +850,9 @@ def translate(self, tgt_expansion, src_expansion, src_coeff_exprs, src_rscale,
 
         translated_coeffs = [
             sum(derivatives[m + j + tgt_expansion.order + src_expansion.order]
-                    * src_coeff_exprs[src_expansion.get_storage_index(m)]
-                for m in src_expansion.get_coefficient_identifiers())
-            for j in tgt_expansion.get_coefficient_identifiers()]
+                    * src_coeff_exprs[src_expansion.get_storage_index((m,))]
+                for m, in src_expansion.get_coefficient_identifiers())
+            for j, in tgt_expansion.get_coefficient_identifiers()]
 
         translated_coeffs = self.postprocess_local_exprs(tgt_expansion,
                 src_expansion, translated_coeffs, src_rscale, tgt_rscale,

sumpy/expansion/multipole.py

@@ -44,6 +44,7 @@
     from collections.abc import Sequence
 
     from sumpy.assignment_collection import SymbolicAssignmentCollection
+    from sumpy.expansion.diff_op import MultiIndex
     from sumpy.kernel import Kernel
 
 
@@ -405,11 +406,14 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase):
     def get_bessel_arg_scaling(self):
         return
 
-    def get_storage_index(self, k):
-        return self.order+k
+    @override
+    def get_storage_index(self, mi: MultiIndex):
+        ind, = mi
+        return self.order+ind
 
+    @override
     def get_coefficient_identifiers(self):
-        return list(range(-self.order, self.order+1))
+        return [(i,) for i in range(-self.order, self.order+1)]
 
     def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None):
         if not self.use_rscale:
@@ -431,7 +435,7 @@ def coefficients_from_source(self, kernel, avec, bvec, rscale, sac=None):
                     / rscale ** abs(c)
                     * sym.exp(sym.I * c * -source_angle_rel_center),
                     avec)
-                for c in self.get_coefficient_identifiers()]
+                for c, in self.get_coefficient_identifiers()]
 
     def evaluate(self, kernel, coeffs, bvec, rscale, sac=None):
         if not self.use_rscale:
@@ -443,12 +447,12 @@ def evaluate(self, kernel, coeffs, bvec, rscale, sac=None):
 
         arg_scale = self.get_bessel_arg_scaling()
 
-        return sum(coeffs[self.get_storage_index(c)]
+        return sum(coeffs[self.get_storage_index((c,))]
                    * kernel.postprocess_at_target(
                        Hankel1(c, arg_scale * bvec_len, 0)
                        * rscale ** abs(c)
                        * sym.exp(sym.I * c * target_angle_rel_center), bvec)
-                for c in self.get_coefficient_identifiers())
+                for c, in self.get_coefficient_identifiers())
 
     def translate_from(self, src_expansion, src_coeff_exprs, src_rscale,
             dvec, tgt_rscale, sac=None):
@@ -468,14 +472,14 @@ def translate_from(self, src_expansion, src_coeff_exprs, src_rscale,
         arg_scale = self.get_bessel_arg_scaling()
 
         translated_coeffs = []
-        for j in self.get_coefficient_identifiers():
+        for j, in self.get_coefficient_identifiers():
             translated_coeffs.append(
-                sum(src_coeff_exprs[src_expansion.get_storage_index(m)]
+                sum(src_coeff_exprs[src_expansion.get_storage_index((m,))]
                     * BesselJ(m - j, arg_scale * dvec_len, 0)
                     * src_rscale ** abs(m)
                     / tgt_rscale ** abs(j)
                     * sym.exp(sym.I * (m - j) * new_center_angle_rel_old_center)
-                for m in src_expansion.get_coefficient_identifiers()))
+                for m, in src_expansion.get_coefficient_identifiers()))
         return translated_coeffs