pystencils/kernelcreation.py → src/pystencils/kernelcreation.py

-import functools
 import itertools
-from types import MappingProxyType
+import warnings
+from typing import Union, List
 
 import sympy as sp
+from pystencils.config import CreateKernelConfig
 
-from pystencils.assignment import Assignment
-from pystencils.astnodes import Block, Conditional, LoopOverCoordinate, SympyAssignment
+from pystencils.assignment import Assignment, AddAugmentedAssignment
+from pystencils.astnodes import Node, Block, Conditional, LoopOverCoordinate, SympyAssignment
 from pystencils.cpu.vectorization import vectorize
+from pystencils.enums import Target, Backend
 from pystencils.field import Field, FieldType
-from pystencils.gpucuda.indexing import indexing_creator_from_params
+from pystencils.node_collection import NodeCollection
 from pystencils.simp.assignment_collection import AssignmentCollection
+from pystencils.kernel_contrains_check import KernelConstraintsCheck
+from pystencils.simplificationfactory import create_simplification_strategy
 from pystencils.stencil import direction_string_to_offset, inverse_direction_string
 from pystencils.transformations import (
     loop_blocking, move_constants_before_loop, remove_conditionals_in_staggered_kernel)
 
 
-def create_kernel(assignments,
-                  target='cpu',
-                  data_type="double",
-                  iteration_slice=None,
-                  ghost_layers=None,
-                  skip_independence_check=False,
-                  cpu_openmp=False,
-                  cpu_vectorize_info=None,
-                  cpu_blocking=None,
-                  omp_single_loop=True,
-                  gpu_indexing='block',
-                  gpu_indexing_params=MappingProxyType({}),
-                  use_textures_for_interpolation=True,
-                  cpu_prepend_optimizations=[],
-                  use_auto_for_assignments=False,
-                  opencl_queue=None,
-                  opencl_ctx=None):
+def create_kernel(assignments: Union[Assignment, List[Assignment],
+                                     AddAugmentedAssignment, List[AddAugmentedAssignment],
+                                     AssignmentCollection, List[Node], NodeCollection],
+                  *,
+                  config: CreateKernelConfig = None, **kwargs):
     """
     Creates abstract syntax tree (AST) of kernel, using a list of update equations.
-
+    This function forms the general API and delegates the kernel creation to others depending on the CreateKernelConfig.
     Args:
         assignments: can be a single assignment, sequence of assignments or an `AssignmentCollection`
-        target: 'cpu', 'llvm', 'gpu' or 'opencl'
-        data_type: data type used for all untyped symbols (i.e. non-fields), can also be a dict from symbol name
-                  to type
-        iteration_slice: rectangular subset to iterate over, if not specified the complete non-ghost layer \
-                         part of the field is iterated over
-        ghost_layers: a single integer specifies the ghost layer count at all borders, can also be a sequence of
-                      pairs ``[(x_lower_gl, x_upper_gl), .... ]``. These layers are excluded from the iteration.
-                      If left to default, the number of ghost layers is determined automatically.
-        skip_independence_check: don't check that loop iterations are independent. This is needed e.g. for
-                                 periodicity kernel, that access the field outside the iteration bounds. Use with care!
-        cpu_openmp: True or number of threads for OpenMP parallelization, False for no OpenMP
-        omp_single_loop: if OpenMP is active: whether multiple outer loops are permitted
-        cpu_vectorize_info: a dictionary with keys, 'vector_instruction_set', 'assume_aligned' and 'nontemporal'
-                            for documentation of these parameters see vectorize function. Example:
-                            '{'instruction_set': 'avx512', 'assume_aligned': True, 'nontemporal':True}'
-        cpu_blocking: a tuple of block sizes or None if no blocking should be applied
-        gpu_indexing: either 'block' or 'line' , or custom indexing class, see `AbstractIndexing`
-        gpu_indexing_params: dict with indexing parameters (constructor parameters of indexing class)
-                             e.g. for 'block' one can specify '{'block_size': (20, 20, 10) }'
-        cpu_prepend_optimizations: list of extra optimizations to perform first on the AST
+        config: CreateKernelConfig which includes the needed configuration
+        kwargs: Arguments for updating the config
 
     Returns:
         abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
@@ -67,8 +41,8 @@ def create_kernel(assignments,
         >>> import numpy as np
         >>> s, d = ps.fields('s, d: [2D]')
         >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
-        >>> ast = ps.create_kernel(assignment, target='cpu', cpu_openmp=True)
-        >>> kernel = ast.compile()
+        >>> kernel_ast = ps.create_kernel(assignment, config=ps.CreateKernelConfig(cpu_openmp=True))
+        >>> kernel = kernel_ast.compile()
         >>> d_arr = np.zeros([5, 5])
         >>> kernel(d=d_arr, s=np.ones([5, 5]))
         >>> d_arr
@@ -78,82 +52,132 @@ def create_kernel(assignments,
                [0., 4., 4., 4., 0.],
                [0., 0., 0., 0., 0.]])
     """
+    # ----  Updating configuration from kwargs
+    if not config:
+        config = CreateKernelConfig(**kwargs)
+    else:
+        for k, v in kwargs.items():
+            if not hasattr(config, k):
+                raise KeyError(f'{v} is not a valid kwarg. Please look in CreateKernelConfig for valid settings')
+            setattr(config, k, v)
+
     # ----  Normalizing parameters
-    if isinstance(assignments, Assignment):
+    if isinstance(assignments, (Assignment, AddAugmentedAssignment)):
         assignments = [assignments]
     assert assignments, "Assignments must not be empty!"
-    split_groups = ()
-    if isinstance(assignments, AssignmentCollection):
-        if 'split_groups' in assignments.simplification_hints:
-            split_groups = assignments.simplification_hints['split_groups']
-        assignments = assignments.all_assignments
-
-    # ----  Creating ast
-    if target == 'cpu':
-        from pystencils.cpu import create_kernel
-        from pystencils.cpu import add_openmp
-        ast = create_kernel(assignments, type_info=data_type, split_groups=split_groups,
-                            iteration_slice=iteration_slice, ghost_layers=ghost_layers,
-                            skip_independence_check=skip_independence_check)
-        for optimization in cpu_prepend_optimizations:
-            optimization(ast)
-        omp_collapse = None
-        if cpu_blocking:
-            omp_collapse = loop_blocking(ast, cpu_blocking)
-        if cpu_openmp:
-            add_openmp(ast, num_threads=cpu_openmp, collapse=omp_collapse, assume_single_outer_loop=omp_single_loop)
-        if cpu_vectorize_info:
-            if cpu_vectorize_info is True:
-                vectorize(ast)
-            elif isinstance(cpu_vectorize_info, dict):
-                vectorize(ast, **cpu_vectorize_info)
-                if cpu_openmp and cpu_blocking and 'nontemporal' in cpu_vectorize_info and \
-                        cpu_vectorize_info['nontemporal'] and 'cachelineZero' in ast.instruction_set:
-                    # This condition is stricter than it needs to be: if blocks along the fastest axis start on a
-                    # cache line boundary, it's okay. But we cannot determine that here.
-                    # We don't need to disallow OpenMP collapsing because it is never applied to the inner loop.
-                    raise ValueError("Blocking cannot be combined with cacheline-zeroing")
-            else:
-                raise ValueError("Invalid value for cpu_vectorize_info")
-    elif target == 'llvm':
-        from pystencils.llvm import create_kernel
-        ast = create_kernel(assignments, type_info=data_type, split_groups=split_groups,
-                            iteration_slice=iteration_slice, ghost_layers=ghost_layers)
-    elif target == 'gpu' or target == 'opencl':
-        from pystencils.gpucuda import create_cuda_kernel
-        ast = create_cuda_kernel(assignments, type_info=data_type,
-                                 indexing_creator=indexing_creator_from_params(gpu_indexing, gpu_indexing_params),
-                                 iteration_slice=iteration_slice, ghost_layers=ghost_layers,
-                                 skip_independence_check=skip_independence_check,
-                                 use_textures_for_interpolation=use_textures_for_interpolation)
-        if target == 'opencl':
-            from pystencils.opencl.opencljit import make_python_function
-            ast._backend = 'opencl'
-            ast.compile = functools.partial(make_python_function, ast, opencl_queue, opencl_ctx)
-            ast._target = 'opencl'
-            ast._backend = 'opencl'
-        return ast
+    if isinstance(assignments, list):
+        assignments = NodeCollection(assignments)
+    elif isinstance(assignments, AssignmentCollection):
+        # TODO Markus check and doku
+        # --- applying first default simplifications
+        try:
+            if config.default_assignment_simplifications:
+                simplification = create_simplification_strategy()
+                assignments = simplification(assignments)
+        except Exception as e:
+            warnings.warn(f"It was not possible to apply the default pystencils optimisations to the "
+                          f"AssignmentCollection due to the following problem :{e}")
+        simplification_hints = assignments.simplification_hints
+        assignments = NodeCollection.from_assignment_collection(assignments)
+        assignments.simplification_hints = simplification_hints
+
+    if config.index_fields:
+        return create_indexed_kernel(assignments, config=config)
     else:
-        raise ValueError(f"Unknown target {target}. Has to be one of 'cpu', 'gpu' or 'llvm' ")
+        return create_domain_kernel(assignments, config=config)
+
+
+def create_domain_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
+    """
+    Creates abstract syntax tree (AST) of kernel, using a NodeCollection.
+
+    Note that `create_domain_kernel` is a lower level function which shoul be accessed by not providing `index_fields`
+    to create_kernel
+
+    Args:
+        assignments: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
+        config: CreateKernelConfig which includes the needed configuration
 
-    if use_auto_for_assignments:
+    Returns:
+        abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+        can be compiled with through its 'compile()' member
+
+    Example:
+        >>> import pystencils as ps
+        >>> import numpy as np
+        >>> from pystencils.kernelcreation import create_domain_kernel
+        >>> from pystencils.node_collection import NodeCollection
+        >>> s, d = ps.fields('s, d: [2D]')
+        >>> assignment = ps.Assignment(d[0,0], s[0, 1] + s[0, -1] + s[1, 0] + s[-1, 0])
+        >>> kernel_config = ps.CreateKernelConfig(cpu_openmp=True)
+        >>> kernel_ast = create_domain_kernel(NodeCollection([assignment]), config=kernel_config)
+        >>> kernel = kernel_ast.compile()
+        >>> d_arr = np.zeros([5, 5])
+        >>> kernel(d=d_arr, s=np.ones([5, 5]))
+        >>> d_arr
+        array([[0., 0., 0., 0., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 4., 4., 4., 0.],
+               [0., 0., 0., 0., 0.]])
+    """
+    # --- eval
+    assignments.evaluate_terms()
+
+    # FUTURE WORK from here we shouldn't NEED sympy
+    # --- check constrains
+    check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
+                                   check_double_write_condition=not config.allow_double_writes)
+
+    check.visit(assignments)
+
+    assignments.bound_fields = check.fields_written
+    assignments.rhs_fields = check.fields_read
+
+    # ----  Creating ast
+    ast = None
+    if config.target == Target.CPU:
+        if config.backend == Backend.C:
+            from pystencils.cpu import add_openmp, create_kernel
+            ast = create_kernel(assignments, config=config)
+            for optimization in config.cpu_prepend_optimizations:
+                optimization(ast)
+            omp_collapse = None
+            if config.cpu_blocking:
+                omp_collapse = loop_blocking(ast, config.cpu_blocking)
+            if config.cpu_openmp:
+                add_openmp(ast, num_threads=config.cpu_openmp, collapse=omp_collapse,
+                           assume_single_outer_loop=config.omp_single_loop)
+            if config.cpu_vectorize_info:
+                if config.cpu_vectorize_info is True:
+                    vectorize(ast)
+                elif isinstance(config.cpu_vectorize_info, dict):
+                    vectorize(ast, **config.cpu_vectorize_info)
+                    if config.cpu_openmp and config.cpu_blocking and 'nontemporal' in config.cpu_vectorize_info and \
+                            config.cpu_vectorize_info['nontemporal'] and 'cachelineZero' in ast.instruction_set:
+                        # This condition is stricter than it needs to be: if blocks along the fastest axis start on a
+                        # cache line boundary, it's okay. But we cannot determine that here.
+                        # We don't need to disallow OpenMP collapsing because it is never applied to the inner loop.
+                        raise ValueError("Blocking cannot be combined with cacheline-zeroing")
+                else:
+                    raise ValueError("Invalid value for cpu_vectorize_info")
+    elif config.target == Target.GPU:
+        if config.backend == Backend.CUDA:
+            from pystencils.gpu import create_cuda_kernel
+            ast = create_cuda_kernel(assignments, config=config)
+
+    if not ast:
+        raise NotImplementedError(
+            f'{config.target} together with {config.backend} is not supported by `create_domain_kernel`')
+
+    if config.use_auto_for_assignments:
         for a in ast.atoms(SympyAssignment):
             a.use_auto = True
 
     return ast
 
 
-def create_indexed_kernel(assignments,
-                          index_fields,
-                          target='cpu',
-                          data_type="double",
-                          coordinate_names=('x', 'y', 'z'),
-                          cpu_openmp=True,
-                          gpu_indexing='block',
-                          gpu_indexing_params=MappingProxyType({}),
-                          use_textures_for_interpolation=True,
-                          opencl_queue=None,
-                          opencl_ctx=None):
+def create_indexed_kernel(assignments: NodeCollection, *, config: CreateKernelConfig):
     """
     Similar to :func:`create_kernel`, but here not all cells of a field are updated but only cells with
     coordinates which are stored in an index field. This traversal method can e.g. be used for boundary handling.
@@ -163,12 +187,22 @@ def create_indexed_kernel(assignments,
     'coordinate_names' parameter. The struct can have also other fields that can be read and written in the kernel, for
     example boundary parameters.
 
-    index_fields: list of index fields, i.e. 1D fields with struct data type
-    coordinate_names: name of the coordinate fields in the struct data type
+    Note that `create_indexed_kernel` is a lower level function which shoul be accessed by providing `index_fields`
+    to create_kernel
+
+    Args:
+        assignments: `pystencils.node_collection.NodeCollection` containing all assignements and nodes to be processed
+        config: CreateKernelConfig which includes the needed configuration
+
+    Returns:
+        abstract syntax tree (AST) object, that can either be printed as source code with `show_code` or
+        can be compiled with through its 'compile()' member
 
     Example:
         >>> import pystencils as ps
+        >>> from pystencils.node_collection import NodeCollection
         >>> import numpy as np
+        >>> from pystencils.kernelcreation import create_indexed_kernel
         >>>
         >>> # Index field stores the indices of the cell to visit together with optional values
         >>> index_arr_dtype = np.dtype([('x', np.int32), ('y', np.int32), ('val', np.double)])
@@ -177,9 +211,10 @@ def create_indexed_kernel(assignments,
         >>>
         >>> # Additional values  stored in index field can be accessed in the kernel as well
         >>> s, d = ps.fields('s, d: [2D]')
-        >>> assignment = ps.Assignment(d[0,0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val'))
-        >>> ast = create_indexed_kernel(assignment, [idx_field], coordinate_names=('x', 'y'))
-        >>> kernel = ast.compile()
+        >>> assignment = ps.Assignment(d[0, 0], 2 * s[0, 1] + 2 * s[1, 0] + idx_field('val'))
+        >>> kernel_config = ps.CreateKernelConfig(index_fields=[idx_field], coordinate_names=('x', 'y'))
+        >>> kernel_ast = create_indexed_kernel(NodeCollection([assignment]), config=kernel_config)
+        >>> kernel = kernel_ast.compile()
         >>> d_arr = np.zeros([5, 5])
         >>> kernel(s=np.ones([5, 5]), d=d_arr, idx=index_arr)
         >>> d_arr
@@ -188,42 +223,37 @@ def create_indexed_kernel(assignments,
                [0. , 0. , 4.2, 0. , 0. ],
                [0. , 0. , 0. , 4.3, 0. ],
                [0. , 0. , 0. , 0. , 0. ]])
+
     """
-    if isinstance(assignments, Assignment):
-        assignments = [assignments]
-    elif isinstance(assignments, AssignmentCollection):
-        assignments = assignments.all_assignments
-    if target == 'cpu':
-        from pystencils.cpu import create_indexed_kernel
-        from pystencils.cpu import add_openmp
-        ast = create_indexed_kernel(assignments, index_fields=index_fields, type_info=data_type,
-                                    coordinate_names=coordinate_names)
-        if cpu_openmp:
-            add_openmp(ast, num_threads=cpu_openmp)
-        return ast
-    elif target == 'llvm':
-        raise NotImplementedError("Indexed kernels are not yet supported in LLVM backend")
-    elif target == 'gpu' or target == 'opencl':
-        from pystencils.gpucuda import created_indexed_cuda_kernel
-        idx_creator = indexing_creator_from_params(gpu_indexing, gpu_indexing_params)
-        ast = created_indexed_cuda_kernel(assignments,
-                                          index_fields,
-                                          type_info=data_type,
-                                          coordinate_names=coordinate_names,
-                                          indexing_creator=idx_creator,
-                                          use_textures_for_interpolation=use_textures_for_interpolation)
-        if target == 'opencl':
-            from pystencils.opencl.opencljit import make_python_function
-            ast._backend = 'opencl'
-            ast.compile = functools.partial(make_python_function, ast, opencl_queue, opencl_ctx)
-            ast._target = 'opencl'
-            ast._backend = 'opencl'
-        return ast
-    else:
-        raise ValueError(f"Unknown target {target}. Has to be either 'cpu' or 'gpu'")
+    # --- eval
+    assignments.evaluate_terms()
+
+    # FUTURE WORK from here we shouldn't NEED sympy
+    # --- check constrains
+    check = KernelConstraintsCheck(check_independence_condition=not config.skip_independence_check,
+                                   check_double_write_condition=not config.allow_double_writes)
+    check.visit(assignments)
+
+    assignments.bound_fields = check.fields_written
+    assignments.rhs_fields = check.fields_read
+
+    ast = None
+    if config.target == Target.CPU and config.backend == Backend.C:
+        from pystencils.cpu import add_openmp, create_indexed_kernel
+        ast = create_indexed_kernel(assignments, config=config)
+        if config.cpu_openmp:
+            add_openmp(ast, num_threads=config.cpu_openmp)
+    elif config.target == Target.GPU:
+        if config.backend == Backend.CUDA:
+            from pystencils.gpu import created_indexed_cuda_kernel
+            ast = created_indexed_cuda_kernel(assignments, config=config)
+
+    if not ast:
+        raise NotImplementedError(f'Indexed kernels are not yet supported for {config.target} with {config.backend}')
+    return ast
 
 
-def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=False, **kwargs):
+def create_staggered_kernel(assignments, target: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs):
     """Kernel that updates a staggered field.
 
     .. image:: /img/staggered_grid.svg
@@ -237,7 +267,7 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=
                      regular fields are passed through to `create_kernel`. Multiple different staggered fields can be
                      used, but they all need to use the same stencil (i.e. the same number of staggered points) and
                      shape.
-        target: 'cpu', 'llvm' or 'gpu'
+        target: 'CPU' or 'GPU'
         gpu_exclusive_conditions: disable the use of multiple conditionals inside the loop. The outer layers are then
                                   handled in an else branch.
         kwargs: passed directly to create_kernel, iteration_slice and ghost_layers parameters are not allowed
@@ -245,7 +275,16 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=
     Returns:
         AST, see `create_kernel`
     """
-    assert 'iteration_slice' not in kwargs and 'ghost_layers' not in kwargs and 'omp_single_loop' not in kwargs
+    # TODO: Add doku like in the other kernels
+    if 'ghost_layers' in kwargs:
+        assert kwargs['ghost_layers'] is None
+        del kwargs['ghost_layers']
+    if 'iteration_slice' in kwargs:
+        assert kwargs['iteration_slice'] is None
+        del kwargs['iteration_slice']
+    if 'omp_single_loop' in kwargs:
+        assert kwargs['omp_single_loop'] is False
+        del kwargs['omp_single_loop']
 
     if isinstance(assignments, AssignmentCollection):
         subexpressions = assignments.subexpressions + [a for a in assignments.main_assignments
@@ -323,7 +362,6 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=
 
         inner_assignment = []
         for assignment in assignments:
-            direction = stencil[assignment.lhs.index[0]]
             inner_assignment.append(SympyAssignment(assignment.lhs, assignment.rhs))
         last_conditional = Conditional(sp.And(*[condition(d) for d in stencil]),
                                        Block(inner_assignment), outer_assignment)
@@ -331,11 +369,8 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=
                             [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
                             [last_conditional]
 
-        if target == 'cpu':
-            from pystencils.cpu import create_kernel as create_kernel_cpu
-            ast = create_kernel_cpu(final_assignments, ghost_layers=ghost_layers, omp_single_loop=False, **kwargs)
-        else:
-            ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, **kwargs)
+        config = CreateKernelConfig(target=target, ghost_layers=ghost_layers, omp_single_loop=False, **kwargs)
+        ast = create_kernel(final_assignments, config=config)
         return ast
 
     for assignment in assignments:
@@ -349,6 +384,11 @@ def create_staggered_kernel(assignments, target='cpu', gpu_exclusive_conditions=
     remove_start_conditional = any([gl[0] == 0 for gl in ghost_layers])
     prepend_optimizations = [lambda ast: remove_conditionals_in_staggered_kernel(ast, remove_start_conditional),
                              move_constants_before_loop]
-    ast = create_kernel(final_assignments, ghost_layers=ghost_layers, target=target, omp_single_loop=False,
-                        cpu_prepend_optimizations=prepend_optimizations, **kwargs)
+    if 'cpu_prepend_optimizations' in kwargs:
+        prepend_optimizations += kwargs['cpu_prepend_optimizations']
+        del kwargs['cpu_prepend_optimizations']
+
+    config = CreateKernelConfig(ghost_layers=ghost_layers, target=target, omp_single_loop=False,
+                                cpu_prepend_optimizations=prepend_optimizations, **kwargs)
+    ast = create_kernel(final_assignments, config=config)
     return ast

src/pystencils/node_collection.py

+from typing import Any, Dict, List, Union, Optional, Set
+
+import sympy
+import sympy as sp
+from sympy.codegen.rewriting import ReplaceOptim, optimize
+
+from pystencils.assignment import Assignment, AddAugmentedAssignment
+import pystencils.astnodes as ast
+from pystencils.backends.cbackend import CustomCodeNode
+from pystencils.functions import DivFunc
+from pystencils.simp import AssignmentCollection
+from pystencils.typing import FieldPointerSymbol
+
+
+class NodeCollection:
+    def __init__(self, assignments: List[Union[ast.Node, Assignment]],
+                 simplification_hints: Optional[Dict[str, Any]] = None,
+                 bound_fields: Set[sp.Symbol] = None, rhs_fields: Set[sp.Symbol] = None):
+
+        def visit(obj):
+            if isinstance(obj, (list, tuple)):
+                return [visit(e) for e in obj]
+            if isinstance(obj, Assignment):
+                if isinstance(obj.lhs, FieldPointerSymbol):
+                    return ast.SympyAssignment(obj.lhs, obj.rhs, is_const=obj.lhs.dtype.const)
+                return ast.SympyAssignment(obj.lhs, obj.rhs)
+            elif isinstance(obj, AddAugmentedAssignment):
+                return ast.SympyAssignment(obj.lhs, obj.lhs + obj.rhs)
+            elif isinstance(obj, ast.SympyAssignment):
+                return obj
+            elif isinstance(obj, ast.Conditional):
+                true_block = visit(obj.true_block)
+                false_block = None if obj.false_block is None else visit(obj.false_block)
+                return ast.Conditional(obj.condition_expr, true_block=true_block, false_block=false_block)
+            elif isinstance(obj, ast.Block):
+                return ast.Block([visit(e) for e in obj.args])
+            elif isinstance(obj, ast.Node) and not isinstance(obj, ast.LoopOverCoordinate):
+                return obj
+            else:
+                raise ValueError("Invalid object in the List of Assignments " + str(type(obj)))
+
+        self.all_assignments = visit(assignments)
+        self.simplification_hints = simplification_hints if simplification_hints else {}
+        self.bound_fields = bound_fields if bound_fields else {}
+        self.rhs_fields = rhs_fields if rhs_fields else {}
+
+    @staticmethod
+    def from_assignment_collection(assignment_collection: AssignmentCollection):
+        return NodeCollection(assignments=assignment_collection.all_assignments,
+                              simplification_hints=assignment_collection.simplification_hints,
+                              bound_fields=assignment_collection.bound_fields,
+                              rhs_fields=assignment_collection.rhs_fields)
+
+    def evaluate_terms(self):
+        evaluate_constant_terms = ReplaceOptim(
+            lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
+            lambda p: p.evalf()
+        )
+
+        evaluate_pow = ReplaceOptim(
+            lambda e: e.is_Pow and e.exp.is_Integer and abs(e.exp) <= 8,
+            lambda p: sp.UnevaluatedExpr(sp.Mul(*([p.base] * +p.exp), evaluate=False)) if p.exp > 0 else
+            (DivFunc(sp.Integer(1), p.base) if p.exp == -1 else
+             DivFunc(sp.Integer(1), sp.UnevaluatedExpr(sp.Mul(*([p.base] * -p.exp), evaluate=False))))
+        )
+        sympy_optimisations = [evaluate_constant_terms, evaluate_pow]
+
+        def visitor(node):
+            if isinstance(node, CustomCodeNode):
+                return node
+            elif isinstance(node, ast.Block):
+                return node.func([visitor(child) for child in node.args])
+            elif isinstance(node, ast.SympyAssignment):
+                new_lhs = visitor(node.lhs)
+                new_rhs = visitor(node.rhs)
+                return node.func(new_lhs, new_rhs, node.is_const, node.use_auto)
+            elif isinstance(node, ast.Node):
+                return node.func(*[visitor(child) for child in node.args])
+            elif isinstance(node, sympy.Basic):
+                return optimize(node, sympy_optimisations)
+            else:
+                raise NotImplementedError(f'{node} {type(node)} has no valid visitor')
+
+        self.all_assignments = [visitor(assignment) for assignment in self.all_assignments]

pystencils/rng.py → src/pystencils/rng.py

@@ -2,7 +2,7 @@ import copy
 import numpy as np
 import sympy as sp
 
-from pystencils.data_types import TypedSymbol, cast_func
+from pystencils.typing import TypedSymbol, CastFunc
 from pystencils.astnodes import LoopOverCoordinate
 from pystencils.backends.cbackend import CustomCodeNode
 from pystencils.sympyextensions import fast_subs
@@ -47,14 +47,13 @@ class RNGBase(CustomCodeNode):
     def get_code(self, dialect, vector_instruction_set, print_arg):
         code = "\n"
         for r in self.result_symbols:
-            if vector_instruction_set and not self.args[1].atoms(cast_func):
+            if vector_instruction_set and not self.args[1].atoms(CastFunc):
                 # this vector RNG has become scalar through substitution
                 code += f"{r.dtype} {r.name};\n"
             else:
-                code += f"{vector_instruction_set[r.dtype.base_name] if vector_instruction_set else r.dtype} " + \
+                code += f"{vector_instruction_set[r.dtype.c_name] if vector_instruction_set else r.dtype} " + \
                         f"{r.name};\n"
-        args = [print_arg(a) for a in self.args] + \
-               [('&' if dialect == 'opencl' else '') + r.name for r in self.result_symbols]
+        args = [print_arg(a) for a in self.args] + ['' + r.name for r in self.result_symbols]
         code += (self._name + "(" + ", ".join(args) + ");\n")
         return code
 
@@ -62,6 +61,15 @@ class RNGBase(CustomCodeNode):
         return ", ".join([str(s) for s in self.result_symbols]) + " \\leftarrow " + \
             self._name.capitalize() + "_RNG(" + ", ".join([str(a) for a in self.args]) + ")"
 
+    def _hashable_content(self):
+        return (self._name, *self.result_symbols, *self.args)
+
+    def __eq__(self, other):
+        return type(self) is type(other) and self._hashable_content() == other._hashable_content()
+
+    def __hash__(self):
+        return hash(self._hashable_content())
+
 
 class PhiloxTwoDoubles(RNGBase):
     _name = "philox_double2"

pystencils/runhelper/db.py → src/pystencils/runhelper/db.py

 import socket
 import time
+from types import MappingProxyType
 from typing import Dict, Iterator, Sequence
 
 import blitzdb
+import six
+from blitzdb.backends.file.backend import serializer_classes
+from blitzdb.backends.file.utils import JsonEncoder
 
 from pystencils.cpu.cpujit import get_compiler_config
+from pystencils import CreateKernelConfig, Target, Backend, Field
+
+import json
+import sympy as sp
+
+from pystencils.typing import BasicType
+
+
+class PystencilsJsonEncoder(JsonEncoder):
+
+    def default(self, obj):
+        if isinstance(obj, CreateKernelConfig):
+            return obj.__dict__
+        if isinstance(obj, (sp.Float, sp.Rational)):
+            return float(obj)
+        if isinstance(obj, sp.Integer):
+            return int(obj)
+        if isinstance(obj, (BasicType, MappingProxyType)):
+            return str(obj)
+        if isinstance(obj, (Target, Backend, sp.Symbol)):
+            return obj.name
+        if isinstance(obj, Field):
+            return f"pystencils.Field(name = {obj.name}, field_type = {obj.field_type.name}, " \
+                   f"dtype = {str(obj.dtype)}, layout = {obj.layout}, shape = {obj.shape}, " \
+                   f"strides = {obj.strides})"
+        return JsonEncoder.default(self, obj)
+
+
+class PystencilsJsonSerializer(object):
+
+    @classmethod
+    def serialize(cls, data):
+        if six.PY3:
+            if isinstance(data, bytes):
+                return json.dumps(data.decode('utf-8'), cls=PystencilsJsonEncoder, ensure_ascii=False).encode('utf-8')
+            else:
+                return json.dumps(data, cls=PystencilsJsonEncoder, ensure_ascii=False).encode('utf-8')
+        else:
+            return json.dumps(data, cls=PystencilsJsonEncoder, ensure_ascii=False).encode('utf-8')
+
+    @classmethod
+    def deserialize(cls, data):
+        if six.PY3:
+            return json.loads(data.decode('utf-8'))
+        else:
+            return json.loads(data.decode('utf-8'))
 
 
 class Database:
@@ -46,7 +96,7 @@ class Database:
     class SimulationResult(blitzdb.Document):
         pass
 
-    def __init__(self, file: str) -> None:
+    def __init__(self, file: str, serializer_info: tuple = None) -> None:
         if file.startswith("mongo://"):
             from pymongo import MongoClient
             db_name = file[len("mongo://"):]
@@ -57,6 +107,10 @@ class Database:
 
         self.backend.autocommit = True
 
+        if serializer_info:
+            serializer_classes.update({serializer_info[0]: serializer_info[1]})
+            self.backend.load_config({'serializer_class': serializer_info[0]}, True)
+
     def save(self, params: Dict, result: Dict, env: Dict = None, **kwargs) -> None:
         """Stores a simulation result in the database.
 
@@ -146,10 +200,15 @@ class Database:
             'cpuCompilerConfig': get_compiler_config(),
         }
         try:
-            from git import Repo, InvalidGitRepositoryError
+            from git import Repo
+        except ImportError:
+            return result
+
+        try:
+            from git import InvalidGitRepositoryError
             repo = Repo(search_parent_directories=True)
             result['git_hash'] = str(repo.head.commit)
-        except (ImportError, InvalidGitRepositoryError):
+        except InvalidGitRepositoryError:
             pass
 
         return result

pystencils/runhelper/parameterstudy.py → src/pystencils/runhelper/parameterstudy.py

@@ -9,6 +9,7 @@ from time import sleep
 from typing import Any, Callable, Dict, Optional, Sequence, Tuple
 
 from pystencils.runhelper import Database
+from pystencils.runhelper.db import PystencilsJsonSerializer
 from pystencils.utils import DotDict
 
 ParameterDict = Dict[str, Any]
@@ -54,10 +55,11 @@ class ParameterStudy:
     Run = namedtuple("Run", ['parameter_dict', 'weight'])
 
     def __init__(self, run_function: Callable[..., Dict], runs: Sequence = (),
-                 database_connector: str = './db') -> None:
+                 database_connector: str = './db',
+                 serializer_info: tuple = ('pystencils_serializer', PystencilsJsonSerializer)) -> None:
         self.runs = list(runs)
         self.run_function = run_function
-        self.db = Database(database_connector)
+        self.db = Database(database_connector, serializer_info)
 
     def add_run(self, parameter_dict: ParameterDict, weight: int = 1) -> None:
         """Schedule a dictionary of parameters to run in this parameter study.

pystencils/session.py → src/pystencils/session.py

@@ -2,7 +2,7 @@ import numpy as np
 import sympy as sp
 
 import pystencils as ps
-import pystencils.jupyter
+from pystencils.jupyter import make_imshow_animation, display_animation, set_display_mode
 import pystencils.plot as plt
 
-__all__ = ['sp', 'np', 'ps', 'plt']
+__all__ = ['sp', 'np', 'ps', 'plt', 'make_imshow_animation', 'display_animation', 'set_display_mode']

pystencils/simp/__init__.py → src/pystencils/simp/__init__.py

@@ -5,10 +5,17 @@ from .simplifications import (
     add_subexpressions_for_sums, apply_on_all_subexpressions, apply_to_all_assignments,
     subexpression_substitution_in_existing_subexpressions,
     subexpression_substitution_in_main_assignments, sympy_cse, sympy_cse_on_assignment_list)
+from .subexpression_insertion import (
+    insert_aliases, insert_zeros, insert_constants,
+    insert_constant_additions, insert_constant_multiples,
+    insert_squares, insert_symbol_times_minus_one)
 from .simplificationstrategy import SimplificationStrategy
 
 __all__ = ['AssignmentCollection', 'SimplificationStrategy',
            'sympy_cse', 'sympy_cse_on_assignment_list', 'apply_to_all_assignments',
            'apply_on_all_subexpressions', 'subexpression_substitution_in_existing_subexpressions',
            'subexpression_substitution_in_main_assignments', 'add_subexpressions_for_constants',
-           'add_subexpressions_for_divisions', 'add_subexpressions_for_sums', 'add_subexpressions_for_field_reads']
+           'add_subexpressions_for_divisions', 'add_subexpressions_for_sums', 'add_subexpressions_for_field_reads',
+           'insert_aliases', 'insert_zeros', 'insert_constants',
+           'insert_constant_additions', 'insert_constant_multiples',
+           'insert_squares', 'insert_symbol_times_minus_one']

pystencils/simp/assignment_collection.py → src/pystencils/simp/assignment_collection.py

@@ -19,13 +19,14 @@ class AssignmentCollection:
     Additionally a dictionary of simplification hints is stored, which are set by the functions that create
     assignment collections to transport information to the simplification system.
 
-    Attributes:
-        main_assignments: list of assignments
-        subexpressions: list of assignments defining subexpressions used in main equations
-        simplification_hints: dict that is used to annotate the assignment collection with hints that are
+    Args:
+        main_assignments: List of assignments. Main assignments are characterised, that the right hand side of each
+                          assignment is a field access. Thus the generated equations write on arrays.
+        subexpressions: List of assignments defining subexpressions used in main equations
+        simplification_hints: Dict that is used to annotate the assignment collection with hints that are
                               used by the simplification system. See documentation of the simplification rules for
                               potentially required hints and their meaning.
-        subexpression_symbol_generator: generator for new symbols that are used when new subexpressions are added
+        subexpression_symbol_generator: Generator for new symbols that are used when new subexpressions are added
                                         used to get new symbols that are unique for this AssignmentCollection
 
     """
@@ -33,9 +34,13 @@ class AssignmentCollection:
     # ------------------------------- Creation & Inplace Manipulation --------------------------------------------------
 
     def __init__(self, main_assignments: Union[List[Assignment], Dict[sp.Expr, sp.Expr]],
-                 subexpressions: Union[List[Assignment], Dict[sp.Expr, sp.Expr]] = {},
+                 subexpressions: Union[List[Assignment], Dict[sp.Expr, sp.Expr]] = None,
                  simplification_hints: Optional[Dict[str, Any]] = None,
                  subexpression_symbol_generator: Iterator[sp.Symbol] = None) -> None:
+
+        if subexpressions is None:
+            subexpressions = {}
+
         if isinstance(main_assignments, Dict):
             main_assignments = [Assignment(k, v)
                                 for k, v in main_assignments.items()]
@@ -56,8 +61,11 @@ class AssignmentCollection:
 
         self.simplification_hints = simplification_hints
 
+        ctrs = [int(n.name[3:])for n in self.rhs_symbols if "xi_" in n.name]
+        max_ctr = max(ctrs) + 1 if len(ctrs) > 0 else 0
+
         if subexpression_symbol_generator is None:
-            self.subexpression_symbol_generator = SymbolGen()
+            self.subexpression_symbol_generator = SymbolGen(ctr=max_ctr)
         else:
             self.subexpression_symbol_generator = subexpression_symbol_generator
 
@@ -102,16 +110,21 @@ class AssignmentCollection:
         return self.subexpressions + self.main_assignments
 
     @property
-    def free_symbols(self) -> Set[sp.Symbol]:
-        """All symbols used in the assignment collection, which do not occur as left hand sides in any assignment."""
-        free_symbols = set()
+    def rhs_symbols(self) -> Set[sp.Symbol]:
+        """All symbols used in the assignment collection, which occur on the rhs of any assignment."""
+        rhs_symbols = set()
         for eq in self.all_assignments:
             if isinstance(eq, Assignment):
-                free_symbols.update(eq.rhs.atoms(sp.Symbol))
+                rhs_symbols.update(eq.rhs.atoms(sp.Symbol))
             elif isinstance(eq, pystencils.astnodes.Node):
-                free_symbols.update(eq.undefined_symbols)
+                rhs_symbols.update(eq.undefined_symbols)
 
-        return free_symbols - self.bound_symbols
+        return rhs_symbols
+
+    @property
+    def free_symbols(self) -> Set[sp.Symbol]:
+        """All symbols used in the assignment collection, which do not occur as left hand sides in any assignment."""
+        return self.rhs_symbols - self.bound_symbols
 
     @property
     def bound_symbols(self) -> Set[sp.Symbol]:
@@ -126,11 +139,15 @@ class AssignmentCollection:
         bound_symbols_set = bound_symbols_set.union(*[
             assignment.symbols_defined for assignment in self.all_assignments
             if isinstance(assignment, pystencils.astnodes.Node)
-        ]
-        )
+        ])
 
         return bound_symbols_set
 
+    @property
+    def rhs_fields(self):
+        """All fields accessed in the assignment collection, which do not occur as left hand sides in any assignment."""
+        return {s.field for s in self.rhs_symbols if hasattr(s, 'field')}
+
     @property
     def free_fields(self):
         """All fields accessed in the assignment collection, which do not occur as left hand sides in any assignment."""
@@ -144,11 +161,9 @@ class AssignmentCollection:
     @property
     def defined_symbols(self) -> Set[sp.Symbol]:
         """All symbols which occur as left-hand-sides of one of the main equations"""
-        return (set(
-            [assignment.lhs for assignment in self.main_assignments if isinstance(assignment, Assignment)]
-        ).union(*[assignment.symbols_defined for assignment in self.main_assignments if isinstance(
-                assignment, pystencils.astnodes.Node)]
-                ))
+        lhs_set = set([assignment.lhs for assignment in self.main_assignments if isinstance(assignment, Assignment)])
+        return (lhs_set.union(*[assignment.symbols_defined for assignment in self.main_assignments
+                                if isinstance(assignment, pystencils.astnodes.Node)]))
 
     @property
     def operation_count(self):
@@ -209,6 +224,7 @@ class AssignmentCollection:
             return {s: func(*args, **kwargs) for s, func in lambdas.items()}
 
         return f
+
     # ---------------------------- Creating new modified collections ---------------------------------------------------
 
     def copy(self,
@@ -270,12 +286,13 @@ class AssignmentCollection:
         processed_other_subexpression_equations = []
         for other_subexpression_eq in other.subexpressions:
             if other_subexpression_eq.lhs in own_subexpression_symbols:
-                if other_subexpression_eq.rhs == own_subexpression_symbols[other_subexpression_eq.lhs]:
+                new_rhs = fast_subs(other_subexpression_eq.rhs, substitution_dict)
+                if new_rhs == own_subexpression_symbols[other_subexpression_eq.lhs]:
                     continue  # exact the same subexpression equation exists already
                 else:
                     # different definition - a new name has to be introduced
                     new_lhs = next(self.subexpression_symbol_generator)
-                    new_eq = Assignment(new_lhs, fast_subs(other_subexpression_eq.rhs, substitution_dict))
+                    new_eq = Assignment(new_lhs, new_rhs)
                     processed_other_subexpression_equations.append(new_eq)
                     substitution_dict[other_subexpression_eq.lhs] = new_lhs
             else:
@@ -298,7 +315,7 @@ class AssignmentCollection:
             if eq.lhs in symbols_to_extract:
                 new_assignments.append(eq)
 
-        new_sub_expr = [eq for eq in self.subexpressions
+        new_sub_expr = [eq for eq in self.all_assignments
                         if eq.lhs in dependent_symbols and eq.lhs not in symbols_to_extract]
         return self.copy(new_assignments, new_sub_expr)
 
@@ -323,8 +340,10 @@ class AssignmentCollection:
         new_eqs = [Assignment(eq.lhs, fast_subs(eq.rhs, subs_dict)) for eq in self.main_assignments]
         return self.copy(new_eqs, new_subexpressions)
 
-    def new_without_subexpressions(self, subexpressions_to_keep: Set[sp.Symbol] = set()) -> 'AssignmentCollection':
+    def new_without_subexpressions(self, subexpressions_to_keep=None) -> 'AssignmentCollection':
         """Returns a new collection where all subexpressions have been inserted."""
+        if subexpressions_to_keep is None:
+            subexpressions_to_keep = set()
         if len(self.subexpressions) == 0:
             return self.copy()
 
@@ -352,6 +371,7 @@ class AssignmentCollection:
 
     def _repr_html_(self):
         """Interface to Jupyter notebook, to display as a nicely formatted HTML table"""
+
         def make_html_equation_table(equations):
             no_border = 'style="border:none"'
             html_table = '<table style="border:none; width: 100%; ">'
@@ -437,8 +457,8 @@ class AssignmentCollection:
 class SymbolGen:
     """Default symbol generator producing number symbols ζ_0, ζ_1, ..."""
 
-    def __init__(self, symbol="xi", dtype=None):
-        self._ctr = 0
+    def __init__(self, symbol="xi", dtype=None, ctr=0):
+        self._ctr = ctr
         self._symbol = symbol
         self._dtype = dtype
 

pystencils/simp/simplifications.py → src/pystencils/simp/simplifications.py

@@ -6,8 +6,9 @@ import sympy as sp
 
 from pystencils.assignment import Assignment
 from pystencils.astnodes import Node
-from pystencils.field import AbstractField, Field
+from pystencils.field import Field
 from pystencils.sympyextensions import subs_additive, is_constant, recursive_collect
+from pystencils.typing import TypedSymbol
 
 
 def sort_assignments_topologically(assignments: Sequence[Union[Assignment, Node]]) -> List[Union[Assignment, Node]]:
@@ -162,18 +163,20 @@ def add_subexpressions_for_sums(ac):
     for eq in ac.all_assignments:
         search_addends(eq.rhs)
 
-    addends = [a for a in addends if not isinstance(a, sp.Symbol) or isinstance(a, AbstractField.AbstractAccess)]
+    addends = [a for a in addends if not isinstance(a, sp.Symbol) or isinstance(a, Field.Access)]
     new_symbol_gen = ac.subexpression_symbol_generator
     substitutions = {addend: new_symbol for new_symbol, addend in zip(new_symbol_gen, addends)}
     return ac.new_with_substitutions(substitutions, True, substitute_on_lhs=False)
 
 
-def add_subexpressions_for_field_reads(ac, subexpressions=True, main_assignments=True):
+def add_subexpressions_for_field_reads(ac, subexpressions=True, main_assignments=True, data_type=None):
     r"""Substitutes field accesses on rhs of assignments with subexpressions
 
     Can change semantics of the update rule (which is the goal of this transformation)
     This is useful if a field should be update in place - all values are loaded before into subexpression variables,
     then the new values are computed and written to the same field in-place.
+    Additionally, if a datatype is given to the function the rhs symbol of the new isolated field read will have
+    this data type. This is useful for mixed precision kernels
     """
     field_reads = set()
     to_iterate = []
@@ -185,7 +188,17 @@ def add_subexpressions_for_field_reads(ac, subexpressions=True, main_assignments
     for assignment in to_iterate:
         if hasattr(assignment, 'lhs') and hasattr(assignment, 'rhs'):
             field_reads.update(assignment.rhs.atoms(Field.Access))
-    substitutions = {fa: next(ac.subexpression_symbol_generator) for fa in field_reads}
+
+    if not field_reads:
+        return ac
+
+    substitutions = dict()
+    for fa in field_reads:
+        lhs = next(ac.subexpression_symbol_generator)
+        if data_type is not None:
+            substitutions.update({fa: TypedSymbol(lhs.name, data_type)})
+        else:
+            substitutions.update({fa: lhs})
     return ac.new_with_substitutions(substitutions, add_substitutions_as_subexpressions=True,
                                      substitute_on_lhs=False, sort_topologically=False)
 
@@ -223,3 +236,31 @@ def apply_on_all_subexpressions(operation: Callable[[sp.Expr], sp.Expr]):
 
     f.__name__ = operation.__name__
     return f
+
+# TODO Markus
+# make this really work for Assignmentcollections
+# this function should ONLY evaluate
+# do the optims_c99 elsewhere optionally
+
+# def apply_sympy_optimisations(ac: AssignmentCollection):
+#     """ Evaluates constant expressions (e.g. :math:`\\sqrt{3}` will be replaced by its floating point representation)
+#         and applies the default sympy optimisations. See sympy.codegen.rewriting
+#     """
+#
+#     # Evaluates all constant terms
+#
+#     assignments = ac.all_assignments
+#
+#     evaluate_constant_terms = ReplaceOptim(lambda e: hasattr(e, 'is_constant') and e.is_constant and not e.is_integer,
+#                                            lambda p: p.evalf())
+#
+#     sympy_optimisations = [evaluate_constant_terms] + list(optims_c99)
+#
+#     assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
+#                    if hasattr(a, 'lhs')
+#                    else a for a in assignments]
+#     assignments_nodes = [a.atoms(SympyAssignment) for a in assignments]
+#     for a in chain.from_iterable(assignments_nodes):
+#         a.optimize(sympy_optimisations)
+#
+#     return AssignmentCollection(assignments)

src/pystencils/simp/subexpression_insertion.py

+import sympy as sp
+from pystencils.sympyextensions import is_constant
+
+#   Subexpression Insertion
+
+
+def insert_subexpressions(ac, selection_callback, skip=None):
+    """
+        Removes a number of subexpressions from an assignment collection by
+        inserting their right-hand side wherever they occur.
+
+        Args:
+         - selection_callback: Function that is called to qualify subexpressions
+            for insertion. Should return `True` for any subexpression that is to be
+            inserted, and `False` otherwise.
+         - skip: Set of symbols (left-hand sides of subexpressions) that should be 
+            ignored even if qualified by the callback.
+    """
+    if skip is None:
+        skip = set()
+    i = 0
+    while i < len(ac.subexpressions):
+        exp = ac.subexpressions[i]
+        if exp.lhs not in skip and selection_callback(exp):
+            ac = ac.new_with_inserted_subexpression(exp.lhs)
+        else:
+            i += 1
+
+    return ac
+
+
+def insert_aliases(ac, **kwargs):
+    """Inserts subexpressions that are aliases of other symbols, 
+    i.e. their right-hand side is only another symbol."""
+    return insert_subexpressions(ac, lambda x: isinstance(x.rhs, sp.Symbol), **kwargs)
+
+
+def insert_zeros(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is zero."""
+    zero = sp.Integer(0)
+    return insert_subexpressions(ac, lambda x: x.rhs == zero, **kwargs)
+
+
+def insert_constants(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is constant, 
+    i.e. contains no symbols."""
+    return insert_subexpressions(ac, lambda x: is_constant(x.rhs), **kwargs)
+
+
+def insert_symbol_times_minus_one(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is just a 
+    negation of another symbol."""
+    def callback(exp):
+        rhs = exp.rhs
+        minus_one = sp.Integer(-1)
+        atoms = rhs.atoms(sp.Symbol)
+        return len(atoms) == 1 and rhs == minus_one * atoms.pop()
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def insert_constant_multiples(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is a constant 
+    multiplied with another symbol."""
+    def callback(exp):
+        rhs = exp.rhs
+        symbols = rhs.atoms(sp.Symbol)
+        numbers = rhs.atoms(sp.Number)
+        return len(symbols) == 1 and len(numbers) == 1 and \
+            rhs == numbers.pop() * symbols.pop()
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def insert_constant_additions(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is a sum of a 
+    constant and another symbol."""
+    def callback(exp):
+        rhs = exp.rhs
+        symbols = rhs.atoms(sp.Symbol)
+        numbers = rhs.atoms(sp.Number)
+        return len(symbols) == 1 and len(numbers) == 1 and \
+            rhs == numbers.pop() + symbols.pop()
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def insert_squares(ac, **kwargs):
+    """Inserts subexpressions whose right-hand side is another symbol squared."""
+    def callback(exp):
+        rhs = exp.rhs
+        symbols = rhs.atoms(sp.Symbol)
+        return len(symbols) == 1 and rhs == symbols.pop() ** 2
+    return insert_subexpressions(ac, callback, **kwargs)
+
+
+def bind_symbols_to_skip(insertion_function, skip):
+    return lambda ac: insertion_function(ac, skip=skip)

src/pystencils/simplificationfactory.py

+from pystencils.simp import (SimplificationStrategy, insert_constants, insert_symbol_times_minus_one,
+                             insert_constant_multiples, insert_constant_additions, insert_squares, insert_zeros)
+
+
+def create_simplification_strategy():
+    """
+    Creates a default simplification `ps.simp.SimplificationStrategy`. The idea behind the default simplification
+    strategy is to reduce the number of subexpressions by inserting single constants and to evaluate constant
+    terms beforehand.
+    """
+    s = SimplificationStrategy()
+    s.add(insert_symbol_times_minus_one)
+    s.add(insert_constant_multiples)
+    s.add(insert_constant_additions)
+    s.add(insert_squares)
+    s.add(insert_zeros)
+    s.add(insert_constants)
+    s.add(lambda ac: ac.new_without_unused_subexpressions())

pystencils/stencil.py → src/pystencils/stencil.py

@@ -5,6 +5,8 @@ from typing import Sequence
 import numpy as np
 import sympy as sp
 
+from pystencils.utils import binary_numbers
+
 
 def inverse_direction(direction):
     """Returns inverse i.e. negative of given direction tuple
@@ -293,6 +295,38 @@ def direction_string_to_offset(direction: str, dim: int = 3):
     return offset[:dim]
 
 
+def adjacent_directions(direction):
+    """
+    Returns all adjacent directions for a direction as tuple of tuples. This is useful for exmple to find all directions
+    relevant for neighbour communication.
+
+    Args:
+        direction: tuple representing a direction. For example (0, 1, 0) for the northern side
+
+    Examples:
+        >>> adjacent_directions((0, 0, 0))
+        ((0, 0, 0),)
+        >>> adjacent_directions((0, 1, 0))
+        ((0, 1, 0),)
+        >>> adjacent_directions((0, 1, 1))
+        ((0, 0, 1), (0, 1, 0), (0, 1, 1))
+        >>> adjacent_directions((-1, -1))
+        ((-1, -1), (-1, 0), (0, -1))
+    """
+    result = set()
+    if all(e == 0 for e in direction):
+        result.add(direction)
+        return tuple(result)
+    binary_numbers_list = binary_numbers(len(direction))
+    for adjacent_direction in binary_numbers_list:
+        for i, entry in enumerate(direction):
+            if entry == 0:
+                adjacent_direction[i] = 0
+            if entry == -1 and adjacent_direction[i] == 1:
+                adjacent_direction[i] = -1
+        if not all(e == 0 for e in adjacent_direction):
+            result.add(tuple(adjacent_direction))
+    return tuple(sorted(result))
 # -------------------------------------- Visualization -----------------------------------------------------------------
 
 
@@ -341,7 +375,7 @@ def plot_2d(stencil, axes=None, figure=None, data=None, textsize='12', **kwargs)
     for direction, annotation in zip(stencil, data):
         assert len(direction) == 2, "Works only for 2D stencils"
         direction = tuple(int(i) for i in direction)
-        if not(direction[0] == 0 and direction[1] == 0):
+        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):
@@ -421,16 +455,17 @@ def plot_3d(stencil, figure=None, axes=None, data=None, textsize='8'):
             FancyArrowPatch.__init__(self, (0, 0), (0, 0), *args, **kwargs)
             self._verts3d = xs, ys, zs
 
-        def draw(self, renderer):
+        def do_3d_projection(self, *_):
             xs3d, ys3d, zs3d = self._verts3d
-            xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, renderer.M)
+            xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
             self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
-            FancyArrowPatch.draw(self, renderer)
+
+            return np.min(zs)
 
     if axes is None:
         if figure is None:
             figure = plt.figure()
-        axes = figure.gca(projection='3d')
+        axes = figure.add_subplot(projection='3d')
         try:
             axes.set_aspect("equal")
         except NotImplementedError:
@@ -446,7 +481,7 @@ def plot_3d(stencil, figure=None, axes=None, data=None, textsize='8'):
     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)
+            axes.plot(*zip(s, e), color="k", alpha=0.5)
 
     for d, annotation in zip(stencil, data):
         assert len(d) == 3, "Works only for 3D stencils"

pystencils/sympyextensions.py → src/pystencils/sympyextensions.py

@@ -6,12 +6,14 @@ from functools import partial, reduce
 from typing import Callable, Dict, Iterable, List, Optional, Sequence, Tuple, TypeVar, Union
 
 import sympy as sp
+from sympy import PolynomialError
 from sympy.functions import Abs
 from sympy.core.numbers import Zero
 
 from pystencils.assignment import Assignment
-from pystencils.data_types import cast_func, get_type_of_expression, PointerType, VectorType
-from pystencils.kernelparameters import FieldPointerSymbol
+from pystencils.functions import DivFunc
+from pystencils.typing import CastFunc, get_type_of_expression, PointerType, VectorType
+from pystencils.typing.typed_sympy import FieldPointerSymbol
 
 T = TypeVar('T')
 
@@ -158,17 +160,23 @@ def fast_subs(expression: T, substitutions: Dict,
     if type(expression) is sp.Matrix:
         return expression.copy().applyfunc(partial(fast_subs, substitutions=substitutions))
 
-    def visit(expr):
+    def visit(expr, evaluate=True):
         if skip and skip(expr):
             return expr
-        if hasattr(expr, "fast_subs"):
+        elif hasattr(expr, "fast_subs"):
             return expr.fast_subs(substitutions, skip)
-        if expr in substitutions:
+        elif expr in substitutions:
             return substitutions[expr]
-        if not hasattr(expr, 'args'):
+        elif not hasattr(expr, 'args'):
             return expr
-        param_list = [visit(a) for a in expr.args]
-        return expr if not param_list else expr.func(*param_list)
+        elif isinstance(expr, (sp.UnevaluatedExpr, DivFunc)):
+            args = [visit(a, False) for a in expr.args]
+            return expr.func(*args)
+        else:
+            param_list = [visit(a, evaluate) for a in expr.args]
+            if isinstance(expr, (sp.Mul, sp.Add)):
+                return expr if not param_list else expr.func(*param_list, evaluate=evaluate)
+            return expr if not param_list else expr.func(*param_list)
 
     if len(substitutions) == 0:
         return expression
@@ -235,6 +243,9 @@ def subs_additive(expr: sp.Expr, replacement: sp.Expr, subexpression: sp.Expr,
 
     normalized_replacement_match = normalize_match_parameter(required_match_replacement, len(subexpression.args))
 
+    if isinstance(subexpression, sp.Number):
+        return expr.subs({replacement: subexpression})
+
     def visit(current_expr):
         if current_expr.is_Add:
             expr_max_length = max(len(current_expr.args), len(subexpression.args))
@@ -263,7 +274,7 @@ def subs_additive(expr: sp.Expr, replacement: sp.Expr, subexpression: sp.Expr,
             return current_expr
         else:
             if current_expr.func == sp.Mul and Zero() in param_list:
-                return Zero()
+                return sp.simplify(current_expr)
             else:
                 return current_expr.func(*param_list, evaluate=False)
 
@@ -345,7 +356,7 @@ def remove_higher_order_terms(expr: sp.Expr, symbols: Sequence[sp.Symbol], order
         factor_count = 0
         if type(product) is Mul:
             for factor in product.args:
-                if type(factor) == Pow:
+                if type(factor) is Pow:
                     if factor.args[0] in symbols:
                         factor_count += factor.args[1]
                 if factor in symbols:
@@ -355,13 +366,13 @@ def remove_higher_order_terms(expr: sp.Expr, symbols: Sequence[sp.Symbol], order
                 factor_count += product.args[1]
         return factor_count
 
-    if type(expr) == Mul or type(expr) == Pow:
+    if type(expr) is Mul or type(expr) is Pow:
         if velocity_factors_in_product(expr) <= order:
             return expr
         else:
-            return sp.Rational(0, 1)
+            return Zero()
 
-    if type(expr) != Add:
+    if type(expr) is not Add:
         return expr
 
     for sum_term in expr.args:
@@ -432,11 +443,14 @@ def extract_most_common_factor(term):
 
 
 def recursive_collect(expr, symbols, order_by_occurences=False):
-    """Applies sympy.collect recursively for a list of symbols, collecting symbol 2 in the coefficients of symbol 1, 
+    """Applies sympy.collect recursively for a list of symbols, collecting symbol 2 in the coefficients of symbol 1,
     and so on.
+    
+    ``expr`` must be rewritable as a polynomial in the given ``symbols``.
+    It it is not, ``recursive_collect`` will fail quietly, returning the original expression.
 
     Args:
-        expr: A sympy expression
+        expr: A sympy expression.
         symbols: A sequence of symbols
         order_by_occurences: If True, during recursive descent, always collect the symbol occuring 
                              most often in the expression.
@@ -447,13 +461,85 @@ def recursive_collect(expr, symbols, order_by_occurences=False):
     if len(symbols) == 0:
         return expr
     symbol = symbols[0]
-    collected_poly = sp.Poly(expr.collect(symbol), symbol)
+    collected = expr.collect(symbol)
+    
+    try:
+        collected_poly = sp.Poly(collected, symbol)
+    except PolynomialError:
+        return expr
+
     coeffs = collected_poly.all_coeffs()[::-1]
     rec_sum = sum(symbol**i * recursive_collect(c, symbols[1:], order_by_occurences) for i, c in enumerate(coeffs))
     return rec_sum
 
 
-def count_operations(term: Union[sp.Expr, List[sp.Expr]],
+def summands(expr):
+    return set(expr.args) if isinstance(expr, sp.Add) else {expr}
+
+
+def simplify_by_equality(expr, a, b, c):
+    """
+    Uses the equality a = b + c, where a and b must be symbols, to simplify expr 
+    by attempting to express additive combinations of two quantities by the third.
+
+    This works on expressions that are reducible to the form 
+    :math:`a * (...) + b * (...) + c * (...)`,
+    without any mixed terms of a, b and c.
+    """
+    if not isinstance(a, sp.Symbol) or not isinstance(b, sp.Symbol):
+        raise ValueError("a and b must be symbols.")
+
+    c = sp.sympify(c)
+
+    if not (isinstance(c, sp.Symbol) or is_constant(c)):
+        raise ValueError("c must be either a symbol or a constant!")
+
+    expr = sp.sympify(expr)
+
+    expr_expanded = sp.expand(expr)
+    a_coeff = expr_expanded.coeff(a, 1)
+    expr_expanded -= (a * a_coeff).expand()
+    b_coeff = expr_expanded.coeff(b, 1)
+    expr_expanded -= (b * b_coeff).expand()
+    if isinstance(c, sp.Symbol):
+        c_coeff = expr_expanded.coeff(c, 1)
+        rest = expr_expanded - (c * c_coeff).expand()
+    else:
+        c_coeff = expr_expanded / c
+        rest = 0
+
+    a_summands = summands(a_coeff)
+    b_summands = summands(b_coeff)
+    c_summands = summands(c_coeff)
+
+    # replace b + c by a
+    b_plus_c_coeffs = b_summands & c_summands
+    for coeff in b_plus_c_coeffs:
+        rest += a * coeff
+    b_summands -= b_plus_c_coeffs
+    c_summands -= b_plus_c_coeffs
+
+    # replace a - b by c
+    neg_b_summands = {-x for x in b_summands}
+    a_minus_b_coeffs = a_summands & neg_b_summands
+    for coeff in a_minus_b_coeffs:
+        rest += c * coeff
+    a_summands -= a_minus_b_coeffs
+    b_summands -= {-x for x in a_minus_b_coeffs}
+
+    # replace a - c by b
+    neg_c_summands = {-x for x in c_summands}
+    a_minus_c_coeffs = a_summands & neg_c_summands
+    for coeff in a_minus_c_coeffs:
+        rest += b * coeff
+    a_summands -= a_minus_c_coeffs
+    c_summands -= {-x for x in a_minus_c_coeffs}
+
+    # put it back together
+    return (rest + a * sum(a_summands) + b * sum(b_summands) + c * sum(c_summands)).expand()
+
+
+def count_operations(term: Union[sp.Expr, List[sp.Expr], List[Assignment]],
                      only_type: Optional[str] = 'real') -> Dict[str, int]:
     """Counts the number of additions, multiplications and division.
 
@@ -519,7 +605,7 @@ def count_operations(term: Union[sp.Expr, List[sp.Expr]],
             visit_children = False
         elif t.is_integer:
             pass
-        elif isinstance(t, cast_func):
+        elif isinstance(t, CastFunc):
             visit_children = False
             visit(t.args[0])
         elif t.func is fast_sqrt:
@@ -553,8 +639,10 @@ def count_operations(term: Union[sp.Expr, List[sp.Expr]],
             for child_term, condition in t.args:
                 visit(child_term)
             visit_children = False
-        elif isinstance(t, sp.Rel):
+        elif isinstance(t, (sp.Rel, sp.UnevaluatedExpr)):
             pass
+        elif isinstance(t, DivFunc):
+            result["divs"] += 1
         else:
             warnings.warn(f"Unknown sympy node of type {str(t.func)} counting will be inaccurate")