misc: Refactor derivative

Author: JDBetteridge

devito/finite_differences/derivative.py

@@ -1,4 +1,4 @@
-from collections import OrderedDict
+from collections import defaultdict
 from collections.abc import Iterable
 from functools import cached_property
 from itertools import chain
@@ -9,7 +9,7 @@
 from .differentiable import Differentiable, diffify, interp_for_fd, Add
 from .tools import direct, transpose
 from .rsfd import d45
-from devito.tools import (as_mapper, as_tuple, filter_ordered, frozendict, is_integer,
+from devito.tools import (as_mapper, as_tuple, frozendict, is_integer,
                           Pickable)
 from devito.types.utils import DimensionTuple
 
@@ -29,7 +29,7 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
         Expression for which the Derivative is produced.
     dims : Dimension or tuple of Dimension
         Dimensions w.r.t. which to differentiate.
-    fd_order : int or tuple of int, optional, default=1
+    fd_order : int, tuple of int or dict of {Dimension: int}, optional, default=1
         Coefficient discretization order. Note: this impacts the width of
         the resulting stencil.
     deriv_order: int or tuple of int, optional, default=1
@@ -93,6 +93,11 @@ class Derivative(sympy.Derivative, Differentiable, Pickable):
                    'x0', 'method', 'weights')
 
     def __new__(cls, expr, *dims, **kwargs):
+        # TODO: Delete this
+        if kwargs.get('preprocessed', False):
+            from warnings import warn
+            warn('I removed the `preprocessed` kwarg')
+
         if type(expr) is sympy.Derivative:
             raise ValueError("Cannot nest sympy.Derivative with devito.Derivative")
         if not isinstance(expr, Differentiable):
@@ -103,14 +108,62 @@ def __new__(cls, expr, *dims, **kwargs):
         if isinstance(expr, sympy.Number):
             return 0
 
-        new_dims, orders, fd_o, var_count = cls._process_kwargs(expr, *dims, **kwargs)
+        # Validate `dims`. It can be:
+        # - a single Dimension ie: x
+        # - an iterable of Dimensions ie: (x, y)
+        # - a single tuple of Dimension and order ie: (x, 2)
+        # - or an iterable of Dimension, order ie: ((x, 2), (y, 2))
+        if len(dims) == 0:
+            raise ValueError('Expected Dimension w.r.t. which to differentiate')
+        elif len(dims) == 1 and isinstance(dims[0], Iterable) and len(dims[0]) != 2:
+            # Iterable of Dimensions
+            raise ValueError(f'Expected `(dim, deriv_order)`, got {dims[0]}')
+        elif len(dims) == 2 and not isinstance(dims[1], Iterable) and is_integer(dims[1]):
+            # special case of single dimension and order
+            dims = (dims, )
+
+        # Use `deriv_order` if specified
+        deriv_order = kwargs.get('deriv_order', (1,)*len(dims))
+        if not isinstance(deriv_order, Iterable):
+            deriv_order = as_tuple(deriv_order)
+        if len(deriv_order) != len(dims):
+            raise ValueError(
+                'Length of `deriv_order` does not match the length of dimensions'
+            )
+
+        # Count the number of derivatives for each dimension
+        dcounter = defaultdict(int)
+        for d, o in zip(dims, deriv_order):
+            if isinstance(d, Iterable):
+                dcounter[d[0]] += d[1]
+            else:
+                dcounter[d] += o
 
-        # Construct the actual Derivative object
-        obj = Differentiable.__new__(cls, expr, *var_count)
-        obj._dims = tuple(OrderedDict.fromkeys(new_dims))
+        # Default finite difference orders depending on input dimension (.dt or .dx)
+        default_fdo = tuple([
+            expr.time_order
+            if getattr(d, 'is_Time', False)
+            else expr.space_order
+            for d in dcounter.keys()
+        ])
 
-        obj._fd_order = DimensionTuple(*as_tuple(fd_o), getters=obj._dims)
-        obj._deriv_order = DimensionTuple(*as_tuple(orders), getters=obj._dims)
+        # SymPy expects the list of variable w.r.t. which we differentiate to be a list
+        # of 2-tuple `(s, count)` where s is the entity to diff wrt and count is the order
+        # of the derivative
+        derivatives = [sympy.Tuple(d, o) for d, o in dcounter.items()]
+
+        # Construct the actual Derivative object
+        obj = Differentiable.__new__(cls, expr, *derivatives)
+        obj._dims = tuple(dcounter.keys())
+
+        obj._fd_order = DimensionTuple(
+            *as_tuple(kwargs.get('fd_order', default_fdo)),
+            getters=obj._dims
+        )
+        obj._deriv_order = DimensionTuple(
+            *as_tuple(dcounter.values()),
+            getters=obj._dims
+        )
         obj._side = kwargs.get("side")
         obj._transpose = kwargs.get("transpose", direct)
         obj._method = kwargs.get("method", 'FD')
@@ -131,78 +184,6 @@ def __new__(cls, expr, *dims, **kwargs):
 
         return obj
 
-    @classmethod
-    def _process_kwargs(cls, expr, *dims, **kwargs):
-        """
-        Process arguments for the construction of a Derivative
-        """
-        # Skip costly processing if constructiong from preprocessed
-        if kwargs.get('preprocessed', False):
-            fd_orders = kwargs.get('fd_order')
-            deriv_orders = kwargs.get('deriv_order')
-            if len(dims) == 1:
-                if isinstance(dims[0], Iterable):
-                    assert dims[0][1] == deriv_orders[0]
-                    dims = tuple([dims[0][0]]*max(1, deriv_orders[0]))
-                else:
-                    dims = tuple([dims[0]]*max(1, deriv_orders[0]))
-            variable_count = [sympy.Tuple(s, dims.count(s))
-                              for s in filter_ordered(dims)]
-            return dims, deriv_orders, fd_orders, variable_count
-
-        # Check `dims`. It can be a single Dimension, an iterable of Dimensions, or even
-        # an iterable of 2-tuple (Dimension, deriv_order)
-        if len(dims) == 0:
-            raise ValueError("Expected Dimension w.r.t. which to differentiate")
-        elif len(dims) == 1:
-            if isinstance(dims[0], Iterable):
-                # Iterable of Dimensions
-                if len(dims[0]) != 2:
-                    raise ValueError("Expected `(dim, deriv_order)`, got %s" % dims[0])
-                orders = kwargs.get('deriv_order', dims[0][1])
-                if dims[0][1] != orders:
-                    raise ValueError("Two different values of `deriv_order`")
-                new_dims = tuple([dims[0][0]]*max(1, dims[0][1]))
-            else:
-                # Single Dimension
-                orders = kwargs.get('deriv_order', 1)
-                if isinstance(orders, Iterable):
-                    orders = orders[0]
-                new_dims = tuple([dims[0]]*max(1, orders))
-        elif len(dims) == 2 and not isinstance(dims[1], Iterable) and is_integer(dims[1]):
-            # special case of single dimension and order
-            orders = dims[1]
-            new_dims = tuple([dims[0]]*max(1, orders))
-        else:
-            # Iterable of 2-tuple, e.g. ((x, 2), (y, 3))
-            new_dims = []
-            orders = []
-            d_ord = kwargs.get('deriv_order', tuple([1]*len(dims)))
-            for d, o in zip(dims, d_ord):
-                if isinstance(d, Iterable):
-                    new_dims.extend([d[0]]*max(1, d[1]))
-                    orders.append(d[1])
-                else:
-                    new_dims.extend([d]*max(1, o))
-                    orders.append(o)
-            new_dims = as_tuple(new_dims)
-            orders = as_tuple(orders)
-
-        # Finite difference orders depending on input dimension (.dt or .dx)
-        odims = filter_ordered(new_dims)
-        fd_orders = kwargs.get('fd_order', tuple([expr.time_order if
-                                                  getattr(d, 'is_Time', False) else
-                                                  expr.space_order for d in odims]))
-        if len(odims) == 1 and isinstance(fd_orders, Iterable):
-            fd_orders = fd_orders[0]
-
-        # SymPy expects the list of variable w.r.t. which we differentiate to be a list
-        # of 2-tuple `(s, count)` where s is the entity to diff wrt and count is the order
-        # of the derivative
-        variable_count = [sympy.Tuple(s, new_dims.count(s))
-                          for s in odims]
-        return new_dims, orders, fd_orders, variable_count
-
     @classmethod
     def _process_x0(cls, dims, **kwargs):
         try:
@@ -268,7 +249,7 @@ def __call__(self, x0=None, fd_order=None, side=None, method=None, **kwargs):
         if fd_order is not None:
             rkw['fd_order'] = self._filter_dims(_fd_order, as_tuple=True)
 
-        return self._rebuild(**rkw)
+        return self._rebuild(*self.args, **rkw)
 
     def _rebuild(self, *args, **kwargs):
         if not args:
@@ -331,7 +312,7 @@ def _xreplace(self, subs):
         expr = self.expr.xreplace(dsubs)
 
         subs = self._ppsubs + (subs,)  # Postponed substitutions
-        return self._rebuild(subs=subs, expr=expr), True
+        return self._rebuild(expr, *self.args[1:], subs=subs), True
 
     @cached_property
     def _metadata(self):
@@ -405,7 +386,7 @@ def T(self):
         else:
             adjoint = direct
 
-        return self._rebuild(transpose=adjoint)
+        return self._rebuild(*self.args, transpose=adjoint)
 
     def _eval_at(self, func):
         """
@@ -438,12 +419,12 @@ def _eval_at(self, func):
             # in most equation with div(a * u) for example. The expression is re-centered
             # at the highest priority index (see _gather_for_diff) to compute the
             # derivative at x0.
-            return self._rebuild(expr=self.expr._gather_for_diff, x0=x0)
+            return self._rebuild(self.expr._gather_for_diff, *self.args[1:], x0=x0)
         else:
             # For every other cases, that has more functions or more complexe arithmetic,
             # there is not actual way to decide what to do so it’s as safe to use
             # the expression as is.
-            return self._rebuild(x0=x0)
+            return self._rebuild(*self.args, x0=x0)
 
     def _evaluate(self, **kwargs):
         # Evaluate finite-difference.

devito/finite_differences/differentiable.py

@@ -1161,7 +1161,7 @@ def interp_for_fd(expr, x0, **kwargs):
 @interp_for_fd.register(sympy.Derivative)
 def _(expr, x0, **kwargs):
     x0_expr = {d: v for d, v in x0.items() if d not in expr.dims}
-    return expr.func(expr=interp_for_fd(expr.expr, x0_expr, **kwargs))
+    return expr.func(interp_for_fd(expr.expr, x0_expr, **kwargs), *expr.args[1:])
 
 
 @interp_for_fd.register(sympy.Expr)