Merge pull request #6 from ZoeLeibowitz/PETScArray

Add liveness to PETSc Array

devito/types/__init__.py

@@ -6,7 +6,6 @@
 from .object import *  # noqa
 from .lazy import *  # noqa
 from .misc import *  # noqa
-from .petsc import *  # noqa
 
 # Needed both within and outside Devito
 from .dimension import *  # noqa
@@ -16,6 +15,7 @@
 
 # Some more internal types which depend on some of the types above
 from .parallel import *  # noqa
+from .petsc import *  # noqa
 
 # Needed only outside Devito
 from .grid import *  # noqa

devito/types/array.py

@@ -17,9 +17,11 @@ class ArrayBasic(AbstractFunction):
 
     is_ArrayBasic = True
 
-    __rkwargs__ = AbstractFunction.__rkwargs__ + ('is_const',)
+    __rkwargs__ = AbstractFunction.__rkwargs__ + ('is_const', 'liveness')
 
     def __init_finalize__(self, *args, **kwargs):
+
+        self._liveness = kwargs.setdefault('liveness', 'lazy')
         super().__init_finalize__(*args, **kwargs)
         self._is_const = kwargs.get('is_const', False)
 
@@ -56,6 +58,18 @@ def shape_allocated(self):
     def is_const(self):
         return self._is_const
 
+    @property
+    def liveness(self):
+        return self._liveness
+
+    @property
+    def _mem_internal_eager(self):
+        return self._liveness == 'eager'
+
+    @property
+    def _mem_internal_lazy(self):
+        return self._liveness == 'lazy'
+
 
 class Array(ArrayBasic):
 
@@ -126,9 +140,6 @@ def __new__(cls, *args, **kwargs):
     def __init_finalize__(self, *args, **kwargs):
         super().__init_finalize__(*args, **kwargs)
 
-        self._liveness = kwargs.get('liveness', 'lazy')
-        assert self._liveness in ['eager', 'lazy']
-
         self._scope = kwargs.get('scope', 'heap')
         assert self._scope in ['heap', 'stack', 'static', 'constant', 'shared']
 
@@ -153,10 +164,6 @@ def __padding_setup__(self, **kwargs):
             raise TypeError("`padding` must be int or %d-tuple of ints" % self.ndim)
         return DimensionTuple(*padding, getters=self.dimensions)
 
-    @property
-    def liveness(self):
-        return self._liveness
-
     @property
     def scope(self):
         return self._scope
@@ -165,14 +172,6 @@ def scope(self):
     def _C_ctype(self):
         return POINTER(dtype_to_ctype(self.dtype))
 
-    @property
-    def _mem_internal_eager(self):
-        return self._liveness == 'eager'
-
-    @property
-    def _mem_internal_lazy(self):
-        return self._liveness == 'lazy'
-
     @property
     def _mem_stack(self):
         return self._scope in ('stack', 'shared')

devito/types/petsc.py

@@ -3,6 +3,8 @@
 from devito.types.array import ArrayBasic
 import numpy as np
 from cached_property import cached_property
+from devito.finite_differences import Differentiable
+from devito.types.basic import AbstractFunction
 
 
 class DM(LocalObject):
@@ -57,15 +59,22 @@ class KSPConvergedReason(LocalObject):
     dtype = CustomDtype('KSPConvergedReason')
 
 
-class PETScArray(ArrayBasic):
+class PETScArray(ArrayBasic, Differentiable):
     """
     PETScArrays are generated by the compiler only and represent
     a customised variant of ArrayBasic. They are designed to
     avoid generating a cast in the low-level code.
+    Differentiable enables compatability with standard Function objects,
+    allowing for the use of the `subs` method.
+    TODO: Potentially re-evaluate and separate into PETScFunction(Differentiable)
+    and then PETScArray(ArrayBasic).
     """
 
     _data_alignment = False
 
+    __rkwargs__ = (AbstractFunction.__rkwargs__ +
+                   ('dimensions', 'liveness'))
+
     @classmethod
     def __dtype_setup__(cls, **kwargs):
         return kwargs.get('dtype', np.float32)

tests/test_petsc.py

@@ -1,4 +1,4 @@
-from devito import Grid
+from devito import Grid, Function, Eq
 from devito.ir.iet import Call, ElementalFunction, Definition, DummyExpr
 from devito.passes.iet.languages.C import CDataManager
 from devito.types import (DM, Mat, Vec, PetscMPIInt, KSP,
@@ -63,3 +63,24 @@ def test_petsc_functions():
     assert str(defn3) == 'PetscInt**restrict ptr3;'
     assert str(defn4) == 'const PetscInt**restrict ptr4;'
     assert str(expr) == 'ptr0[x][y] = ptr1[x][y] + 1;'
+
+
+def test_petsc_subs():
+    """
+    Test support for PETScArrays in substitutions.
+    """
+    grid = Grid((2, 2))
+
+    f1 = Function(name='f1', grid=grid, space_order=2)
+    f2 = Function(name='f2', grid=grid, space_order=2)
+
+    arr = PETScArray(name='arr', dimensions=f2.dimensions, dtype=f2.dtype)
+
+    eqn = Eq(f1, f2.laplace)
+    eqn_subs = eqn.subs(f2, arr)
+
+    assert str(eqn) == 'Eq(f1(x, y), Derivative(f2(x, y), (x, 2))' +  \
+        ' + Derivative(f2(x, y), (y, 2)))'
+
+    assert str(eqn_subs) == 'Eq(f1(x, y), Derivative(arr(x, y), (x, 2))' +  \
+        ' + Derivative(arr(x, y), (y, 2)))'