Merge pull request #2373 from devitocodes/corner-stagg-fd

api: fix corner case staggered fd for centered x0

.github/workflows/pytest-core-nompi.yml

@@ -72,7 +72,7 @@ jobs:
 
         - name: pytest-ubuntu-py312-gcc13-omp
           python-version: '3.12'
-          os: ubuntu-20.04
+          os: ubuntu-24.04
           arch: "gcc-13"
           language: "openmp"
           sympy: "1.11"

devito/finite_differences/coefficients.py

@@ -2,7 +2,7 @@
 from cached_property import cached_property
 
 from devito.finite_differences import Weights, generate_indices
-from devito.finite_differences.tools import numeric_weights, symbolic_weights
+from devito.finite_differences.tools import numeric_weights
 from devito.tools import filter_ordered, as_tuple
 
 __all__ = ['Coefficient', 'Substitutions', 'default_rules']
@@ -131,7 +131,7 @@ class Substitutions(object):
 
     Now define some partial d/dx FD coefficients of the Function u:
 
-    >>> u_x_coeffs = Coefficient(1, u, x, np.array([-0.6, 0.1, 0.6]))
+    >>> u_x_coeffs = Coefficient(2, u, x, np.array([-0.6, 0.1, 0.6]))
 
     And now create our Substitutions object to pass to equation:
 
@@ -141,7 +141,7 @@ class Substitutions(object):
     Now create a Devito equation and pass to it 'subs'
 
     >>> from devito import Eq
-    >>> eq = Eq(u.dt+u.dx, coefficients=subs)
+    >>> eq = Eq(u.dt+u.dx2, coefficients=subs)
 
     When evaluated, the derivatives will use the custom coefficients. We can
     check that by
@@ -231,7 +231,7 @@ def default_rules(obj, functions):
 
     from devito.symbolics.search import retrieve_dimensions
 
-    def generate_subs(deriv_order, function, index):
+    def generate_subs(deriv_order, function, index, indices):
         dim = retrieve_dimensions(index)[0]
 
         if dim.is_Time:
@@ -244,27 +244,15 @@ def generate_subs(deriv_order, function, index):
 
         subs = {}
 
-        mapper = {dim: index}
-        # Get full range of indices and weights
-        indices, x0 = generate_indices(function, dim,
-                                       fd_order, side=None, x0=mapper)
-        sweights = symbolic_weights(function, deriv_order, indices, x0)
-
         # Actual FD used indices and weights
         if deriv_order == 1 and fd_order == 2:
             fd_order = 1
 
-        indices, x0 = generate_indices(function, dim,
-                                       fd_order, side=None, x0=mapper)
-
-        coeffs = numeric_weights(function, deriv_order, indices, x0)
+        coeffs = numeric_weights(function, deriv_order, indices, index)
 
         for (c, i) in zip(coeffs, indices):
             subs.update({function._coeff_symbol(i, deriv_order, function, index): c})
 
-        # Set all unused weights to zero
-        subs.update({w: 0 for w in sweights if w not in subs})
-
         return subs
 
     # Determine which 'rules' are missing
@@ -274,7 +262,11 @@ def generate_subs(deriv_order, function, index):
         for w in i.weights:
             terms.update(w.find(sym))
 
-    args_present = filter_ordered(term.args[1:] for term in terms)
+    args_present = {}
+    for term in terms:
+        key = term.args[1:]
+        idx = term.args[0]
+        args_present.setdefault(key, []).append(idx)
 
     subs = obj.substitutions
     if subs:
@@ -288,7 +280,7 @@ def generate_subs(deriv_order, function, index):
     args_provided = list(set(args_provided))
 
     rules = {}
-    not_provided = []
+    not_provided = {}
     for i0 in args_present:
         if any(i0 == i1 for i1 in args_provided):
             # Perfect match, as expected by the legacy custom coeffs API
@@ -313,10 +305,10 @@ def generate_subs(deriv_order, function, index):
         if mapper:
             rules.update(mapper)
         else:
-            not_provided.append(i0)
+            not_provided.update({i0: args_present[i0]})
 
-    for i in not_provided:
-        rules = {**rules, **generate_subs(*i)}
+    for i, indices in not_provided.items():
+        rules = {**rules, **generate_subs(*i, indices)}
 
     return rules
 

devito/finite_differences/derivative.py

@@ -356,8 +356,12 @@ def _eval_at(self, func):
         setup where one could have Eq(u(x + h_x/2), v(x).dx)) in which case v(x).dx
         has to be computed at x=x + h_x/2.
         """
-        # If an x0 already exists do not overwrite it
-        x0 = self.x0 or func.indices_ref._getters
+        # If an x0 already exists or evaluating at the same function (i.e u = u.dx)
+        # do not overwrite it
+        if self.x0 or self.side is not None or func.function is self.expr.function:
+            return self
+
+        x0 = func.indices_ref._getters
         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

devito/finite_differences/differentiable.py

@@ -121,6 +121,13 @@ def _fd(self):
     def _symbolic_functions(self):
         return frozenset([i for i in self._functions if i.coefficients == 'symbolic'])
 
+    @cached_property
+    def function(self):
+        if len(self._functions) == 1:
+            return self._functions.pop()
+        else:
+            return None
+
     @cached_property
     def _uses_symbolic_coefficients(self):
         return bool(self._symbolic_functions)

devito/finite_differences/finite_difference.py

@@ -213,7 +213,7 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
     side = None
     # First order derivative with 2nd order FD is strongly discouraged so taking
     # first order fd that is a lot better
-    if deriv_order == 1 and fd_order == 2 and coefficients != 'symbolic':
+    if deriv_order == 1 and fd_order == 2:
         fd_order = 1
 
     # Zeroth order derivative is just the expression itself if not shifted

devito/finite_differences/tools.py

@@ -275,6 +275,7 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None):
     -------
     An IndexSet, representing an ordered list of indices.
     """
+    stagg_fd = expr.is_Staggered and x0
     # Evaluation point
     x0 = sympify(((x0 or {}).get(dim) or expr.indices_ref[dim]))
 
@@ -284,12 +285,17 @@ def generate_indices(expr, dim, order, side=None, matvec=None, x0=None):
         iexpr = x0 + d * dim.spacing
         return IndexSet(dim, expr=iexpr), x0
 
-    # Shift for side
-    side = side or centered
-
     # Evaluation point relative to the expression's grid
     mid = (x0 - expr.indices_ref[dim]).subs({dim: 0, dim.spacing: 1})
 
+    # Centered scheme for staggered field create artifacts if not shifted
+    if stagg_fd and mid == 0 and not dim.is_Time and side is not centered:
+        mid = 0.5 if x0 is dim else -0.5
+        x0 = x0 + mid * dim.spacing
+
+    # Shift for side
+    side = side or centered
+
     # Indices range
     o_min = int(np.ceil(mid - order/2)) + side.val
     o_max = int(np.floor(mid + order/2)) + side.val

examples/seismic/tti/operators.py

@@ -299,10 +299,10 @@ def kernel_staggered_2d(model, u, v, **kwargs):
     else:
         # Stencils
         phdx = ((costheta*epsilon*u).dx - (sintheta*epsilon*u).dyc +
-                (costheta*delta*v).dxc - (sintheta*delta*v).dy)
+                (costheta*delta*v).dx - (sintheta*delta*v).dyc)
         u_vx = Eq(vx.backward, dampl * vx + dampl * s * phdx)
 
-        pvdz = ((sintheta*delta*u).dx + (costheta*delta*u).dyc +
+        pvdz = ((sintheta*delta*u).dxc + (costheta*delta*u).dy +
                 (sintheta*v).dxc + (costheta*v).dy)
         u_vz = Eq(vz.backward, dampl * vz + dampl * s * pvdz)
 

examples/seismic/tutorials/06_elastic_varying_parameters.ipynb

@@ -392,22 +392,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Operator `Kernel` ran in 0.22 s\n"
+      "Operator `Kernel` ran in 0.23 s\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "  PerfEntry(time=0.19680499999999987, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.20279700000000003, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section1', rank=None),\n",
-       "  PerfEntry(time=0.005080999999999991, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.005569999999999994, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section2', rank=None),\n",
-       "  PerfEntry(time=0.005789000000000039, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.006224000000000026, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section3', rank=None),\n",
-       "  PerfEntry(time=0.005510000000000015, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.006004000000000021, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section4', rank=None),\n",
-       "  PerfEntry(time=0.005339000000000027, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "  PerfEntry(time=0.00578600000000002, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
      "execution_count": 14,
@@ -511,22 +511,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Operator `Kernel` ran in 0.24 s\n"
+      "Operator `Kernel` ran in 0.23 s\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "  PerfEntry(time=0.20575800000000005, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.20478399999999966, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section1', rank=None),\n",
-       "  PerfEntry(time=0.006976999999999987, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.005923000000000004, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section2', rank=None),\n",
-       "  PerfEntry(time=0.007339000000000036, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.00643000000000003, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section3', rank=None),\n",
-       "  PerfEntry(time=0.0069080000000000166, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.006100000000000016, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section4', rank=None),\n",
-       "  PerfEntry(time=0.006813000000000026, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "  PerfEntry(time=0.006256000000000025, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
      "execution_count": 17,
@@ -699,20 +699,20 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Operator `Kernel` ran in 0.21 s\n"
+      "Operator `Kernel` ran in 0.26 s\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "  PerfEntry(time=0.1906600000000002, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.22333499999999995, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section1', rank=None),\n",
-       "  PerfEntry(time=0.005354999999999988, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.008422999999999998, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section2', rank=None),\n",
-       "  PerfEntry(time=0.005681000000000019, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.009065000000000017, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section3', rank=None),\n",
-       "  PerfEntry(time=0.005547000000000023, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "  PerfEntry(time=0.009086000000000014, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
      "execution_count": 24,
@@ -796,20 +796,20 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Operator `Kernel` ran in 0.27 s\n"
+      "Operator `Kernel` ran in 0.22 s\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "  PerfEntry(time=0.22429000000000004, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.19835300000000017, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section1', rank=None),\n",
-       "  PerfEntry(time=0.013601000000000016, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.005520999999999994, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section2', rank=None),\n",
-       "  PerfEntry(time=0.012073000000000028, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.005916000000000023, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section3', rank=None),\n",
-       "  PerfEntry(time=0.012027000000000043, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "  PerfEntry(time=0.006079000000000015, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
      "execution_count": 26,
@@ -1016,22 +1016,22 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Operator `Kernel` ran in 0.57 s\n"
+      "Operator `Kernel` ran in 0.71 s\n"
      ]
     },
     {
      "data": {
       "text/plain": [
        "PerformanceSummary([(PerfKey(name='section0', rank=None),\n",
-       "  PerfEntry(time=0.5206050000000007, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.6343909999999997, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section1', rank=None),\n",
-       "  PerfEntry(time=0.010398000000000041, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.017495000000000028, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section2', rank=None),\n",
-       "  PerfEntry(time=0.012117999999999964, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.018152000000000022, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section3', rank=None),\n",
-       "  PerfEntry(time=0.011993000000000005, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
+       "  PerfEntry(time=0.017839, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),\n",
        " (PerfKey(name='section4', rank=None),\n",
-       "  PerfEntry(time=0.011515999999999974, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
+       "  PerfEntry(time=0.017724000000000007, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])"
       ]
      },
      "execution_count": 34,

tests/test_derivatives.py

@@ -440,6 +440,20 @@ def test_new_x0_eval_at(self):
         v = Function(name="v", grid=grid, space_order=2)
         assert u.dx(x0=x - x.spacing/2)._eval_at(v).x0 == {x: x - x.spacing/2}
 
+    @pytest.mark.parametrize('stagg', [True, False])
+    def test_eval_at_centered(self, stagg):
+        grid = Grid((10,))
+        x = grid.dimensions[0]
+        stagg = NODE if stagg else x
+        x0 = x if stagg else x + .5 * x.spacing
+
+        u = Function(name="u", grid=grid, space_order=2, staggered=stagg)
+        v = Function(name="v", grid=grid, space_order=2, staggered=stagg)
+
+        assert u.dx._eval_at(v).evaluate == u.dx(x0=x0).evaluate
+        assert v.dx._eval_at(u).evaluate == v.dx(x0=x0).evaluate
+        assert u.dx._eval_at(u).evaluate == u.dx.evaluate
+
     def test_fd_new_lo(self):
         grid = Grid((10,))
         x = grid.dimensions[0]

tests/test_mpi.py

@@ -2334,9 +2334,9 @@ def test_staggering(self, mode):
         op(time_M=2)
 
         # Expected norms computed "manually" from sequential runs
-        assert np.isclose(norm(ux), 6253.4349, rtol=1.e-4)
-        assert np.isclose(norm(uxx), 80001.0304, rtol=1.e-4)
-        assert np.isclose(norm(uxy), 61427.853, rtol=1.e-4)
+        assert np.isclose(norm(ux), 6042.554, rtol=1.e-4)
+        assert np.isclose(norm(uxx), 64632.75, rtol=1.e-4)
+        assert np.isclose(norm(uxy), 59737.77, rtol=1.e-4)
 
     @pytest.mark.parallel(mode=2)
     def test_op_new_dist(self, mode):

tests/test_staggered_utils.py

@@ -99,7 +99,7 @@ def test_gather_for_diff(expr, expected):
 
 @pytest.mark.parametrize('expr, expected', [
     ('((a + b).dx._eval_at(a)).is_Add', 'True'),
-    ('(a + b).dx._eval_at(a)', 'a.dx._eval_at(a) + b.dx._eval_at(a)'),
+    ('(a + b).dx._eval_at(a)', 'a.dx(x0=a.indices_ref._getters) + b.dx._eval_at(a)'),
     ('(a*b).dx._eval_at(a).expr', 'a.subs({x: x0}) * b'),
     ('(a * b.dx).dx._eval_at(b).expr._eval_deriv ',
      'a.subs({x: x0}) * b.dx.evaluate')])
@@ -110,7 +110,6 @@ def test_stagg_fd_composite(expr, expected):
     y0 = y + y.spacing/2  # noqa
     a = Function(name="a", grid=grid, staggered=NODE)  # noqa
     b = Function(name="b", grid=grid, staggered=x)  # noqa
-
     assert eval(expr) == eval(expected)