test_symbolics.py

from ctypes import c_void_p

import sympy
import pytest
import numpy as np

from sympy import Expr, Symbol
from devito import (Constant, Dimension, Grid, Function, solve, TimeFunction, Eq,  # noqa
                    Operator, SubDimension, norm, Le, Ge, Gt, Lt, Abs, sin, cos,
                    Min, Max)
from devito.finite_differences.differentiable import SafeInv, Weights
from devito.ir import Expression, FindNodes
from devito.symbolics import (retrieve_functions, retrieve_indexed, evalrel,  # noqa
                              CallFromPointer, Cast, DefFunction, FieldFromPointer,
                              INT, FieldFromComposite, IntDiv, Namespace, Rvalue,
                              ReservedWord, ListInitializer, ccode, uxreplace,
                              retrieve_derivatives)
from devito.symbolics.extraction import separate_eqn, centre_stencil
from devito.tools import as_tuple
from devito.types import (Array, Bundle, FIndexed, LocalObject, Object,
                          ComponentAccess, StencilDimension, Symbol as dSymbol)
from devito.types.basic import AbstractSymbol


def test_float_indices():
    """
    Test that indices only contain Integers.
    """
    grid = Grid((10,))
    x = grid.dimensions[0]
    x0 = x + 1.0 * x.spacing
    u = Function(name="u", grid=grid, space_order=2)
    indices = u.subs({x: x0}).indexify().indices[0]
    assert len(indices.atoms(sympy.Float)) == 0
    assert indices == x + 1

    indices = u.subs({x: 1.0}).indexify().indices[0]
    assert len(indices.atoms(sympy.Float)) == 0
    assert indices == 1


def test_func_of_indices():
    """
    Test that origin is correctly processed with functions
    """
    grid = Grid((10,))
    x = grid.dimensions[0]
    u = Function(name="u", grid=grid, space_order=2, staggered=x)
    us = u.subs({u.indices[0]: INT(Abs(u.indices[0]))})
    assert us.indices[0] == INT(Abs(x + x.spacing/2))
    assert us.indexify().indices[0] == INT(Abs(x))


@pytest.mark.parametrize('dtype,expected', [
    (np.float32, "float r0 = 1.0F/h_x;"),
    (np.float64, "double r0 = 1.0/h_x;")
])
def test_floatification_issue_1627(dtype, expected):
    """
    MFE for issue #1627.
    """
    grid = Grid(shape=(10, 10), dtype=dtype)
    x, y = grid.dimensions

    u = TimeFunction(name='u', grid=grid)

    eq = Eq(u.forward, ((u/x.spacing) + 2.0)/x.spacing)

    op = Operator(eq)

    exprs = FindNodes(Expression).visit(op)
    assert len(exprs) == 2
    assert str(exprs[0]) == expected


def test_sympy_assumptions():
    """
    Ensure that AbstractSymbol assumptions are set correctly and
    preserved during rebuild.
    """
    s0 = AbstractSymbol('s')
    s1 = AbstractSymbol('s', nonnegative=True, integer=False, real=True)

    assert s0.is_negative is None
    assert s0.is_positive is None
    assert s0.is_integer is None
    assert s0.is_real is True
    assert s1.is_negative is False
    assert s1.is_positive is True
    assert s1.is_integer is False
    assert s1.is_real is True

    s0r = s0._rebuild()
    s1r = s1._rebuild()

    assert s0.assumptions0 == s0r.assumptions0
    assert s0 == s0r

    assert s1.assumptions0 == s1r.assumptions0
    assert s1 == s1r


def test_modified_sympy_assumptions():
    """
    Check that sympy assumptions can be changed during a rebuild.
    """
    s0 = AbstractSymbol('s')
    s1 = AbstractSymbol('s', nonnegative=True, integer=False, real=True)

    s2 = s0._rebuild(nonnegative=True, integer=False, real=True)

    assert s2.assumptions0 == s1.assumptions0
    assert s2 == s1


def test_constant():
    c = Constant(name='c')

    assert c.free_symbols == {c}
    assert c.bound_symbols == set()


def test_dimension():
    d = Dimension(name='d')

    assert d.free_symbols == {d}
    assert d.bound_symbols == {d.symbolic_min, d.symbolic_max, d.symbolic_size}


def test_subdimension():
    d = Dimension(name='d')

    di = SubDimension.middle(name='di', parent=d, thickness_left=4, thickness_right=4)
    assert di.free_symbols == {di}
    assert di.bound_symbols == {d.symbolic_min, d.symbolic_max} | set(di.thickness)

    dl = SubDimension.left(name='dl', parent=d, thickness=4)
    assert dl.free_symbols == {dl}
    assert dl.bound_symbols == {d.symbolic_min, dl.thickness.left}

    dr = SubDimension.right(name='dr', parent=d, thickness=4)
    assert dr.free_symbols == {dr}
    assert dr.bound_symbols == {d.symbolic_max, dr.thickness.right}


def test_timefunction():
    grid = Grid(shape=(4, 4))
    x, y = grid.dimensions
    t = grid.stepping_dim

    f = TimeFunction(name='f', grid=grid)

    assert f.free_symbols == {t, x, y}
    assert f.bound_symbols == ({f.indexed, f._C_symbol} |
                               t.bound_symbols | x.bound_symbols | y.bound_symbols)


def test_indexed():
    grid = Grid(shape=(10, 10))
    x, y = grid.dimensions

    u = Function(name='u', grid=grid)

    assert u.free_symbols == {x, y}
    assert u.indexed.free_symbols == {u.indexed}

    # Test reconstruction
    assert u.indexed.func() == u.indexed
    ru = u.indexed.func(*u.indexed.args)
    assert ru is not u.indexed
    assert ru == u.indexed
    assert ru.function is u

    ub = Array(name='ub', dtype=u.dtype, dimensions=u.dimensions)

    assert ub.free_symbols == {x, y}
    assert ub.indexed.free_symbols == {ub.indexed}


def test_indexed_staggered():
    grid = Grid(shape=(10, 10))
    x, y = grid.dimensions
    hx, hy = x.spacing, y.spacing

    u = Function(name='u', grid=grid, staggered=(x, y))
    u0 = u.subs({x: 1, y: 2})
    assert u0.indices == (1 + hx / 2, 2 + hy / 2)
    assert u0.indexify().indices == (1, 2)


def test_bundle():
    grid = Grid(shape=(4, 4))

    f = Function(name='f', grid=grid)
    g = Function(name='g', grid=grid)

    fg = Bundle(name='fg', components=(f, g))

    # Test reconstruction
    fg._rebuild().components == fg.components


def test_call_from_pointer():
    s = Symbol('s')

    # Test construction
    cfp0 = CallFromPointer('foo', 's')
    cfp1 = CallFromPointer('foo', s)
    assert str(cfp0) == str(cfp1) == 's->foo()'
    assert cfp0 != cfp1  # As `cfc0`'s underlying Symbol is a types.Symbol
    cfp2 = CallFromPointer('foo', s, [0])
    assert str(cfp2) == 's->foo(0)'
    cfp3 = CallFromPointer('foo', s, [0, 'a'])
    assert str(cfp3) == 's->foo(0, a)'

    # Test hashing
    assert hash(cfp0) != hash(cfp1)  # Same reason as above
    assert hash(cfp0) != hash(cfp2)
    assert hash(cfp2) != hash(cfp3)

    # Test reconstruction
    cfp4 = cfp3.func(*cfp3.args)
    assert cfp3 == cfp4

    # Free symbols
    a = cfp3.args[2][1]
    assert str(a) == 'a'
    assert cfp3.free_symbols == {s, a}


def test_field_from_pointer():
    s = Symbol('s')

    # Test construction
    ffp0 = FieldFromPointer('foo', 's')
    ffp1 = FieldFromPointer('foo', s)
    assert str(ffp0) == str(ffp1) == 's->foo'
    assert ffp0 != ffp1  # As `ffc0`'s underlying Symbol is a types.Symbol
    ffp2 = FieldFromPointer('bar', 's')
    assert ffp0 != ffp2

    # Test hashing
    assert hash(ffp0) != hash(ffp1)  # Same reason as above
    assert hash(ffp0) != hash(ffp2)

    # Test reconstruction
    ffp3 = ffp0.func(*ffp0.args)
    assert ffp0 == ffp3

    # Free symbols
    assert ffp1.free_symbols == {s}


def test_field_from_composite():
    s = Symbol('s')

    # Test construction
    ffc0 = FieldFromComposite('foo', 's')
    ffc1 = FieldFromComposite('foo', s)
    assert str(ffc0) == str(ffc1) == 's.foo'
    assert ffc0 != ffc1  # As `ffc0`'s underlying Symbol is a types.Symbol
    ffc2 = FieldFromComposite('bar', 's')
    assert ffc0 != ffc2

    # Test hashing
    assert hash(ffc0) != hash(ffc1)  # Same reason as above
    assert hash(ffc0) != hash(ffc2)

    # Test reconstruction
    ffc3 = ffc0.func(*ffc0.args)
    assert ffc0 == ffc3

    # Free symbols
    assert ffc1.free_symbols == {s}


def test_extended_sympy_arithmetic():
    # NOTE: `s` gets turned into a devito.Symbol, whose dtype
    # defaults to np.int32
    cfp = CallFromPointer('foo', 's')
    ffp = FieldFromPointer('foo', 's')
    ffc = FieldFromComposite('foo', 's')

    assert ccode(cfp + 1) == 's->foo() + 1'
    assert ccode(ffp + 1) == 's->foo + 1'
    assert ccode(ffc + 1) == 's.foo + 1'

    grid = Grid(shape=(4, 4))
    u = Function(name='u', grid=grid)
    foo = FieldFromPointer('foo', u._C_symbol)
    assert ccode(-1 + foo) == 'u_vec->foo - 1'

    # Now wrapping an object whose dtype is not numeric, hence
    # noncommutative
    o = Object(name='o', dtype=c_void_p)
    bar = FieldFromPointer('bar', o)
    assert ccode(-1 + bar) == '-1 + o->bar'


def test_integer_abs():
    i1 = Dimension(name="i1")
    assert ccode(Abs(i1 - 1)) == "abs(i1 - 1)"
    assert ccode(Abs(i1 - .5)) == "fabsf(i1 - 5.0e-1F)"
    assert ccode(
        Abs(i1 - Constant('half', dtype=np.float64, default_value=0.5))
    ) == "fabs(i1 - half)"


def test_cos_vs_cosf():
    a = dSymbol('a', dtype=np.float32)
    assert ccode(cos(a)) == "cosf(a)"

    b = dSymbol('b', dtype=np.float64)
    assert ccode(cos(b)) == "cos(b)"

    # Doesn't make much sense, but it's legal
    c = dSymbol('c', dtype=np.int32)
    assert ccode(cos(c)) == "cos(c)"


def test_intdiv():
    a = Symbol('a')
    b = Symbol('b')

    # IntDiv by 1 automatically simplified
    v = IntDiv(a, 1)
    assert v is a

    v = IntDiv(a, 2)
    assert ccode(v) == 'a / 2'

    # Within larger expressions -> parentheses
    v = b*IntDiv(a, 2)
    assert ccode(v) == 'b*(a / 2)'
    v = 3*IntDiv(a, 2)
    assert ccode(v) == '3*(a / 2)'
    v = b*IntDiv(a, 2) + 3
    assert ccode(v) == 'b*(a / 2) + 3'

    # IntDiv by zero or non-integer fails
    with pytest.raises(ValueError):
        IntDiv(a, 0)
    with pytest.raises(ValueError):
        IntDiv(a, 3.5)

    v = b*IntDiv(a + b, 2) + 3
    assert ccode(v) == 'b*((a + b) / 2) + 3'


def test_safeinv():
    grid = Grid(shape=(11, 11))
    x, y = grid.dimensions

    u1 = Function(name='u', grid=grid)
    u2 = Function(name='u', grid=grid, dtype=np.float64)

    op1 = Operator(Eq(u1, SafeInv(u1, u1)))
    op2 = Operator(Eq(u2, SafeInv(u2, u2)))

    assert 'SAFEINV' in str(op1)
    assert 'SAFEINV' in str(op2)


def test_def_function():
    foo0 = DefFunction('foo', arguments=['a', 'b'], template=['int'])
    foo1 = DefFunction('foo', arguments=['a', 'b'], template=['int'])
    foo2 = DefFunction('foo', arguments=['a', 'b'])
    foo3 = DefFunction('foo', arguments=['a'])

    # Code generation
    assert str(foo0) == 'foo<int>(a, b)'
    assert str(foo3) == 'foo(a)'

    # Hashing and equality
    assert hash(foo0) == hash(foo1)
    assert foo0 == foo1
    assert hash(foo0) != hash(foo2)
    assert hash(foo2) != hash(foo3)

    # Reconstruction
    assert foo0 == foo0._rebuild()
    assert str(foo0._rebuild('bar', template=['float'])) == 'bar<float>(a, b)'


def test_namespace():
    ns0 = Namespace(['std', 'algorithms', 'parallel'])
    assert str(ns0) == 'std::algorithms::parallel'

    ns1 = Namespace(['std'])
    ns2 = Namespace(['std', 'algorithms', 'parallel'])

    # Test hashing and equality
    assert hash(ns0) != hash(ns1)  # Same reason as above
    assert ns0 != ns1
    assert hash(ns0) == hash(ns2)
    assert ns0 == ns2

    # Free symbols
    assert not ns0.free_symbols


def test_rvalue():
    ctype = ReservedWord('dummytype')
    ns = Namespace(['my', 'namespace'])
    init = ListInitializer(())

    assert str(Rvalue(ctype, ns, init)) == 'my::namespace::dummytype{}'


def test_cast():
    s = Symbol(name='s', dtype=np.float32)

    class BarCast(Cast):
        _base_typ = 'bar'

    v = BarCast(s, '**')
    assert ccode(v) == '(bar**)s'

    # Reconstruction
    assert ccode(v.func(*v.args)) == '(bar**)s'

    v1 = BarCast(s, '****')
    assert v != v1


def test_findexed():
    grid = Grid(shape=(3, 3, 3))
    x, y, z = grid.dimensions

    f = Function(name='f', grid=grid)

    strides_map = {x: 1, y: 2, z: 3}
    fi = FIndexed(f.base, x+1, y, z-2, strides_map=strides_map)
    assert ccode(fi) == 'f(x + 1, y, z - 2)'

    # Binding
    _, fi1 = fi.bind('fL')
    assert fi1.base is f.base
    assert ccode(fi1) == 'fL(x + 1, y, z - 2)'

    # Reconstruction
    strides_map = {x: 3, y: 2, z: 1}
    new_fi = fi.func(strides_map=strides_map, accessor=fi1.accessor)

    assert new_fi.name == fi.name == 'f'
    assert new_fi.accessor == fi1.accessor
    assert new_fi.accessor.name == 'fL'
    assert new_fi.indices == fi.indices
    assert new_fi.strides_map == strides_map


def test_canonical_ordering_of_weights():
    grid = Grid(shape=(3, 3, 3))
    x, y, z = grid.dimensions

    f = Function(name='f', grid=grid)

    i = StencilDimension('i0', 0, 2)
    w = Weights(name='w0', dimensions=i, initvalue=[1.0, 2.0, 3.0])

    fi = f[x, y + i, z]
    wi = w[i]
    cf = ComponentAccess(fi, 0)

    assert (ccode(1.0*f[x, y, z] + 2.0*f[x, y + 1, z] + 3.0*f[x, y + 2, z]) ==
            '1.0F*f[x][y][z] + 2.0F*f[x][y + 1][z] + 3.0F*f[x][y + 2][z]')
    assert ccode(fi*wi) == 'w0[i0]*f[x][y + i0][z]'
    assert ccode(cf*wi) == 'w0[i0]*f[x][y + i0][z].x'


def test_symbolic_printing():
    b = Symbol('b')

    v = CallFromPointer('foo', 's') + b
    assert str(v) == 'b + s->foo()'

    class MyLocalObject(LocalObject, Expr):
        pass

    lo = MyLocalObject(name='lo')
    assert str(lo + 2) == '2 + lo'

    grid = Grid((10,))
    f = Function(name="f", grid=grid)
    fi = FIndexed(f.base, grid.dimensions[0])
    df = DefFunction('aaa', arguments=[fi])
    assert ccode(df) == 'aaa(f(x))'


def test_is_on_grid():
    grid = Grid((10,))
    x = grid.dimensions[0]
    x0 = x + .5 * x.spacing
    u = Function(name="u", grid=grid, space_order=2)

    assert u._grid_map == {}
    assert u.subs({x: x0})._grid_map == {x: x0}
    assert all(uu._grid_map == {} for uu in retrieve_functions(u.subs({x: x0}).evaluate))


@pytest.mark.parametrize('expr,expected', [
    ('f[x+2]*g[x+4] + f[x+3]*g[x+5] + f[x+4] + f[x+1]',
     ['f[x+2]', 'g[x+4]', 'f[x+3]', 'g[x+5]', 'f[x+1]', 'f[x+4]']),
    ('f[x]*g[x+2] + f[x+1]*g[x+3]', ['f[x]', 'g[x+2]', 'f[x+1]', 'g[x+3]']),
])
def test_canonical_ordering(expr, expected):
    """
    Test that the `expr.args` are stored in canonical ordering.
    """
    grid = Grid(shape=(10,))
    x, = grid.dimensions  # noqa

    f = Function(name='f', grid=grid)  # noqa
    g = Function(name='g', grid=grid)  # noqa

    expr = eval(expr)
    for n, i in enumerate(list(expected)):
        expected[n] = eval(i)

    assert retrieve_indexed(expr) == expected


def test_solve_time():
    """
    Tests that solves only evaluate the time derivative.
    """
    grid = Grid(shape=(11, 11))
    u = TimeFunction(name="u", grid=grid, time_order=2, space_order=4)
    m = Function(name="m", grid=grid, space_order=4)
    dt = grid.time_dim.spacing
    eq = m * u.dt2 + u.dx
    sol = solve(eq, u.forward)
    # Check u.dx is not evaluated. Need to simplify because the solution
    # contains some Dummy in the Derivative subs that make equality break.
    assert len(retrieve_derivatives(sol)) == 1
    assert sympy.simplify(u.dx - retrieve_derivatives(sol)[0]) == 0
    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

    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():

    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

    grid = Grid(shape=(4, 4))

    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_minmax():
    grid = Grid(shape=(5, 5))
    x, y = grid.dimensions

    f = Function(name="f", grid=grid)
    s = dSymbol(name="s")

    eqns = [Eq(s, 2),
            Eq(f, Max(y, s, 4))]

    op = Operator(eqns)

    op.apply()
    assert np.all(f.data == 4)


@pytest.mark.parametrize('dtype,expected', [
    (np.float32, ("fmaxf(", "fminf(")),
    (np.float64, ("fmax(", "fmin(")),
])
def test_minmax_precision(dtype, expected):
    grid = Grid(shape=(5, 5), dtype=dtype)

    f = Function(name="f", grid=grid)
    g = Function(name="g", grid=grid)

    eqn = Eq(f, Min(g, 4.0) + Max(g, 2.0))

    op = Operator(eqn)

    g.data[:] = 3.0

    op.apply()

    # Check generated code -- ensure it's using the fp64 versions of min/max,
    # that is fminf/fmaxf
    assert all(i in str(op) for i in expected)

    assert np.all(f.data == 6.0)


@pytest.mark.parametrize('dtype,expected', [
    (np.float32, "powf("),
    (np.float64, "pow("),
])
def test_pow_precision(dtype, expected):
    grid = Grid(shape=(5, 5), dtype=dtype)

    f = Function(name="f", grid=grid)
    g = Function(name="g", grid=grid)

    eqn = Eq(f, g**1.5)

    op = Operator(eqn)

    g.data[:] = 4.0

    op.apply()

    assert expected in str(op)
    assert np.all(f.data == 8.0)


@pytest.mark.parametrize('dtype,expected', [
    (np.float32, "fabsf("),
    (np.float64, "fabs("),
])
def test_abs_precision(dtype, expected):
    grid = Grid(shape=(5, 5), dtype=dtype)

    f = Function(name="f", grid=grid)
    g = Function(name="g", grid=grid)

    eqn = Eq(f, abs(g))

    op = Operator(eqn)

    g.data[:] = -1.0

    op.apply()

    assert expected in str(op)
    assert np.all(f.data == 1.0)


class TestRelationsWithAssumptions:

    def test_multibounds_op(self):
        """
        Tests evalrel function on a simple example.
        """

        grid = Grid(shape=(16, 16, 16))

        a = Function(name='a', grid=grid)
        b = Function(name='b', grid=grid)
        c = Function(name='c', grid=grid)
        d = Function(name='d', grid=grid)
        a.data[:] = 5

        b = pow(a, 2)
        c = a + 10
        d = 2*a

        f = TimeFunction(name='f', grid=grid, space_order=2)

        f.data[:] = 0.1
        eqns = [Eq(f.forward, f.laplace + f * evalrel(min, [f, b, c, d]))]

        op = Operator(eqns, opt=('advanced'))
        op.apply(time_M=5)
        fnorm = norm(f)

        c2 = Function(name='c2', grid=grid)
        d2 = Function(name='d2', grid=grid)

        f.data[:] = 0.1
        eqns = [Eq(f.forward, f.laplace + f * evalrel(min, [f, b, c2, d2], [Ge(d, c)]))]
        op = Operator(eqns, opt=('advanced'))
        op.apply(time_M=5)
        fnorm2 = norm(f)

        assert fnorm == fnorm2

    @pytest.mark.parametrize('op, expr, assumptions, expected', [
        ([min, '[a, b, c, d]', '[]', 'Min(a, Min(b, Min(c, d)))']),
        ([max, '[a, b, c, d]', '[]', 'Max(a, Max(b, Max(c, d)))']),
        ([min, '[a]', '[]', 'a']),
        ([min, '[a, b]', '[Le(d, a), Ge(c, b)]', 'Min(a, b)']),
        ([min, '[a, b, c]', '[]', 'Min(a, Min(b, c))']),
        ([min, '[a, b, c, d]', '[Le(d, a), Ge(c, b)]', 'Min(b, d)']),
        ([min, '[a, b, c, d]', '[Ge(a, b), Ge(d, a), Ge(b, c)]', 'c']),
        ([max, '[a]', '[Le(a, a)]', 'a']),
        ([max, '[a, b]', '[Le(a, b)]', 'b']),
        ([max, '[a, b, c]', '[Le(c, b), Le(c, a)]', 'Max(a, b)']),
        ([max, '[a, b, c, d]', '[Ge(a, b), Ge(d, a), Ge(b, c)]', 'd']),
        ([max, '[a, b, c, d]', '[Ge(a, b), Le(b, c)]', 'Max(a, Max(c, d))']),
        ([max, '[a, b, c, d]', '[Ge(a, b), Le(c, b)]', 'Max(a, d)']),
        ([max, '[a, b, c, d]', '[Ge(b, a), Ge(a, b)]', 'Max(a, Max(c, d))']),
        ([min, '[a, b, c, d]', '[Ge(b, a), Ge(a, b), Le(c, b), Le(b, a)]', 'Min(c, d)']),
        ([min, '[a, b, c, d]', '[Ge(b, a), Ge(a, b), Le(c, b), Le(b, d)]', 'c']),
        ([min, '[a, b, c, d]', '[Ge(b, a + d)]', 'Min(a, Min(c, d))']),
        ([min, '[a, b, c, d]', '[Lt(b + a, d)]', 'Min(a, Min(b, c))']),
        ([max, '[a, b, c, d]', '[Lt(b + a, d)]', 'Max(c, d)']),
        ([max, '[a, b, c, d]', '[Gt(a, b + c + d)]', 'a']),
    ])
    def test_relations_w_complex_assumptions(self, op, expr, assumptions, expected):
        """
        Tests evalmin/evalmax with multiple args and assumtpions"""
        a = Symbol('a', positive=True)  # noqa
        b = Symbol('b', positive=True)  # noqa
        c = Symbol('c', positive=True)  # noqa
        d = Symbol('d', positive=True)  # noqa

        eqn = eval(expr)
        assumptions = eval(assumptions)
        assert str(evalrel(op, eqn, assumptions)) == expected

    @pytest.mark.parametrize('op, expr, assumptions, expected', [
        ([min, '[a, b, c, d]', '[Ge(b, a), Ge(a, b), Le(c, b), Le(b, d)]', 'c']),
        ([min, '[a, b, c, d]', '[Ge(b, a + d)]', 'Min(a, Min(b, Min(c, d)))']),
        ([min, '[a, b, c, d]', '[Ge(c, a + d)]', 'Min(a, b)']),
        ([max, '[a, b, c, d]', '[Ge(c, a + d), Gt(b, a + d)]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Ge(a + d, b), Gt(b, a + d)]',
         'Max(a, Max(b, Max(c, d)))']),
        ([max, '[a, b, c, d]', '[Le(c, a + d)]', 'Max(a, Max(b, Max(c, d)))']),
        ([max, '[a, b, c, d]', '[Le(c, d), Le(a, b)]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Le(c, d), Le(d, c)]', 'Max(a, Max(b, c))']),
        ([min, '[a, b, c, d]', '[Le(c, d).negated, Le(a, b).negated]', 'Min(b, d)']),
        ([min, '[a, b, c, d]', '[Le(c, d).negated, Le(a, b).negated]', 'Min(b, d)']),
        ([min, '[a, b, c, d]', '[Gt(c, d).negated, Ge(a, b).negated]', 'Min(a, c)']),
        ([min, '[a, b, c, d]', '[Gt(c, d).reversed, Ge(a, b).reversed]', 'Min(b, d)']),
        ([min, '[a, b, c, d]', '[Le(c, d).negated, Le(a, b).negated]', 'Min(b, d)']),
        ([max, '[a, b, c, d]', '[Le(c, d).negated, Le(a, b).negated]', 'Max(a, c)']),
        ([max, '[a, b, c, d]', '[Gt(c, d).negated, Ge(a, b).negated]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Gt(c, d).reversed, Ge(a, b).reversed]', 'Max(a, c)']),
        ([max, '[a, b, c, d]', '[Lt(c, d).reversed, Le(a, b).reversed]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Gt(c, d + a).negated]', 'Max(a, Max(b, Max(c, d)))']),
        ([max, '[a, b, c, d]', '[Lt(c, d + a).negated]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Le(c, d + a).negated]', 'Max(b, d)']),
        ([max, '[a, b, c, d]', '[Le(c + b, d + a).negated]',
          'Max(a, Max(b, Max(c, d)))']),
        ([max, '[a, b, c, d, e]', '[Gt(a, b + c + e)]',
          'Max(a, Max(b, Max(c, Max(d, e))))']),
        ([max, '[a, b, c, d, e]', '[Ge(c, a), Ge(b, a), Ge(a, c), Ge(e, c), Ge(d, e)]',
          'Max(b, d)']),
    ])
    def test_relations_w_complex_assumptions_II(self, op, expr, assumptions, expected):
        """
        Tests evalmin/evalmax with multiple args and assumtpions"""
        a = Symbol('a', positive=False)  # noqa
        b = Symbol('b', positive=False)  # noqa
        c = Symbol('c', positive=True)  # noqa
        d = Symbol('d', positive=True)  # noqa
        e = Symbol('e', positive=True)  # noqa

        eqn = eval(expr)
        assumptions = eval(assumptions)
        assert str(evalrel(op, eqn, assumptions)) == expected

    @pytest.mark.parametrize('op, expr, assumptions, expected', [
        ([min, '[a, b, c, d]', '[Ge(b, a)]', 'a']),
        ([min, '[a, b, c, d]', '[Ge(b, d)]', 'Min(a, d)']),
        ([min, '[a, b, c, d]', '[Ge(c, a + d)]', 'Min(a, b)']),
        ([max, '[a, b, c, d, e]', 'None', 'Max(e, d)']),
    ])
    def test_assumptions(self, op, expr, assumptions, expected):
        """
        Tests evalmin/evalmax with multiple args and assumtpions"""
        a = Symbol('a', positive=False)
        b = Symbol('b', positive=False)
        c = Symbol('c', positive=True)
        d = Symbol('d', positive=True)
        e = Symbol('e', positive=True)

        c = a + 9  # noqa
        d = b + 10  # noqa
        e = c + 1  # noqa
        eqn = eval(expr)
        assumptions = eval(assumptions)
        expected = eval(expected)
        assert evalrel(op, eqn, assumptions) == expected


@pytest.mark.parametrize('eqn, target, expected', [
    ('Eq(f1.laplace, g1)',
     'f1', ('g1(x, y)', 'Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')),
    ('Eq(g1, f1.laplace)',
     'f1', ('-g1(x, y)', '-Derivative(f1(x, y), (x, 2)) - Derivative(f1(x, y), (y, 2))')),
    ('Eq(g1, f1.laplace)', 'g1',
     ('Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))', 'g1(x, y)')),
    ('Eq(f1 + f1.laplace, g1)', 'f1', ('g1(x, y)',
     'f1(x, y) + Derivative(f1(x, y), (x, 2)) + Derivative(f1(x, y), (y, 2))')),
    ('Eq(g1.dx + f1.dx, g1)', 'f1',
     ('g1(x, y) - Derivative(g1(x, y), x)', 'Derivative(f1(x, y), x)')),
    ('Eq(g1.dx + f1.dx, g1)', 'g1',
     ('-Derivative(f1(x, y), x)', '-g1(x, y) + Derivative(g1(x, y), x)')),
    ('Eq(f1 * g1.dx, g1)', 'g1', ('0', 'f1(x, y)*Derivative(g1(x, y), x) - g1(x, y)')),
    ('Eq(f1 * g1.dx, g1)', 'f1', ('g1(x, y)', 'f1(x, y)*Derivative(g1(x, y), x)')),
    ('Eq((f1 * g1.dx).dy, f1)', 'f1',
     ('0', '-f1(x, y) + Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')),
    ('Eq((f1 * g1.dx).dy, f1)', 'g1',
     ('f1(x, y)', 'Derivative(f1(x, y)*Derivative(g1(x, y), x), y)')),
    ('Eq(f2.laplace, g2)', 'g2',
     ('-Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))',
      '-g2(t, x, y)')),
    ('Eq(f2.laplace, g2)', 'f2', ('g2(t, x, y)',
     'Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')),
    ('Eq(f2.laplace, f2)', 'f2', ('0',
     '-f2(t, x, y) + Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2))')),
    ('Eq(f2*g2, f2)', 'f2', ('0', 'f2(t, x, y)*g2(t, x, y) - f2(t, x, y)')),
    ('Eq(f2*g2, f2)', 'g2', ('f2(t, x, y)', 'f2(t, x, y)*g2(t, x, y)')),
    ('Eq(g2*f2.laplace, f2)', 'g2', ('f2(t, x, y)',
     '(Derivative(f2(t, x, y), (x, 2)) + Derivative(f2(t, x, y), (y, 2)))*g2(t, x, y)')),
    ('Eq(f2.forward, f2)', 'f2.forward', ('f2(t, x, y)', 'f2(t + dt, x, y)')),
    ('Eq(f2.forward, f2)', 'f2', ('-f2(t + dt, x, y)', '-f2(t, x, y)')),
    ('Eq(f2.forward.laplace, f2)', 'f2.forward', ('f2(t, x, y)',
     'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')),
    ('Eq(f2.forward.laplace, f2)', 'f2',
     ('-Derivative(f2(t + dt, x, y), (x, 2)) - Derivative(f2(t + dt, x, y), (y, 2))',
      '-f2(t, x, y)')),
    ('Eq(f2.laplace + f2.forward.laplace, g2)', 'f2.forward',
     ('g2(t, x, y) - Derivative(f2(t, x, y), (x, 2)) - Derivative(f2(t, x, y), (y, 2))',
      'Derivative(f2(t + dt, x, y), (x, 2)) + Derivative(f2(t + dt, x, y), (y, 2))')),
    ('Eq(g2.laplace, f2 + g2.forward)', 'g2.forward',
     ('f2(t, x, y) - Derivative(g2(t, x, y), (x, 2)) - Derivative(g2(t, x, y), (y, 2))',
      '-g2(t + dt, x, y)'))
])
def test_separate_eqn(eqn, target, expected):
    """
    Test the separate_eqn function.
    """
    grid = Grid((2, 2))

    so = 2

    f1 = Function(name='f1', grid=grid, space_order=so)  # noqa
    g1 = Function(name='g1', grid=grid, space_order=so)  # noqa

    f2 = TimeFunction(name='f2', grid=grid, space_order=so)  # noqa
    g2 = TimeFunction(name='g2', grid=grid, space_order=so)  # noqa

    b, F, _ = separate_eqn(eval(eqn), eval(target))
    expected_b, expected_F = expected

    assert str(b) == expected_b
    assert str(F) == expected_F


@pytest.mark.parametrize('eqn, target, expected', [
    ('Eq(f1.laplace, g1).evaluate', 'f1',
        (
            'g1(x, y)',
            '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 '
            '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2'
        )),
    ('Eq(g1, f1.laplace).evaluate', 'f1',
        (
            '-g1(x, y)',
            '-(-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2) '
            '- (-2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2)'
        )),
    ('Eq(g1, f1.laplace).evaluate', 'g1',
        (
            '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 '
            '- 2.0*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2',
            'g1(x, y)'
        )),
    ('Eq(f1 + f1.laplace, g1).evaluate', 'f1',
        (
            'g1(x, y)',
            '-2.0*f1(x, y)/h_x**2 + f1(x - h_x, y)/h_x**2 + f1(x + h_x, y)/h_x**2 - 2.0'
            '*f1(x, y)/h_y**2 + f1(x, y - h_y)/h_y**2 + f1(x, y + h_y)/h_y**2 + f1(x, y)'
        )),
    ('Eq(g1.dx + f1.dx, g1).evaluate', 'f1',
        (
            '-(-g1(x, y)/h_x + g1(x + h_x, y)/h_x) + g1(x, y)',
            '-f1(x, y)/h_x + f1(x + h_x, y)/h_x'
        )),
    ('Eq(g1.dx + f1.dx, g1).evaluate', 'g1',
        (
            '-(-f1(x, y)/h_x + f1(x + h_x, y)/h_x)',
            '-g1(x, y)/h_x + g1(x + h_x, y)/h_x - g1(x, y)'
        )),
    ('Eq(f1 * g1.dx, g1).evaluate', 'g1',
        (
            '0', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) - g1(x, y)'
        )),
    ('Eq(f1 * g1.dx, g1).evaluate', 'f1',
        (
            'g1(x, y)', '(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)'
        )),
    ('Eq((f1 * g1.dx).dy, f1).evaluate', 'f1',
        (
            '0', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) '
            '+ (-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y '
            '- f1(x, y)'
        )),
    ('Eq((f1 * g1.dx).dy, f1).evaluate', 'g1',
        (
            'f1(x, y)', '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y) + '
            '(-g1(x, y + h_y)/h_x + g1(x + h_x, y + h_y)/h_x)*f1(x, y + h_y)/h_y'
        )),
    ('Eq(f2.laplace, g2).evaluate', 'g2',
        (
            '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)'
            '/h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + '
            'f2(t, x, y + h_y)/h_y**2)', '-g2(t, x, y)'
        )),
    ('Eq(f2.laplace, g2).evaluate', 'f2',
        (
            'g2(t, x, y)', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + '
            'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)'
            '/h_y**2 + f2(t, x, y + h_y)/h_y**2'
        )),
    ('Eq(f2.laplace, f2).evaluate', 'f2',
        (
            '0', '-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + '
            'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2'
            ' + f2(t, x, y + h_y)/h_y**2 - f2(t, x, y)'
        )),
    ('Eq(g2*f2.laplace, f2).evaluate', 'g2',
        (
            'f2(t, x, y)', '(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + '
            'f2(t, x + h_x, y)/h_x**2 - 2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2'
            ' + f2(t, x, y + h_y)/h_y**2)*g2(t, x, y)'
        )),
    ('Eq(f2.forward.laplace, f2).evaluate', 'f2.forward',
        (
            'f2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2'
            ' + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*f2(t + dt, x, y)/h_y**2 + '
            'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2'
        )),
    ('Eq(f2.forward.laplace, f2).evaluate', 'f2',
        (
            '-(-2.0*f2(t + dt, x, y)/h_x**2 + f2(t + dt, x - h_x, y)/h_x**2 + '
            'f2(t + dt, x + h_x, y)/h_x**2) - (-2.0*f2(t + dt, x, y)/h_y**2 + '
            'f2(t + dt, x, y - h_y)/h_y**2 + f2(t + dt, x, y + h_y)/h_y**2)',
            '-f2(t, x, y)'
        )),
    ('Eq(f2.laplace + f2.forward.laplace, g2).evaluate', 'f2.forward',
        (
            '-(-2.0*f2(t, x, y)/h_x**2 + f2(t, x - h_x, y)/h_x**2 + f2(t, x + h_x, y)/'
            'h_x**2) - (-2.0*f2(t, x, y)/h_y**2 + f2(t, x, y - h_y)/h_y**2 + '
            'f2(t, x, y + h_y)/h_y**2) + g2(t, x, y)', '-2.0*f2(t + dt, x, y)/h_x**2 + '
            'f2(t + dt, x - h_x, y)/h_x**2 + f2(t + dt, x + h_x, y)/h_x**2 - 2.0*'
            'f2(t + dt, x, y)/h_y**2 + f2(t + dt, x, y - h_y)/h_y**2 + '
            'f2(t + dt, x, y + h_y)/h_y**2'
        )),
    ('Eq(g2.laplace, f2 + g2.forward).evaluate', 'g2.forward',
        (
            '-(-2.0*g2(t, x, y)/h_x**2 + g2(t, x - h_x, y)/h_x**2 + '
            'g2(t, x + h_x, y)/h_x**2) - (-2.0*g2(t, x, y)/h_y**2 + g2(t, x, y - h_y)'
            '/h_y**2 + g2(t, x, y + h_y)/h_y**2) + f2(t, x, y)', '-g2(t + dt, x, y)'
        ))
])
def test_separate_eval_eqn(eqn, target, expected):
    """
    Test the separate_eqn function on pre-evaluated equations.
    """
    grid = Grid((2, 2))

    so = 2

    f1 = Function(name='f1', grid=grid, space_order=so)  # noqa
    g1 = Function(name='g1', grid=grid, space_order=so)  # noqa

    f2 = TimeFunction(name='f2', grid=grid, space_order=so)  # noqa
    g2 = TimeFunction(name='g2', grid=grid, space_order=so)  # noqa

    b, F, _ = separate_eqn(eval(eqn), eval(target))
    expected_b, expected_F = expected

    assert str(b) == expected_b
    assert str(F) == expected_F


@pytest.mark.parametrize('expr, so, target, expected', [
    ('f1.laplace', 2, 'f1', '-2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'),
    ('f1 + f1.laplace', 2, 'f1',
     'f1(x, y) - 2.0*f1(x, y)/h_y**2 - 2.0*f1(x, y)/h_x**2'),
    ('g1.dx + f1.dx', 2, 'f1', '-f1(x, y)/h_x'),
    ('10 + f1.dx2', 2, 'g1', '0'),
    ('(f1 * g1.dx).dy', 2, 'f1',
     '(-1/h_y)*(-g1(x, y)/h_x + g1(x + h_x, y)/h_x)*f1(x, y)'),
    ('(f1 * g1.dx).dy', 2, 'g1', '-(-1/h_y)*f1(x, y)*g1(x, y)/h_x'),
    ('f2.laplace', 2, 'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'),
    ('f2*g2', 2, 'f2', 'f2(t, x, y)*g2(t, x, y)'),
    ('g2*f2.laplace', 2, 'f2',
     '(-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2)*g2(t, x, y)'),
    ('f2.forward', 2, 'f2.forward', 'f2(t + dt, x, y)'),
    ('f2.forward.laplace', 2, 'f2.forward',
     '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'),
    ('f2.laplace + f2.forward.laplace', 2, 'f2.forward',
     '-2.0*f2(t + dt, x, y)/h_y**2 - 2.0*f2(t + dt, x, y)/h_x**2'),
    ('f2.laplace + f2.forward.laplace', 2,
     'f2', '-2.0*f2(t, x, y)/h_y**2 - 2.0*f2(t, x, y)/h_x**2'),
    ('f2.laplace', 4, 'f2', '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'),
    ('f2.laplace + f2.forward.laplace', 4, 'f2.forward',
     '-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2'),
    ('f2.laplace + f2.forward.laplace', 4, 'f2',
     '-2.5*f2(t, x, y)/h_y**2 - 2.5*f2(t, x, y)/h_x**2'),
    ('f2.forward*f2.forward.laplace', 4, 'f2.forward',
     '(-2.5*f2(t + dt, x, y)/h_y**2 - 2.5*f2(t + dt, x, y)/h_x**2)*f2(t + dt, x, y)')
])
def test_centre_stencil(expr, so, target, expected):
    """
    Test extraction of centre stencil from an equation.
    """
    grid = Grid((2, 2))

    f1 = Function(name='f1', grid=grid, space_order=so)  # noqa
    g1 = Function(name='g1', grid=grid, space_order=so)  # noqa

    f2 = TimeFunction(name='f2', grid=grid, space_order=so)  # noqa
    g2 = TimeFunction(name='g2', grid=grid, space_order=so)  # noqa

    centre = centre_stencil(eval(expr), eval(target))

    assert str(centre) == expected