Merge pull request #1786 from devitocodes/sympy19

reqs: Extend SymPy support to 1.9

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

@@ -19,7 +19,6 @@ jobs:
       DEVITO_ARCH: "${{ matrix.arch }}"
       DEVITO_LANGUAGE: ${{ matrix.language }}
       OMP_NUM_THREADS: 2
-      PYTHON_VERSION: "${{ matrix.python-version }}"
 
     strategy:
       # Prevent all build to stop if a single one fails
@@ -29,7 +28,7 @@ jobs:
         name: [
            pytest-ubuntu-py37-gcc5-omp,
            pytest-ubuntu-py38-gcc6-omp,
-           pytest-ubuntu-py36-gcc7-omp-sympy17,
+           pytest-ubuntu-py36-gcc7-omp,
            pytest-ubuntu-py37-gcc7-noomp,
            pytest-ubuntu-py37-gcc8-omp,
            pytest-ubuntu-py38-gcc9-omp,
@@ -43,48 +42,56 @@ jobs:
           os: ubuntu-18.04
           arch: "gcc-5"
           language: "openmp"
+          sympy: 1.8
 
         - name: pytest-ubuntu-py38-gcc6-omp
           python-version: 3.8
           os: ubuntu-18.04
           arch: "gcc-6"
           language: "openmp"
+          sympy: 1.7
 
-        - name: pytest-ubuntu-py36-gcc7-omp-sympy17
+        - name: pytest-ubuntu-py36-gcc7-omp
           python-version: 3.6
           os: ubuntu-18.04
           arch: "gcc-7"
           language: "openmp"
+          sympy: 1.8
 
         - name: pytest-ubuntu-py37-gcc7-noomp
           python-version: 3.7
           os: ubuntu-18.04
           arch: "gcc-7"
           language: "C"
+          sympy: 1.8
 
         - name: pytest-ubuntu-py37-gcc8-omp
           python-version: 3.7
           os: ubuntu-18.04
           arch: "gcc-8"
           language: "openmp"
+          sympy: 1.9
 
         - name: pytest-ubuntu-py38-gcc9-omp
           python-version: 3.8
           os: ubuntu-20.04
           arch: "gcc-9"
           language: "openmp"
+          sympy: 1.8
 
         - name: pytest-osx-py37-clang-omp
           python-version: 3.7
           os: macos-latest
           arch: "osx"
           language: "C"
+          sympy: 1.8
 
         - name: pytest-docker-py36-gcc-omp
           python-version: 3.6
           os: ubuntu-18.04
           arch: "gcc"
           language: "openmp"
+          sympy: 1.8
 
         - set: base
           test-set: 'not adjoint'
@@ -139,10 +146,7 @@ jobs:
       run: |
         pip install --upgrade pip
         pip install -e .
-
-    - name: Alternative sympy version
-      if: matrix.name == 'pytest-ubuntu-py36-gcc7-omp-sympy17'
-      run: pip install sympy==1.7
+        pip install sympy==${{matrix.sympy}}
 
     - name: Test with pytest
       run: |

devito/tools/abc.py

@@ -147,7 +147,7 @@ def __reduce_ex__(self, proto):
 
     def __getnewargs_ex__(self):
         return (tuple(getattr(self, i) for i in self._pickle_args),
-                {i.lstrip('_'): getattr(self, i) for i in self._pickle_kwargs})
+                {i: getattr(self, i) for i in self._pickle_kwargs})
 
 
 class Singleton(type):

devito/types/basic.py

@@ -220,6 +220,10 @@ class AbstractSymbol(sympy.Symbol, Basic, Pickable, Evaluable):
     def _filter_assumptions(cls, **kwargs):
         """Extract sympy.Symbol-specific kwargs."""
         assumptions = {}
+        # pop predefined assumptions
+        for key in ('real', 'imaginary', 'commutative'):
+            kwargs.pop(key, None)
+        # extract sympy.Symbol-specific kwargs
         for i in list(kwargs):
             if i in _assume_rules.defined_facts:
                 assumptions[i] = kwargs.pop(i)
@@ -333,6 +337,11 @@ def _subs(self, old, new, **hints):
     _pickle_kwargs = ['name', 'dtype', 'is_const']
     __reduce_ex__ = Pickable.__reduce_ex__
 
+    def __getnewargs_ex__(self):
+        args, kwargs = Pickable.__getnewargs_ex__(self)
+        kwargs.update(self.assumptions0)
+        return args, kwargs
+
 
 class Symbol(AbstractSymbol, Cached):
 
@@ -368,22 +377,23 @@ def _cache_key(cls, *args, **kwargs):
         return frozendict(key)
 
     def __new__(cls, *args, **kwargs):
-        key = cls._cache_key(*args, **kwargs)
+        assumptions, kwargs = cls._filter_assumptions(**kwargs)
+        key = cls._cache_key(*args, **{**assumptions, **kwargs})
         obj = cls._cache_get(key)
 
         if obj is not None:
             return obj
 
         # Not in cache. Create a new Symbol via sympy.Symbol
-        name = kwargs.get('name') or args[0]
-        assumptions, kwargs = cls._filter_assumptions(**kwargs)
+        args = list(args)
+        name = kwargs.pop('name', None) or args.pop(0)
 
         # Note: use __xnew__ to bypass sympy caching
         newobj = sympy.Symbol.__xnew__(cls, name, **assumptions)
 
         # Initialization
         newobj._dtype = cls.__dtype_setup__(**kwargs)
-        newobj.__init_finalize__(*args, **kwargs)
+        newobj.__init_finalize__(name, *args, **kwargs)
 
         # Store new instance in symbol cache
         Cached.__init__(newobj, key)
@@ -536,22 +546,11 @@ def _new(cls, *args, **kwargs):
                 newobj = super(AbstractTensor, cls)._new(args[2])
 
             # Filter grid and dimensions
-            grids = {getattr(c, 'grid', None) for c in newobj._mat} - {None}
-            dimensions = {d for c in newobj._mat
-                          for d in getattr(c, 'dimensions', ())} - {None}
-            # If none of the components are devito objects, returns a sympy Matrix
-            if len(grids) == 0 and len(dimensions) == 0:
+            grid, dimensions = newobj._infer_dims()
+            if grid is None and dimensions is None:
                 return sympy.ImmutableDenseMatrix(*args)
-            elif len(grids) > 0:
-                dimensions = None
-                assert len(grids) == 1
-                grid = grids.pop()
-            else:
-                grid = None
-                dimensions = tuple(dimensions)
-
             # Initialized with constructed object
-            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj._mat,
+            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
                                      grid=grid, dimensions=dimensions)
         else:
             # Initialize components and create new Matrix from standard
@@ -562,6 +561,52 @@ def _new(cls, *args, **kwargs):
 
         return newobj
 
+    @classmethod
+    def _fromrep(cls, rep):
+        """
+        This the new constructor mechanism for matrices in sympy 1.9.
+        Standard new object go through `_new` but arithmetic operations directly use
+        the representation based one.
+        This class method is only accessible from an existing AbstractTensor
+        that contains a grid or dimensions.
+        """
+        newobj = super(AbstractTensor, cls)._fromrep(rep)
+        grid, dimensions = newobj._infer_dims()
+        try:
+            # This is needed when `_fromrep` is called directly in 1.9
+            # for example with mul.
+            newobj.__init_finalize__(newobj.rows, newobj.cols, newobj.flat(),
+                                     grid=grid, dimensions=dimensions)
+        except TypeError:
+            # We can end up here when `_fromrep` is called through the default _new
+            # when input `comps` don't have grid or dimensions. For example
+            # `test_non_devito_tens` in `test_tensor.py`.
+            pass
+        return newobj
+
+    def _infer_dims(self):
+        grids = {getattr(c, 'grid', None) for c in self.flat()} - {None}
+        dimensions = {d for c in self.flat()
+                      for d in getattr(c, 'dimensions', ())} - {None}
+        # If none of the components are devito objects, returns a sympy Matrix
+        if len(grids) == 0 and len(dimensions) == 0:
+            return None, None
+        elif len(grids) > 0:
+            dimensions = None
+            assert len(grids) == 1
+            grid = grids.pop()
+        else:
+            grid = None
+            dimensions = tuple(dimensions)
+
+        return grid, dimensions
+
+    def flat(self):
+        try:
+            return super().flat()
+        except AttributeError:
+            return self._mat
+
     def __init_finalize__(self, *args, **kwargs):
         pass
 
@@ -583,8 +628,11 @@ def _eval_matrix_mul(self, other):
         # expected behavior is to produce an n x m matrix of zeros
         if self.cols != 0 and other.rows != 0:
             self_cols = self.cols
-            mat = self._mat
-            other_mat = other._mat
+            mat = self.flat()
+            try:
+                other_mat = other.flat()
+            except AttributeError:
+                other_mat = other._mat
             for i in range(new_len):
                 row, col = i // other.cols, i % other.cols
                 row_indices = range(self_cols*row, self_cols*(row+1))

devito/types/constant.py

@@ -55,6 +55,10 @@ def __dtype_setup__(cls, **kwargs):
     def is_const(self):
         return True
 
+    @property
+    def value(self):
+        return self.data
+
     @property
     def data(self):
         """The value of the data object, as a scalar (int, float, ...)."""
@@ -108,4 +112,4 @@ def _arg_check(self, args, intervals):
         except AttributeError:
             pass
 
-    _pickle_kwargs = DataSymbol._pickle_kwargs + ['_value']
+    _pickle_kwargs = DataSymbol._pickle_kwargs + ['value']

devito/types/dimension.py

@@ -310,6 +310,7 @@ def _arg_check(self, args, size, interval):
     _pickle_args = ['name']
     _pickle_kwargs = ['spacing']
     __reduce_ex__ = Pickable.__reduce_ex__
+    __getnewargs_ex__ = Pickable.__getnewargs_ex__
 
 
 class BasicDimension(Dimension, Symbol):

devito/types/parallel.py

@@ -33,7 +33,7 @@ class NThreadsBase(Scalar):
     is_Input = True
     is_PerfKnob = True
 
-    def __new__(cls, **kwargs):
+    def __new__(cls, *args, **kwargs):
         kwargs.setdefault('name', cls.name)
         kwargs['is_const'] = True
         return super().__new__(cls, **kwargs)
@@ -71,7 +71,7 @@ class NPThreads(NThreadsBase):
 
     name = 'npthreads'
 
-    def __new__(cls, **kwargs):
+    def __new__(cls, *args, **kwargs):
         obj = super().__new__(cls, **kwargs)
 
         # Size of the thread pool
@@ -388,7 +388,7 @@ class DeviceSymbol(Scalar):
     is_Input = True
     is_PerfKnob = True
 
-    def __new__(cls, **kwargs):
+    def __new__(cls, *args, **kwargs):
         kwargs['name'] = cls.name
         kwargs['is_const'] = True
         return super().__new__(cls, **kwargs)

requirements.txt

@@ -1,6 +1,6 @@
 pip>=9.0.1
 numpy>1.16
-sympy>=1.7,<=1.8
+sympy>=1.7,<=1.9
 scipy
 flake8>=2.1.0
 nbval