Fourier/Hankel/LineTaylor: Use tuples as coeff identifiers

Author: inducer

sumpy/expansion/local.py

@@ -48,6 +48,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 (
         HankelBased2DMultipoleExpansion,
@@ -118,12 +119,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,
@@ -165,15 +168,15 @@ def coefficients_from_source(self,
             return [kernel.postprocess_at_source(
                             line_kernel.diff(tau, i), avec)
                     .subs(tau, 0)
-                    for i in self.get_coefficient_identifiers()]
+                    for i, in self.get_coefficient_identifiers()]
 
     def evaluate(self, tgt_kernel, coeffs, bvec, rscale, sac=None):
         # no point in heeding rscale here--just ignore it
 
         # 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
@@ -557,11 +560,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,
@@ -585,7 +591,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,
@@ -604,12 +610,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,
@@ -634,14 +640,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

@@ -892,9 +892,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,
@@ -930,10 +930,10 @@ def translation_classes_dependent_data(self,
         # evaluations. https://dlmf.nist.gov/10.23.F1
         # This loop computes the first row and the last column vector sufficient
         # to specify the matrix entries.
-        for j in tgt_expansion.get_coefficient_identifiers():
-            idx_j = tgt_expansion.get_storage_index(j)
-            for m in src_expansion.get_coefficient_identifiers():
-                idx_m = src_expansion.get_storage_index(m)
+        for j, in tgt_expansion.get_coefficient_identifiers():
+            idx_j = tgt_expansion.get_storage_index((j,))
+            for m, in src_expansion.get_coefficient_identifiers():
+                idx_m = src_expansion.get_storage_index((m,))
                 translation_classes_dependent_data[idx_j + idx_m] = (
                     Hankel1(m + j, arg_scale * dvec_len, 0)
                     * sym.exp(sym.I * (m + j) * new_center_angle_rel_old_center))
@@ -944,8 +944,8 @@ def preprocess_multipole_exprs(self, tgt_expansion, src_expansion,
             src_coeff_exprs, sac, src_rscale):
 
         src_coeff_exprs = list(src_coeff_exprs)
-        for m in src_expansion.get_coefficient_identifiers():
-            src_coeff_exprs[src_expansion.get_storage_index(m)] *= src_rscale**abs(m)
+        for m, in src_expansion.get_coefficient_identifiers():
+            src_coeff_exprs[src_expansion.get_storage_index((m,))] *= src_rscale**abs(m)
         return src_coeff_exprs
 
     def preprocess_multipole_nexprs(self, tgt_expansion, src_expansion):
@@ -956,8 +956,8 @@ def postprocess_local_exprs(self, tgt_expansion, src_expansion,
 
         # Filter out the dummy rows and scale them for target
         result = []
-        for j in tgt_expansion.get_coefficient_identifiers():
-            result.append(m2l_result[tgt_expansion.get_storage_index(j)]
+        for j, in tgt_expansion.get_coefficient_identifiers():
+            result.append(m2l_result[tgt_expansion.get_storage_index((j,))]
                     * tgt_rscale**(abs(j)) * sym.Integer(-1)**j)
 
         return result

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