Moved stencils from lbmpy->pystencils, New coeff extraction funcs

- new functions to easily extract stencil coefficients and visualize them
- Moved stencil functions from lbmpy to pystencils

__init__.py

@@ -10,6 +10,7 @@ from .assignment import Assignment
 from .sympyextensions import SymbolCreator
 from .datahandling import create_data_handling
 from .kernel_decorator import kernel
+from .stencils import visualize_stencil_expression
 from . import fd
 
 __all__ = ['Field', 'FieldType', 'fields',
@@ -22,4 +23,5 @@ __all__ = ['Field', 'FieldType', 'fields',
            'SymbolCreator',
            'create_data_handling',
            'kernel',
-           'fd']
+           'fd',
+           'visualize_stencil_expression']

datahandling/serial_datahandling.py

@@ -381,3 +381,16 @@ class SerialDataHandling(DataHandling):
 
     def save_all(self, file):
         np.savez_compressed(file, **self.cpu_arrays)
+
+    def load_all(self, file):
+        file_contents = np.load(file)
+        for arr_name, arr_contents in self.cpu_arrays.items():
+            if arr_name not in file_contents:
+                print("Skipping read data {} because there is no data with this name in data handling".format(arr_name))
+                continue
+            if file_contents[arr_name].shape != arr_contents.shape:
+                print("Skipping read data {} because shapes don't match. "
+                      "Read array shape {}, exising array shape {}".format(arr_name, file_contents[arr_name].shape,
+                                                                           arr_contents.shape))
+                continue
+            np.copyto(arr_contents, file_contents[arr_name])

field.py

@@ -602,6 +602,16 @@ class Field:
             else:
                 return "{{%s}_{%s}}" % (n, offset_str)
 
+        def __str__(self):
+            n = self._field.latex_name if self._field.latex_name else self._field.name
+            offset_str = ",".join([sp.latex(o) for o in self.offsets])
+            if self.is_absolute_access:
+                offset_str = "[abs]{}".format(offset_str)
+            if self.index and self.index != (0,):
+                return "%s[%s](%s)" % (n, offset_str, self.index if len(self.index) > 1 else self.index[0])
+            else:
+                return "%s[%s]" % (n, offset_str)
+
 
 def get_layout_from_strides(strides: Sequence[int], index_dimension_ids: Optional[List[int]] = None):
     index_dimension_ids = [] if index_dimension_ids is None else index_dimension_ids

stencils.py

+import sympy as sp
+from collections import defaultdict
+from pystencils import Field
+
+
+def inverse_direction(direction):
+    """Returns inverse i.e. negative of given direction tuple"""
+    return tuple([-i for i in direction])
+
+
+def is_valid_stencil(stencil, max_neighborhood=None):
+    """
+    Tests if a nested sequence is a valid stencil i.e. all the inner sequences have the same length.
+    If max_neighborhood is specified, it is also verified that the stencil does not contain any direction components
+    with absolute value greater than the maximal neighborhood.
+    """
+    expected_dim = len(stencil[0])
+    for d in stencil:
+        if len(d) != expected_dim:
+            return False
+        if max_neighborhood is not None:
+            for d_i in d:
+                if abs(d_i) > max_neighborhood:
+                    return False
+    return True
+
+
+def is_symmetric_stencil(stencil):
+    """Tests for every direction d, that -d is also in the stencil"""
+    for d in stencil:
+        if inverse_direction(d) not in stencil:
+            return False
+    return True
+
+
+def stencils_have_same_entries(s1, s2):
+    if len(s1) != len(s2):
+        return False
+    return len(set(s1) - set(s2)) == 0
+
+
+# -------------------------------------Expression - Coefficient Form Conversion ----------------------------------------
+
+
+def stencil_coefficient_dict(expr):
+    """Extracts coefficients in front of field accesses in a expression.
+
+    Expression may only access a single field at a single index.
+
+    Returns:
+        center, coefficient dict, nonlinear part
+        where center is the single field that is accessed in expression accessed at center
+        and coefficient dict maps offsets to coefficients. The nonlinear part is everything that is not in the form of
+        coefficient times field access.
+
+    Examples:
+        >>> import pystencils as ps
+        >>> f = ps.fields("f(3) : double[2D]")
+        >>> field, coeffs, nonlinear_part = stencil_coefficient_dict(2 * f[0, 1](1) + 3 * f[-1, 0](1) + 123)
+        >>> assert nonlinear_part == 123 and field == f(1)
+        >>> sorted(coeffs.items())
+        [((-1, 0), 3), ((0, 1), 2)]
+    """
+    expr = expr.expand()
+    field_accesses = expr.atoms(Field.Access)
+    fields = set(fa.field for fa in field_accesses)
+    accessed_indices = set(fa.index for fa in field_accesses)
+
+    if len(fields) != 1:
+        raise ValueError("Could not extract stencil coefficients. "
+                         "Expression has to be a linear function of exactly one field.")
+    if len(accessed_indices) != 1:
+        raise ValueError("Could not extract stencil coefficients. Field is accessed at multiple indices")
+
+    field = fields.pop()
+    idx = accessed_indices.pop()
+
+    coefficients = defaultdict(lambda: 0)
+    coefficients.update({fa.offsets: expr.coeff(fa) for fa in field_accesses})
+
+    linear_part = sum(c * field[off](*idx) for off, c in coefficients.items())
+    nonlinear_part = expr - linear_part
+    return field(*idx), coefficients, nonlinear_part
+
+
+def stencil_coefficients(expr):
+    """Returns two lists - one with accessed offsets and one with their coefficients.
+
+    Same restrictions as `stencil_coefficient_dict` apply. Expression must not have any nonlinear part
+
+    >>> import pystencils as ps
+    >>> f = ps.fields("f(3) : double[2D]")
+    >>> stencil_coefficients(2 * f[0, 1](1) + 3 * f[-1, 0](1))
+    ([(-1, 0), (0, 1)], [3, 2])
+    """
+    field_center, coefficients, nonlinear_part = stencil_coefficient_dict(expr)
+    assert nonlinear_part == 0
+    stencil = list(coefficients.keys())
+    entries = [coefficients[c] for c in stencil]
+    return stencil, entries
+
+
+def stencil_coefficient_list(expr, matrix_form=False):
+    """Returns stencil coefficients in the form of nested lists
+
+    Same restrictions as `stencil_coefficient_dict` apply. Expression must not have any nonlinear part
+
+    Examples:
+        >>> import pystencils as ps
+        >>> f = ps.fields("f: double[2D]")
+        >>> stencil_coefficient_list(2 * f[0, 1] + 3 * f[-1, 0])
+        [[0, 0, 0], [3, 0, 0], [0, 2, 0]]
+        >>> stencil_coefficient_list(2 * f[0, 1] + 3 * f[-1, 0], matrix_form=True)
+        Matrix([
+        [0, 2, 0],
+        [3, 0, 0],
+        [0, 0, 0]])
+    """
+    field_center, coefficients, nonlinear_part = stencil_coefficient_dict(expr)
+    assert nonlinear_part == 0
+    field = field_center.field
+
+    dim = field.spatial_dimensions
+    max_offsets = defaultdict(lambda: 0)
+    for offset in coefficients.keys():
+        for d, off in enumerate(offset):
+            max_offsets[d] = max(max_offsets[d], abs(off))
+
+    if dim == 1:
+        result = [coefficients[(i,)] for i in range(-max_offsets[0], max_offsets[0] + 1)]
+        return sp.Matrix(result) if matrix_form else result
+    else:
+        y_range = list(range(-max_offsets[1], max_offsets[1] + 1))
+        if matrix_form:
+            y_range.reverse()
+        if dim == 2:
+            result = [[coefficients[(i, j)]
+                       for i in range(-max_offsets[0], max_offsets[0] + 1)]
+                      for j in y_range]
+            return sp.Matrix(result) if matrix_form else result
+        elif dim == 3:
+            result = [[[coefficients[(i, j, k)]
+                        for i in range(-max_offsets[0], max_offsets[0] + 1)]
+                       for j in y_range]
+                      for k in range(-max_offsets[2], max_offsets[2] + 1)]
+            return [sp.Matrix(l) for l in result] if matrix_form else result
+        else:
+            raise ValueError("Can only handle fields with 1,2 or 3 spatial dimensions")
+
+
+# -------------------------------------- Visualization -----------------------------------------------------------------
+
+
+def visualize_stencil(stencil, **kwargs):
+    dim = len(stencil[0])
+    if dim == 2:
+        visualize_stencil_2d(stencil, **kwargs)
+    else:
+        slicing = False
+        if 'slice' in kwargs:
+            slicing = kwargs['slice']
+            del kwargs['slice']
+
+        if slicing:
+            visualize_stencil_3d_by_slicing(stencil, **kwargs)
+        else:
+            visualize_stencil_3d(stencil, **kwargs)
+
+
+def visualize_stencil_2d(stencil, axes=None, figure=None, data=None, textsize='12', **kwargs):
+    """
+    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
+    """
+    from matplotlib.patches import BoxStyle
+    import matplotlib.pyplot as plt
+
+    if axes is None:
+        if figure is None:
+            figure = plt.gcf()
+        axes = figure.gca()
+
+    text_box_style = BoxStyle("Round", pad=0.3)
+    head_length = 0.1
+    max_offsets = [max(abs(d[c]) for d in stencil) for c in (0, 1)]
+
+    if data is None:
+        data = list(range(len(stencil)))
+
+    for direction, annotation in zip(stencil, data):
+        assert len(direction) == 2, "Works only for 2D stencils"
+
+        if not(direction[0] == 0 and direction[1] == 0):
+            axes.arrow(0, 0, direction[0], direction[1], head_width=0.08, head_length=head_length, color='k')
+
+        if isinstance(annotation, sp.Basic):
+            annotation = "$" + sp.latex(annotation) + "$"
+        else:
+            annotation = str(annotation)
+
+        def position_correction(d, magnitude=0.18):
+            if d < 0:
+                return -magnitude
+            elif d > 0:
+                return +magnitude
+            else:
+                return 0
+        text_position = [direction[c] + position_correction(direction[c]) for c in (0, 1)]
+        axes.text(*text_position, annotation, verticalalignment='center',
+                  zorder=30, horizontalalignment='center', size=textsize,
+                  bbox=dict(boxstyle=text_box_style, facecolor='#00b6eb', alpha=0.85, linewidth=0))
+
+    axes.set_axis_off()
+    axes.set_aspect('equal')
+    max_offsets = [m if m > 0 else 0.1 for m in max_offsets]
+    border = 0.1
+    axes.set_xlim([-border - max_offsets[0], border + max_offsets[0]])
+    axes.set_ylim([-border - max_offsets[1], border + max_offsets[1]])
+
+
+def visualize_stencil_3d_by_slicing(stencil, slice_axis=2, figure=None, data=None, **kwargs):
+    """Visualizes a 3D, first-neighborhood stencil by plotting 3 slices along a given axis.
+
+    Args:
+        stencil: stencil as sequence of directions
+        slice_axis: 0, 1, or 2 indicating the axis to slice through
+        data: optional data to print as text besides the arrows
+    """
+    import matplotlib.pyplot as plt
+
+    for d in stencil:
+        for element in d:
+            assert element == -1 or element == 0 or element == 1, "This function can only first neighborhood stencils"
+
+    if figure is None:
+        figure = plt.gcf()
+
+    axes = [figure.add_subplot(1, 3, i + 1) for i in range(3)]
+    splitted_directions = [[], [], []]
+    splitted_data = [[], [], []]
+    axes_names = ['x', 'y', 'z']
+
+    for i, d in enumerate(stencil):
+        split_idx = d[slice_axis] + 1
+        reduced_dir = tuple([element for j, element in enumerate(d) if j != slice_axis])
+        splitted_directions[split_idx].append(reduced_dir)
+        splitted_data[split_idx].append(i if data is None else data[i])
+
+    for i in range(3):
+        visualize_stencil_2d(splitted_directions[i], axes=axes[i], data=splitted_data[i], **kwargs)
+    for i in [-1, 0, 1]:
+        axes[i + 1].set_title("Cut at %s=%d" % (axes_names[slice_axis], i))
+
+
+def visualize_stencil_3d(stencil, figure=None, axes=None, data=None, textsize='8'):
+    """
+    Draws 3D stencil into a 3D coordinate system, parameters are similar to :func:`visualize_stencil_2d`
+    If data is None, no labels are drawn. To draw the labels as in the 2D case, use ``data=list(range(len(stencil)))``
+    """
+    from matplotlib.patches import FancyArrowPatch
+    from mpl_toolkits.mplot3d import proj3d
+    import matplotlib.pyplot as plt
+    from matplotlib.patches import BoxStyle
+    from itertools import product, combinations
+    import numpy as np
+
+    class Arrow3D(FancyArrowPatch):
+        def __init__(self, xs, ys, zs, *args, **kwargs):
+            FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
+            self._verts3d = xs, ys, zs
+
+        def draw(self, renderer):
+            xs3d, ys3d, zs3d = self._verts3d
+            xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
+            self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
+            FancyArrowPatch.draw(self, renderer)
+
+    if axes is None:
+        if figure is None:
+            figure = plt.figure()
+        axes = figure.gca(projection='3d')
+        axes.set_aspect("equal")
+
+    if data is None:
+        data = [None] * len(stencil)
+
+    text_offset = 1.25
+    text_box_style = BoxStyle("Round", pad=0.3)
+
+    # Draw cell (cube)
+    r = [-1, 1]
+    for s, e in combinations(np.array(list(product(r, r, r))), 2):
+        if np.sum(np.abs(s - e)) == r[1] - r[0]:
+            axes.plot3D(*zip(s, e), color="k", alpha=0.5)
+
+    for d, annotation in zip(stencil, data):
+        assert len(d) == 3, "Works only for 3D stencils"
+        if not (d[0] == 0 and d[1] == 0 and d[2] == 0):
+            if d[0] == 0:
+                color = '#348abd'
+            elif d[1] == 0:
+                color = '#fac364'
+            elif sum([abs(d) for d in d]) == 2:
+                color = '#95bd50'
+            else:
+                color = '#808080'
+
+            a = Arrow3D([0, d[0]], [0, d[1]], [0, d[2]], mutation_scale=20, lw=2, arrowstyle="-|>", color=color)
+            axes.add_artist(a)
+
+        if annotation:
+            if isinstance(annotation, sp.Basic):
+                annotation = "$" + sp.latex(annotation) + "$"
+            else:
+                annotation = str(annotation)
+
+            axes.text(d[0] * text_offset, d[1] * text_offset, d[2] * text_offset,
+                      annotation, verticalalignment='center', zorder=30,
+                      size=textsize, bbox=dict(boxstyle=text_box_style, facecolor='#777777', alpha=0.6, linewidth=0))
+
+    axes.set_xlim([-text_offset * 1.1, text_offset * 1.1])
+    axes.set_ylim([-text_offset * 1.1, text_offset * 1.1])
+    axes.set_zlim([-text_offset * 1.1, text_offset * 1.1])
+    axes.set_axis_off()
+
+
+def visualize_stencil_expression(expr, **kwargs):
+    stencil, coefficients = stencil_coefficients(expr)
+    dim = len(stencil[0])
+    assert 0 < dim <= 3
+    if dim == 1:
+        return stencil_coefficient_list(expr, matrix_form=True)
+    elif dim == 2:
+        return visualize_stencil_2d(stencil, data=coefficients, **kwargs)
+    elif dim == 3:
+        return visualize_stencil_3d_by_slicing(stencil, data=coefficients, **kwargs)

utils.py

@@ -2,6 +2,7 @@ import os
 from tempfile import NamedTemporaryFile
 from contextlib import contextmanager
 from typing import Mapping
+from collections import Counter
 
 
 class DotDict(dict):
@@ -66,3 +67,19 @@ def atomic_file_write(file_path):
         f.file.close()
         yield f.name
     os.rename(f.name, file_path)
+
+
+def fully_contains(l1, l2):
+    """Tests if elements of sequence 1 are in sequence 2 in same or higher number.
+
+    >>> fully_contains([1, 1, 2], [1, 2])  # 1 is only present once in second list
+    False
+    >>> fully_contains([1, 1, 2], [1, 1, 4, 2])
+    True
+    """
+    l1_counter = Counter(l1)
+    l2_counter = Counter(l2)
+    for element, count in l1_counter.items():
+        if l2_counter[element] < count:
+            return False
+    return True