From fed0d48824a1da3bda8236db1f04d7e872025bea Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Tue, 30 May 2023 18:25:59 -0500 Subject: [PATCH 1/5] Use the fact that helmholtz k is real for better codegen --- sumpy/kernel.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 1be1dec93..b79d40ad3 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -519,6 +519,10 @@ def __init__(self, dim, helmholtz_k_name="k", scaling = var("I")/4 elif dim == 3: r = pymbolic_real_norm_2(make_sym_vector("d", dim)) + 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("I")*k*r)/r scaling = 1/(4*var("pi")) else: From 154bbe6930ce0101ebe3f7caa4ac53fbcb844c2f Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Wed, 31 May 2023 01:01:11 -0500 Subject: [PATCH 2/5] Fix expr Co-authored-by: Alex Fikl --- sumpy/kernel.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sumpy/kernel.py b/sumpy/kernel.py index b79d40ad3..9c1d434f8 100644 --- a/sumpy/kernel.py +++ b/sumpy/kernel.py @@ -523,7 +523,6 @@ def __init__(self, dim, helmholtz_k_name="k", 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("I")*k*r)/r scaling = 1/(4*var("pi")) else: raise RuntimeError("unsupported dimensionality") From 495dfc5951cb8098407e58a6cc2bf718ee3d1bc9 Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Wed, 31 May 2023 14:36:31 -0500 Subject: [PATCH 3/5] rewrite at codegen time --- sumpy/codegen.py | 33 +++++++++++++++++++++++++++++++++ sumpy/kernel.py | 11 ++++++++++- 2 files changed, 43 insertions(+), 1 deletion(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 3b19fcb9e..28861f197 100644 --- a/sumpy/codegen.py +++ b/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": diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 9c1d434f8..72b227617 100644 --- a/sumpy/kernel.py +++ b/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") From abf6db989e94ca7bfd3ce09dca18134dd4af74a4 Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Thu, 1 Jun 2023 12:37:02 -0500 Subject: [PATCH 4/5] Use a symbolic function --- sumpy/codegen.py | 33 --------------------------------- sumpy/kernel.py | 16 ++++++---------- sumpy/symbolic.py | 25 +++++++++++++++++++++++++ 3 files changed, 31 insertions(+), 43 deletions(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 28861f197..3b19fcb9e 100644 --- a/sumpy/codegen.py +++ b/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": diff --git a/sumpy/kernel.py b/sumpy/kernel.py index 72b227617..444eb4293 100644 --- a/sumpy/kernel.py +++ b/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") diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index 48d43a5ef..a0bf815f5 100644 --- a/sumpy/symbolic.py +++ b/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 From 51eb520ca3348311e706e3f13520285cbe44888d Mon Sep 17 00:00:00 2001 From: Isuru Fernando Date: Thu, 1 Jun 2023 13:12:56 -0500 Subject: [PATCH 5/5] Fix conversions --- sumpy/codegen.py | 23 +---------------------- sumpy/symbolic.py | 14 +++++++++----- 2 files changed, 10 insertions(+), 27 deletions(-) diff --git a/sumpy/codegen.py b/sumpy/codegen.py index 3b19fcb9e..2f9be4bc8 100644 --- a/sumpy/codegen.py +++ b/sumpy/codegen.py @@ -32,7 +32,7 @@ from pytools import memoize_method -from sumpy.symbolic import (SympyToPymbolicMapper as SympyToPymbolicMapperBase) +from sumpy.symbolic import SympyToPymbolicMapper import logging logger = logging.getLogger(__name__) @@ -49,27 +49,6 @@ """ -# {{{ sympy -> pymbolic mapper - -import sumpy.symbolic as sym -_SPECIAL_FUNCTION_NAMES = frozenset(dir(sym.functions)) - - -class SympyToPymbolicMapper(SympyToPymbolicMapperBase): - - def not_supported(self, expr): - if isinstance(expr, int): - return expr - elif getattr(expr, "is_Function", False): - func_name = SympyToPymbolicMapperBase.function_name(self, expr) - return prim.Variable(func_name)( - *tuple(self.rec(arg) for arg in expr.args)) - else: - return SympyToPymbolicMapperBase.not_supported(self, expr) - -# }}} - - # {{{ bessel -> loopy codegen BESSEL_PREAMBLE = """//CL// diff --git a/sumpy/symbolic.py b/sumpy/symbolic.py index a0bf815f5..90a80fa06 100644 --- a/sumpy/symbolic.py +++ b/sumpy/symbolic.py @@ -313,12 +313,16 @@ 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) + def not_supported(self, expr): + if getattr(expr, "is_Function", False): + if self.function_name(expr) == "ExpI": + arg = self.rec(expr.args[0]) + return prim.Variable("cos")(arg) + 1j * prim.Variable("sin")(arg) + else: + return prim.Variable(self.function_name(expr))( + *[self.rec(arg) for arg in expr.args]) else: - return SympyToPymbolicMapperBase.map_FunctionSymbol(self, expr) + return SympyToPymbolicMapperBase.not_supported(self, expr) class PymbolicToSympyMapperWithSymbols(PymbolicToSympyMapper):