Refactored finite difference spatial discretization

- added isotropic version

fd/derivative.py

 import sympy as sp
 from collections import namedtuple, defaultdict
+
+from pystencils import Field
 from pystencils.sympyextensions import normalize_product, prod
 
 
@@ -22,6 +24,8 @@ class Diff(sp.Expr):
     def __new__(cls, argument, target=-1, superscript=-1):
         if argument == 0:
             return sp.Rational(0, 1)
+        if isinstance(argument, Field):
+            argument = argument.center
         return sp.Expr.__new__(cls, argument.expand(), sp.sympify(target), sp.sympify(superscript))
 
     @property
@@ -176,6 +180,35 @@ class DiffOperator(sp.Expr):
 # ----------------------------------------------------------------------------------------------------------------------
 
 
+def diff(expr, *args):
+    """Shortcut function to create nested derivatives
+
+    >>> f = sp.Symbol("f")
+    >>> diff(f, 0, 0, 1) == Diff(Diff( Diff(f, 1), 0), 0)
+    True
+    """
+    if len(args) == 0:
+        return expr
+    result = expr
+    for index in reversed(args):
+        result = Diff(result, index)
+    return result
+
+
+def diff_args(expr):
+    """Extracts the indices and argument of possibly nested derivative - inverse of diff function
+
+    >>> args = (sp.Symbol("x"), 0, 1, 2, 5, 1)
+    >>> e = diff(*args)
+    >>> assert diff_args(e) == args
+    """
+    if not isinstance(expr, Diff):
+        return expr,
+    else:
+        inner_res = diff_args(expr.args[0])
+        return (inner_res[0], expr.args[1], *inner_res[1:])
+
+
 def diff_terms(expr):
     """Returns set of all derivatives in an expression.
 
@@ -200,16 +233,6 @@ def collect_diffs(expr):
     return expr.collect(diff_terms(expr))
 
 
-def create_nested_diff(arg, *args):
-    """Shortcut to create nested derivatives"""
-    assert arg is not None
-    args = sorted(args, reverse=True, key=lambda e: e.name if isinstance(e, sp.Symbol) else e)
-    res = arg
-    for i in args:
-        res = Diff(res, i)
-    return res
-
-
 def replace_diff(expr, replacement_dict):
     """replacement_dict: maps variable (target) to a new Differential operator"""
 
@@ -464,6 +487,27 @@ def combine_diff_products(expr):
     return combine(expr)
 
 
+def replace_generic_laplacian(expr, dim=None):
+    """Laplacian can be written as Diff(Diff(term)) without explicitly giving the dimensions.
+
+    This function replaces these constructs by diff(term, 0, 0) + diff(term, 1, 1) + ...
+    For this to work, the arguments of the derivative have to be field or field accesses such that the spatial
+    dimension can be determined.
+    """
+    if isinstance(expr, Diff):
+        arg, *indices = diff_args(expr)
+        if isinstance(arg, Field.Access):
+            dim = arg.field.spatial_dimensions
+        assert dim is not None
+        if len(indices) == 2 and all(i == -1 for i in indices):
+            return sum(diff(arg, i, i) for i in range(dim))
+        else:
+            return expr
+    else:
+        new_args = [replace_generic_laplacian(a, dim) for a in expr.args]
+        return expr.func(*new_args) if new_args else expr
+
+
 def functional_derivative(functional, v):
     r"""Computes functional derivative of functional with respect to v using Euler-Lagrange equation
 

fd/finitedifferences.py

@@ -73,11 +73,11 @@ class Discretization2ndOrder:
         self.dt = dt
 
     @staticmethod
-    def __diff_order(e):
+    def _diff_order(e):
         if not isinstance(e, Diff):
             return 0
         else:
-            return 1 + Discretization2ndOrder.__diff_order(e.args[0])
+            return 1 + Discretization2ndOrder._diff_order(e.args[0])
 
     def _discretize_diffusion(self, expr):
         result = 0
@@ -110,7 +110,7 @@ class Discretization2ndOrder:
             return e.func(*new_args) if new_args else e
 
     def _discretize_diff(self, e):
-        order = self.__diff_order(e)
+        order = self._diff_order(e)
         if order == 1:
             fa = e.args[0]
             index = e.target

fd/spatial.py

+import sympy as sp
+from functools import partial
+from pystencils import AssignmentCollection, Field
+from pystencils.fd import Diff
+from .derivative import diff_args
+
+
+def fd_stencils_standard(indices, dx, fa):
+    order = len(indices)
+    if order == 1:
+        idx = indices[0]
+        return (fa.neighbor(idx, 1) - fa.neighbor(idx, -1)) / (2 * dx)
+    elif order == 2:
+        if indices[0] == indices[1]:
+            return (-2 * fa + fa.neighbor(indices[0], -1) + fa.neighbor(indices[0], +1)) / (dx ** 2)
+        else:
+            offsets = [(1, 1), [-1, 1], [1, -1], [-1, -1]]
+            return sum(o1 * o2 * fa.neighbor(indices[0], o1).neighbor(indices[1], o2)
+                       for o1, o2 in offsets) / (4 * dx ** 2)
+    raise NotImplementedError("Supports only derivatives up to order 2")
+
+
+def fd_stencils_isotropic(indices, dx, fa):
+    dim = fa.field.spatial_dimensions
+    if dim == 1:
+        return fd_stencils_standard(indices, dx, fa)
+    elif dim == 2:
+        order = len(indices)
+
+        if order == 1:
+            idx = indices[0]
+            assert 0 <= idx < 2
+            other_idx = 1 if indices[0] == 0 else 0
+            weights = {-1: sp.Rational(1, 12) / dx,
+                       0: sp.Rational(1, 3) / dx,
+                       1: sp.Rational(1, 12) / dx}
+            upper_terms = sum(fa.neighbor(idx, +1).neighbor(other_idx, off) * w for off, w in weights.items())
+            lower_terms = sum(fa.neighbor(idx, -1).neighbor(other_idx, off) * w for off, w in weights.items())
+            return upper_terms - lower_terms
+        elif order == 2:
+            if indices[0] == indices[1]:
+                idx = indices[0]
+                other_idx = 1 if idx == 0 else 0
+                diagonals = sp.Rational(1, 12) * sum(fa.neighbor(0, i).neighbor(1, j) for i in (-1, 1) for j in (-1, 1))
+                div_direction = sp.Rational(5, 6) * sum(fa.neighbor(idx, i) for i in (-1, 1))
+                other_direction = - sp.Rational(1, 6) * sum(fa.neighbor(other_idx, i) for i in (-1, 1))
+                center = - sp.Rational(5, 3) * fa
+                return (diagonals + div_direction + other_direction + center) / (dx ** 2)
+            else:
+                return fd_stencils_standard(indices, dx, fa)
+    raise NotImplementedError("Supports only derivatives up to order 2 for 1D and 2D setups")
+
+
+def discretize_spatial(expr, dx, stencil=fd_stencils_standard):
+    if isinstance(stencil, str):
+        if stencil == 'standard':
+            stencil = fd_stencils_standard
+        elif stencil == 'isotropic':
+            stencil = fd_stencils_isotropic
+        else:
+            raise ValueError("Unknown stencil. Supported 'standard' and 'isotropic'")
+
+    if isinstance(expr, list):
+        return [discretize_spatial(e, dx, stencil) for e in expr]
+    elif isinstance(expr, sp.Matrix):
+        return expr.applyfunc(partial(discretize_spatial, dx=dx, stencil=stencil))
+    elif isinstance(expr, AssignmentCollection):
+        return expr.copy(main_assignments=[e for e in expr.main_assignments],
+                         subexpressions=[e for e in expr.subexpressions])
+    elif isinstance(expr, Diff):
+        arg, *indices = diff_args(expr)
+        if not isinstance(arg, Field.Access):
+            raise ValueError("Only derivatives with field or field accesses as arguments can be discretized")
+        return stencil(indices, dx, arg)
+    else:
+        new_args = [discretize_spatial(a, dx, stencil) for a in expr.args]
+        return expr.func(*new_args) if new_args else expr

field.py

@@ -386,6 +386,8 @@ class Field:
         return hash((self._layout, self.shape, self.strides, self._dtype, self.field_type, self._field_name))
 
     def __eq__(self, other):
+        if not isinstance(other, Field):
+            return False
         self_tuple = (self.shape, self.strides, self.name, self.dtype, self.field_type)
         other_tuple = (other.shape, other.strides, other.name, other.dtype, other.field_type)
         return self_tuple == other_tuple

stencils.py

@@ -171,10 +171,11 @@ def visualize_stencil_2d(stencil, axes=None, figure=None, data=None, textsize='1
     """
     Creates a matplotlib 2D plot of the stencil
 
-    :param stencil: sequence of directions
-    :param axes: optional matplotlib axes
-    :param data: data to annotate the directions with, if none given, the indices are used
-    :param textsize: size of annotation text
+    Args:
+        stencil: sequence of directions
+        axes: optional matplotlib axes
+        data: data to annotate the directions with, if none given, the indices are used
+        textsize: size of annotation text
     """
     from matplotlib.patches import BoxStyle
     import matplotlib.pyplot as plt
@@ -329,6 +330,7 @@ def visualize_stencil_3d(stencil, figure=None, axes=None, data=None, textsize='8
 
 
 def visualize_stencil_expression(expr, **kwargs):
+    """Displays coefficients of a linear update expression of a single field as matplotlib arrow drawing."""
     stencil, coefficients = stencil_coefficients(expr)
     dim = len(stencil[0])
     assert 0 < dim <= 3