Merge pull request #2323 from devitocodes/stagg-fd-custom

api: Fix custom fd for staggered

devito/finite_differences/coefficients.py

@@ -89,7 +89,7 @@ def index(self):
         The dimension to which the coefficients will be applied plus the offset
         in that dimension if the function is staggered.
         """
-        return self._function.indices_ref[self._dimension]
+        return self._dimension
 
     @property
     def weights(self):
@@ -268,7 +268,6 @@ def generate_subs(deriv_order, function, index):
         return subs
 
     # Determine which 'rules' are missing
-
     sym = get_sym(functions)
     terms = obj.find(sym)
     for i in obj.find(Weights):
@@ -304,10 +303,12 @@ def generate_subs(deriv_order, function, index):
         try:
             for k, v in subs.rules.items():
                 ofs, do, f, d = k.args
-                if deriv_order == do and dim is d and f in expr._functions:
-                    mapper[k.func(ofs, do, expr, d)] = v
+                is_dim = dim == expr.indices_ref[d]
+                if deriv_order == do and is_dim and f in expr._functions:
+                    mapper[k.func(ofs, do, expr, expr.indices_ref[d])] = v
         except AttributeError:
-            assert subs is None
+            # assert subs is None
+            pass
 
         if mapper:
             rules.update(mapper)

tests/test_symbolic_coefficients.py

@@ -267,7 +267,8 @@ def test_staggered_array(self, order):
         eq_f = Eq(f, f.dx2, coefficients=Substitutions(coeffs_f))
         eq_g = Eq(g, g.dx2, coefficients=Substitutions(coeffs_g))
 
-        Operator([eq_f, eq_g])()
+        op = Operator([eq_f, eq_g])
+        op()
 
         assert np.allclose(f.data, g.data, atol=1e-7)
 
@@ -303,6 +304,25 @@ def test_staggered_function(self, order):
 
         assert np.allclose(f.data, g.data, atol=1e-7)
 
+    def test_staggered_function_evalat(self):
+        grid = Grid(shape=(11,), extent=(10.,))
+        x = grid.dimensions[0]
+
+        f = Function(name='f', grid=grid, space_order=2,
+                     coefficients='symbolic')
+        g = Function(name='g', grid=grid, space_order=2,
+                     coefficients='symbolic', staggered=x)
+
+        w = Function(name='w', space_order=0, shape=(3,),
+                     dimensions=(Dimension("p"),))
+        coeffs_g = Coefficient(2, g, x, w)
+
+        eq_fg = Eq(f, g.dx2, coefficients=Substitutions(coeffs_g))
+
+        expected = 'Eq(f(x), g(x - h_x/2)*w[0] + g(x + h_x/2)*w[1])'
+
+        assert str(eq_fg.evaluate) == expected
+
     def test_staggered_equation(self):
         """
         Check that expressions with substitutions are consistent with