Merge pull request #2811 from devitocodes/grad-kwargs

api: Extend side API to all derivative operator shortcuts

Author: FabioLuporini

devito/finite_differences/derivative.py

@@ -38,7 +38,7 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
         Derivative order.
     side : Side or tuple of Side, optional, default=centered
         Side of the finite difference location, centered (at x), left (at x - 1)
-        or right (at x +1).
+        or right (at x + 1).
     transpose : Transpose, optional, default=direct
         Forward (matvec=direct) or transpose (matvec=transpose) mode of the
         finite difference.

devito/finite_differences/differentiable.py

@@ -310,7 +310,7 @@ def laplace(self):
         """
         return self.laplacian()
 
-    def laplacian(self, shift=None, order=None, method='FD', **kwargs):
+    def laplacian(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Laplacian of the Differentiable with shifted derivatives and custom
         FD order.
@@ -329,19 +329,22 @@ def laplacian(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
-            Custom weights for the finite differences.
+            Custom weights for the finite difference coefficients.
         """
         w = kwargs.get('weights', kwargs.get('w'))
         order = order or self.space_order
         space_dims = self.root_dimensions
         shift_x0 = make_shift_x0(shift, (len(space_dims),))
         derivs = tuple(f'd{d.name}2' for d in space_dims)
         return Add(*[getattr(self, d)(x0=shift_x0(shift, space_dims[i], None, i),
-                                      method=method, fd_order=order, w=w)
+                                      method=method, fd_order=order, side=side, w=w)
                      for i, d in enumerate(derivs)])
 
-    def div(self, shift=None, order=None, method='FD', **kwargs):
+    def div(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Divergence of the input Function.
 
@@ -357,6 +360,9 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite difference coefficients.
         """
@@ -365,10 +371,11 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         shift_x0 = make_shift_x0(shift, (len(space_dims),))
         order = order or self.space_order
         return Add(*[getattr(self, f'd{d.name}')(x0=shift_x0(shift, d, None, i),
-                                                 fd_order=order, method=method, w=w)
+                                                 fd_order=order, method=method, side=side,
+                                                 w=w)
                      for i, d in enumerate(space_dims)])
 
-    def grad(self, shift=None, order=None, method='FD', **kwargs):
+    def grad(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Gradient of the input Function.
 
@@ -384,16 +391,21 @@ def grad(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
-            Custom weights for the finite
+            Custom weights for the finite difference coefficients.
         """
         from devito.types.tensor import VectorFunction, VectorTimeFunction
         space_dims = self.root_dimensions
         shift_x0 = make_shift_x0(shift, (len(space_dims),))
         order = order or self.space_order
+
         w = kwargs.get('weights', kwargs.get('w'))
         comps = [getattr(self, f'd{d.name}')(x0=shift_x0(shift, d, None, i),
-                                             fd_order=order, method=method, w=w)
+                                             fd_order=order, method=method, side=side,
+                                             w=w)
                  for i, d in enumerate(space_dims)]
         vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
         return vec_func(name=f'grad_{self.name}', time_order=self.time_order,

devito/finite_differences/operators.py

@@ -1,4 +1,4 @@
-def div(func, shift=None, order=None, method='FD', **kwargs):
+def div(func, shift=None, order=None, method='FD', side=None, **kwargs):
     """
     Divergence of the input Function.
 
@@ -14,12 +14,15 @@ def div(func, shift=None, order=None, method='FD', **kwargs):
     method: str, optional, default='FD'
         Discretization method. Options are 'FD' (default) and
         'RSFD' (rotated staggered grid finite-difference).
+    side : Side or tuple of Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x + 1).
     weights/w: list, tuple, or dict, optional, default=None
         Custom weights for the finite difference coefficients.
     """
     w = kwargs.get('weights', kwargs.get('w'))
     try:
-        return func.div(shift=shift, order=order, method=method, w=w)
+        return func.div(shift=shift, order=order, method=method, side=side, w=w)
     except AttributeError:
         return 0
 
@@ -41,7 +44,7 @@ def div45(func, shift=None, order=None):
     return div(func, shift=shift, order=order, method='RSFD')
 
 
-def grad(func, shift=None, order=None, method='FD', **kwargs):
+def grad(func, shift=None, order=None, method='FD', side=None, **kwargs):
     """
     Gradient of the input Function.
 
@@ -57,12 +60,15 @@ def grad(func, shift=None, order=None, method='FD', **kwargs):
     method: str, optional, default='FD'
         Discretization method. Options are 'FD' (default) and
         'RSFD' (rotated staggered grid finite-difference).
+    side : Side or tuple of Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x + 1).
     weights/w: list, tuple, or dict, optional, default=None
         Custom weights for the finite difference coefficients.
     """
     w = kwargs.get('weights', kwargs.get('w'))
     try:
-        return func.grad(shift=shift, order=order, method=method, w=w)
+        return func.grad(shift=shift, order=order, method=method, side=side, w=w)
     except AttributeError:
         raise AttributeError("Gradient not supported for class %s" % func.__class__)
 
@@ -84,7 +90,7 @@ def grad45(func, shift=None, order=None):
     return grad(func, shift=shift, order=order, method='RSFD')
 
 
-def curl(func, shift=None, order=None, method='FD', **kwargs):
+def curl(func, shift=None, order=None, method='FD', side=None, **kwargs):
     """
     Curl of the input Function. Only supported for VectorFunction
 
@@ -100,12 +106,15 @@ def curl(func, shift=None, order=None, method='FD', **kwargs):
     method: str, optional, default='FD'
         Discretization method. Options are 'FD' (default) and
         'RSFD' (rotated staggered grid finite-difference).
+    side : Side or tuple of Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x + 1).
     weights/w: list, tuple, or dict, optional, default=None
         Custom weights for the finite difference coefficients.
     """
     w = kwargs.get('weights', kwargs.get('w'))
     try:
-        return func.curl(shift=shift, order=order, method=method, w=w)
+        return func.curl(shift=shift, order=order, method=method, side=side, w=w)
     except AttributeError:
         raise AttributeError("Curl only supported for 3D VectorFunction")
 
@@ -128,7 +137,7 @@ def curl45(func, shift=None, order=None):
     return curl(func, shift=shift, order=order, method='RSFD')
 
 
-def laplace(func, shift=None, order=None, method='FD', **kwargs):
+def laplace(func, shift=None, order=None, method='FD', side=None, **kwargs):
     """
     Laplacian of the input Function.
 
@@ -143,12 +152,15 @@ def laplace(func, shift=None, order=None, method='FD', **kwargs):
         Uses `func.space_order` when not specified
     method: str, optional, default='FD'
         Discretization method. Options are 'FD' (default) and 'RSFD'
+    side : Side or tuple of Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x + 1).
     weights/w: list, tuple, or dict, optional, default=None
         Custom weights for the finite difference coefficients.
     """
     w = kwargs.get('weights', kwargs.get('w'))
     try:
-        return func.laplacian(shift=shift, order=order, method=method, w=w)
+        return func.laplacian(shift=shift, order=order, method=method, side=side, w=w)
     except AttributeError:
         return 0
 

devito/types/tensor.py

@@ -220,7 +220,7 @@ def values(self):
         else:
             return super().values()
 
-    def div(self, shift=None, order=None, method='FD', **kwargs):
+    def div(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Divergence of the TensorFunction (is a VectorFunction).
 
@@ -234,6 +234,9 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite differences.
         """
@@ -247,7 +250,7 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         for i in range(len(self.space_dimensions)):
             comps.append(sum([getattr(self[j, i], 'd%s' % d.name)
                               (x0=shift_x0(shift, d, i, j), fd_order=order,
-                               method=method, w=w)
+                               method=method, side=side, w=w)
                               for j, d in enumerate(space_dims)]))
         return func._new(comps)
 
@@ -258,7 +261,7 @@ def laplace(self):
         """
         return self.laplacian()
 
-    def laplacian(self, shift=None, order=None, method='FD', **kwargs):
+    def laplacian(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Laplacian of the TensorFunction with shifted derivatives and custom
         FD order.
@@ -277,6 +280,9 @@ def laplacian(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite
         """
@@ -290,7 +296,7 @@ def laplacian(self, shift=None, order=None, method='FD', **kwargs):
         for j in range(ndim):
             comps.append(sum([getattr(self[j, i], 'd%s2' % d.name)
                               (x0=shift_x0(shift, d, j, i), fd_order=order,
-                               method=method, w=w)
+                               method=method, side=side, w=w)
                               for i, d in enumerate(space_dims)]))
         return func._new(comps)
 
@@ -358,7 +364,7 @@ def __str__(self):
 
     __repr__ = __str__
 
-    def div(self, shift=None, order=None, method='FD', **kwargs):
+    def div(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Divergence of the VectorFunction, creates the divergence Function.
 
@@ -372,6 +378,9 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite difference coefficients.
         """
@@ -380,7 +389,8 @@ def div(self, shift=None, order=None, method='FD', **kwargs):
         order = order or self.space_order
         space_dims = self.root_dimensions
         return sum([getattr(self[i], 'd%s' % d.name)(x0=shift_x0(shift, d, None, i),
-                                                     fd_order=order, method=method, w=w)
+                                                     fd_order=order, method=method,
+                                                     side=side, w=w)
                     for i, d in enumerate(space_dims)])
 
     @property
@@ -390,7 +400,7 @@ def laplace(self):
         """
         return self.laplacian()
 
-    def laplacian(self, shift=None, order=None, method='FD', **kwargs):
+    def laplacian(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Laplacian of the VectorFunction, creates the Laplacian VectorFunction.
 
@@ -404,6 +414,9 @@ def laplacian(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite
         """
@@ -413,12 +426,13 @@ def laplacian(self, shift=None, order=None, method='FD', **kwargs):
         order = order or self.space_order
         space_dims = self.root_dimensions
         comps = [sum([getattr(s, 'd%s2' % d.name)(x0=shift_x0(shift, d, None, i),
-                                                  fd_order=order, w=w, method=method)
+                                                  fd_order=order, method=method,
+                                                  side=side, w=w)
                       for i, d in enumerate(space_dims)])
                  for s in self]
         return func._new(comps)
 
-    def curl(self, shift=None, order=None, method='FD', **kwargs):
+    def curl(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Gradient of the (3D) VectorFunction, creates the curl VectorFunction.
 
@@ -432,6 +446,9 @@ def curl(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite difference coefficients.
         """
@@ -444,21 +461,27 @@ def curl(self, shift=None, order=None, method='FD', **kwargs):
         shift_x0 = make_shift_x0(shift, (len(dims), len(dims)))
         order = order or self.space_order
         comp1 = (getattr(self[2], derivs[1])(x0=shift_x0(shift, dims[1], 2, 1),
-                                             fd_order=order, method=method, w=w) -
+                                             fd_order=order, method=method,
+                                             side=side, w=w) -
                  getattr(self[1], derivs[2])(x0=shift_x0(shift, dims[2], 1, 2),
-                                             fd_order=order, method=method, w=w))
+                                             fd_order=order, method=method,
+                                             side=side, w=w))
         comp2 = (getattr(self[0], derivs[2])(x0=shift_x0(shift, dims[2], 0, 2),
-                                             fd_order=order, method=method, w=w) -
+                                             fd_order=order, method=method,
+                                             side=side, w=w) -
                  getattr(self[2], derivs[0])(x0=shift_x0(shift, dims[0], 2, 0),
-                                             fd_order=order, method=method, w=w))
+                                             fd_order=order, method=method,
+                                             side=side, w=w))
         comp3 = (getattr(self[1], derivs[0])(x0=shift_x0(shift, dims[0], 1, 0),
-                                             fd_order=order, method=method, w=w) -
+                                             fd_order=order, method=method,
+                                             side=side, w=w) -
                  getattr(self[0], derivs[1])(x0=shift_x0(shift, dims[1], 0, 1),
-                                             fd_order=order, method=method, w=w))
+                                             fd_order=order, method=method,
+                                             side=side, w=w))
         func = vec_func(self)
         return func._new(3, 1, [comp1, comp2, comp3])
 
-    def grad(self, shift=None, order=None, method='FD', **kwargs):
+    def grad(self, shift=None, order=None, method='FD', side=None, **kwargs):
         """
         Gradient of the VectorFunction, creates the gradient TensorFunction.
 
@@ -472,6 +495,9 @@ def grad(self, shift=None, order=None, method='FD', **kwargs):
         method: str, optional, default='FD'
             Discretization method. Options are 'FD' (default) and
             'RSFD' (rotated staggered grid finite-difference).
+        side : Side or tuple of Side, optional, default=centered
+            Side of the finite difference location, centered (at x), left (at x - 1)
+            or right (at x + 1).
         weights/w: list, tuple, or dict, optional, default=None
             Custom weights for the finite difference coefficients.
         """
@@ -481,8 +507,9 @@ def grad(self, shift=None, order=None, method='FD', **kwargs):
         shift_x0 = make_shift_x0(shift, (ndim, ndim))
         order = order or self.space_order
         space_dims = self.root_dimensions
-        comps = [[getattr(f, 'd%s' % d.name)(x0=shift_x0(shift, d, i, j), w=w,
-                                             fd_order=order, method=method)
+        comps = [[getattr(f, 'd%s' % d.name)(x0=shift_x0(shift, d, i, j),
+                                             fd_order=order, method=method,
+                                             side=side, w=w)
                   for j, d in enumerate(space_dims)]
                  for i, f in enumerate(self)]
         return func._new(comps)