rewrite at codegen time

Author: isuruf

sumpy/codegen.py

@@ -592,6 +592,39 @@ def map_sum(self, expr, *args):
 # }}}
 
 
+# {{{ helmholtz rewrite
+class HelmholtzRewriter(CSECachingIdentityMapper, CallExternalRecMapper):
+    def __init__(self, k, ik):
+        self.k = k
+        self.ik = ik
+
+    def map_variable(self, expr, *args):
+        if expr.name == self.ik.name:
+            return 1j*self.k
+        else:
+            return expr
+
+    def map_call(self, expr, *args):
+        if isinstance(expr.function, prim.Variable) \
+                and expr.function.name == "exp":
+            params = expr.parameters
+            assert len(params) == 1
+            param = self.rec(params[0])
+            if isinstance(param, prim.Product) and 1j in param.children:
+                children = list(param.children)
+                del children[children.index(1j)]
+                params = (prim.Product(tuple(children)),)
+                return prim.Call(prim.Variable("cos"), params) + \
+                        1j * prim.Call(prim.Variable("sin"), params)
+
+        return super().map_call(expr, *args)
+
+    map_common_subexpression_uncached = IdentityMapper.map_common_subexpression
+
+
+# }}}
+
+
 class MathConstantRewriter(CSECachingIdentityMapper, CallExternalRecMapper):
     def map_variable(self, expr, *args):
         if expr.name == "pi":

sumpy/kernel.py

@@ -522,7 +522,7 @@ def __init__(self, dim, helmholtz_k_name="k",
             if allow_evanescent:
                 expr = var("exp")(var("I")*k*r)/r
             else:
-                expr = (var("cos")(k*r) + var("I")*var("sin")(k*r))/r
+                expr = var("exp")(var("Ik")*r)/r
             scaling = 1/(4*var("pi"))
         else:
             raise RuntimeError("unsupported dimensionality")
@@ -579,6 +579,15 @@ def get_pde_as_diff_op(self):
         k = sym.Symbol(self.helmholtz_k_name)
         return (laplacian(w) + k**2 * w)
 
+    def get_code_transformer(self):
+        k = SpatialConstant(self.helmholtz_k_name)
+
+        if self.allow_evanescent:
+            return lambda expr: expr
+        else:
+            from sumpy.codegen import HelmholtzRewriter
+            return HelmholtzRewriter(k, var("Ik"))
+
 
 class YukawaKernel(ExpressionKernel):
     init_arg_names = ("dim", "yukawa_lambda_name")