Merge pull request #2738 from devitocodes/canonicalize-weights-position

compiler: Enforce indexed weights as rightmost arg

Author: FabioLuporini

devito/finite_differences/differentiable.py

@@ -19,7 +19,7 @@
 from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
                           infer_dtype, extract_dtype, is_integer, split, is_number)
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
-from devito.types.basic import AbstractFunction
+from devito.types.basic import AbstractFunction, Indexed
 
 __all__ = ['Differentiable', 'DiffDerivative', 'IndexDerivative', 'EvalDerivative',
            'Weights', 'Real', 'Imag', 'Conj']
@@ -769,12 +769,22 @@ def free_symbols(self):
     func = DifferentiableOp._rebuild
 
 
+class WeightsIndexed(Indexed):
+    pass
+
+
 class Weights(Array):
 
     """
     The weights (or coefficients) of a finite-difference expansion.
     """
 
+    # Use IndexedWeights for the underlying Indexed objects because they
+    # are guaranteed to appear at the end on an expression's .args.
+    # This makes it dramatically easier to implement substutions. It also makes
+    # it easier to visually parse IndexDerivatives when looking at them
+    _indexed_cls = WeightsIndexed
+
     def __init_finalize__(self, *args, **kwargs):
         dimensions = as_tuple(kwargs.get('dimensions'))
         weights = kwargs.get('initvalue')

devito/types/basic.py

@@ -1663,7 +1663,13 @@ def sort_key(self, order=None):
 
     def __getitem__(self, indices, **kwargs):
         """Produce a types.Indexed, rather than a sympy.Indexed."""
-        return Indexed(self, *as_tuple(indices))
+        # Is there a specific Indexed class to use?
+        try:
+            cls = self.function._indexed_cls
+        except AttributeError:
+            cls = Indexed
+
+        return cls(self, *as_tuple(indices))
 
     def _hashable_content(self):
         return super()._hashable_content() + (self.function,)

tests/test_symbolics.py

@@ -558,8 +558,8 @@ def test_canonical_ordering_of_weights():
 
     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'
+    assert ccode(fi*wi) == 'f[x][y + i0][z]*w0[i0]'
+    assert ccode(cf*wi) == 'f[x][y + i0][z].x*w0[i0]'
 
 
 def test_symbolic_printing():