From 73bedc5f0642cb1a656916caae5b726100ba176a Mon Sep 17 00:00:00 2001
From: Andrew <ac1721@ic.ac.uk>
Date: Thu, 14 Sep 2023 14:33:04 +0100
Subject: [PATCH] tests: Add new test to test_ir and flake 8

---
 devito/ir/support/basic.py | 29 +++++++++++-------------
 tests/test_operator.py     | 46 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 16 deletions(-)

diff --git a/devito/ir/support/basic.py b/devito/ir/support/basic.py
index 3740d3c423..60de20ed05 100644
--- a/devito/ir/support/basic.py
+++ b/devito/ir/support/basic.py
@@ -123,7 +123,7 @@ def index_mode(self):
     def aindices(self):
         retval = []
         for i, fi in zip(self, self.findices):
-            dims = {j for j in i.free_symbols if isinstance(j, Dimension)}
+            dims = set(d.root if d.indirect else d for d in i.atoms(Dimension))
             sdims = {d for d in dims if d.is_Stencil}
             candidates = dims - sdims
 
@@ -425,27 +425,24 @@ def distance(self, other):
                 ret.append(S.Zero)
             else:
                 v = i - j
-                # If (i-j) is a finite number, then we potentially have an Imaginary dependence
                 if v.is_Number and v.is_finite:
                     # If i and j are numbers, we append an (I) distance
                     if i.is_Number and j.is_Number:
                         return Vector(S.ImaginaryUnit)
-                    # If i and j are not numbers, there may be dimension-dependent dependencies
-                    # so we append the distance
+                    # If both i and j are not numbers, there may be dimension-dependent
+                    # dependencies so we append the distance
                     else:
                         ret.append(v)
 
                 # We are writing over an entire dimension but reading from one point.
                 # If there are overlaps between the two then we would have a dependency
-                # This is a conservative estimation as there are cases (example below) where
-                # we potentially don't actually have a dependency given that we don't write
-                # over the entire dimension
-                # But we would need the exact iteration intervals to compute this,
-                # which is only known at compile time
-
-                # No Dependency Example:
-                # Eq(u[0,y], 1)
-                # Eq(u[1, y+1], u[0,1])
+                # This is a conservative estimation as there are cases (example below)
+                # where we may or may not have a dependency given that we don't write
+                # depending on domain size, which is not compilation-time known
+
+                # For example:
+                # Eq(u[0, y], 1)
+                # Eq(u[1, y+1], u[0, 1])
                 elif i.is_Number and not j.is_Number:
                     ret.append(S.Infinity)
 
@@ -681,9 +678,9 @@ def is_const(self, dim):
         """
         True if a constant dependence, that is no Dimensions involved, False otherwise.
         """
-        return (self.source.aindices[dim] is None and
-                self.sink.aindices[dim] is None and
-                self.distance_mapper[dim] == 0)
+        return (self.source.aindices.get(dim) is None and
+                self.sink.aindices.get(dim) is None and
+                self.distance_mapper.get(dim, 0) == 0)
 
     @memoized_meth
     def is_carried(self, dim=None):
diff --git a/tests/test_operator.py b/tests/test_operator.py
index f38ac01942..ed39bb7ddf 100644
--- a/tests/test_operator.py
+++ b/tests/test_operator.py
@@ -1934,6 +1934,52 @@ def test_topofuse_w_numeric_dim(self):
 
         assert_structure(op, ['r,i', 'r'], 'r,i')
 
+    @pytest.mark.parametrize('eqns, expected, exp_trees, exp_iters', [
+        (['Eq(u[0, x], 1)',
+            'Eq(u[1, x], u[0, x + h_x] + u[0, x - h_x] - 2*u[0, x])'],
+            np.array([[1., 1., 1.],
+                      [-1., 0., -1.]]),
+            ['x', 'x'], 'x,x')
+    ])
+    def test_2194(self, eqns, expected, exp_trees, exp_iters):
+        grid = Grid(shape=(3, ))
+        u = TimeFunction(name='u', grid=grid)
+        x = grid.dimensions[0]
+        h_x = x.spacing  # noqa: F841
+
+        for i, e in enumerate(list(eqns)):
+            eqns[i] = eval(e)
+
+        op = Operator(eqns)
+        assert_structure(op, exp_trees, exp_iters)
+
+        op.apply()
+        assert(np.all(u.data[:] == expected[:]))
+
+    @pytest.mark.parametrize('eqns, expected, exp_trees, exp_iters', [
+        (['Eq(u[0, y], 1)', 'Eq(u[1, y], u[0, y + 1])'],
+            np.array([[1., 1.],
+                      [1., 0.]]),
+            ['y', 'y'], 'y,y'),
+        (['Eq(u[0, y], 1)', 'Eq(u[1, y], u[0, 2])'],
+            np.array([[1., 1.],
+                      [0., 0.]]),
+            ['y', 'y'], 'y,y')
+    ])
+    def test_2194_v2(self, eqns, expected, exp_trees, exp_iters):
+        grid = Grid(shape=(2, 2))
+        u = Function(name='u', grid=grid)
+        x, y = grid.dimensions
+
+        for i, e in enumerate(list(eqns)):
+            eqns[i] = eval(e)
+
+        op = Operator(eqns)
+        assert_structure(op, exp_trees, exp_iters)
+
+        op.apply()
+        assert(np.all(u.data[:] == expected[:]))
+
 
 class TestInternals(object):