Merge pull request #3002 from devitocodes/patch-sympy-tensor-args

dsl: misc patches from recent updates (interp, sympy args)

Author: mloubout

devito/types/basic.py

@@ -4,6 +4,7 @@
 from contextlib import contextmanager, suppress
 from ctypes import POINTER, Structure, _Pointer, c_char, c_char_p
 from functools import cached_property, reduce
+from numbers import Number
 from operator import mul
 
 import numpy as np
@@ -1539,7 +1540,10 @@ def _new(cls, *args, **kwargs):
             # Filter grid and dimensions
             grid, dimensions = newobj._infer_dims()
             if grid is None and dimensions is None:
-                return sympy.ImmutableDenseMatrix(*args)
+                # Downgrade to a plain Matrix, reusing the representation rather
+                # than rebuilding from `args`, as the latter would sympify the
+                # entries with sympy's `sympify` instead of `cls._sympify`
+                return sympy.ImmutableDenseMatrix._fromrep(newobj._rep)
             # Initialized with constructed object
             newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
                                      grid=grid, dimensions=dimensions)
@@ -1581,13 +1585,18 @@ def __subfunc_setup__(cls, *args, **kwargs):
     @classmethod
     def _sympify(cls, arg):
         # This is used internally by sympy to process arguments at rebuilt. And since
-        # some of our properties are non-sympyfiable we need to have a fallback.
-        # `strict` so that strings are left alone rather than parsed into Symbols,
-        # while plain numbers are turned into `Expr` as sympy expects (a Matrix
-        # holding non-`Expr` entries, such as a plain `int` 0, is deprecated)
+        # some of our properties are non-sympyfiable we need to have a fallback
+        if isinstance(arg, Number):
+            # Plain numbers must be sympified, as sympy assigns the `EXRAW` domain
+            # to a Matrix holding non-`Expr` entries such as a plain `int` 0
+            return sympy.sympify(arg)
         try:
-            return sympy.sympify(arg, strict=True)
-        except sympy.SympifyError:
+            # Pure sympy object
+            return arg._sympy_()
+        except AttributeError:
+            # Anything else, such as a `Staggering`, is passed through untouched.
+            # Note that sympifying is not an option here, as it would convert
+            # away the type, `Staggering` being a `tuple` for example
             return arg
 
     @classmethod

devito/types/sparse.py

@@ -1008,20 +1008,24 @@ def _arg_defaults(self, alias=None, estimate_memory=False):
         if estimate_memory:
             return defaults
         key = alias or self
-        coords = defaults.get(key.coordinates.name, key.coordinates.data)
+        coords = defaults.get(key.coordinates.name, self.coordinates.data)
         defaults.update(key.interpolator._arg_defaults(coords=coords,
-                                                       sfunc=key))
+                                                       sfunc=self))
         return defaults
 
     def _arg_values(self, estimate_memory=False, **kwargs):
         values = super()._arg_values(estimate_memory=estimate_memory, **kwargs)
         if estimate_memory:
             return values
 
-        # Resolve the runtime grid origin (honours `o_x`/`o_y`/... overrides)
-        # and hand it to the interpolator so tables reflect the actual frame
-        # of reference used by the kernel.
+        # `super` has already tabulated through `_arg_defaults`, in the frame
+        # of whichever object supplied the runtime values.  Only an explicit
+        # `o_x`/`o_y`/... override moves that frame again, and the tables then
+        # have to be rebuilt against it.
         onames = [o.name for o in self.grid.origin_symbols]
+        if not any(n in kwargs for n in onames):
+            return values
+
         origin = tuple(kwargs.get(n, o) for n, o in
                        zip(onames, self.grid.origin, strict=True))
         coords = values.get(self.coordinates.name, self.coordinates.data)

devito/types/tensor.py

@@ -24,12 +24,15 @@ def staggering(stagg, i, j, d, dims):
     if stagg is None:
         # No input
         return NODE if i == j else (d, dims[j])
+    elif isinstance(stagg, MatrixBase):
+        # From rebuild/tensor property. Indexed as a sympy Matrix. Note that this
+        # may be a plain Matrix rather than an AbstractTensor, as rebuilding a
+        # tensor component-wise downgrades it when the components aren't Devito
+        # objects, which is the case for a Matrix of `Staggering`
+        return stagg[i, j]
     elif isinstance(stagg, (tuple, list)):
         # User input as list or tuple
         return stagg[i][j]
-    elif isinstance(stagg, AbstractTensor):
-        # From rebuild/tensor property. Indexed as a sympy Matrix
-        return stagg[i, j]
 
 
 class TensorFunction(AbstractTensor):

examples/seismic/tti/operators.py

@@ -62,7 +62,7 @@ def trig_func(model):
     return costheta, sintheta
 
 
-def Gzz_centered(model, field):
+def Gzz_centered(model, field, b=None):
     """
     3D rotated second order derivative in the direction z.
 
@@ -72,12 +72,17 @@ def Gzz_centered(model, field):
         Physical parameters model structure.
     field : Function
         Input for which the derivative is computed.
+    b : Function, optional
+        Buoyancy to build the operator with, defaulting to the model's. Since
+        the operator is linear in it, passing a perturbation here gives the
+        derivative of the operator with respect to the buoyancy in that
+        direction.
 
     Returns
     -------
     Rotated second order derivative w.r.t. z.
     """
-    b = getattr(model, 'b', 1)
+    b = getattr(model, 'b', 1) if b is None else b
     costheta, sintheta, cosphi, sinphi = trig_func(model)
 
     order1 = field.space_order // 2
@@ -99,7 +104,7 @@ def Gzz_centered(model, field):
     return Gzz
 
 
-def Gzz_centered_2d(model, field):
+def Gzz_centered_2d(model, field, b=None):
     """
     2D rotated second order derivative in the direction z.
 
@@ -109,12 +114,17 @@ def Gzz_centered_2d(model, field):
         Physical parameters model structure.
     field : Function
         Input for which the derivative is computed.
+    b : Function, optional
+        Buoyancy to build the operator with, defaulting to the model's. Since
+        the operator is linear in it, passing a perturbation here gives the
+        derivative of the operator with respect to the buoyancy in that
+        direction.
 
     Returns
     -------
     Rotated second order derivative w.r.t. z.
     """
-    b = getattr(model, 'b', 1)
+    b = getattr(model, 'b', 1) if b is None else b
     costheta, sintheta = trig_func(model)
 
     order1 = field.space_order // 2
@@ -133,7 +143,7 @@ def Gzz_centered_2d(model, field):
 
 
 # Centered case produces directly Gxx + Gyy
-def Gh_centered(model, field):
+def Gh_centered(model, field, b=None):
     """
     Sum of the 3D rotated second order derivative in the direction x and y.
     As the Laplacian is rotation invariant, it is computed as the conventional
@@ -146,13 +156,19 @@ def Gh_centered(model, field):
         Physical parameters model structure.
     field : Function
         Input field.
+    b : Function, optional
+        Buoyancy to build the operator with, defaulting to the model's. See
+        :func:`Gzz_centered`.
 
     Returns
     -------
     Sum of the 3D rotated second order derivative in the direction x and y.
     """
-    Gzz = Gzz_centered(model, field) if model.dim == 3 else Gzz_centered_2d(model, field)
-    b = getattr(model, 'b', None)
+    b = getattr(model, 'b', None) if b is None else b
+    if model.dim == 3:  # noqa: SIM108
+        Gzz = Gzz_centered(model, field, b=b)
+    else:
+        Gzz = Gzz_centered_2d(model, field, b=b)
     if b is not None:
         _diff = lambda f, d: getattr(f, f'd{d.name}')
         so = field.space_order // 2

tests/test_interpolation.py

@@ -1077,6 +1077,37 @@ def test_position(self, shape):
 
         assert(np.allclose(rec.data, rec1.data, atol=1e-5))
 
+    @pytest.mark.parametrize('interpolation,r', [('linear', 1), ('sinc', 4)])
+    def test_position_override_grid(self, interpolation, r):
+        """
+        Inject through an Operator built on a grid whose origin differs from
+        the one it is applied to, as when an Operator compiled against one
+        model is applied to another. The point must land where the runtime
+        origin puts it, not the compile-time one.
+        """
+        shape, spacing, coord = (41, 41), (10., 10.), 120.
+        extent = tuple((s - 1) * h for s, h in zip(shape, spacing, strict=True))
+        kw = dict(interpolation=interpolation, r=r)
+
+        def setup(origin):
+            grid = Grid(shape=shape, extent=extent, origin=origin)
+            u = TimeFunction(name='u', grid=grid, space_order=8)
+            src = SparseTimeFunction(name='src', grid=grid, npoint=1, nt=2, **kw)
+            src.coordinates.data[0, :] = coord
+            src.data[:] = 1.
+            return u, src
+
+        u_build, src_build = setup((0., 0.))
+        op = Operator(src_build.inject(field=u_build.forward, expr=src_build))
+
+        shift = -100.
+        u, src = setup((shift, shift))
+        op.apply(time_M=0, u=u, src=src)
+
+        expected = tuple(int((coord - shift) / h) for h in spacing)
+        peak = np.unravel_index(np.argmax(np.abs(u.data)), u.data.shape)[1:]
+        assert peak == expected
+
     def test_sparse_first(self):
         """
         Tests custom sprase function with sparse dimension as first index.

tests/test_tensors.py

@@ -12,6 +12,7 @@
 )
 from devito.symbolics import retrieve_derivatives
 from devito.types import NODE
+from devito.types.utils import Staggering
 
 
 def dimify(dimensions):
@@ -528,6 +529,25 @@ def test_diag_sympified_zeros(func1):
     assert all(isinstance(c, sympy.Expr) for c in f2.flat())
 
 
+@pytest.mark.parametrize('func1', [TensorFunction, TensorTimeFunction,
+                                   VectorFunction, VectorTimeFunction])
+def test_staggered_attribute_roundtrip(func1):
+    """
+    Accessing an attribute rebuilds the tensor component-wise, which must not
+    sympify a `Staggering` away, otherwise it can no longer be fed back as the
+    `staggered` kwarg.
+    """
+    grid = Grid(tuple([5]*3))
+    f1 = func1(name="f1", grid=grid, time_order=1)
+
+    stagg = f1.staggered
+    assert all(isinstance(s, Staggering) for s in stagg.flat())
+
+    f2 = func1(name="f2", grid=grid, time_order=1, staggered=stagg)
+    assert all(c1.staggered == c2.staggered
+               for c1, c2 in zip(f1.flat(), f2.flat(), strict=True))
+
+
 def test_non_expr_components():
     """
     A tensor may legitimately hold non-`Expr` components, which sympy deprecates