Use a symbolic function

Author: isuruf

sumpy/codegen.py

@@ -592,39 +592,6 @@ 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,12 @@ def __init__(self, dim, helmholtz_k_name="k",
             if allow_evanescent:
                 expr = var("exp")(var("I")*k*r)/r
             else:
-                expr = var("exp")(var("Ik")*r)/r
+                # expi is a function that takes in a real and returns a
+                # complex number such that
+                # expi(x) = exp(I * x)
+                # Retaining the information that the input is real leads
+                # to better code generation
+                expr = var("expi")(k*r)/r
             scaling = 1/(4*var("pi"))
         else:
             raise RuntimeError("unsupported dimensionality")
@@ -579,15 +584,6 @@ 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")

sumpy/symbolic.py

@@ -313,6 +313,13 @@ def map_Mul(self, expr):  # noqa: N802
 
         return math.prod(num_args) / math.prod(den_args)
 
+    def map_FunctionSymbol(self, expr):
+        if expr.get_name() == "ExpI":
+            arg = self.rec(expr.args[0])
+            return prim.Variable("cos")(arg) + 1j * prim.Variable("sin")(arg)
+        else:
+            return SympyToPymbolicMapperBase.map_FunctionSymbol(self, expr)
+
 
 class PymbolicToSympyMapperWithSymbols(PymbolicToSympyMapper):
     def map_variable(self, expr):
@@ -338,6 +345,9 @@ def map_call(self, expr):
             args = [self.rec(param) for param in expr.parameters]
             args.append(0)
             return BesselJ(*args)
+        elif expr.function.name == "expi":
+            args = [self.rec(param) for param in expr.parameters]
+            return ExpI(*args)
         else:
             return PymbolicToSympyMapper.map_call(self, expr)
 
@@ -369,8 +379,20 @@ class Hankel1(_BesselOrHankel):
     pass
 
 
+class ExpI(sympy.Function):
+    """A symbolic function that takes a real value as an
+    input and returns a complex number such that
+    expi(x) = exp(i*x).
+    """
+    nargs = (1,)
+
+    def fdiff(self, argindex=1):
+        return self.func(self.args[0]) * sympy.I
+
+
 _SympyBesselJ = BesselJ
 _SympyHankel1 = Hankel1
+_SympyExpI = ExpI
 
 if USE_SYMENGINE:
     def BesselJ(*args):   # noqa: N802  # pylint: disable=function-redefined
@@ -379,4 +401,7 @@ def BesselJ(*args):   # noqa: N802  # pylint: disable=function-redefined
     def Hankel1(*args):   # noqa: N802  # pylint: disable=function-redefined
         return sym.sympify(_SympyHankel1(*args))
 
+    def ExpI(*args):   # noqa: N802  # pylint: disable=function-redefined
+        return sym.sympify(_SympyExpI(*args))
+
 # vim: fdm=marker