Add get_storage_index for the wranglers

Author: isuruf

sumpy/expansion/__init__.py

@@ -30,6 +30,7 @@
 import sumpy.symbolic as sym
 from sumpy.kernel import Kernel
 from sumpy.tools import add_mi
+import pymbolic.primitives as prim
 
 import logging
 logger = logging.getLogger(__name__)
@@ -377,6 +378,18 @@ def _split_coeffs_into_hyperplanes(
 
 
 class FullExpansionTermsWrangler(ExpansionTermsWrangler):
+
+    def get_storage_index(self, mi, order=None):
+        if not order:
+            order = sum(mi)
+        if self.dim == 3:
+            return (order*(order + 1)*(order + 2))//6 + \
+                    (order + 2)*mi[2] - (mi[2]*(mi[2] + 1))//2 + mi[1]
+        elif self.dim == 2:
+            return (order*(order + 1))//2 + mi[1]
+        else:
+            raise NotImplementedError
+
     def get_coefficient_identifiers(self):
         return super().get_full_coefficient_identifiers()
 
@@ -584,6 +597,51 @@ def get_full_coefficient_identifiers(self):
         key, _ = self._get_mi_ordering_key_and_axis_permutation()
         return sorted(identifiers, key=key)
 
+    def get_storage_index(self, mi, order=None):
+        if not order:
+            order = sum(mi)
+
+        ordering_key, axis_permutation = \
+                self._get_mi_ordering_key_and_axis_permutation()
+        deriv_id_to_coeff, = self.knl.get_pde_as_diff_op().eqs
+        max_mi = max(deriv_id_to_coeff, key=ordering_key).mi
+
+        if all(m != 0 for m in max_mi):
+            raise NotImplementedError("non-elliptic PDEs")
+
+        c = max_mi[axis_permutation[0]]
+
+        mi = list(mi)
+        mi[axis_permutation[0]], mi[0] = mi[0], mi[axis_permutation[0]]
+
+        if self.dim == 3:
+            if all(isinstance(axis, int) for axis in mi):
+                if order < c - 1:
+                    return (order*(order + 1)*(order + 2))/6 + \
+                        (order + 2)*mi[0] - (mi[0]*(mi[0] + 1))/2 + mi[1]
+                else:
+                    return (c*(c-1)*(c-2))/6 + (c * order * (2 + order - c)
+                        + mi[0]*(3 - mi[0]+2*order))/2 + mi[1]
+            else:
+                return prim.If(prim.Comparison(order, "<", c - 1),
+                    (order*(order + 1)*(order + 2))/6
+                        + (order + 2)*mi[0] - (mi[0]*(mi[0] + 1))/2 + mi[1],
+                    (c*(c-1)*(c-2))/6 + (c * order * (2 + order - c)
+                        + mi[0]*(3 - mi[0]+2*order))/2 + mi[1]
+                )
+        elif self.dim == 2:
+            if all(isinstance(axis, int) for axis in mi):
+                if order < c - 1:
+                    return (order*(order + 1))//2 + mi[0]
+                else:
+                    return (c*(c-1))//2 + c*(order - c + 1) + mi[0]
+            else:
+                return prim.If(prim.Comparison(order, "<", c - 1),
+                    (order*(order + 1))//2 + mi[0],
+                    (c*(c-1))//2 + c*(order - c + 1) + mi[0])
+        else:
+            raise NotImplementedError
+
     @memoize_method
     def get_stored_ids_and_unscaled_projection_matrix(self):
         from pytools import ProcessLogger

test/test_misc.py

@@ -49,6 +49,9 @@
     make_identity_diff_op, concat, as_scalar_pde, diff,
     gradient, divergence, laplacian, curl)
 
+from sumpy.expansion import (FullExpansionTermsWrangler,
+    LinearPDEBasedExpansionTermsWrangler)
+
 import logging
 logger = logging.getLogger(__name__)
 
@@ -516,6 +519,45 @@ def get_pde_as_diff_op(self):
 # }}}
 
 
+# {{{ test_get_storage_index
+
+class TestKernel(ExpressionKernel):
+    def __init__(self, dim, max_mi):
+        super().__init__(dim=dim, expression=1, global_scaling_const=1,
+                is_complex_valued=False)
+        self._max_mi = max_mi
+
+    def get_pde_as_diff_op(self):
+        w = make_identity_diff_op(self.dim)
+        pde = diff(w, tuple(self._max_mi))
+        return pde
+
+
+@pytest.mark.parametrize("order", [6])
+@pytest.mark.parametrize("knl", [
+    LaplaceKernel(2),
+    LaplaceKernel(3),
+    TestKernel(2, (3, 0)),
+    TestKernel(2, (0, 3)),
+    TestKernel(3, (3, 0, 0)),
+    TestKernel(3, (0, 3, 0)),
+    TestKernel(3, (0, 0, 3)),
+    BiharmonicKernel(2),
+    BiharmonicKernel(3),
+])
+@pytest.mark.parametrize("compressed", (True, False))
+def test_get_storage_index(order, knl, compressed):
+    dim = knl.dim
+    if compressed:
+        wrangler = LinearPDEBasedExpansionTermsWrangler(order, dim, knl=knl)
+    else:
+        wrangler = FullExpansionTermsWrangler(order, dim)
+    for i, mi in enumerate(wrangler.get_coefficient_identifiers()):
+        assert i == wrangler.get_storage_index(mi)
+
+# }}}
+
+
 # You can test individual routines by typing
 # $ python test_misc.py 'test_pde_check_kernels(_acf,
 #       KernelInfo(HelmholtzKernel(2), k=5), order=5)'