code: from itertools import product
  
  import numpy as np
  from sympy import And
  import pytest
  
  from conftest import skipif
  from devito import (ConditionalDimension, Grid, Function, TimeFunction, SparseFunction,  # noqa
                      Eq, Operator, Constant, Dimension, SubDimension, switchconfig,
                      SubDomain, Lt, Le, Gt, Ge, Ne, Buffer)
  from devito.ir.iet import Expression, Iteration, FindNodes, retrieve_iteration_tree
  from devito.symbolics import indexify, retrieve_functions
  from devito.types import Array
  
  
  class TestBufferedDimension(object):
  
      def test_multi_buffer(self):
          grid = Grid((3, 3))
          f = TimeFunction(name="f", grid=grid)
          g = TimeFunction(name="g", grid=grid, save=Buffer(7))
  
          op = Operator([Eq(f.forward, 1), Eq(g, f.forward)])
          op(time_M=3)
          # f looped all time_order buffer and is 1 everywhere
          assert np.allclose(f.data, 1)
          # g looped indices 0 to 3, rest is still 0
          assert np.allclose(g.data[0:4], 1)
          assert np.allclose(g.data[4:], 0)
  
      def test_multi_buffer_long_time(self):
          grid = Grid((3, 3))
          time = grid.time_dim
          f = TimeFunction(name="f", grid=grid)
          g = TimeFunction(name="g", grid=grid, save=Buffer(7))
  
          op = Operator([Eq(f.forward, time), Eq(g, time+1)])
          op(time_M=20)
          # f[0] is time=19, f[1] is time=20
          assert np.allclose(f.data[0], 19)
          assert np.allclose(f.data[1], 20)
          # g is time 15 to 21 (loop twice the 7 buffer then 15->21)
          for i in range(7):
              assert np.allclose(g.data[i], 14+i+1)
  
  
  class TestSubDimension(object):
  
      def test_interior(self):
          """
          Tests application of an Operator consisting of a single equation
          over the ``interior`` subdomain.
          """
          grid = Grid(shape=(4, 4, 4))
          x, y, z = grid.dimensions
  
          interior = grid.interior
  
          u = TimeFunction(name='u', grid=grid)
  
          eqn = [Eq(u.forward, u + 2, subdomain=interior)]
  
          op = Operator(eqn)
          op.apply(time_M=2)
          assert np.all(u.data[1, 1:-1, 1:-1, 1:-1] == 6.)
          assert np.all(u.data[1, :, 0] == 0.)
          assert np.all(u.data[1, :, -1] == 0.)
          assert np.all(u.data[1, :, :, 0] == 0.)
          assert np.all(u.data[1, :, :, -1] == 0.)
  
      def test_domain_vs_interior(self):
          """
          Tests application of an Operator consisting of two equations, one
          over the whole domain (default), and one over the ``interior`` subdomain.
          """
          grid = Grid(shape=(4, 4, 4))
          x, y, z = grid.dimensions
          t = grid.stepping_dim  # noqa
  
          interior = grid.interior
  
          u = TimeFunction(name='u', grid=grid)  # noqa
          eqs = [Eq(u.forward, u + 1),
                 Eq(u.forward, u.forward + 2, subdomain=interior)]
  
          op = Operator(eqs, opt='noop')
          trees = retrieve_iteration_tree(op)
          assert len(trees) == 2
  
          op.apply(time_M=1)
          assert np.all(u.data[1, 0, :, :] == 1)
          assert np.all(u.data[1, -1, :, :] == 1)
          assert np.all(u.data[1, :, 0, :] == 1)
          assert np.all(u.data[1, :, -1, :] == 1)
          assert np.all(u.data[1, :, :, 0] == 1)
          assert np.all(u.data[1, :, :, -1] == 1)
          assert np.all(u.data[1, 1:3, 1:3, 1:3] == 3)
  
      def test_subdim_middle(self):
          """
          Tests that instantiating SubDimensions using the classmethod
          constructors works correctly.
          """
          grid = Grid(shape=(4, 4, 4))
          x, y, z = grid.dimensions
          t = grid.stepping_dim  # noqa
  
          u = TimeFunction(name='u', grid=grid)  # noqa
          xi = SubDimension.middle(name='xi', parent=x,
                                   thickness_left=1,
                                   thickness_right=1)
          eqs = [Eq(u.forward, u + 1)]
          eqs = [e.subs(x, xi) for e in eqs]
  
          op = Operator(eqs)
  
          u.data[:] = 1.0
          op.apply(time_M=1)
          assert np.all(u.data[1, 0, :, :] == 1)
          assert np.all(u.data[1, -1, :, :] == 1)
          assert np.all(u.data[1, 1:3, :, :] == 2)
  
      def test_symbolic_size(self):
          """Check the symbolic size of all possible SubDimensions is as expected."""
          grid = Grid(shape=(4,))
          x, = grid.dimensions
          thickness = 4
  
          xleft = SubDimension.left(name='xleft', parent=x, thickness=thickness)
          assert xleft.symbolic_size == xleft.thickness.left[0]
  
          xi = SubDimension.middle(name='xi', parent=x,
                                   thickness_left=thickness, thickness_right=thickness)
          assert xi.symbolic_size == (x.symbolic_max - x.symbolic_min -
                                      xi.thickness.left[0] - xi.thickness.right[0] + 1)
  
          xright = SubDimension.right(name='xright', parent=x, thickness=thickness)
          assert xright.symbolic_size == xright.thickness.right[0]
  
      def test_bcs(self):
          """
          Tests application of an Operator consisting of multiple equations
          defined over different sub-regions, explicitly created through the
          use of SubDimensions.
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
          t = grid.stepping_dim
          thickness = 4
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=0, time_order=1)
  
          xleft = SubDimension.left(name='xleft', parent=x, thickness=thickness)
          xi = SubDimension.middle(name='xi', parent=x,
                                   thickness_left=thickness, thickness_right=thickness)
          xright = SubDimension.right(name='xright', parent=x, thickness=thickness)
  
          yi = SubDimension.middle(name='yi', parent=y,
                                   thickness_left=thickness, thickness_right=thickness)
  
          t_in_centre = Eq(u[t+1, xi, yi], 1)
          leftbc = Eq(u[t+1, xleft, yi], u[t+1, xleft+1, yi] + 1)
          rightbc = Eq(u[t+1, xright, yi], u[t+1, xright-1, yi] + 1)
  
          op = Operator([t_in_centre, leftbc, rightbc])
  
          op.apply(time_m=1, time_M=1)
  
          assert np.all(u.data[0, :, 0:thickness] == 0.)
          assert np.all(u.data[0, :, -thickness:] == 0.)
          assert all(np.all(u.data[0, i, thickness:-thickness] == (thickness+1-i))
                     for i in range(thickness))
          assert all(np.all(u.data[0, -i, thickness:-thickness] == (thickness+2-i))
                     for i in range(1, thickness + 1))
          assert np.all(u.data[0, thickness:-thickness, thickness:-thickness] == 1.)
  
      def test_flow_detection_interior(self):
          """
          Test detection of flow directions when SubDimensions are used
          (in this test they are induced by the ``interior`` subdomain).
  
          Stencil uses values at new timestep as well as those at previous ones
          This forces an evaluation order onto x.
          Weights are:
  
                 x=0     x=1     x=2     x=3
           t=N    2    ---3
                  v   /
           t=N+1  o--+----4
  
          Flow dependency should traverse x in the negative direction
  
                 x=2     x=3     x=4     x=5      x=6
          t=0             0   --- 0     -- 1    -- 0
                          v  /    v    /   v   /
          t=1            44 -+--- 11 -+--- 2--+ -- 0
          """
          grid = Grid(shape=(10, 10))
          x, y = grid.dimensions
  
          interior = grid.interior
  
          u = TimeFunction(name='u', grid=grid, save=10, time_order=1, space_order=0)
  
          step = Eq(u.forward, 2*u
                    + 3*u.subs(x, x+x.spacing)
                    + 4*u.forward.subs(x, x+x.spacing),
                    subdomain=interior)
          op = Operator(step)
  
          u.data[0, 5, 5] = 1.0
          op.apply(time_M=0)
          assert u.data[1, 5, 5] == 2
          assert u.data[1, 4, 5] == 11
          assert u.data[1, 3, 5] == 44
          assert u.data[1, 2, 5] == 4*44
          assert u.data[1, 1, 5] == 4*4*44
  
          # This point isn't updated because of the `interior` selection
          assert u.data[1, 0, 5] == 0
  
          assert np.all(u.data[1, 6:, :] == 0)
          assert np.all(u.data[1, :, 0:5] == 0)
          assert np.all(u.data[1, :, 6:] == 0)
  
      @pytest.mark.parametrize('exprs,expected,', [
          # Carried dependence in both /t/ and /x/
          (['Eq(u[t+1, x, y], u[t+1, x-1, y] + u[t, x, y])'], 'y'),
          (['Eq(u[t+1, x, y], u[t+1, x-1, y] + u[t, x, y], subdomain=interior)'], 'i0y'),
          # Carried dependence in both /t/ and /y/
          (['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y])'], 'x'),
          (['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y], subdomain=interior)'], 'i0x'),
          # Carried dependence in /y/, leading to separate /y/ loops, one
          # going forward, the other backward
          (['Eq(u[t+1, x, y], u[t+1, x, y-1] + u[t, x, y], subdomain=interior)',
            'Eq(u[t+1, x, y], u[t+1, x, y+1] + u[t, x, y], subdomain=interior)'], 'i0x'),
      ])
      def test_iteration_property_parallel(self, exprs, expected):
          """Tests detection of sequental and parallel Iterations when applying
          equations over different subdomains."""
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions  # noqa
          t = grid.time_dim  # noqa
  
          interior = grid.interior  # noqa
  
          u = TimeFunction(name='u', grid=grid, save=10, time_order=1)  # noqa
  
          # List comprehension would need explicit locals/globals mappings to eval
          for i, e in enumerate(list(exprs)):
              exprs[i] = eval(e)
  
          op = Operator(exprs, opt='noop')
          iterations = FindNodes(Iteration).visit(op)
          assert all(i.is_Sequential for i in iterations if i.dim.name != expected)
          assert all(i.is_Parallel for i in iterations if i.dim.name == expected)
  
      @skipif(['device'])
      @pytest.mark.parametrize('exprs,expected,', [
          # All parallel, the innermost Iteration gets vectorized
          (['Eq(u[time, x, yleft], u[time, x, yleft] + 1.)'], ['yleft']),
          # All outers are parallel, carried dependence in `yleft`, so the middle
          # Iteration over `x` gets vectorized
          (['Eq(u[time, x, yleft], u[time, x, yleft+1] + 1.)'], ['x']),
          # Only the middle Iteration is parallel, so no vectorization (the Iteration
          # is left non-vectorised for OpenMP parallelism)
          (['Eq(u[time+1, x, yleft], u[time, x, yleft+1] + u[time+1, x, yleft+1])'], [])
      ])
      def test_iteration_property_vector(self, exprs, expected):
          """Tests detection of vector Iterations when using subdimensions."""
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions  # noqa
          time = grid.time_dim  # noqa
  
          # The leftmost 10 elements
          yleft = SubDimension.left(name='yleft', parent=y, thickness=10) # noqa
  
          u = TimeFunction(name='u', grid=grid, save=10, time_order=0, space_order=1)  # noqa
  
          # List comprehension would need explicit locals/globals mappings to eval
          for i, e in enumerate(list(exprs)):
              exprs[i] = eval(e)
  
          op = Operator(exprs, opt='simd')
          iterations = FindNodes(Iteration).visit(op)
          vectorized = [i.dim.name for i in iterations if i.is_Vectorized]
          assert set(vectorized) == set(expected)
  
      def test_subdimmiddle_parallel(self):
          """
          Tests application of an Operator consisting of a subdimension
          defined over different sub-regions, explicitly created through the
          use of SubDimensions.
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
          t = grid.stepping_dim
          thickness = 4
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=0, time_order=1)
  
          xi = SubDimension.middle(name='xi', parent=x,
                                   thickness_left=thickness, thickness_right=thickness)
  
          yi = SubDimension.middle(name='yi', parent=y,
                                   thickness_left=thickness, thickness_right=thickness)
  
          # a 5 point stencil that can be computed in parallel
          centre = Eq(u[t+1, xi, yi], u[t, xi, yi] + u[t, xi-1, yi]
                                      + u[t, xi+1, yi] + u[t, xi, yi-1] + u[t, xi, yi+1])
  
          u.data[0, 10, 10] = 1.0
  
          op = Operator([centre])
  
          iterations = FindNodes(Iteration).visit(op)
          assert all(i.is_Affine and i.is_Parallel for i in iterations if i.dim in [xi, yi])
  
          op.apply(time_m=0, time_M=0)
  
          assert np.all(u.data[1, 9:12, 10] == 1.0)
          assert np.all(u.data[1, 10, 9:12] == 1.0)
  
          # Other than those, it should all be 0
          u.data[1, 9:12, 10] = 0.0
          u.data[1, 10, 9:12] = 0.0
          assert np.all(u.data[1, :] == 0)
  
      def test_subdimleft_parallel(self):
          """
          Tests application of an Operator consisting of a subdimension
          defined over different sub-regions, explicitly created through the
          use of SubDimensions.
  
          This tests that flow direction is not being automatically inferred
          from whether the subdimension is on the left or right boundary.
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
          t = grid.stepping_dim
          thickness = 4
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=0, time_order=1)
  
          xl = SubDimension.left(name='xl', parent=x, thickness=thickness)
  
          yi = SubDimension.middle(name='yi', parent=y,
                                   thickness_left=thickness, thickness_right=thickness)
  
          # Can be done in parallel
          eq = Eq(u[t+1, xl, yi], u[t, xl, yi] + 1)
  
          op = Operator([eq])
  
          iterations = FindNodes(Iteration).visit(op)
          assert all(i.is_Affine and i.is_Parallel for i in iterations if i.dim in [xl, yi])
  
          op.apply(time_m=0, time_M=0)
  
          assert np.all(u.data[1, 0:thickness, 0:thickness] == 0)
          assert np.all(u.data[1, 0:thickness, -thickness:] == 0)
          assert np.all(u.data[1, 0:thickness, thickness:-thickness] == 1)
          assert np.all(u.data[1, thickness+1:, :] == 0)
  
      def test_subdimmiddle_notparallel(self):
          """
          Tests application of an Operator consisting of a subdimension
          defined over different sub-regions, explicitly created through the
          use of SubDimensions.
  
          Different from ``test_subdimmiddle_parallel`` because an interior
          dimension cannot be evaluated in parallel.
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
          t = grid.stepping_dim
          thickness = 4
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=0, time_order=1)
  
          xi = SubDimension.middle(name='xi', parent=x,
                                   thickness_left=thickness, thickness_right=thickness)
  
          yi = SubDimension.middle(name='yi', parent=y,
                                   thickness_left=thickness, thickness_right=thickness)
  
          # flow dependencies in x and y which should force serial execution
          # in reverse direction
          centre = Eq(u[t+1, xi, yi], u[t, xi, yi] + u[t+1, xi+1, yi+1])
          u.data[0, 10, 10] = 1.0
  
          op = Operator([centre])
  
          iterations = FindNodes(Iteration).visit(op)
          assert all(i.is_Affine and i.is_Sequential for i in iterations if i.dim == xi)
          assert all(i.is_Affine and i.is_Parallel for i in iterations if i.dim == yi)
  
          op.apply(time_m=0, time_M=0)
  
          for i in range(4, 11):
              assert u.data[1, i, i] == 1.0
              u.data[1, i, i] = 0.0
  
          assert np.all(u.data[1, :] == 0)
  
      def test_subdimleft_notparallel(self):
          """
          Tests application of an Operator consisting of a subdimension
          defined over different sub-regions, explicitly created through the
          use of SubDimensions.
  
          This tests that flow direction is not being automatically inferred
          from whether the subdimension is on the left or right boundary.
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
          t = grid.stepping_dim
          thickness = 4
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=1, time_order=0)
  
          xl = SubDimension.left(name='xl', parent=x, thickness=thickness)
  
          yi = SubDimension.middle(name='yi', parent=y,
                                   thickness_left=thickness, thickness_right=thickness)
  
          # Flows inward (i.e. forward) rather than outward
          eq = Eq(u[t+1, xl, yi], u[t+1, xl-1, yi] + 1)
  
          op = Operator([eq])
  
          iterations = FindNodes(Iteration).visit(op)
          assert all(i.is_Affine and i.is_Sequential for i in iterations if i.dim == xl)
          assert all(i.is_Affine and i.is_Parallel for i in iterations if i.dim == yi)
  
          op.apply(time_m=1, time_M=1)
  
          assert all(np.all(u.data[0, :thickness, thickness+i] == [1, 2, 3, 4])
                     for i in range(12))
          assert np.all(u.data[0, thickness:] == 0)
          assert np.all(u.data[0, :, thickness+12:] == 0)
  
      def test_subdim_fd(self):
          """
          Test that the FD shortcuts are handled correctly with SubDimensions
          """
          grid = Grid(shape=(20, 20))
          x, y = grid.dimensions
  
          u = TimeFunction(name='u', save=None, grid=grid, space_order=1, time_order=1)
          u.data[:] = 2.
  
          # Flows inward (i.e. forward) rather than outward
          eq = [Eq(u.forward, u.dx + u.dy, subdomain=grid.interior)]
  
          op = Operator(eq)
  
          op.apply(time_M=0)
  
          assert np.all(u.data[1, -1, :] == 2.)
          assert np.all(u.data[1, :, 0] == 2.)
          assert np.all(u.data[1, :, -1] == 2.)
          assert np.all(u.data[1, 0, :] == 2.)
          assert np.all(u.data[1, 1:18, 1:18] == 0.)
  
      def test_arrays_defined_over_subdims(self):
          """
          Check code generation when an Array uses a SubDimension.
          """
          grid = Grid(shape=(3,))
          x, = grid.dimensions
          xi, = grid.interior.dimensions
  
          f = Function(name='f', grid=grid)
          a = Array(name='a', dimensions=(xi,), dtype=grid.dtype)
          op = Operator([Eq(a[xi], 1), Eq(f, f + a[xi + 1], subdomain=grid.interior)],
                        openmp=False)
          assert len(op.parameters) == 6
          # neither `x_size` nor `xi_size` are expected here
          assert not any(i.name in ('x_size', 'xi_size') for i in op.parameters)
          # Try running it -- regardless of what it will produce, this should run
          # ie, this checks this error isn't raised:
          # "ValueError: No value found for parameter xi_size"
          op()
  
  
  class TestConditionalDimension(object):
  
      """
      A collection of tests to check the correct functioning of ConditionalDimensions.
      """
  
      def test_basic(self):
          nt = 19
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid)
          assert(grid.stepping_dim in u.indices)
  
          u2 = TimeFunction(name='u2', grid=grid, save=nt)
          assert(time in u2.indices)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
                               time_dim=time_subsampled)
          assert(time_subsampled in usave.indices)
  
          eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.), Eq(usave, u)]
          op = Operator(eqns)
          op.apply(t_M=nt-2)
          assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1))
          assert np.all([np.allclose(u2.data[i], i) for i in range(nt)])
          assert np.all([np.allclose(usave.data[i], i*factor)
                        for i in range((nt+factor-1)//factor)])
  
      def test_basic_shuffles(self):
          """
          Like ``test_basic``, but with different equation orderings. Nevertheless,
          we assert against the same exact values as in ``test_basic``, since we
          save `u`, not `u.forward`.
          """
          nt = 19
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid)
  
          u2 = TimeFunction(name='u2', grid=grid, save=nt)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
                               time_dim=time_subsampled)
  
          # Shuffle 1
          eqns = [Eq(usave, u), Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.)]
          op = Operator(eqns)
          op.apply(t_M=nt-2)
          assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1))
          assert np.all([np.allclose(u2.data[i], i) for i in range(nt)])
          assert np.all([np.allclose(usave.data[i], i*factor)
                        for i in range((nt+factor-1)//factor)])
  
          # Shuffle 2
          usave.data[:] = 0.
          u.data[:] = 0.
          u2.data[:] = 0.
          eqns = [Eq(u.forward, u + 1.), Eq(usave, u), Eq(u2.forward, u2 + 1.)]
          op = Operator(eqns)
          op.apply(t_M=nt-2)
          assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1))
          assert np.all([np.allclose(u2.data[i], i) for i in range(nt)])
          assert np.all([np.allclose(usave.data[i], i*factor)
                        for i in range((nt+factor-1)//factor)])
  
      def test_spacial_subsampling(self):
          """
          Test conditional dimension for the spatial ones.
          This test saves u every two grid points :
          u2[x, y] = u[2*x, 2*y]
          """
          nt = 19
          grid = Grid(shape=(12, 12))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid, save=nt)
          assert(grid.time_dim in u.indices)
  
          # Creates subsampled spatial dimensions and accordine grid
          dims = tuple([ConditionalDimension(d.name+'sub', parent=d, factor=2)
                        for d in u.grid.dimensions])
          grid2 = Grid((6, 6), dimensions=dims, time_dimension=time)
          u2 = TimeFunction(name='u2', grid=grid2, save=nt)
          assert(time in u2.indices)
  
          eqns = [Eq(u.forward, u + 1.), Eq(u2, u)]
          op = Operator(eqns)
          op.apply(time_M=nt-2)
          # Verify that u2[x,y]= u[2*x, 2*y]
          assert np.allclose(u.data[:-1, 0:-1:2, 0:-1:2], u2.data[:-1, :, :])
  
      def test_time_subsampling_fd(self):
          nt = 19
          grid = Grid(shape=(11, 11))
          x, y = grid.dimensions
          time = grid.time_dim
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
                               time_dim=time_subsampled, time_order=2)
  
          dx2 = [indexify(i) for i in retrieve_functions(usave.dt2.evaluate)]
          assert dx2 == [usave[time_subsampled - 1, x, y],
                         usave[time_subsampled + 1, x, y],
                         usave[time_subsampled, x, y]]
  
      def test_subsampled_fd(self):
          """
          Test that the FD shortcuts are handled correctly with ConditionalDimensions
          """
          grid = Grid(shape=(21, 21))
          time = grid.time_dim
          # Creates subsampled spatial dimensions and accordine grid
          dims = tuple([ConditionalDimension(d.name+'sub', parent=d, factor=2)
                        for d in grid.dimensions])
          grid2 = Grid((6, 6), dimensions=dims, time_dimension=time)
          u2 = TimeFunction(name='u2', grid=grid2, space_order=2, time_order=1)
          u2.data.fill(2.)
          eqns = [Eq(u2.forward, u2.dx + u2.dy)]
          op = Operator(eqns)
          op.apply(time_M=0, x_M=11, y_M=11)
          # Verify that u2 contains subsampled fd values
          assert np.all(u2.data[0, :, :] == 2.)
          assert np.all(u2.data[1, 0, 0] == 0.)
          assert np.all(u2.data[1, -1, -1] == -40.)
          assert np.all(u2.data[1, 0, -1] == -20.)
          assert np.all(u2.data[1, -1, 0] == -20.)
          assert np.all(u2.data[1, 1:-1, 0] == 0.)
          assert np.all(u2.data[1, 0, 1:-1] == 0.)
          assert np.all(u2.data[1, 1:-1, -1] == -20.)
          assert np.all(u2.data[1, -1, 1:-1] == -20.)
          assert np.all(u2.data[1, 1:4, 1:4] == 0.)
  
      # This test generates an openmp loop form which makes older gccs upset
      @switchconfig(openmp=False)
      def test_nothing_in_negative(self):
          """Test the case where when the condition is false, there is nothing to do."""
          nt = 4
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', save=nt, grid=grid)
          assert(grid.time_dim in u.indices)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
                               time_dim=time_subsampled)
          assert(time_subsampled in usave.indices)
  
          eqns = [Eq(usave, u)]
          op = Operator(eqns)
  
          u.data[:] = 1.0
          usave.data[:] = 0.0
          op.apply(time_m=1, time_M=1)
          assert np.allclose(usave.data, 0.0)
  
          op.apply(time_m=0, time_M=0)
          assert np.allclose(usave.data, 1.0)
  
      def test_laplace(self):
          grid = Grid(shape=(20, 20, 20))
          x, y, z = grid.dimensions
          time = grid.time_dim
          t = grid.stepping_dim
          tsave = ConditionalDimension(name='tsave', parent=time, factor=2)
  
          u = TimeFunction(name='u', grid=grid, save=None, time_order=2)
          usave = TimeFunction(name='usave', grid=grid, time_dim=tsave,
                               time_order=0, space_order=0)
  
          steps = []
          # save of snapshot
          steps.append(Eq(usave, u))
          # standard laplace-like thing
          steps.append(Eq(u[t+1, x, y, z],
                          u[t, x, y, z] - u[t-1, x, y, z]
                          + u[t, x-1, y, z] + u[t, x+1, y, z]
                          + u[t, x, y-1, z] + u[t, x, y+1, z]
                          + u[t, x, y, z-1] + u[t, x, y, z+1]))
  
          op = Operator(steps)
  
          u.data[:] = 0.0
          u.data[0, 10, 10, 10] = 1.0
          op.apply(time_m=0, time_M=0)
          assert np.sum(u.data[0, :, :, :]) == 1.0
          assert np.sum(u.data[1, :, :, :]) == 7.0
          assert np.all(usave.data[0, :, :, :] == u.data[0, :, :, :])
  
      def test_as_expr(self):
          nt = 19
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid)
          assert(grid.stepping_dim in u.indices)
  
          u2 = TimeFunction(name='u2', grid=grid, save=nt)
          assert(time in u2.indices)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
                               time_dim=time_subsampled)
          assert(time_subsampled in usave.indices)
  
          eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.),
                  Eq(usave, time_subsampled * u)]
          op = Operator(eqns)
          op.apply(t=nt-2)
          assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1))
          assert np.all([np.allclose(u2.data[i], i) for i in range(nt)])
          assert np.all([np.allclose(usave.data[i], i*factor*i)
                        for i in range((nt+factor-1)//factor)])
  
      def test_shifted(self):
          nt = 19
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid)
          assert(grid.stepping_dim in u.indices)
  
          u2 = TimeFunction(name='u2', grid=grid, save=nt)
          assert(time in u2.indices)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
          usave = TimeFunction(name='usave', grid=grid, save=2, time_dim=time_subsampled)
          assert(time_subsampled in usave.indices)
  
          t_sub_shift = Constant(name='t_sub_shift', dtype=np.int32)
  
          eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.),
                  Eq(usave.subs(time_subsampled, time_subsampled - t_sub_shift), u)]
          op = Operator(eqns)
  
          # Starting at time_m=10, so time_subsampled - t_sub_shift is in range
          op.apply(time_m=10, time_M=nt-2, t_sub_shift=3)
          assert np.all(np.allclose(u.data[0], 8))
          assert np.all([np.allclose(u2.data[i], i - 10) for i in range(10, nt)])
          assert np.all([np.allclose(usave.data[i], 2+i*factor) for i in range(2)])
  
      def test_no_index(self):
          """Test behaviour when the ConditionalDimension is used as a symbol in
          an expression."""
          nt = 19
          grid = Grid(shape=(11, 11))
          time = grid.time_dim
  
          u = TimeFunction(name='u', grid=grid)
          assert(grid.stepping_dim in u.indices)
  
          v = Function(name='v', grid=grid)
  
          factor = 4
          time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
  
          eqns = [Eq(u.forward, u + 1), Eq(v, v + u*u*time_subsampled)]
          op = Operator(eqns)
          op.apply(t_M=nt-2)
          assert np.all(np.allclose(u.data[(nt-1) % 3], nt-1))
          # expected result is 1024
          # v = u[0]**2 * 0 + u[4]**2 * 1 + u[8]**2 * 2 + u[12]**2 * 3 + u[16]**2 * 4
          # with u[t] = t
          # v = 16 * 1 + 64 * 2 + 144 * 3 + 256 * 4 = 1600
          assert np.all(np.allclose(v.data, 1600))
  
      def test_no_index_sparse(self):
          """Test behaviour when the ConditionalDimension is used as a symbol in
          an expression over sparse data objects."""
          grid = Grid(shape=(4, 4), extent=(3.0, 3.0))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid, save=1)
          f.data[:] = 0.
  
          coordinates = [(0.5, 0.5), (0.5, 2.5), (2.5, 0.5), (2.5, 2.5)]
          sf = SparseFunction(name='sf', grid=grid, npoint=4, coordinates=coordinates)
          sf.data[:] = 1.
          sd = sf.dimensions[sf._sparse_position]
  
          # We want to write to `f` through `sf` so that we obtain the
          # following 4x4 grid (the '*' show the position of the sparse points)
          # We do that by emulating an injection
          #
          # 0 --- 0 --- 0 --- 0
          # |  *  |     |  *  |
          # 0 --- 1 --- 1 --- 0
          # |     |     |     |
          # 0 --- 1 --- 1 --- 0
          # |  *  |     |  *  |
          # 0 --- 0 --- 0 --- 0
  
          radius = 1
          indices = [(i, i+radius) for i in sf._coordinate_indices]
          bounds = [i.symbolic_size - radius for i in grid.dimensions]
  
          eqs = []
          for e, i in enumerate(product(*indices)):
              args = [j > 0 for j in i]
              args.extend([j < k for j, k in zip(i, bounds)])
              condition = And(*args, evaluate=False)
              cd = ConditionalDimension('sfc%d' % e, parent=sd, condition=condition)
              index = [time] + list(i)
              eqs.append(Eq(f[index], f[index] + sf[cd]))
  
          op = Operator(eqs)
          op.apply(time=0)
  
          assert np.all(f.data[0, 1:-1, 1:-1] == 1.)
          assert np.all(f.data[0, 0] == 0.)
          assert np.all(f.data[0, -1] == 0.)
          assert np.all(f.data[0, :, 0] == 0.)
          assert np.all(f.data[0, :, -1] == 0.)
  
      def test_symbolic_factor(self):
          """
          Test ConditionalDimension with symbolic factor (provided as a Constant).
          """
          g = Grid(shape=(4, 4, 4))
  
          u = TimeFunction(name='u', grid=g, time_order=0)
  
          fact = Constant(name='fact', dtype=np.int32, value=4)
          tsub = ConditionalDimension(name='tsub', parent=g.time_dim, factor=fact)
          usave = TimeFunction(name='usave', grid=g, time_dim=tsub, save=4)
  
          op = Operator([Eq(u, u + 1), Eq(usave, u)])
  
          op.apply(time=7)  # Use `fact`'s default value, 4
          assert np.all(usave.data[0] == 1)
          assert np.all(usave.data[1] == 5)
  
          u.data[:] = 0.
          op.apply(time=7, fact=2)
          assert np.all(usave.data[0] == 1)
          assert np.all(usave.data[1] == 3)
          assert np.all(usave.data[2] == 5)
          assert np.all(usave.data[3] == 7)
  
      def test_implicit_dims(self):
          """
          Test ConditionalDimension as an implicit dimension for an equation.
          """
  
          # This test makes an Operator that should create a vector of increasing
          # integers, but stop incrementing when a certain stop value is reached
  
          shape = (50,)
          stop_value = 20
  
          time = Dimension(name='time')
          f = TimeFunction(name='f', shape=shape, dimensions=[time])
  
          # The condition to stop incrementing
          cond = ConditionalDimension(name='cond',
                                      parent=time, condition=f[time] < stop_value)
  
          eqs = [Eq(f.forward, f), Eq(f.forward, f.forward + 1, implicit_dims=[cond])]
          op = Operator(eqs)
          op.apply(time_M=shape[0] - 2)
  
          # Make the same calculation in python to assert the result
          F = np.zeros(shape[0])
          for i in range(shape[0]):
              F[i] = i if i < stop_value else stop_value
  
          assert np.all(f.data == F)
  
      def test_stepping_dim_in_condition_lowering(self):
          """
          Check that the compiler performs lowering on conditions
          with TimeDimensions and generates the expected code::
  
          if (g[t][x + 1][y + 1] <= 10){          if (g[t0][x + 1][y + 1] <= 10){
              ...                          -->       ...
          }                                       }
  
          This test increments a function by one at every timestep until it is
          less-or-equal to 10 (g<=10) while although operator runs for 13 timesteps.
          """
          grid = Grid(shape=(4, 4))
          _, y = grid.dimensions
  
          ths = 10
          g = TimeFunction(name='g', grid=grid)
  
          ci = ConditionalDimension(name='ci', parent=y, condition=Le(g, ths))
  
          op = Operator(Eq(g.forward, g + 1, implicit_dims=ci))
  
          op.apply(time_M=ths+3)
          assert np.all(g.data[0, :, :] == ths)
          assert np.all(g.data[1, :, :] == ths + 1)
          assert 'if (g[t0][x + 1][y + 1] <= 10)\n'
          '{\n g[t1][x + 1][y + 1] = g[t0][x + 1][y + 1] + 1' in str(op.ccode)
  
      def test_expr_like_lowering(self):
          """
          Test the lowering of an expr-like ConditionalDimension's condition.
          This test makes an Operator that should indexify and lower the condition
          passed in the Conditional Dimension
          """
  
          grid = Grid(shape=(3, 3))
          g1 = Function(name='g1', grid=grid)
          g2 = Function(name='g2', grid=grid)
  
          g1.data[:] = 0.49
          g2.data[:] = 0.49
          x, y = grid.dimensions
          ci = ConditionalDimension(name='ci', parent=y, condition=Le((g1 + g2),
                                    1.01*(g1 + g2)))
  
          f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
          Operator(Eq(f, g1+g2)).apply()
  
          assert np.all(f.data[:] == g1.data[:] + g2.data[:])
  
      @pytest.mark.parametrize('setup_rel, rhs, c1, c2, c3, c4', [
          # Relation, RHS, c1 to c4 used as indexes in assert
          (Lt, 3, 2, 4, 4, -1), (Le, 2, 2, 4, 4, -1), (Ge, 3, 4, 6, 1, 4),
          (Gt, 2, 4, 6, 1, 4), (Ne, 5, 2, 6, 1, 2)
      ])
      def test_relational_classes(self, setup_rel, rhs, c1, c2, c3, c4):
          """
          Test ConditionalDimension using conditions based on Relations over SubDomains.
          """
  
          class InnerDomain(SubDomain):
              name = 'inner'
  
              def define(self, dimensions):
                  return {d: ('middle', 2, 2) for d in dimensions}
  
          inner_domain = InnerDomain()
          grid = Grid(shape=(8, 8), subdomains=(inner_domain,))
          g = Function(name='g', grid=grid)
          g2 = Function(name='g2', grid=grid)
  
          for i in [g, g2]:
              i.data[:4, :4] = 1
              i.data[4:, :4] = 2
              i.data[4:, 4:] = 3
              i.data[:4, 4:] = 4
  
          xi, yi = grid.subdomains['inner'].dimensions
  
          cond = setup_rel(0.25*g + 0.75*g2, rhs, subdomain=grid.subdomains['inner'])
          ci = ConditionalDimension(name='ci', parent=yi, condition=cond)
          f = Function(name='f', shape=grid.shape, dimensions=(xi, ci))
  
          eq1 = Eq(f, 0.4*g + 0.6*g2)
          eq2 = Eq(f, 5)
  
          Operator([eq1, eq2]).apply()
          assert np.all(f.data[2:6, c1:c2] == 5.)
          assert np.all(f.data[:, c3:c4] < 5.)
  
      def test_from_cond_to_param(self):
          """
          Test that Functions appearing in the condition of a ConditionalDimension
          but not explicitly in an Eq are actually part of the Operator input
          (stems from issue #1298).
          """
          grid = Grid(shape=(8, 8))
          x, y = grid.dimensions
  
          g = Function(name='g', grid=grid)
          h = Function(name='h', grid=grid)
          ci = ConditionalDimension(name='ci', parent=y, condition=Lt(g, 2 + h))
          f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
  
          for _ in range(5):
              # issue #1298 was non deterministic
              Operator(Eq(f, 5)).apply()
  
      @skipif('device')
      def test_no_fusion_simple(self):
          """
          If ConditionalDimensions are present, then Clusters must not be fused so
          that ultimately Eqs get scheduled to different loop nests.
          """
          grid = Grid(shape=(4, 4, 4))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid)
          g = Function(name='g', grid=grid)
          h = Function(name='h', grid=grid)
  
          # No ConditionalDimensions yet. Will be fused and optimized
          eqns = [Eq(f.forward, f + 1),
                  Eq(h, f + 1),
                  Eq(g, f + 1)]
  
          op = Operator(eqns)
  
          exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
          assert len(exprs) == 4
          assert exprs[1].expr.rhs is exprs[0].output
          assert exprs[2].expr.rhs is exprs[0].output
          assert exprs[3].expr.rhs is exprs[0].output
  
          # Now with a ConditionalDimension. No fusion, no optimization
          ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
  
          eqns = [Eq(f.forward, f + 1),
                  Eq(h, f + 1),
                  Eq(g, f + 1, implicit_dims=[ctime])]
  
          op = Operator(eqns)
          exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
          assert len(exprs) == 3
          assert exprs[1].expr.rhs is exprs[0].output
          assert exprs[2].expr.rhs is exprs[0].output
          exprs = FindNodes(Expression).visit(op._func_table['bf1'].root)
          assert len(exprs) == 1
  
      @skipif('device')
      def test_no_fusion_convoluted(self):
          """
          Conceptually like `test_no_fusion_simple`, but with more expressions
          and non-trivial data flow.
          """
          grid = Grid(shape=(4, 4, 4))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid)
          g = Function(name='g', grid=grid)
          h = Function(name='h', grid=grid)
  
          ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
  
          eqns = [Eq(f.forward, f + 1),
                  Eq(h, f + 1),
                  Eq(g, f + 1, implicit_dims=[ctime]),
                  Eq(f.forward, f + 1, implicit_dims=[ctime]),
                  Eq(f.forward, f + 1),
                  Eq(g, f + 1)]
  
          op = Operator(eqns)
  
          exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
          assert len(exprs) == 3
          assert exprs[1].expr.rhs is exprs[0].output
          assert exprs[2].expr.rhs is exprs[0].output
  
          exprs = FindNodes(Expression).visit(op._func_table['bf1'].root)
          assert len(exprs) == 3
  
          exprs = FindNodes(Expression).visit(op._func_table['bf2'].root)
          assert len(exprs) == 3
          assert exprs[1].expr.rhs is exprs[0].output
          assert exprs[2].expr.rhs is exprs[0].output
  
  
  class TestMashup(object):
  
      """
      Check the correct functioning of the compiler in presence of many Dimension types.
      """
  
      def test_topofusion_w_subdims_conddims(self):
          """
          Check that topological fusion works across guarded Clusters over different
          iteration spaces and in presence of anti-dependences.
  
          This test uses both SubDimensions (via SubDomains) and ConditionalDimensions.
          """
          grid = Grid(shape=(4, 4, 4))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid, time_order=2)
          g = TimeFunction(name='g', grid=grid, time_order=2)
          h = TimeFunction(name='h', grid=grid, time_order=2)
          fsave = TimeFunction(name='fsave', grid=grid, time_order=2, save=5)
          gsave = TimeFunction(name='gsave', grid=grid, time_order=2, save=5)
  
          ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
  
          eqns = [Eq(f.forward, f + 1),
                  Eq(g.forward, g + 1),
                  Eq(fsave, f.dt2, implicit_dims=[ctime]),
                  Eq(h, f + g, subdomain=grid.interior),
                  Eq(gsave, g.dt2, implicit_dims=[ctime])]
  
          op = Operator(eqns)
  
          # Check generated code -- expect the gsave equation to be scheduled together
          # in the same loop nest with the fsave equation
          assert len(op._func_table) == 3
  
          exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
          assert len(exprs) == 2
          assert exprs[0].write is f
          assert exprs[1].write is g
  
          exprs = FindNodes(Expression).visit(op._func_table['bf1'].root)
          assert len(exprs) == 3
          assert exprs[1].write is fsave
          assert exprs[2].write is gsave
  
          exprs = FindNodes(Expression).visit(op._func_table['bf2'].root)
          assert len(exprs) == 1
          assert exprs[0].write is h
  
      def test_topofusion_w_subdims_conddims_v2(self):
          """
          Like `test_topofusion_w_subdims_conddims` but with more SubDomains,
          so we expect fewer loop nests.
          """
          grid = Grid(shape=(4, 4, 4))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid, time_order=2)
          g = TimeFunction(name='g', grid=grid, time_order=2)
          h = TimeFunction(name='h', grid=grid, time_order=2)
          fsave = TimeFunction(name='fsave', grid=grid, time_order=2, save=5)
          gsave = TimeFunction(name='gsave', grid=grid, time_order=2, save=5)
  
          ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
  
          eqns = [Eq(f.forward, f + 1, subdomain=grid.interior),
                  Eq(g.forward, g + 1, subdomain=grid.interior),
                  Eq(fsave, f.dt2, implicit_dims=[ctime]),
                  Eq(h, f + g, subdomain=grid.interior),
                  Eq(gsave, g.dt2, implicit_dims=[ctime])]
  
          op = Operator(eqns)
  
          # Check generated code -- expect the gsave equation to be scheduled together
          # in the same loop nest with the fsave equation
          assert len(op._func_table) == 2
          assert len(FindNodes(Expression).visit(op._func_table['bf0'].root)) == 3
          assert len(FindNodes(Expression).visit(op._func_table['bf1'].root)) == 2 + 1  # r0
  
      def test_topofusion_w_subdims_conddims_v3(self):
          """
          Like `test_topofusion_w_subdims_conddims_v2` but with an extra anti-dependence,
          which causes scheduling over more loop nests.
          """
          grid = Grid(shape=(4, 4, 4))
          time = grid.time_dim
  
          f = TimeFunction(name='f', grid=grid, time_order=2)
          g = TimeFunction(name='g', grid=grid, time_order=2)
          h = TimeFunction(name='h', grid=grid, time_order=2)
          fsave = TimeFunction(name='fsave', grid=grid, time_order=2, save=5)
          gsave = TimeFunction(name='gsave', grid=grid, time_order=2, save=5)
  
          ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
  
          eqns = [Eq(f.forward, f + 1, subdomain=grid.interior),
                  Eq(g.forward, g + 1, subdomain=grid.interior),
                  Eq(fsave, f.dt2, implicit_dims=[ctime]),
                  Eq(h, f.dt2.dx + g, subdomain=grid.interior),
                  Eq(gsave, g.dt2, implicit_dims=[ctime])]
  
          op = Operator(eqns)
  
          # Check generated code -- expect the gsave equation to be scheduled together
          # in the same loop nest with the fsave equation
          assert len(op._func_table) == 3
  
          exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
          assert len(exprs) == 2
          assert exprs[0].write is f
          assert exprs[1].write is g
  
          exprs = FindNodes(Expression).visit(op._func_table['bf1'].root)
          assert len(exprs) == 3
          assert exprs[1].write is fsave
          assert exprs[2].write is gsave
  
          exprs = FindNodes(Expression).visit(op._func_table['bf2'].root)
          assert len(exprs) == 2
          assert exprs[1].write is h
  
apis: numpy.allclose numpy.sum numpy.all numpy.zeros