Merge pull request #2326 from devitocodes/stagg-fd-custom

api: add support for 45 degree FD approx

devito/finite_differences/__init__.py

@@ -5,3 +5,4 @@
 from .tools import generate_fd_shortcuts  # noqa
 from .coefficients import *  # noqa
 from .operators import *  # noqa
+from .rsfd import *  # noqa

devito/finite_differences/coefficients.py

@@ -257,7 +257,7 @@ def generate_subs(deriv_order, function, index):
         indices, x0 = generate_indices(function, dim,
                                        fd_order, side=None, x0=mapper)
 
-        coeffs = numeric_weights(deriv_order, indices, x0)
+        coeffs = numeric_weights(function, deriv_order, indices, x0)
 
         for (c, i) in zip(coeffs, indices):
             subs.update({function._coeff_symbol(i, deriv_order, function, index): c})

devito/finite_differences/derivative.py

@@ -7,6 +7,7 @@
 from .finite_difference import generic_derivative, first_derivative, cross_derivative
 from .differentiable import Differentiable
 from .tools import direct, transpose
+from .rsfd import d45
 from devito.tools import as_mapper, as_tuple, filter_ordered, frozendict
 from devito.types.utils import DimensionTuple
 
@@ -86,7 +87,8 @@ class Derivative(sympy.Derivative, Differentiable):
     _fd_priority = 3
 
     __rargs__ = ('expr', 'dims')
-    __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs', 'x0')
+    __rkwargs__ = ('side', 'deriv_order', 'fd_order', 'transpose', '_ppsubs',
+                   'x0', 'method')
 
     def __new__(cls, expr, *dims, **kwargs):
         if type(expr) is sympy.Derivative:
@@ -106,6 +108,7 @@ def __new__(cls, expr, *dims, **kwargs):
         obj._deriv_order = orders if skip else DimensionTuple(*orders, getters=obj._dims)
         obj._side = kwargs.get("side")
         obj._transpose = kwargs.get("transpose", direct)
+        obj._method = kwargs.get("method", 'FD')
 
         ppsubs = kwargs.get("subs", kwargs.get("_ppsubs", []))
         processed = []
@@ -183,12 +186,14 @@ def _process_kwargs(cls, expr, *dims, **kwargs):
                           for s in filter_ordered(new_dims)]
         return new_dims, orders, fd_orders, variable_count
 
-    def __call__(self, x0=None, fd_order=None, side=None):
+    def __call__(self, x0=None, fd_order=None, side=None, method=None):
         if self.ndims == 1:
-            _fd_order = fd_order or self._fd_order
-            _side = side or self._side
-            new_x0 = {self.dims[0]: x0} if x0 is not None else self.x0
-            return self._new_from_self(fd_order=_fd_order, side=_side, x0=new_x0)
+            fd_order = fd_order or self._fd_order
+            side = side or self._side
+            method = method or self._method
+            x0 = {self.dims[0]: x0} if x0 is not None else self.x0
+            return self._new_from_self(fd_order=fd_order, side=side, x0=x0,
+                                       method=method)
 
         if side is not None:
             raise TypeError("Side only supported for first order single"
@@ -210,7 +215,7 @@ def _new_from_self(self, **kwargs):
         expr = kwargs.pop('expr', self.expr)
         _kwargs = {'deriv_order': self.deriv_order, 'fd_order': self.fd_order,
                    'side': self.side, 'transpose': self.transpose, 'subs': self._ppsubs,
-                   'x0': self.x0, 'preprocessed': True}
+                   'x0': self.x0, 'preprocessed': True, 'method': self.method}
         _kwargs.update(**kwargs)
         return Derivative(expr, *self.dims, **_kwargs)
 
@@ -305,6 +310,10 @@ def transpose(self):
     def is_TimeDependent(self):
         return self.expr.is_TimeDependent
 
+    @property
+    def method(self):
+        return self._method
+
     @property
     def T(self):
         """Transpose of the Derivative.
@@ -384,14 +393,21 @@ def _eval_fd(self, expr, **kwargs):
         expand = kwargs.get('expand', True)
 
         # Step 2: Evaluate FD of the new expression
-        if self.side is not None and self.deriv_order == 1:
+        if self.method == 'RSFD':
+            assert len(self.dims) == 1
+            assert self.deriv_order == 1
+            res = d45(expr, self.dims[0], x0=self.x0, expand=expand)
+        elif self.side is not None and self.deriv_order == 1:
+            assert self.method == 'FD'
             res = first_derivative(expr, self.dims[0], self.fd_order,
                                    side=self.side, matvec=self.transpose,
                                    x0=self.x0, expand=expand)
         elif len(self.dims) > 1:
+            assert self.method == 'FD'
             res = cross_derivative(expr, self.dims, self.fd_order, self.deriv_order,
                                    matvec=self.transpose, x0=self.x0, expand=expand)
         else:
+            assert self.method == 'FD'
             res = generic_derivative(expr, *self.dims, self.fd_order, self.deriv_order,
                                      matvec=self.transpose, x0=self.x0, expand=expand)
 

devito/finite_differences/differentiable.py

@@ -125,6 +125,17 @@ def _symbolic_functions(self):
     def _uses_symbolic_coefficients(self):
         return bool(self._symbolic_functions)
 
+    @cached_property
+    def coefficients(self):
+        coefficients = {f.coefficients for f in self._functions}
+        # If there is multiple ones, we have to revert to the highest priority
+        # i.e we have to remove symbolic
+        if len(coefficients) == 2:
+            return (coefficients - {'symbolic'}).pop()
+        else:
+            assert len(coefficients) == 1
+            return coefficients.pop()
+
     @cached_property
     def _coeff_symbol(self, *args, **kwargs):
         if self._uses_symbolic_coefficients:
@@ -305,7 +316,7 @@ def laplacian(self, shift=None, order=None):
                                       fd_order=order)
                      for i, d in enumerate(derivs)])
 
-    def div(self, shift=None, order=None):
+    def div(self, shift=None, order=None, method='FD'):
         """
         Divergence of the input Function.
 
@@ -318,16 +329,18 @@ def div(self, shift=None, order=None):
         order: int, optional, default=None
             Discretization order for the finite differences.
             Uses `func.space_order` when not specified
-
+        method: str, optional, default='FD'
+            Discretization method. Options are 'FD' (default) and
+            'RSFD' (rotated staggered grid finite-difference).
         """
         space_dims = [d for d in self.dimensions if d.is_Space]
         shift_x0 = make_shift_x0(shift, (len(space_dims),))
         order = order or self.space_order
         return Add(*[getattr(self, 'd%s' % d.name)(x0=shift_x0(shift, d, None, i),
-                                                   fd_order=order)
+                                                   fd_order=order, method=method)
                      for i, d in enumerate(space_dims)])
 
-    def grad(self, shift=None, order=None):
+    def grad(self, shift=None, order=None, method='FD'):
         """
         Gradient of the input Function.
 
@@ -340,14 +353,16 @@ def grad(self, shift=None, order=None):
         order: int, optional, default=None
             Discretization order for the finite differences.
             Uses `func.space_order` when not specified
-
+        method: str, optional, default='FD'
+            Discretization method. Options are 'FD' (default) and
+            'RSFD' (rotated staggered grid finite-difference).
         """
         from devito.types.tensor import VectorFunction, VectorTimeFunction
         space_dims = [d for d in self.dimensions if d.is_Space]
         shift_x0 = make_shift_x0(shift, (len(space_dims),))
         order = order or self.space_order
         comps = [getattr(self, 'd%s' % d.name)(x0=shift_x0(shift, d, None, i),
-                                               fd_order=order)
+                                               fd_order=order, method=method)
                  for i, d in enumerate(space_dims)]
         vec_func = VectorTimeFunction if self.is_TimeDependent else VectorFunction
         return vec_func(name='grad_%s' % self.name, time_order=self.time_order,
@@ -659,7 +674,7 @@ def _evaluate(self, **kwargs):
         expr = self.expr._evaluate(**kwargs)
 
         if not kwargs.get('expand', True):
-            return self.func(expr, self.dimensions)
+            return self._rebuild(expr)
 
         values = product(*[list(d.range) for d in self.dimensions])
         terms = []
@@ -821,7 +836,7 @@ def _evaluate(self, **kwargs):
         mapper = {w.subs(d, i): f.weights[n] for n, i in enumerate(d.range)}
         expr = expr.xreplace(mapper)
 
-        return EvalDerivative(expr, base=self.base)
+        return EvalDerivative(*expr.args, base=self.base)
 
 
 class DiffDerivative(IndexDerivative, DifferentiableOp):

devito/finite_differences/finite_difference.py

@@ -1,9 +1,8 @@
 from sympy import sympify
 
 from .differentiable import EvalDerivative, DiffDerivative, Weights
-from .tools import (numeric_weights, symbolic_weights, left, right,
-                    generate_indices, centered, direct, transpose,
-                    check_input, check_symbolic)
+from .tools import (left, right, generate_indices, centered, direct, transpose,
+                    check_input, check_symbolic, fd_weights_registry)
 
 __all__ = ['first_derivative', 'cross_derivative', 'generic_derivative',
            'left', 'right', 'centered', 'transpose', 'generate_indices']
@@ -16,7 +15,7 @@
 @check_input
 @check_symbolic
 def first_derivative(expr, dim, fd_order=None, side=centered, matvec=direct, x0=None,
-                     symbolic=False, expand=True):
+                     coefficients='taylor', expand=True):
     """
     First-order derivative of a given expression.
 
@@ -26,22 +25,22 @@ def first_derivative(expr, dim, fd_order=None, side=centered, matvec=direct, x0=
         Expression for which the first-order derivative is produced.
     dim : Dimension
         The Dimension w.r.t. which to differentiate.
-    fd_order : int, optional
+    fd_order : int, optional, default=expr.space_order
         Coefficient discretization order. Note: this impacts the width of
-        the resulting stencil. Defaults to `expr.space_order`.
-    side : Side, optional
+        the resulting stencil.
+    side : Side, optional, default=centered
         Side of the finite difference location, centered (at x), left (at x - 1)
-        or right (at x +1). Defaults to `centered`.
-    matvec : Transpose, optional
+        or right (at x +1).
+    matvec : Transpose, optional, default=direct
         Forward (matvec=direct) or transpose (matvec=transpose) mode of the
-        finite difference. Defaults to `direct`.
-    x0 : dict, optional
+        finite difference.
+    x0 : dict, optional, default=None
         Origin of the finite-difference scheme as a map dim: origin_dim.
-    symbolic : bool, optional
-        Use default or custom coefficients (weights). Defaults to False.
-    expand : bool, optional
+    coefficients : string, optional, default='taylor'
+        Use taylor or custom coefficients (weights).
+    expand : bool, optional, default=True
         If True, the derivative is fully expanded as a sum of products,
-        otherwise an IndexSum is returned. Defaults to True.
+        otherwise an IndexSum is returned.
 
     Returns
     -------
@@ -88,8 +87,12 @@ def first_derivative(expr, dim, fd_order=None, side=centered, matvec=direct, x0=
     fd_order = fd_order or expr.space_order
     deriv_order = 1
 
+    # Enforce stable time coefficients
+    if dim.is_Time and coefficients != 'symbolic':
+        coefficients = 'taylor'
+
     return make_derivative(expr, dim, fd_order, deriv_order, side,
-                           matvec, x0, symbolic, expand)
+                           matvec, x0, coefficients, expand)
 
 
 @check_input
@@ -104,21 +107,22 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, **kwargs):
         Expression for which the cross derivative is produced.
     dims : tuple of Dimension
         Dimensions w.r.t. which to differentiate.
-    fd_order : tuple of ints
+    fd_order : int, optional, default=expr.space_order
         Coefficient discretization order. Note: this impacts the width of
         the resulting stencil.
-    deriv_order : tuple of ints
-        Derivative order, e.g. 2 for a second-order derivative.
-    matvec : Transpose, optional
+    side : Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x +1).
+    matvec : Transpose, optional, default=direct
         Forward (matvec=direct) or transpose (matvec=transpose) mode of the
-        finite difference. Defaults to `direct`.
-    x0 : dict, optional
+        finite difference.
+    x0 : dict, optional, default=None
         Origin of the finite-difference scheme as a map dim: origin_dim.
-    symbolic : bool, optional
-        Use default or custom coefficients (weights). Defaults to False.
-    expand : bool, optional
+    coefficients : string, optional, default='taylor'
+        Use taylor or custom coefficients (weights).
+    expand : bool, optional, default=True
         If True, the derivative is fully expanded as a sum of products,
-        otherwise an IndexSum is returned. Defaults to True.
+        otherwise an IndexSum is returned.
 
     Returns
     -------
@@ -166,7 +170,7 @@ def cross_derivative(expr, dims, fd_order, deriv_order, x0=None, **kwargs):
 @check_input
 @check_symbolic
 def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
-                       symbolic=False, expand=True):
+                       coefficients='taylor', expand=True):
     """
     Arbitrary-order derivative of a given expression.
 
@@ -176,54 +180,57 @@ def generic_derivative(expr, dim, fd_order, deriv_order, matvec=direct, x0=None,
         Expression for which the derivative is produced.
     dim : Dimension
         The Dimension w.r.t. which to differentiate.
-    fd_order : int
+    fd_order : int, optional, default=expr.space_order
         Coefficient discretization order. Note: this impacts the width of
         the resulting stencil.
-    deriv_order : int
-        Derivative order, e.g. 2 for a second-order derivative.
-    matvec : Transpose, optional
+    side : Side, optional, default=centered
+        Side of the finite difference location, centered (at x), left (at x - 1)
+        or right (at x +1).
+    matvec : Transpose, optional, default=direct
         Forward (matvec=direct) or transpose (matvec=transpose) mode of the
-        finite difference. Defaults to `direct`.
-    x0 : dict, optional
+        finite difference.
+    x0 : dict, optional, default=None
         Origin of the finite-difference scheme as a map dim: origin_dim.
-    symbolic : bool, optional
-        Use default or custom coefficients (weights). Defaults to False.
-    expand : bool, optional
+    coefficients : string, optional, default='taylor'
+        Use taylor or custom coefficients (weights).
+    expand : bool, optional, default=True
         If True, the derivative is fully expanded as a sum of products,
-        otherwise an IndexSum is returned. Defaults to True.
+        otherwise an IndexSum is returned.
 
     Returns
     -------
     expr-like
         ``deriv-order`` derivative of ``expr``.
     """
     side = None
-    # First order derivative with 2nd order FD is highly non-recommended so taking
+    # 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 not symbolic:
+    if deriv_order == 1 and fd_order == 2 and coefficients != 'symbolic':
         fd_order = 1
 
+    # Enforce stable time coefficients
+    if dim.is_Time and coefficients != 'symbolic':
+        coefficients = 'taylor'
+
     return make_derivative(expr, dim, fd_order, deriv_order, side,
-                           matvec, x0, symbolic, expand)
+                           matvec, x0, coefficients, expand)
 
 
-def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, symbolic,
+def make_derivative(expr, dim, fd_order, deriv_order, side, matvec, x0, coefficients,
                     expand):
     # The stencil indices
     indices, x0 = generate_indices(expr, dim, fd_order, side=side, matvec=matvec,
                                    x0=x0)
-
-    # Finite difference weights from Taylor approximation given these positions
-    if symbolic:
-        weights = symbolic_weights(expr, deriv_order, indices, x0)
-    else:
-        weights = numeric_weights(deriv_order, indices, x0)
+    # Finite difference weights corresponding to the indices. Computed via the
+    # `coefficients` method (`taylor` or `symbolic`)
+    weights = fd_weights_registry[coefficients](expr, deriv_order, indices, x0)
 
     # Enforce fixed precision FD coefficients to avoid variations in results
-    weights = [sympify(w).evalf(_PRECISION) for w in weights][::matvec.val]
+    weights = [sympify(w).evalf(_PRECISION) for w in weights]
 
     # Transpose the FD, if necessary
     if matvec == transpose:
+        weights = weights[::-1]
         indices = indices.transpose()
 
     # Shift index due to staggering, if any