From 41c3bc21a903c7fdb6fb12a20c15a1593c3a35c2 Mon Sep 17 00:00:00 2001
From: Lawrence Mitchell <lawrence@wence.uk>
Date: Sat, 9 May 2020 19:02:33 +0100
Subject: [PATCH] loopy: Refactor mathfunction mapping

Remove need for math_table since loopy does name mapping and the only
names that need translating are Bessel functions and ln. Additionally,
don't forbid mathfunctions that only operator on real values in
complex mode (since their argument may be real!).
---
 tsfc/loopy.py | 93 ++++++++++++++-------------------------------------
 1 file changed, 26 insertions(+), 67 deletions(-)

diff --git a/tsfc/loopy.py b/tsfc/loopy.py
index ca6a5356..e04ba2ba 100644
--- a/tsfc/loopy.py
+++ b/tsfc/loopy.py
@@ -22,37 +22,6 @@
 from contextlib import contextmanager
 
 
-# Table of handled math functions in real and complex modes
-# Note that loopy handles addition of type prefixes and suffixes itself.
-math_table = {
-    'sqrt': ('sqrt', 'sqrt'),
-    'abs': ('abs', 'abs'),
-    'cos': ('cos', 'cos'),
-    'sin': ('sin', 'sin'),
-    'tan': ('tan', 'tan'),
-    'acos': ('acos', 'acos'),
-    'asin': ('asin', 'asin'),
-    'atan': ('atan', 'atan'),
-    'cosh': ('cosh', 'cosh'),
-    'sinh': ('sinh', 'sinh'),
-    'tanh': ('tanh', 'tanh'),
-    'acosh': ('acosh', 'acosh'),
-    'asinh': ('asinh', 'asinh'),
-    'atanh': ('atanh', 'atanh'),
-    'power': ('pow', 'pow'),
-    'exp': ('exp', 'exp'),
-    'ln': ('log', 'log'),
-    'real': (None, 'real'),
-    'imag': (None, 'imag'),
-    'conj': (None, 'conj'),
-    'erf': ('erf', None),
-    'atan_2': ('atan2', None),
-    'atan2': ('atan2', None),
-    'min_value': ('min', None),
-    'max_value': ('max', None)
-}
-
-
 maxtype = partial(numpy.find_common_type, [])
 
 
@@ -419,49 +388,39 @@ def _expression_power(expr, ctx):
 
 @_expression.register(gem.MathFunction)
 def _expression_mathfunction(expr, ctx):
-
-    complex_mode = int(is_complex(ctx.scalar_type))
-
-    # Bessel functions
     if expr.name.startswith('cyl_bessel_'):
-        if complex_mode:
-            msg = "Bessel functions for complex numbers: missing implementation"
-            raise NotImplementedError(msg)
+        # Bessel functions
+        if is_complex(ctx.scalar_type):
+            raise NotImplementedError("Bessel functions for complex numbers: "
+                                      "missing implementation")
         nu, arg = expr.children
-        nu_thunk = lambda: expression(nu, ctx)
-        arg_loopy = expression(arg, ctx)
-        if expr.name == 'cyl_bessel_j':
-            if nu == gem.Zero():
-                return p.Variable("j0")(arg_loopy)
-            elif nu == gem.one:
-                return p.Variable("j1")(arg_loopy)
-            else:
-                return p.Variable("jn")(nu_thunk(), arg_loopy)
-        if expr.name == 'cyl_bessel_y':
-            if nu == gem.Zero():
-                return p.Variable("y0")(arg_loopy)
-            elif nu == gem.one:
-                return p.Variable("y1")(arg_loopy)
-            else:
-                return p.Variable("yn")(nu_thunk(), arg_loopy)
-
+        nu_ = expression(nu, ctx)
+        arg_ = expression(arg, ctx)
         # Modified Bessel functions (C++ only)
         #
         # These mappings work for FEniCS only, and fail with Firedrake
         # since no Boost available.
-        if expr.name in ['cyl_bessel_i', 'cyl_bessel_k']:
+        if expr.name in {'cyl_bessel_i', 'cyl_bessel_k'}:
             name = 'boost::math::' + expr.name
-            return p.Variable(name)(nu_thunk(), arg_loopy)
-
-        assert False, "Unknown Bessel function: {}".format(expr.name)
-
-    # Other math functions
-    name = math_table[expr.name][complex_mode]
-    if name is None:
-        raise RuntimeError("{} not supported in {} mode".format(expr.name,
-                                                                ("real", "complex")[complex_mode]))
-
-    return p.Variable(name)(*[expression(c, ctx) for c in expr.children])
+            return p.Variable(name)(nu_, arg_)
+        else:
+            # cyl_bessel_{jy} -> {jy}
+            name = expr.name[-1:]
+            if nu == gem.Zero():
+                return p.Variable(f"{name}0")(arg_)
+            elif nu == gem.one:
+                return p.Variable(f"{name}1")(arg_)
+            else:
+                return p.Variable(f"{name}n")(nu_, arg_)
+    else:
+        if expr.name == "ln":
+            name = "log"
+        else:
+            name = expr.name
+        # Not all mathfunctions apply to complex numbers, but this
+        # will be picked up in loopy. This way we allow erf(real(...))
+        # in complex mode (say).
+        return p.Variable(name)(*(expression(c, ctx) for c in expr.children))
 
 
 @_expression.register(gem.MinValue)