Merge pull request #2571 from devitocodes/hotfix-weight-eval

compiler: Fix uxreplace behaviour with numeric expressions

Author: mloubout

devito/symbolics/manipulation.py

@@ -14,7 +14,7 @@
 from devito.symbolics.queries import q_leaf
 from devito.symbolics.search import retrieve_indexed, retrieve_functions
 from devito.symbolics.unevaluation import (
-    Add as UnevalAdd, Mul as UnevalMul, Pow as UnevalPow
+    Add as UnevalAdd, Mul as UnevalMul, Pow as UnevalPow, UnevaluableMixin
 )
 from devito.tools import as_list, as_tuple, flatten, split, transitive_closure
 from devito.types.basic import Basic, Indexed
@@ -300,7 +300,11 @@ def _eval_numbers(expr, args):
     """
     numbers, others = split(args, lambda i: i.is_Number)
     if len(numbers) > 1:
-        args[:] = [expr.func(*numbers)] + others
+        if isinstance(expr, UnevaluableMixin):
+            cls = expr.func.__base__
+        else:
+            cls = expr.func
+        args[:] = [cls(*numbers)] + others
 
 
 def flatten_args(args, op, ignore=None):

tests/test_symbolics.py

@@ -4,7 +4,7 @@
 import pytest
 import numpy as np
 
-from sympy import Expr, Symbol
+from sympy import Expr, Number, Symbol
 from devito import (Constant, Dimension, Grid, Function, solve, TimeFunction, Eq,  # noqa
                     Operator, SubDimension, norm, Le, Ge, Gt, Lt, Abs, sin, cos,
                     Min, Max)
@@ -13,7 +13,7 @@
 from devito.symbolics import (retrieve_functions, retrieve_indexed, evalrel,  # noqa
                               CallFromPointer, Cast, DefFunction, FieldFromPointer,
                               INT, FieldFromComposite, IntDiv, Namespace, Rvalue,
-                              ReservedWord, ListInitializer, uxreplace,
+                              ReservedWord, ListInitializer, uxreplace, pow_to_mul,
                               retrieve_derivatives, BaseCast)
 from devito.tools import as_tuple
 from devito.types import (Array, Bundle, FIndexed, LocalObject, Object,
@@ -589,51 +589,67 @@ def test_solve_time():
     assert sympy.simplify(sympy.expand(sol - (-dt**2*u.dx/m + 2.0*u - u.backward))) == 0
 
 
-@pytest.mark.parametrize('expr,subs,expected', [
-    ('f', '{f: g}', 'g'),
-    ('f[x, y+1]', '{f.indexed: g.indexed}', 'g[x, y+1]'),
-    ('cos(f)', '{cos: sin}', 'sin(f)'),
-    ('cos(f + sin(g))', '{cos: sin, sin: cos}', 'sin(f + cos(g))'),
-    ('FIndexed(f.indexed, x, y)', '{x: 0}', 'FIndexed(f.indexed, 0, y)'),
-])
-def test_uxreplace(expr, subs, expected):
-    grid = Grid(shape=(4, 4))
-    x, y = grid.dimensions  # noqa
+class TestUxreplace:
 
-    f = Function(name='f', grid=grid)  # noqa
-    g = Function(name='g', grid=grid)  # noqa
+    @pytest.mark.parametrize('expr,subs,expected', [
+        ('f', '{f: g}', 'g'),
+        ('f[x, y+1]', '{f.indexed: g.indexed}', 'g[x, y+1]'),
+        ('cos(f)', '{cos: sin}', 'sin(f)'),
+        ('cos(f + sin(g))', '{cos: sin, sin: cos}', 'sin(f + cos(g))'),
+        ('FIndexed(f.indexed, x, y)', '{x: 0}', 'FIndexed(f.indexed, 0, y)'),
+    ])
+    def test_expressions(self, expr, subs, expected):
+        grid = Grid(shape=(4, 4))
+        x, y = grid.dimensions  # noqa
 
-    assert uxreplace(eval(expr), eval(subs)) == eval(expected)
+        f = Function(name='f', grid=grid)  # noqa
+        g = Function(name='g', grid=grid)  # noqa
 
+        assert uxreplace(eval(expr), eval(subs)) == eval(expected)
 
-def test_uxreplace_custom_reconstructable():
+    def test_custom_reconstructable(self):
 
-    class MyDefFunction(DefFunction):
-        __rargs__ = ('name', 'arguments')
-        __rkwargs__ = ('p0', 'p1', 'p2')
+        class MyDefFunction(DefFunction):
+            __rargs__ = ('name', 'arguments')
+            __rkwargs__ = ('p0', 'p1', 'p2')
 
-        def __new__(cls, name=None, arguments=None, p0=None, p1=None, p2=None):
-            obj = super().__new__(cls, name=name, arguments=arguments)
-            obj.p0 = p0
-            obj.p1 = as_tuple(p1)
-            obj.p2 = p2
-            return obj
+            def __new__(cls, name=None, arguments=None, p0=None, p1=None, p2=None):
+                obj = super().__new__(cls, name=name, arguments=arguments)
+                obj.p0 = p0
+                obj.p1 = as_tuple(p1)
+                obj.p2 = p2
+                return obj
 
-    grid = Grid(shape=(4, 4))
+        grid = Grid(shape=(4, 4))
 
-    f = Function(name='f', grid=grid)
-    g = Function(name='g', grid=grid)
+        f = Function(name='f', grid=grid)
+        g = Function(name='g', grid=grid)
+
+        func = MyDefFunction(name='foo', arguments=f.indexify(),
+                             p0=f, p1=f, p2='bar')
+
+        mapper = {f: g, f.indexify(): g.indexify()}
+        func1 = uxreplace(func, mapper)
+
+        assert func1.arguments == (g.indexify(),)
+        assert func1.p0 is g
+        assert func1.p1 == (g,)
+        assert func1.p2 == 'bar'
+
+    def test_reduce_to_number(self):
+        grid = Grid(shape=(4, 4))
+        x, _ = grid.dimensions
+        h_x = x.spacing
 
-    func = MyDefFunction(name='foo', arguments=f.indexify(),
-                         p0=f, p1=f, p2='bar')
+        # Emulate lowered coefficient
+        w = -0.0354212/(h_x*h_x)
+        w_lowered = pow_to_mul(w)
 
-    mapper = {f: g, f.indexify(): g.indexify()}
-    func1 = uxreplace(func, mapper)
+        w_sub = uxreplace(w_lowered, {h_x: Number(3)})
 
-    assert func1.arguments == (g.indexify(),)
-    assert func1.p0 is g
-    assert func1.p1 == (g,)
-    assert func1.p2 == 'bar'
+        assert np.isclose(w_sub, -0.003935689)
+        assert not w_sub.is_Mul
+        assert w_sub.is_Number
 
 
 def test_minmax():