tests: add staggered origin in Abs MFE

devito/operations/interpolators.py

@@ -188,16 +188,14 @@ def _interpolation_indices(self, variables, offset=0, field_offset=0,
                 condition = sympy.And(lb, ub, evaluate=False)
                 mapper[d] = ConditionalDimension(p.name, self.sfunction._sparse_dim,
                                                  condition=condition, indirect=True)
+
             # Apply mapper to each variable with origin correction before the
-            # dimension get replaced.
-            vmapper = {}
-            for v in variables:
-                try:
-                    vmapper[v] = v.subs({k: c - v.origin[k] for k, c in mapper.items()})
-                except KeyError:
-                    vmapper[v] = v.subs(mapper)
+            # Dimensions get replaced
+            subs = {v: v.subs({k: c - v.origin.get(k, 0) for k, c in mapper.items()})
+                    for v in variables}
+
             # Track Indexed substitutions
-            idx_subs.append(vmapper)
+            idx_subs.append(subs)
 
         # Temporaries for the position
         temps = [Eq(v, k, implicit_dims=implicit_dims)

devito/tools/data_structures.py

@@ -66,6 +66,9 @@ def __getnewargs_ex__(self):
         # objects with varying number of attributes
         return (tuple(self), dict(self.__dict__))
 
+    def get(self, key, val):
+        return self._getters.get(key, val)
+
 
 class ReducerMap(MultiDict):
 

tests/test_symbolics.py

@@ -9,7 +9,7 @@
                     Operator, SubDimension, norm, Le, Ge, Gt, Lt, Abs, sin, cos, Min, Max)
 from devito.ir import Expression, FindNodes
 from devito.symbolics import (retrieve_functions, retrieve_indexed, evalrel,  # noqa
-                              CallFromPointer, Cast, FieldFromPointer,
+                              CallFromPointer, Cast, FieldFromPointer, INT,
                               FieldFromComposite, IntDiv, ccode, uxreplace)
 from devito.types import Array, Bundle, LocalObject, Object, Symbol as dSymbol
 
@@ -31,6 +31,18 @@ def test_float_indices():
     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;")