Merge branch 'devitocodes:main' into master

Author: ziyiyin97

devito/finite_differences/derivative.py

@@ -503,7 +503,27 @@ def _eval_at(self, func):
         if self.expr.staggered == func.staggered and self.expr.is_Function:
             return self
 
+        # Check if x0's keys come from a DerivedDimension
         x0 = func.indices_ref.getters
+        psubs = {}
+        nx0 = x0.copy()
+        for d, d0 in x0.items():
+            if d in self.dims:
+                # d is a valid Derivative dimension
+                continue
+            for sd in self.dims:
+                if sd in d._defines:
+                    # x0 key is a DerivedDimension of the derivative dimension
+                    # e.g f.dx(x0={ix: ix + h_x/2}) for a subdomain
+                    # Set x0 to the derivative dimension and add a substitution
+                    # to the parent
+                    # e.g f.dx(x0={x: x + h_x/2}).subs({x: ix})
+                    psubs[sd] = d
+                    nx0[sd] = nx0.pop(d)._subs(d, sd)
+        rkw = {'x0': nx0}
+        if psubs:
+            rkw['subs'] = (psubs,)
+
         if self.expr.is_Add:
             # If `expr` has both staggered and non-staggered terms such as
             # `(u(x + h_x/2) + v(x)).dx` then we exploit linearity of FD to split
@@ -512,19 +532,19 @@ def _eval_at(self, func):
             mapper = as_mapper(self.expr._args_diff, lambda i: i.staggered)
             args = [self.expr.func(*v) for v in mapper.values()]
             args.extend([a for a in self.expr.args if a not in self.expr._args_diff])
-            args = [self._rebuild(expr=a, x0=x0) for a in args]
+            args = [self._rebuild(a, **rkw) for a in args]
             return self.expr.func(*args)
         elif self.expr.is_Mul:
             # For Mul, We treat the basic case `u(x + h_x/2) * v(x) which is what appear
             # in most equation with div(a * u) for example. The expression is re-centered
             # at the highest priority index (see _gather_for_diff) to compute the
             # derivative at x0.
-            return self._rebuild(self.expr._gather_for_diff, x0=x0)
+            return self._rebuild(self.expr._gather_for_diff, **rkw)
         else:
             # For every other cases, that has more functions or more complexe arithmetic,
             # there is not actual way to decide what to do so it’s as safe to use
             # the expression as is.
-            return self._rebuild(x0=x0)
+            return self._rebuild(self.expr, **rkw)
 
     def _evaluate(self, **kwargs):
         # Evaluate finite-difference.

devito/finite_differences/differentiable.py

@@ -750,7 +750,12 @@ def dimensions(self):
         return self._dimensions
 
     def _evaluate(self, **kwargs):
-        expr = self.expr._evaluate(**kwargs)
+        try:
+            expr = self.expr._evaluate(**kwargs)
+        except AttributeError:
+            # There are rare circumstances in which `self.expr` is a plain
+            # SymPy object rather than an Evaluable
+            expr = Evaluable._evaluate_maybe_nested(self.expr, **kwargs)
 
         if not kwargs.get('expand', True):
             return self._rebuild(expr)
@@ -770,7 +775,10 @@ def free_symbols(self):
 
 
 class WeightsIndexed(Indexed):
-    pass
+
+    @property
+    def dimension(self):
+        return self.function.dimension
 
 
 class Weights(Array):

tests/test_interpolation.py

@@ -1098,6 +1098,7 @@ def test_inject_subdomain_sinc(self):
                          ['p_sr0rsr0xrsr0y'],
                          'p_sr0rsr0xrsr0y')
 
+    @pytest.mark.xfail(reason="OOB issue")
     @pytest.mark.parallel(mode=4)
     def test_interpolate_subdomain_mpi(self, mode):
         """

tests/test_subdomains.py

@@ -1778,3 +1778,24 @@ def define(self, dimensions):
         if grid.distributor.myrank == 0:
             assert np.all(np.isclose(fdata[:], gdata[:, 2:-2, 2:-2]))
             assert np.isclose(np.linalg.norm(fdata[:]), norm)
+
+    def test_mixed_domain_fd_staggered(self):
+        grid = Grid(shape=(20, 20, 20))
+        x = grid.dimensions[0]
+
+        class SD1(SubDomain):
+            name = 'sd1'
+
+            def define(self, dimensions):
+                x, y, z = dimensions
+                return {x: ('left', 2), y: y, z: z}
+
+        sd1 = SD1(grid=grid)
+        f = TimeFunction(name='f', grid=grid, time_order=1, space_order=2,
+                         staggered=(x,))
+        g = TimeFunction(name='g', grid=sd1, time_order=1, space_order=2,
+                         staggered=(x,))
+
+        eq = Eq(g, g + f.dx)
+        eqe = eq.evaluate
+        assert eqe.rhs == g + f.dx(x0=x).evaluate._subs(x, g.dimensions[1])