Allow multiple kernels in p2e

They are aggregated using a strength expression

sumpy/expansion/local.py

@@ -55,18 +55,21 @@ class LineTaylorLocalExpansion(LocalExpansionBase):
     def get_coefficient_identifiers(self):
         return list(range(self.order+1))
 
-    def coefficients_from_source(self, avec, bvec, rscale, sac=None):
+    def coefficients_from_source(self, avec, bvec, rscale, sac=None, kernel=None):
         # no point in heeding rscale here--just ignore it
         if bvec is None:
             raise RuntimeError("cannot use line-Taylor expansions in a setting "
                     "where the center-target vector is not known at coefficient "
                     "formation")
 
+        if kernel is None:
+            kernel = self.kernel
+
         tau = sym.Symbol("tau")
 
         avec_line = avec + tau*bvec
 
-        line_kernel = self.kernel.get_expression(avec_line)
+        line_kernel = kernel.get_expression(avec_line)
 
         from sumpy.symbolic import USE_SYMENGINE
 
@@ -75,8 +78,8 @@ class LineTaylorLocalExpansion(LocalExpansionBase):
             deriv_taker = MiDerivativeTaker(line_kernel, (tau,))
 
             return [my_syntactic_subs(
-                        self.kernel.postprocess_at_target(
-                            self.kernel.postprocess_at_source(
+                        kernel.postprocess_at_target(
+                            kernel.postprocess_at_source(
                                 deriv_taker.diff(i),
                                 avec), bvec),
                         {tau: 0})
@@ -89,8 +92,8 @@ class LineTaylorLocalExpansion(LocalExpansionBase):
             #
             # See also https://gitlab.tiker.net/inducer/pytential/merge_requests/12
 
-            return [self.kernel.postprocess_at_target(
-                        self.kernel.postprocess_at_source(
+            return [kernel.postprocess_at_target(
+                        kernel.postprocess_at_source(
                             line_kernel.diff("tau", i), avec),
                         bvec)
                     .subs("tau", 0)
@@ -113,10 +116,11 @@ class VolumeTaylorLocalExpansionBase(LocalExpansionBase):
     Coefficients represent derivative values of the kernel.
     """
 
-    def coefficients_from_source(self, avec, bvec, rscale, sac=None):
+    def coefficients_from_source(self, avec, bvec, rscale, sac=None, kernel=None):
         from sumpy.tools import MiDerivativeTaker
-        ppkernel = self.kernel.postprocess_at_source(
-                self.kernel.get_expression(avec), avec)
+        if kernel is None:
+            kernel = self.kernel
+        ppkernel = kernel.postprocess_at_source(kernel.get_expression(avec), avec)
 
         taker = MiDerivativeTaker(ppkernel, avec)
         return [
@@ -372,9 +376,11 @@ class _FourierBesselLocalExpansion(LocalExpansionBase):
     def get_coefficient_identifiers(self):
         return list(range(-self.order, self.order+1))
 
-    def coefficients_from_source(self, avec, bvec, rscale, sac=None):
+    def coefficients_from_source(self, avec, bvec, rscale, sac=None, kernel=None):
         if not self.use_rscale:
             rscale = 1
+        if kernel is None:
+            kernel = self.kernel
 
         from sumpy.symbolic import sym_real_norm_2
         hankel_1 = sym.Function("hankel_1")
@@ -384,7 +390,7 @@ class _FourierBesselLocalExpansion(LocalExpansionBase):
         # The coordinates are negated since avec points from source to center.
         source_angle_rel_center = sym.atan2(-avec[1], -avec[0])
         avec_len = sym_real_norm_2(avec)
-        return [self.kernel.postprocess_at_source(
+        return [kernel.postprocess_at_source(
                     hankel_1(c, arg_scale * avec_len)
                     * rscale ** abs(c)
                     * sym.exp(sym.I * c * source_angle_rel_center), avec)

sumpy/expansion/multipole.py

@@ -53,9 +53,10 @@ class VolumeTaylorMultipoleExpansionBase(MultipoleExpansionBase):
     Coefficients represent the terms in front of the kernel derivatives.
     """
 
-    def coefficients_from_source(self, avec, bvec, rscale, sac=None):
+    def coefficients_from_source(self, avec, bvec, rscale, sac=None, kernel=None):
         from sumpy.kernel import DirectionalSourceDerivative
-        kernel = self.kernel
+        if kernel is None:
+            kernel = self.kernel
 
         from sumpy.tools import mi_power, mi_factorial
 
@@ -284,10 +285,13 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase):
     def get_coefficient_identifiers(self):
         return list(range(-self.order, self.order+1))
 
-    def coefficients_from_source(self, avec, bvec, rscale, sac=None):
+    def coefficients_from_source(self, avec, bvec, rscale, sac=None, kernel=None):
         if not self.use_rscale:
             rscale = 1
 
+        if kernel is None:
+            kernel = self.kernel
+
         from sumpy.symbolic import sym_real_norm_2
         bessel_j = sym.Function("bessel_j")
         avec_len = sym_real_norm_2(avec)
@@ -297,7 +301,7 @@ class _HankelBased2DMultipoleExpansion(MultipoleExpansionBase):
         # The coordinates are negated since avec points from source to center.
         source_angle_rel_center = sym.atan2(-avec[1], -avec[0])
         return [
-                self.kernel.postprocess_at_source(
+                kernel.postprocess_at_source(
                     bessel_j(c, arg_scale * avec_len)
                     / rscale ** abs(c)
                     * sym.exp(sym.I * c * -source_angle_rel_center),

sumpy/p2e.py

@@ -23,6 +23,8 @@ THE SOFTWARE.
 import numpy as np
 import loopy as lp
 from loopy.version import MOST_RECENT_LANGUAGE_VERSION
+import pymbolic
+from pymbolic.mapper import WalkMapper
 
 from sumpy.tools import KernelCacheWrapper
 
@@ -45,8 +47,8 @@ Particle-to-expansion
 # {{{ P2E base class
 
 class P2EBase(KernelCacheWrapper):
-    def __init__(self, ctx, expansion,
-            options=[], name=None, device=None):
+    def __init__(self, ctx, expansion, kernels=None,
+            options=[], name=None, device=None, strength_expr=None, nstrengths=1):
         """
         :arg expansion: a subclass of :class:`sympy.expansion.ExpansionBase`
         :arg strength_usage: A list of integers indicating which expression
@@ -58,6 +60,13 @@ class P2EBase(KernelCacheWrapper):
         if device is None:
             device = ctx.devices[0]
 
+        if kernels is None:
+            kernels = [expansion.kernel]
+
+        if strength_expr is None:
+            import pymbolic
+            strength_expr = pymbolic.parse("strength0 * kernel0")
+
         from sumpy.kernel import TargetDerivativeRemover
         expansion = expansion.with_kernel(
                 TargetDerivativeRemover()(expansion.kernel))
@@ -67,9 +76,20 @@ class P2EBase(KernelCacheWrapper):
         self.options = options
         self.name = name or self.default_name
         self.device = device
+        self.kernels = kernels
+        self.strength_expr, self.extra_source_variables = \
+            process_strength_expr(strength_expr, nstrengths, len(kernels))
+        self.nstrengths = nstrengths
 
         self.dim = expansion.dim
 
+    def get_result_expr(self, icoeff):
+        subst_dict = dict((
+            pymbolic.var(f"kernel{iknl}"), pymbolic.var(f"coeff{icoeff}_{iknl}")
+            ) for iknl in range(len(self.kernels)))
+        res = pymbolic.substitute(self.strength_expr, subst_dict)
+        return res
+
     def get_loopy_instructions(self):
         from sumpy.symbolic import make_sym_vector
         avec = make_sym_vector("a", self.dim)
@@ -80,10 +100,15 @@ class P2EBase(KernelCacheWrapper):
         from sumpy.assignment_collection import SymbolicAssignmentCollection
         sac = SymbolicAssignmentCollection()
 
-        coeff_names = [
-            sac.assign_unique("coeff%d" % i, coeff_i)
+        coeff_names = []
+        for knl_idx, kernel in enumerate(self.kernels):
             for i, coeff_i in enumerate(
-                self.expansion.coefficients_from_source(avec, None, rscale, sac))]
+                self.expansion.coefficients_from_source(avec, None, rscale,
+                     sac, kernel=kernel)
+            ):
+                coeff_names.append(
+                    sac.assign_unique(f"coeff{i}_{knl_idx}", coeff_i)
+                )
 
         sac.run_global_cse()
 
@@ -127,20 +152,21 @@ class P2EFromSingleBox(P2EBase):
                     for isrc
                         <> a[idim] = center[idim] - sources[idim, isrc] {dup=idim}
 
-                        <> strength = strengths[isrc]
+                        <> strength = strengths[0, isrc]
                         """] + self.get_loopy_instructions() + ["""
                     end
-                    """] + ["""
+                    """] + [f"""
                     tgt_expansions[src_ibox-tgt_base_ibox, {coeffidx}] = \
-                            simul_reduce(sum, isrc, strength*coeff{coeffidx}) \
+                        simul_reduce(sum, isrc, {self.get_result_expr(coeffidx)}) \
                             {{id_prefix=write_expn}}
-                    """.format(coeffidx=i) for i in range(ncoeffs)] + ["""
+                    """ for coeffidx in range(ncoeffs)] + ["""
                 end
                 """],
                 [
                     lp.GlobalArg("sources", None, shape=(self.dim, "nsources"),
                         dim_tags="sep,c"),
-                    lp.GlobalArg("strengths", None, shape="nsources"),
+                    lp.GlobalArg("strengths", None, shape="nstrengths, nsources",
+                        dim_tags="sep,C"),
                     lp.GlobalArg("box_source_starts,box_source_counts_nonchild",
                         None, shape=None),
                     lp.GlobalArg("centers", None, shape="dim, aligned_nboxes"),
@@ -155,7 +181,7 @@ class P2EFromSingleBox(P2EBase):
                 assumptions="nsrc_boxes>=1",
                 silenced_warnings="write_race(write_expn*)",
                 default_offset=lp.auto,
-                fixed_parameters=dict(dim=self.dim),
+                fixed_parameters=dict(dim=self.dim, nstrengths=self.nstrengths),
                 lang_version=MOST_RECENT_LANGUAGE_VERSION)
 
         loopy_knl = self.expansion.prepare_loopy_kernel(loopy_knl)
@@ -206,7 +232,8 @@ class P2EFromCSR(P2EBase):
                 [
                     lp.GlobalArg("sources", None, shape=(self.dim, "nsources"),
                         dim_tags="sep,c"),
-                    lp.GlobalArg("strengths", None, shape="nsources"),
+                    lp.GlobalArg("strengths", None, shape="nstrengths, nsources",
+                        dim_tags="sep,C"),
                     lp.GlobalArg("source_box_starts,source_box_lists",
                         None, shape=None, offset=lp.auto),
                     lp.GlobalArg("box_source_starts,box_source_counts_nonchild",
@@ -245,14 +272,14 @@ class P2EFromCSR(P2EBase):
                             <> a[idim] = center[idim] - sources[idim, isrc] \
                                     {dup=idim}
 
-                            <> strength = strengths[isrc]
+                            <> strength = strengths[0, isrc]
                             """] + self.get_loopy_instructions() + ["""
                         end
                     end
                     """] + ["""
                     tgt_expansions[tgt_ibox - tgt_base_ibox, {coeffidx}] = \
                             simul_reduce(sum, (isrc_box, isrc),
-                                strength*coeff{coeffidx}) \
+                                strength*coeff{coeffidx}_0) \
                             {{id_prefix=write_expn}}
                     """.format(coeffidx=i) for i in range(ncoeffs)] + ["""
                 end
@@ -262,7 +289,7 @@ class P2EFromCSR(P2EBase):
                 assumptions="ntgt_boxes>=1",
                 silenced_warnings="write_race(write_expn*)",
                 default_offset=lp.auto,
-                fixed_parameters=dict(dim=self.dim),
+                fixed_parameters=dict(dim=self.dim, nstrengths=self.nstrengths),
                 lang_version=MOST_RECENT_LANGUAGE_VERSION)
 
         loopy_knl = self.expansion.prepare_loopy_kernel(loopy_knl)
@@ -298,4 +325,39 @@ class P2EFromCSR(P2EBase):
 
 # }}}
 
+
+# {{{ helper functions
+
+class CollectVariableMapper(WalkMapper):
+    def __init__(self):
+        self.variables = set()
+
+    def post_visit(self, expr, *args, **kwargs):
+        if isinstance(expr, pymbolic.primitives.Variable):
+            self.variables.add(expr)
+
+
+def process_strength_expr(expr, nstrengths, nkernels):
+    collect_variable_mapper = CollectVariableMapper()
+    collect_variable_mapper(expr)
+    variables = collect_variable_mapper.variables
+
+    # Get variables that are not strengths nor kernels
+    source_variables = [var for var in variables if
+        not (var.name.startswith("strength") or var.name.startswith("kernel"))]
+
+    # Use strengths array and index with isrc
+    subst = {}
+    for i in range(nstrengths):
+        old = pymbolic.var(f"strength{i}")
+        new = pymbolic.var("strengths")[0, pymbolic.var("isrc")]
+        subst[old] = new
+    print(subst)
+    expr = pymbolic.substitute(expr, subst)
+
+    return expr, source_variables
+
+
+#}}}
+
 # vim: foldmethod=marker

test/test_kernels.py

@@ -229,7 +229,7 @@ def test_p2e2p(ctx_factory, base_knl, expn_class, order, with_source_derivative)
                 box_source_counts_nonchild=box_source_counts_nonchild,
                 centers=centers,
                 sources=sources,
-                strengths=strengths,
+                strengths=(strengths,),
                 nboxes=1,
                 tgt_base_ibox=0,
                 rscale=rscale,
@@ -495,7 +495,7 @@ def test_translations(ctx_factory, knl, local_expn_class, mpole_expn_class):
                 box_source_counts_nonchild=p2m_box_source_counts_nonchild,
                 centers=centers,
                 sources=sources,
-                strengths=strengths,
+                strengths=(strengths,),
                 nboxes=nboxes,
                 rscale=m1_rscale,