src/pystencils/include/ppc_altivec_helpers.h

+/*
+Copyright 2021, Michael Kuron.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+* Redistributions of source code must retain the above copyright
+  notice, this list of conditions, and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright
+  notice, this list of conditions, and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+
+* Neither the name of of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+#include <altivec.h>
+#undef vector
+#undef bool
+
+inline void cachelineZero(void * p) {
+#ifdef __xlC__
+	__dcbz(p);
+#else
+	__asm__ volatile("dcbz 0, %0"::"r"(p):"memory");
+#endif
+}
+
+inline size_t _cachelineSize() {
+	// allocate and fill with ones
+	const size_t max_size = 0x100000;
+	uint8_t data[2*max_size];
+	for (size_t i = 0; i < 2*max_size; ++i) {
+		data[i] = 0xff;
+	}
+	
+	// find alignment offset
+	size_t offset = max_size - ((uintptr_t) data) % max_size;
+
+	// zero a cacheline
+	cachelineZero((void*) (data + offset));
+
+	// make sure that at least one byte was zeroed
+	if (data[offset] != 0) {
+		return SIZE_MAX;
+	}
+
+	// make sure that nothing was zeroed before the pointer
+	if (data[offset-1] == 0) {
+		return SIZE_MAX;
+	}
+
+	// find the last byte that was zeroed
+	for (size_t size = 1; size < max_size; ++size) {
+		if (data[offset + size] != 0) {
+			return size;
+		}
+	}
+	
+	// too much was zeroed
+	return SIZE_MAX;
+}
+
+inline size_t cachelineSize() {
+	static size_t size = _cachelineSize();
+	return size;
+}

src/pystencils/include/riscv_v_helpers.h

+/*
+Copyright 2023, Michael Kuron.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are
+met:
+
+* Redistributions of source code must retain the above copyright
+  notice, this list of conditions, and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright
+  notice, this list of conditions, and the following disclaimer in the
+  documentation and/or other materials provided with the distribution.
+
+* Neither the name of of the copyright holder nor the names of its
+  contributors may be used to endorse or promote products derived from
+  this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+*/
+
+inline void cachelineZero(void * p) {
+#ifdef __riscv_zicboz
+	__asm__ volatile("cbo.zero (%0)"::"r"(p):"memory");
+#endif
+}
+
+inline size_t _cachelineSize() {
+	// allocate and fill with ones
+	const size_t max_size = 0x100000;
+	uint8_t data[2*max_size];
+	for (size_t i = 0; i < 2*max_size; ++i) {
+		data[i] = 0xff;
+	}
+	
+	// find alignment offset
+	size_t offset = max_size - ((uintptr_t) data) % max_size;
+
+	// zero a cacheline
+	cachelineZero((void*) (data + offset));
+
+	// make sure that at least one byte was zeroed
+	if (data[offset] != 0) {
+		return SIZE_MAX;
+	}
+
+	// make sure that nothing was zeroed before the pointer
+	if (data[offset-1] == 0) {
+		return SIZE_MAX;
+	}
+
+	// find the last byte that was zeroed
+	for (size_t size = 1; size < max_size; ++size) {
+		if (data[offset + size] != 0) {
+			return size;
+		}
+	}
+	
+	// too much was zeroed
+	return SIZE_MAX;
+}
+
+inline size_t cachelineSize() {
+#ifdef __riscv_zicboz
+	static size_t size = _cachelineSize();
+	return size;
+#else
+	return SIZE_MAX;
+#endif
+}

src/pystencils/integer_functions.py

+# TODO #47 move to a module functions
+import numpy as np
+import sympy as sp
+
+from pystencils.typing import CastFunc, collate_types, create_type, get_type_of_expression
+from pystencils.sympyextensions import is_integer_sequence
+
+
+class IntegerFunctionTwoArgsMixIn(sp.Function):
+    is_integer = True
+
+    def __new__(cls, arg1, arg2):
+        args = []
+        for a in (arg1, arg2):
+            if isinstance(a, sp.Number) or isinstance(a, int):
+                args.append(CastFunc(a, create_type("int")))
+            elif isinstance(a, np.generic):
+                args.append(CastFunc(a, a.dtype))
+            else:
+                args.append(a)
+
+        for a in args:
+            try:
+                dtype = get_type_of_expression(a)
+                if not dtype.is_int():
+                    raise ValueError("Argument to integer function is not an int but " + str(dtype))
+            except NotImplementedError:
+                raise ValueError("Integer functions can only be constructed with typed expressions")
+        return super().__new__(cls, *args)
+
+    def _eval_evalf(self, *pargs, **kwargs):
+        arg1 = self.args[0].evalf(*pargs, **kwargs) if hasattr(self.args[0], 'evalf') else self.args[0]
+        arg2 = self.args[1].evalf(*pargs, **kwargs) if hasattr(self.args[1], 'evalf') else self.args[1]
+        return self._eval_op(arg1, arg2)
+
+    def _eval_op(self, arg1, arg2):
+        return self
+
+
+# noinspection PyPep8Naming
+class bitwise_xor(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class bit_shift_right(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class bit_shift_left(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class bitwise_and(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class bitwise_or(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class int_div(IntegerFunctionTwoArgsMixIn):
+    
+    def _eval_op(self, arg1, arg2):
+        return int(arg1 // arg2)
+
+
+# noinspection PyPep8Naming
+class int_power_of_2(IntegerFunctionTwoArgsMixIn):
+    pass
+
+
+# noinspection PyPep8Naming
+class modulo_floor(sp.Function):
+    """Returns the next smaller integer divisible by given divisor.
+
+    Examples:
+        >>> modulo_floor(9, 4)
+        8
+        >>> modulo_floor(11, 4)
+        8
+        >>> modulo_floor(12, 4)
+        12
+        >>> from pystencils import TypedSymbol
+        >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32")
+        >>> modulo_floor(a, b).to_c(str)
+        '(int64_t)((a) / (b)) * (b)'
+    """
+    nargs = 2
+    is_integer = True
+
+    def __new__(cls, integer, divisor):
+        if is_integer_sequence((integer, divisor)):
+            return (int(integer) // int(divisor)) * divisor
+        else:
+            return super().__new__(cls, integer, divisor)
+
+    def to_c(self, print_func):
+        dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1])))
+        assert dtype.is_int()
+        return "({dtype})(({0}) / ({1})) * ({1})".format(print_func(self.args[0]),
+                                                         print_func(self.args[1]), dtype=dtype)
+
+
+# noinspection PyPep8Naming
+class modulo_ceil(sp.Function):
+    """Returns the next bigger integer divisible by given divisor.
+
+    Examples:
+        >>> modulo_ceil(9, 4)
+        12
+        >>> modulo_ceil(11, 4)
+        12
+        >>> modulo_ceil(12, 4)
+        12
+        >>> from pystencils import TypedSymbol
+        >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32")
+        >>> modulo_ceil(a, b).to_c(str)
+        '((a) % (b) == 0 ? a : ((int64_t)((a) / (b))+1) * (b))'
+    """
+    nargs = 2
+    is_integer = True
+
+    def __new__(cls, integer, divisor):
+        if is_integer_sequence((integer, divisor)):
+            return integer if integer % divisor == 0 else ((integer // divisor) + 1) * divisor
+        else:
+            return super().__new__(cls, integer, divisor)
+
+    def to_c(self, print_func):
+        dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1])))
+        assert dtype.is_int()
+        code = "(({0}) % ({1}) == 0 ? {0} : (({dtype})(({0}) / ({1}))+1) * ({1}))"
+        return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype)
+
+
+# noinspection PyPep8Naming
+class div_ceil(sp.Function):
+    """Integer division that is always rounded up
+
+    Examples:
+        >>> div_ceil(9, 4)
+        3
+        >>> div_ceil(8, 4)
+        2
+        >>> from pystencils import TypedSymbol
+        >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32")
+        >>> div_ceil(a, b).to_c(str)
+        '( (a) % (b) == 0 ? (int64_t)(a) / (int64_t)(b) : ( (int64_t)(a) / (int64_t)(b) ) +1 )'
+    """
+    nargs = 2
+    is_integer = True
+
+    def __new__(cls, integer, divisor):
+        if is_integer_sequence((integer, divisor)):
+            return integer // divisor if integer % divisor == 0 else (integer // divisor) + 1
+        else:
+            return super().__new__(cls, integer, divisor)
+
+    def to_c(self, print_func):
+        dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1])))
+        assert dtype.is_int()
+        code = "( ({0}) % ({1}) == 0 ? ({dtype})({0}) / ({dtype})({1}) : ( ({dtype})({0}) / ({dtype})({1}) ) +1 )"
+        return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype)
+
+
+# noinspection PyPep8Naming
+class div_floor(sp.Function):
+    """Integer division
+
+    Examples:
+        >>> div_floor(9, 4)
+        2
+        >>> div_floor(8, 4)
+        2
+        >>> from pystencils import TypedSymbol
+        >>> a, b = TypedSymbol("a", "int64"), TypedSymbol("b", "int32")
+        >>> div_floor(a, b).to_c(str)
+        '((int64_t)(a) / (int64_t)(b))'
+    """
+    nargs = 2
+    is_integer = True
+
+    def __new__(cls, integer, divisor):
+        if is_integer_sequence((integer, divisor)):
+            return integer // divisor
+        else:
+            return super().__new__(cls, integer, divisor)
+
+    def to_c(self, print_func):
+        dtype = collate_types((get_type_of_expression(self.args[0]), get_type_of_expression(self.args[1])))
+        assert dtype.is_int()
+        code = "(({dtype})({0}) / ({dtype})({1}))"
+        return code.format(print_func(self.args[0]), print_func(self.args[1]), dtype=dtype)

src/pystencils/integer_set_analysis.py

+"""Transformations using integer sets based on ISL library"""
+
+import islpy as isl
+import sympy as sp
+
+import pystencils.astnodes as ast
+from pystencils.typing import parents_of_type
+from pystencils.backends.cbackend import CustomSympyPrinter
+
+
+def remove_brackets(s):
+    return s.replace('[', '').replace(']', '')
+
+
+def _degrees_of_freedom_as_string(expr):
+    expr = sp.sympify(expr)
+    indexed = expr.atoms(sp.Indexed)
+    symbols = expr.atoms(sp.Symbol)
+    symbols_without_indexed_base = symbols - {ind.base.args[0] for ind in indexed}
+    symbols_without_indexed_base.update(indexed)
+    return {remove_brackets(str(s)) for s in symbols_without_indexed_base}
+
+
+def isl_iteration_set(node: ast.Node):
+    """Builds up an ISL set describing the iteration space by analysing the enclosing loops of the given node. """
+    conditions = []
+    degrees_of_freedom = set()
+
+    for loop in parents_of_type(node, ast.LoopOverCoordinate):
+        if loop.step != 1:
+            raise NotImplementedError("Loops with strides != 1 are not yet supported.")
+
+        degrees_of_freedom.update(_degrees_of_freedom_as_string(loop.loop_counter_symbol))
+        degrees_of_freedom.update(_degrees_of_freedom_as_string(loop.start))
+        degrees_of_freedom.update(_degrees_of_freedom_as_string(loop.stop))
+
+        loop_start_str = remove_brackets(str(loop.start))
+        loop_stop_str = remove_brackets(str(loop.stop))
+        ctr_name = loop.loop_counter_name
+        set_string_description = f"{ctr_name} >= {loop_start_str} and {ctr_name} < {loop_stop_str}"
+        conditions.append(remove_brackets(set_string_description))
+
+    symbol_names = ','.join(degrees_of_freedom)
+    condition_str = ' and '.join(conditions)
+    set_description = f"{{ [{symbol_names}] : {condition_str} }}"
+    return degrees_of_freedom, isl.BasicSet(set_description)
+
+
+def simplify_loop_counter_dependent_conditional(conditional):
+    """Removes conditionals that depend on the loop counter or iteration limits if they are always true/false."""
+    dofs_in_condition = _degrees_of_freedom_as_string(conditional.condition_expr)
+    dofs_in_loops, iteration_set = isl_iteration_set(conditional)
+    if dofs_in_condition.issubset(dofs_in_loops):
+        symbol_names = ','.join(dofs_in_loops)
+        condition_str = CustomSympyPrinter().doprint(conditional.condition_expr)
+        condition_str = remove_brackets(condition_str)
+        condition_set = isl.BasicSet(f"{{ [{symbol_names}] : {condition_str} }}")
+
+        if condition_set.is_empty():
+            conditional.replace_by_false_block()
+            return
+
+        intersection = iteration_set.intersect(condition_set)
+        if intersection.is_empty():
+            conditional.replace_by_false_block()
+        elif intersection == iteration_set:
+            conditional.replace_by_true_block()

src/pystencils/jupyter.py

+import base64
+from tempfile import NamedTemporaryFile
+
+import matplotlib.animation as animation
+import sympy as sp
+from IPython.display import HTML
+
+import pystencils.plot as plt
+
+VIDEO_TAG = """<video controls width="80%">
+ <source src="data:video/x-m4v;base64,{0}" type="video/mp4">
+ Your browser does not support the video tag.
+</video>"""
+
+
+def __animation_to_html(animation, fps):
+    if not hasattr(animation, 'encoded_video'):
+        with NamedTemporaryFile(suffix='.mp4') as f:
+            animation.save(f.name, fps=fps, extra_args=['-vcodec', 'libx264', '-pix_fmt',
+                                                        'yuv420p', '-profile:v', 'baseline', '-level', '3.0'])
+            video = open(f.name, "rb").read()
+        animation.encoded_video = base64.b64encode(video).decode('ascii')
+
+    return VIDEO_TAG.format(animation.encoded_video)
+
+
+def make_imshow_animation(grid, grid_update_function, frames=90, **_):
+    from functools import partial
+    fig = plt.figure()
+    im = plt.imshow(grid, interpolation='none')
+
+    def update_figure(*_, **kwargs):
+        image = kwargs['image']
+        image = grid_update_function(image)
+        im.set_array(image)
+        return im,
+
+    return animation.FuncAnimation(fig, partial(update_figure, image=grid), frames=frames)
+
+
+# -------   Version 1: Embed the animation as HTML5 video --------- ----------------------------------
+
+def display_as_html_video(animation, fps=30, show=True, **_):
+    try:
+        plt.close(animation._fig)
+        res = __animation_to_html(animation, fps)
+        if show:
+            return HTML(res)
+        else:
+            return HTML("")
+    except KeyboardInterrupt:
+        pass
+
+
+# -------   Version 2: Animation is shown in extra matplotlib window ----------------------------------
+
+
+def display_in_extra_window(*_, **__):
+    fig = plt.gcf()
+    try:
+        fig.canvas.manager.window.raise_()
+    except Exception:
+        pass
+    plt.show()
+
+
+# -------   Version 3: Animation is shown in images that are updated directly in website --------------
+
+def display_as_html_image(animation, show=True, *args, **kwargs):
+    from IPython import display
+
+    try:
+        if show:
+            animation._init_draw()
+        for _ in animation.frame_seq:
+            if show:
+                fig = plt.gcf()
+                display.display(fig)
+            animation._step()
+            if show:
+                display.clear_output(wait=True)
+    except KeyboardInterrupt:
+        display.clear_output(wait=False)
+
+
+# Dispatcher
+
+animation_display_mode = 'image_update'
+display_animation_func = None
+
+
+def display_animation(*args, **kwargs):
+    from IPython import get_ipython
+    ipython = get_ipython()
+    if not ipython:
+        return
+
+    if not display_animation_func:
+        raise Exception("Call set_display_mode first")
+    return display_animation_func(*args, **kwargs)
+
+
+def set_display_mode(mode):
+    from IPython import get_ipython
+    ipython = get_ipython()
+    if not ipython:
+        return
+    global animation_display_mode
+    global display_animation_func
+    animation_display_mode = mode
+    if animation_display_mode == 'video':
+        ipython.magic("matplotlib inline")
+        display_animation_func = display_as_html_video
+    elif animation_display_mode == 'window':
+        ipython.magic("matplotlib qt")
+        display_animation_func = display_in_extra_window
+    elif animation_display_mode == 'image_update':
+        ipython.magic("matplotlib inline")
+        display_animation_func = display_as_html_image
+    else:
+        raise Exception("Unknown mode. Available modes 'image_update', 'video' and 'window' ")
+
+
+def activate_ipython():
+    from IPython import get_ipython
+    ipython = get_ipython()
+    if ipython:
+        set_display_mode('image_update')
+        ipython.magic("config InlineBackend.rc = { }")
+        ipython.magic("matplotlib inline")
+        plt.rc('figure', figsize=(16, 6))
+        sp.init_printing()
+
+
+activate_ipython()

src/pystencils/kernel_contrains_check.py

+from collections import namedtuple, defaultdict
+from typing import Union
+
+import sympy as sp
+from sympy.codegen import Assignment
+
+from pystencils.simp import AssignmentCollection
+from pystencils import astnodes as ast, TypedSymbol
+from pystencils.field import Field
+from pystencils.node_collection import NodeCollection
+from pystencils.transformations import NestedScopes
+
+# TODO use this in Constraint Checker
+accepted_functions = [
+    sp.Pow,
+    sp.sqrt,
+    sp.log,
+    # TODO trigonometric functions (and whatever tests will fail)
+]
+
+
+class KernelConstraintsCheck:
+    # TODO: proper specification
+    # TODO: More checks :)
+    """Checks if the input to create_kernel is valid.
+
+    Test the following conditions:
+
+    - SSA Form for pure symbols:
+        -  Every pure symbol may occur only once as left-hand-side of an assignment
+        -  Every pure symbol that is read, may not be written to later
+    - Independence / Parallelization condition:
+        - a field that is written may only be read at exact the same spatial position
+
+    (Pure symbols are symbols that are not Field.Accesses)
+    """
+    FieldAndIndex = namedtuple('FieldAndIndex', ['field', 'index'])
+
+    def __init__(self, check_independence_condition=True, check_double_write_condition=True):
+        self.scopes = NestedScopes()
+        self.field_reads = defaultdict(set)
+        self.field_writes = defaultdict(set)
+        self.fields_read = set()
+        self.check_independence_condition = check_independence_condition
+        self.check_double_write_condition = check_double_write_condition
+
+    def visit(self, obj):
+        if isinstance(obj, (AssignmentCollection, NodeCollection)):
+            [self.visit(e) for e in obj.all_assignments]
+        elif isinstance(obj, list) or isinstance(obj, tuple):
+            [self.visit(e) for e in obj]
+        elif isinstance(obj, (sp.Eq, ast.SympyAssignment, Assignment)):
+            self.process_assignment(obj)
+        elif isinstance(obj, ast.Conditional):
+            self.scopes.push()
+            # Disable double write check inside conditionals
+            # would be triggered by e.g. in-kernel boundaries
+            old_double_write = self.check_double_write_condition
+            old_independence_condition = self.check_independence_condition
+            self.check_double_write_condition = False
+            self.check_independence_condition = False
+            if obj.false_block:
+                self.visit(obj.false_block)
+            self.process_expression(obj.condition_expr)
+            self.process_expression(obj.true_block)
+            self.check_double_write_condition = old_double_write
+            self.check_independence_condition = old_independence_condition
+            self.scopes.pop()
+        elif isinstance(obj, ast.Block):
+            self.scopes.push()
+            [self.visit(e) for e in obj.args]
+            self.scopes.pop()
+        elif isinstance(obj, ast.Node) and not isinstance(obj, ast.LoopOverCoordinate):
+            pass
+        else:
+            raise ValueError(f'Invalid object in kernel {type(obj)}')
+
+    def process_assignment(self, assignment: Union[sp.Eq, ast.SympyAssignment, Assignment]):
+        # for checks it is crucial to process rhs before lhs to catch e.g. a = a + 1
+        self.process_expression(assignment.rhs)
+        self.process_lhs(assignment.lhs)
+
+    def process_expression(self, rhs):
+        # TODO constraint for accepted functions, see TODO above
+        self.update_accesses_rhs(rhs)
+        if isinstance(rhs, Field.Access):
+            self.fields_read.add(rhs.field)
+            self.fields_read.update(rhs.indirect_addressing_fields)
+        else:
+            for arg in rhs.args:
+                self.process_expression(arg)
+
+    @property
+    def fields_written(self):
+        """
+        Return all rhs fields
+        """
+        return set(k.field for k, v in self.field_writes.items() if len(v))
+
+    def process_lhs(self, lhs: Union[Field.Access, TypedSymbol, sp.Symbol]):
+        assert isinstance(lhs, sp.Symbol)
+        self.update_accesses_lhs(lhs)
+
+    def update_accesses_lhs(self, lhs):
+        if isinstance(lhs, Field.Access):
+            fai = self.FieldAndIndex(lhs.field, lhs.index)
+            if self.check_double_write_condition and lhs.offsets in self.field_writes[fai]:
+                raise ValueError(f"Field {lhs.field.name} is written twice at the same location")
+
+            self.field_writes[fai].add(lhs.offsets)
+
+            if self.check_double_write_condition and len(self.field_writes[fai]) > 1:
+                raise ValueError(
+                    f"Field {lhs.field.name} is written at two different locations")
+
+            if fai in self.field_reads:
+                reads = tuple(self.field_reads[fai])
+                if len(reads) > 1 or lhs.offsets != reads[0]:
+                    if self.check_independence_condition:
+                        raise ValueError(f"Field {lhs.field.name} is written at different location than it was read. "
+                                         f"This means the resulting kernel would not be thread safe")
+        elif isinstance(lhs, sp.Symbol):
+            if self.scopes.is_defined_locally(lhs):
+                raise ValueError(f"Assignments not in SSA form, multiple assignments to {lhs.name}")
+            if lhs in self.scopes.free_parameters:
+                raise ValueError(f"Symbol {lhs.name} is written, after it has been read")
+            self.scopes.define_symbol(lhs)
+
+    def update_accesses_rhs(self, rhs):
+        if isinstance(rhs, Field.Access) and self.check_independence_condition:
+            fai = self.FieldAndIndex(rhs.field, rhs.index)
+            writes = self.field_writes[fai]
+            self.field_reads[fai].add(rhs.offsets)
+            for write_offset in writes:
+                assert len(writes) == 1
+                if write_offset != rhs.offsets:
+                    raise ValueError(f"Violation of loop independence condition. Field "
+                                     f"{rhs.field} is read at {rhs.offsets} and written at {write_offset}")
+            self.fields_read.add(rhs.field)
+        elif isinstance(rhs, sp.Symbol):
+            self.scopes.access_symbol(rhs)

src/pystencils/kernel_decorator.py

+import ast
+import inspect
+import textwrap
+from typing import Callable, Union, List, Dict, Tuple
+
+import sympy as sp
+
+from pystencils.assignment import Assignment
+from pystencils.sympyextensions import SymbolCreator
+from pystencils.config import CreateKernelConfig
+
+__all__ = ['kernel', 'kernel_config']
+
+
+def _kernel(func: Callable[..., None], **kwargs) -> Tuple[List[Assignment], str]:
+    """
+    Convenient function for kernel decorator to prevent code duplication
+    Args:
+        func: decorated function
+        **kwargs: kwargs for the function
+    Returns:
+        assignments, function_name
+    """
+    source = inspect.getsource(func)
+    source = textwrap.dedent(source)
+    a = ast.parse(source)
+    KernelFunctionRewrite().visit(a)
+    ast.fix_missing_locations(a)
+    gl = func.__globals__.copy()
+
+    assignments = []
+
+    def assignment_adder(lhs, rhs):
+        assignments.append(Assignment(lhs, rhs))
+
+    gl['_add_assignment'] = assignment_adder
+    gl['_Piecewise'] = sp.Piecewise
+    gl.update(inspect.getclosurevars(func).nonlocals)
+    exec(compile(a, filename="<ast>", mode="exec"), gl)
+    func = gl[func.__name__]
+    args = inspect.getfullargspec(func).args
+    if 's' in args and 's' not in kwargs:
+        kwargs['s'] = SymbolCreator()
+    func(**kwargs)
+    return assignments, func.__name__
+
+
+def kernel(func: Callable[..., None], **kwargs) -> List[Assignment]:
+    """Decorator to simplify generation of pystencils Assignments.
+
+    Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
+    in the result list. Furthermore the meaning of the ternary inline 'if-else' changes meaning to denote a
+    sympy Piecewise.
+
+    The decorated function may not receive any arguments, with exception of an argument called 's' that specifies
+    a SymbolCreator()
+    Args:
+        func: decorated function
+        **kwargs: kwargs for the function
+
+    Examples:
+        >>> import pystencils as ps
+        >>> @kernel
+        ... def my_kernel(s):
+        ...     f, g = ps.fields('f, g: [2D]')
+        ...     s.neighbors @= f[0,1] + f[1,0]
+        ...     g[0,0]      @= s.neighbors + f[0,0] if f[0,0] > 0 else 0
+        >>> f, g = ps.fields('f, g: [2D]')
+        >>> assert my_kernel[0].rhs == f[0,1] + f[1,0]
+    """
+    assignments, _ = _kernel(func, **kwargs)
+    return assignments
+
+
+def kernel_config(config: CreateKernelConfig, **kwargs) -> Callable[..., Dict]:
+    """Decorator to simplify generation of pystencils Assignments, which takes a configuration
+    and updates the function name accordingly.
+
+    Changes the meaning of the '@=' operator. Each line containing this operator gives a symbolic assignment
+    in the result list. Furthermore, the meaning of the ternary inline 'if-else' changes meaning to denote a
+    sympy Piecewise.
+
+    The decorated function may not receive any arguments, with exception to an argument called 's' that specifies
+    a SymbolCreator()
+    Args:
+        config: Specify whether to return the list with assignments, or a dictionary containing additional settings
+                like func_name
+    Returns:
+        decorator with config
+
+    Examples:
+        >>> import pystencils as ps
+        >>> kernel_configuration = ps.CreateKernelConfig()
+        >>> @kernel_config(kernel_configuration)
+        ... def my_kernel(s):
+        ...     src, dst = ps.fields('src, dst: [2D]')
+        ...     s.neighbors @= src[0, 1] + src[1, 0]
+        ...     dst[0, 0]      @= s.neighbors + src[0, 0] if src[0, 0] > 0 else 0
+        >>> f, g = ps.fields('src, dst: [2D]')
+        >>> assert my_kernel['assignments'][0].rhs == f[0, 1] + f[1, 0]
+    """
+    def decorator(func: Callable[..., None]) -> Union[List[Assignment], Dict]:
+        """
+        Args:
+            func: decorated function
+        Returns:
+            Dict for unpacking into create_kernel
+        """
+        assignments, func_name = _kernel(func, **kwargs)
+        config.function_name = func_name
+        return {'assignments': assignments, 'config': config}
+    return decorator
+
+
+# noinspection PyMethodMayBeStatic
+class KernelFunctionRewrite(ast.NodeTransformer):
+
+    def visit_IfExp(self, node):
+        piecewise_func = ast.Name(id='_Piecewise', ctx=ast.Load())
+        piecewise_func = ast.copy_location(piecewise_func, node)
+        piecewise_args = [ast.Tuple(elts=[node.body, node.test], ctx=ast.Load()),
+                          ast.Tuple(elts=[node.orelse, ast.NameConstant(value=True)], ctx=ast.Load())]
+        result = ast.Call(func=piecewise_func, args=piecewise_args, keywords=[])
+
+        return ast.copy_location(result, node)
+
+    def visit_AugAssign(self, node):
+        self.generic_visit(node)
+        node.target.ctx = ast.Load()
+        new_node = ast.Expr(ast.Call(func=ast.Name(id='_add_assignment', ctx=ast.Load()),
+                                     args=[node.target, node.value],
+                                     keywords=[]))
+        return ast.copy_location(new_node, node)
+
+    def visit_FunctionDef(self, node):
+        self.generic_visit(node)
+        node.decorator_list = []
+        return node

src/pystencils/kernel_wrapper.py

+import pystencils
+
+
+class KernelWrapper:
+    """
+    Light-weight wrapper around a compiled kernel.
+
+    Can be called while still providing access to underlying AST.
+    """
+
+    def __init__(self, kernel, parameters, ast_node: pystencils.astnodes.KernelFunction):
+        self.kernel = kernel
+        self.parameters = parameters
+        self.ast = ast_node
+        self.num_regs = None
+
+    def __call__(self, **kwargs):
+        return self.kernel(**kwargs)
+
+    @property
+    def code(self):
+        return pystencils.get_code_str(self.ast)

src/pystencils/kernelcreation.py

+import itertools
+import warnings
+from typing import Union, List
+
+import sympy as sp
+from pystencils.config import CreateKernelConfig
+
+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.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: 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`
+        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
+        can be compiled with through its 'compile()' member
+
+    Example:
+        >>> import pystencils as ps
+        >>> 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])
+        >>> 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
+        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.]])
+    """
+    # ----  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, AddAugmentedAssignment)):
+        assignments = [assignments]
+    assert assignments, "Assignments must not be empty!"
+    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:
+        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
+
+    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: 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.
+
+    The coordinates are stored in a separated index_field, which is a one dimensional array with struct data type.
+    This struct has to contain fields named 'x', 'y' and for 3D fields ('z'). These names are configurable with the
+    'coordinate_names' parameter. The struct can have also other fields that can be read and written in the kernel, for
+    example boundary parameters.
+
+    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)])
+        >>> index_arr = np.array([(1, 1, 0.1), (2, 2, 0.2), (3, 3, 0.3)], dtype=index_arr_dtype)
+        >>> idx_field = ps.fields(idx=index_arr)
+        >>>
+        >>> # 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'))
+        >>> 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
+        array([[0. , 0. , 0. , 0. , 0. ],
+               [0. , 4.1, 0. , 0. , 0. ],
+               [0. , 0. , 4.2, 0. , 0. ],
+               [0. , 0. , 0. , 4.3, 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
+
+    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: Target = Target.CPU, gpu_exclusive_conditions=False, **kwargs):
+    """Kernel that updates a staggered field.
+
+    .. image:: /img/staggered_grid.svg
+
+    For a staggered field, the first index coordinate defines the location of the staggered value.
+    Further index coordinates can be used to store vectors/tensors at each point.
+
+    Args:
+        assignments: a sequence of assignments or an AssignmentCollection.
+                     Assignments to staggered field are processed specially, while subexpressions and assignments to
+                     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' 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
+
+    Returns:
+        AST, see `create_kernel`
+    """
+    # 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
+                                                       if not hasattr(a, 'lhs')
+                                                       or type(a.lhs) is not Field.Access
+                                                       or not FieldType.is_staggered(a.lhs.field)]
+        assignments = [a for a in assignments.main_assignments if hasattr(a, 'lhs')
+                       and type(a.lhs) is Field.Access
+                       and FieldType.is_staggered(a.lhs.field)]
+    else:
+        subexpressions = [a for a in assignments if not hasattr(a, 'lhs')
+                          or type(a.lhs) is not Field.Access
+                          or not FieldType.is_staggered(a.lhs.field)]
+        assignments = [a for a in assignments if hasattr(a, 'lhs')
+                       and type(a.lhs) is Field.Access
+                       and FieldType.is_staggered(a.lhs.field)]
+    if len(set([tuple(a.lhs.field.staggered_stencil) for a in assignments])) != 1:
+        raise ValueError("All assignments need to be made to staggered fields with the same stencil")
+    if len(set([a.lhs.field.shape for a in assignments])) != 1:
+        raise ValueError("All assignments need to be made to staggered fields with the same shape")
+
+    staggered_field = assignments[0].lhs.field
+    stencil = staggered_field.staggered_stencil
+    dim = staggered_field.spatial_dimensions
+    shape = staggered_field.shape
+
+    counters = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(dim)]
+
+    final_assignments = []
+
+    # find out whether any of the ghost layers is not needed
+    common_exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+    for direction in stencil:
+        exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+        for elementary_direction in direction:
+            exclusions.remove(inverse_direction_string(elementary_direction))
+        common_exclusions.intersection_update(exclusions)
+    ghost_layers = [[0, 0] for d in range(dim)]
+    for direction in common_exclusions:
+        direction = direction_string_to_offset(direction)
+        for d, s in enumerate(direction):
+            if s == 1:
+                ghost_layers[d][1] = 1
+            elif s == -1:
+                ghost_layers[d][0] = 1
+
+    def condition(direction):
+        """exclude those staggered points that correspond to fluxes between ghost cells"""
+        exclusions = set(["E", "W", "N", "S", "T", "B"][:2 * dim])
+
+        for elementary_direction in direction:
+            exclusions.remove(inverse_direction_string(elementary_direction))
+        conditions = []
+        for e in exclusions:
+            if e in common_exclusions:
+                continue
+            offset = direction_string_to_offset(e)
+            for i, o in enumerate(offset):
+                if o == 1:
+                    conditions.append(counters[i] < shape[i] - 1)
+                elif o == -1:
+                    conditions.append(counters[i] > 0)
+        return sp.And(*conditions)
+
+    if gpu_exclusive_conditions:
+        outer_assignment = None
+        conditions = {direction: condition(direction) for direction in stencil}
+        for num_conditions in range(len(stencil)):
+            for combination in itertools.combinations(conditions.values(), num_conditions):
+                for assignment in assignments:
+                    direction = stencil[assignment.lhs.index[0]]
+                    if conditions[direction] in combination:
+                        assignment = SympyAssignment(assignment.lhs, assignment.rhs)
+                        outer_assignment = Conditional(sp.And(*combination), Block([assignment]), outer_assignment)
+
+        inner_assignment = []
+        for assignment in assignments:
+            inner_assignment.append(SympyAssignment(assignment.lhs, assignment.rhs))
+        last_conditional = Conditional(sp.And(*[condition(d) for d in stencil]),
+                                       Block(inner_assignment), outer_assignment)
+        final_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
+                            [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
+                            [last_conditional]
+
+        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:
+        direction = stencil[assignment.lhs.index[0]]
+        sp_assignments = [s for s in subexpressions if not hasattr(s, 'lhs')] + \
+                         [SympyAssignment(s.lhs, s.rhs) for s in subexpressions if hasattr(s, 'lhs')] + \
+                         [SympyAssignment(assignment.lhs, assignment.rhs)]
+        last_conditional = Conditional(condition(direction), Block(sp_assignments))
+        final_assignments.append(last_conditional)
+
+    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]
+    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]

src/pystencils/placeholder_function.py

+from typing import List
+
+import sympy as sp
+
+from pystencils.assignment import Assignment
+from pystencils.astnodes import Node
+from pystencils.sympyextensions import is_constant
+from pystencils.transformations import generic_visit
+
+
+class PlaceholderFunction:
+    pass
+
+
+def to_placeholder_function(expr, name):
+    """Replaces an expression by a sympy function.
+
+    - replacing an expression with just a symbol would lead to problem when calculating derivatives
+    - placeholder functions get rid of this problem
+
+    Examples:
+        >>> x, t = sp.symbols("x, t")
+        >>> temperature = x**2 + t**4 # some 'complicated' dependency
+        >>> temperature_placeholder = to_placeholder_function(temperature, 'T')
+        >>> diffusivity = temperature_placeholder + 42 * t
+        >>> sp.diff(diffusivity, t)  # returns a symbol instead of the computed derivative
+        _dT_dt + 42
+        >>> result, subexpr = remove_placeholder_functions(diffusivity)
+        >>> result
+        T + 42*t
+        >>> subexpr
+        [Assignment(T, t**4 + x**2), Assignment(_dT_dt, 4*t**3), Assignment(_dT_dx, 2*x)]
+
+    """
+    symbols = list(expr.atoms(sp.Symbol))
+    symbols.sort(key=lambda e: e.name)
+    derivative_symbols = [sp.Symbol(f"_d{name}_d{s.name}") for s in symbols]
+    derivatives = [sp.diff(expr, s) for s in symbols]
+
+    assignments = [Assignment(sp.Symbol(name), expr)]
+    assignments += [Assignment(symbol, derivative)
+                    for symbol, derivative in zip(derivative_symbols, derivatives)
+                    if not is_constant(derivative)]
+
+    def fdiff(_, index):
+        result = derivatives[index - 1]
+        return result if is_constant(result) else derivative_symbols[index - 1]
+
+    func = type(name, (sp.Function, PlaceholderFunction),
+                {'fdiff': fdiff,
+                 'value': sp.Symbol(name),
+                 'subexpressions': assignments,
+                 'nargs': len(symbols)})
+    return func(*symbols)
+
+
+def remove_placeholder_functions(expr):
+    subexpressions = []
+
+    def visit(e):
+        if isinstance(e, Node):
+            return e
+        elif isinstance(e, PlaceholderFunction):
+            for se in e.subexpressions:
+                if se.lhs not in {a.lhs for a in subexpressions}:
+                    subexpressions.append(se)
+            return e.value
+        else:
+            new_args = [visit(a) for a in e.args]
+            return e.func(*new_args) if new_args else e
+
+    return generic_visit(expr, visit), subexpressions
+
+
+def prepend_placeholder_functions(assignments: List[Assignment]):
+    result, subexpressions = remove_placeholder_functions(assignments)
+    return subexpressions + result

src/pystencils/plot.py

+"""
+This module extends the pyplot module with functions to show scalar and vector fields in the usual
+simulation coordinate system (y-axis goes up), instead of the "image coordinate system" (y axis goes down) that
+matplotlib normally uses.
+"""
+import warnings
+from itertools import cycle
+
+from matplotlib.pyplot import *
+
+
+def vector_field(array, step=2, **kwargs):
+    """Plots given vector field as quiver (arrow) plot.
+
+    Args:
+        array: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+               coordinate should have shape 2 and stores the 2 velocity components
+        step: plots only every steps's cell, increase the step for high resolution arrays
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.quiver`
+
+    Returns:
+        quiver plot object
+    """
+    assert len(array.shape) == 3, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    assert array.shape[2] == 2, "Last array dimension is expected to store 2D vectors"
+    vel_n = array.swapaxes(0, 1)
+    res = quiver(vel_n[::step, ::step, 0], vel_n[::step, ::step, 1], **kwargs)
+    axis('equal')
+    return res
+
+
+def vector_field_magnitude(array, **kwargs):
+    """Plots the magnitude of a vector field as colormap.
+
+    Args:
+        array: numpy array with 3 dimensions, first two are spatial x,y coordinate, the last
+               coordinate should have shape 2 and stores the 2 velocity components
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+
+    Returns:
+        imshow object
+    """
+    assert len(array.shape) == 3, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    assert array.shape[2] in (2, 3), "Wrong size of the last coordinate. Has to be a 2D or 3D vector field."
+    from numpy.linalg import norm
+    norm = norm(array, axis=2, ord=2)
+    if hasattr(array, 'mask'):
+        norm = np.ma.masked_array(norm, mask=array.mask[:, :, 0])
+    return scalar_field(norm, **kwargs)
+
+
+def scalar_field(array, **kwargs):
+    """Plots field values as colormap.
+
+    Works just as imshow, but uses coordinate system where second coordinate (y) points upwards.
+
+    Args:
+        array: two dimensional numpy array
+        kwargs: keyword arguments passed to :func:`matplotlib.pyplot.imshow`
+
+    Returns:
+        imshow object
+    """
+    import numpy
+    array = numpy.swapaxes(array, 0, 1)
+    res = imshow(array, origin='lower', **kwargs)
+    axis('equal')
+    return res
+
+
+def scalar_field_surface(array, **kwargs):
+    """Plots scalar field as 3D surface
+
+    Args:
+        array: the two dimensional numpy array to plot
+        kwargs: keyword arguments passed to :func:`mpl_toolkits.mplot3d.Axes3D.plot_surface`
+    """
+    from mpl_toolkits.mplot3d import Axes3D
+    from matplotlib import cm
+
+    fig = gcf()
+    ax = fig.add_subplot(111, projection='3d')
+    x, y = np.meshgrid(np.arange(array.shape[0]), np.arange(array.shape[1]), indexing='ij')
+    kwargs.setdefault('rstride', 2)
+    kwargs.setdefault('cstride', 2)
+    kwargs.setdefault('color', 'b')
+    kwargs.setdefault('cmap', cm.coolwarm)
+    return ax.plot_surface(x, y, array, **kwargs)
+
+
+def scalar_field_alpha_value(array, color, clip=False, **kwargs):
+    """Plots an image with same color everywhere, using the array values as transparency.
+
+    Array is supposed to have values between 0 and 1 (if this is not the case it is normalized).
+    An image is plotted that has the same color everywhere, the passed array determines the transparency.
+    Regions where the array is 1  are fully opaque, areas with 0 are fully transparent.
+
+    Args:
+        array: 2D array with alpha values
+        color: fill color
+        clip: if True, all values in the array larger than 1 are set to 1, all values smaller than 0 are set to zero
+              if False, the array is linearly scaled to the [0, 1] interval
+        **kwargs: arguments passed to imshow
+
+    Returns:
+        imshow object
+    """
+    import numpy
+    import matplotlib
+    assert len(array.shape) == 2, "Wrong shape of array - did you forget to slice your 3D domain first?"
+    array = numpy.swapaxes(array, 0, 1)
+
+    if clip:
+        normalized_field = array.copy()
+        normalized_field[normalized_field < 0] = 0
+        normalized_field[normalized_field > 1] = 1
+    else:
+        minimum, maximum = numpy.min(array), numpy.max(array)
+        normalized_field = (array - minimum) / (maximum - minimum)
+
+    color = matplotlib.colors.to_rgba(color)
+    field_to_plot = numpy.empty(array.shape + (4,))
+    # set the complete array to the color
+    for i in range(3):
+        field_to_plot[:, :, i] = color[i]
+    # only the alpha channel varies using the array values
+    field_to_plot[:, :, 3] = normalized_field
+
+    res = imshow(field_to_plot, origin='lower', **kwargs)
+    axis('equal')
+    return res
+
+
+def scalar_field_contour(array, **kwargs):
+    """Small wrapper around contour to transform the coordinate system.
+
+    For details see  :func:`matplotlib.pyplot.imshow`
+    """
+    array = np.swapaxes(array, 0, 1)
+    res = contour(array, **kwargs)
+    axis('equal')
+    return res
+
+
+def multiple_scalar_fields(array, **kwargs):
+    """Plots a 3D array by slicing the last dimension and creates on plot for each entry of the last dimension.
+
+    Args:
+        array: 3D array to plot.
+        **kwargs: passed along to imshow
+    """
+    assert len(array.shape) == 3
+    sub_plots = array.shape[-1]
+    for i in range(sub_plots):
+        subplot(1, sub_plots, i + 1)
+        title(str(i))
+        scalar_field(array[..., i], **kwargs)
+        colorbar()
+
+
+def phase_plot(phase_field: np.ndarray, linewidth=1.0, clip=True) -> None:
+    """Plots a phase field array using the phase variables as alpha channel.
+
+    Args:
+        phase_field: array with len(shape) == 3, first two dimensions are spatial, the last one indexes the phase
+                     components.
+        linewidth: line width of the 0.5 contour lines that are drawn over the alpha blended phase images
+        clip: see scalar_field_alpha_value function
+    """
+    color_cycle = cycle(['#fe0002', '#00fe00', '#0000ff', '#ffa800', '#f600ff'])
+
+    assert len(phase_field.shape) == 3
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        for i in range(phase_field.shape[-1]):
+            scalar_field_alpha_value(phase_field[..., i], next(color_cycle), clip=clip, interpolation='bilinear')
+        if linewidth:
+            for i in range(phase_field.shape[-1]):
+                scalar_field_contour(phase_field[..., i], levels=[0.5], colors='k', linewidths=[linewidth])
+
+
+def sympy_function(expr, x_values=None, **kwargs):
+    """Plots the graph of a sympy term that depends on one symbol only.
+
+    Args:
+        expr: sympy term that depends on one symbol only, which is plotted on the x axis
+        x_values: describes sampling of x axis. Possible values are:
+                     * tuple of (start, stop) or (start, stop, nr_of_steps)
+                     * None, then start=0, stop=1, nr_of_steps=100
+                     * 1D numpy array with x values
+        **kwargs: passed on to :func:`matplotlib.pyplot.plot`
+
+    Returns:
+        plot object
+    """
+    import sympy as sp
+    if x_values is None:
+        x_arr = np.linspace(0, 1, 100)
+    elif type(x_values) is tuple:
+        x_arr = np.linspace(*x_values)
+    elif isinstance(x_values, np.ndarray):
+        assert len(x_values.shape) == 1
+        x_arr = x_values
+    else:
+        raise ValueError("Invalid value for parameter x_values")
+    symbols = expr.atoms(sp.Symbol)
+    assert len(symbols) == 1, "Sympy expression may only depend on one variable only. Depends on " + str(symbols)
+    y_arr = sp.lambdify(symbols.pop(), expr)(x_arr)
+    return plot(x_arr, y_arr, **kwargs)
+
+
+# ------------------------------------------- Animations ---------------------------------------------------------------
+
+def __scale_array(arr):
+    from numpy.linalg import norm
+    norm_arr = norm(arr, axis=2, ord=2)
+    if isinstance(arr, np.ma.MaskedArray):
+        norm_arr = np.ma.masked_array(norm_arr, arr.mask[..., 0])
+    return arr / norm_arr.max()
+
+
+def vector_field_animation(run_function, step=2, rescale=True, plot_setup_function=lambda *_: None,
+                           plot_update_function=lambda *_: None, interval=200, frames=180, **kwargs):
+    """Creates a matplotlib animation of a vector field using a quiver plot.
+
+    Args:
+        run_function: callable without arguments, returning a 2D vector field i.e. numpy array with len(shape)==3
+        step: see documentation of vector_field function
+        rescale: if True, the length of the arrows is rescaled in every time step
+        plot_setup_function: optional callable with the quiver object as argument,
+                             that can be used to set up the plot (title, legend,..)
+        plot_update_function: optional callable with the quiver object as argument
+                              that is called of the quiver object was updated
+        interval: delay between frames in milliseconds (see matplotlib.FuncAnimation)
+        frames: how many frames should be generated, see matplotlib.FuncAnimation
+        **kwargs: passed to quiver plot
+
+    Returns:
+        matplotlib animation object
+    """
+    import matplotlib.animation as animation
+
+    fig = gcf()
+    im = None
+    field = run_function()
+    if rescale:
+        field = __scale_array(field)
+        kwargs.setdefault('scale', 0.6)
+        kwargs.setdefault('angles', 'xy')
+        kwargs.setdefault('scale_units', 'xy')
+
+    quiver_plot = vector_field(field, step=step, **kwargs)
+    plot_setup_function(quiver_plot)
+
+    def update_figure(*_):
+        f = run_function()
+        f = np.swapaxes(f, 0, 1)
+        if rescale:
+            f = __scale_array(f)
+        u, v = f[::step, ::step, 0], f[::step, ::step, 1]
+        quiver_plot.set_UVC(u, v)
+        plot_update_function(quiver_plot)
+        return im,
+
+    return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames)
+
+
+def vector_field_magnitude_animation(run_function, plot_setup_function=lambda *_: None, rescale=False,
+                                     plot_update_function=lambda *_: None, interval=30, frames=180, **kwargs):
+    """Animation of a vector field, showing the magnitude as colormap.
+
+    For arguments, see vector_field_animation
+    """
+    import matplotlib.animation as animation
+    from numpy.linalg import norm
+
+    fig = gcf()
+    im = None
+    field = run_function()
+    if rescale:
+        field = __scale_array(field)
+    im = vector_field_magnitude(field, **kwargs)
+    plot_setup_function(im)
+
+    def update_figure(*_):
+        f = run_function()
+        if rescale:
+            f = __scale_array(f)
+        normed = norm(f, axis=2, ord=2)
+        if hasattr(f, 'mask'):
+            normed = np.ma.masked_array(normed, mask=f.mask[:, :, 0])
+        normed = np.swapaxes(normed, 0, 1)
+        im.set_array(normed)
+        plot_update_function(im)
+        return im,
+
+    return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames)
+
+
+def scalar_field_animation(run_function, plot_setup_function=lambda *_: None, rescale=True,
+                           plot_update_function=lambda *_: None, interval=30, frames=180, **kwargs):
+    """Animation of scalar field as colored image, see `scalar_field`."""
+    import matplotlib.animation as animation
+
+    fig = gcf()
+    im = None
+    field = run_function()
+    if rescale:
+        f_min, f_max = np.min(field), np.max(field)
+        field = (field - f_min) / (f_max - f_min)
+        im = scalar_field(field, vmin=0.0, vmax=1.0, **kwargs)
+    else:
+        im = scalar_field(field, **kwargs)
+    plot_setup_function(im)
+
+    def update_figure(*_):
+        f = run_function()
+        if rescale:
+            f_min, f_max = np.min(f), np.max(f)
+            f = (f - f_min) / (f_max - f_min)
+        if hasattr(f, 'mask'):
+            f = np.ma.masked_array(f, mask=f.mask[:, :])
+        f = np.swapaxes(f, 0, 1)
+        im.set_array(f)
+        plot_update_function(im)
+        return im,
+
+    return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames)
+
+
+def surface_plot_animation(run_function, frames=90, interval=30, zlim=None,  **kwargs):
+    """Animation of scalar field as 3D plot."""
+    from mpl_toolkits.mplot3d import Axes3D
+    import matplotlib.animation as animation
+    from matplotlib import cm
+    fig = gcf()
+    ax = fig.add_subplot(111, projection='3d')
+    data = run_function()
+    x, y = np.meshgrid(np.arange(data.shape[0]), np.arange(data.shape[1]), indexing='ij')
+    kwargs.setdefault('rstride', 2)
+    kwargs.setdefault('cstride', 2)
+    kwargs.setdefault('color', 'b')
+    kwargs.setdefault('cmap', cm.coolwarm)
+    ax.plot_surface(x, y, data, **kwargs)
+    if zlim is not None:
+        ax.set_zlim(*zlim)
+
+    def update_figure(*_):
+        d = run_function()
+        ax.clear()
+        plot = ax.plot_surface(x, y, d, **kwargs)
+        if zlim is not None:
+            ax.set_zlim(*zlim)
+        return plot,
+
+    return animation.FuncAnimation(fig, update_figure, interval=interval, frames=frames, blit=False)

src/pystencils/rng.py

+import copy
+import numpy as np
+import sympy as sp
+
+from pystencils.typing import TypedSymbol, CastFunc
+from pystencils.astnodes import LoopOverCoordinate
+from pystencils.backends.cbackend import CustomCodeNode
+from pystencils.sympyextensions import fast_subs
+
+
+class RNGBase(CustomCodeNode):
+
+    id = 0
+
+    def __init__(self, dim, time_step=TypedSymbol("time_step", np.uint32), offsets=None, keys=None):
+        if keys is None:
+            keys = (0,) * self._num_keys
+        if offsets is None:
+            offsets = (0,) * dim
+        if len(keys) != self._num_keys:
+            raise ValueError(f"Provided {len(keys)} keys but need {self._num_keys}")
+        if len(offsets) != dim:
+            raise ValueError(f"Provided {len(offsets)} offsets but need {dim}")
+        coordinates = [LoopOverCoordinate.get_loop_counter_symbol(i) + offsets[i] for i in range(dim)]
+        if dim < 3:
+            coordinates.append(0)
+
+        self._args = sp.sympify([time_step, *coordinates, *keys])
+        self.result_symbols = tuple(TypedSymbol(f'random_{self.id}_{i}', self._data_type)
+                                    for i in range(self._num_vars))
+        symbols_read = set.union(*[s.atoms(sp.Symbol) for s in self.args])
+        super().__init__("", symbols_read=symbols_read, symbols_defined=self.result_symbols)
+
+        self.headers = [f'"{self._name.split("_")[0]}_rand.h"']
+
+        RNGBase.id += 1
+
+    @property
+    def args(self):
+        return self._args
+
+    def fast_subs(self, subs_dict, skip):
+        rng = copy.deepcopy(self)
+        rng._args = [fast_subs(a, subs_dict, skip) for a in rng._args]
+        return rng
+
+    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(CastFunc):
+                # this vector RNG has become scalar through substitution
+                code += f"{r.dtype} {r.name};\n"
+            else:
+                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] + ['' + r.name for r in self.result_symbols]
+        code += (self._name + "(" + ", ".join(args) + ");\n")
+        return code
+
+    def __repr__(self):
+        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"
+    _data_type = np.float64
+    _num_vars = 2
+    _num_keys = 2
+
+
+class PhiloxFourFloats(RNGBase):
+    _name = "philox_float4"
+    _data_type = np.float32
+    _num_vars = 4
+    _num_keys = 2
+
+
+class AESNITwoDoubles(RNGBase):
+    _name = "aesni_double2"
+    _data_type = np.float64
+    _num_vars = 2
+    _num_keys = 4
+
+
+class AESNIFourFloats(RNGBase):
+    _name = "aesni_float4"
+    _data_type = np.float32
+    _num_vars = 4
+    _num_keys = 4
+
+
+def random_symbol(assignment_list, dim, seed=TypedSymbol("seed", np.uint32), rng_node=PhiloxTwoDoubles,
+                  time_step=TypedSymbol("time_step", np.uint32), offsets=None):
+    """Return a symbol generator for random numbers
+    
+    Args:
+        assignment_list: the subexpressions member of an AssignmentCollection, into which helper variables assignments
+                         will be inserted
+        dim: 2 or 3 for two or three spatial dimensions
+        seed: an integer or TypedSymbol(..., np.uint32) to seed the random number generator. If you create multiple
+              symbol generators, please pass them different seeds so you don't get the same stream of random numbers!
+        rng_node: which random number generator to use (PhiloxTwoDoubles, PhiloxFourFloats, AESNITwoDoubles,
+                  AESNIFourFloats).
+        time_step: TypedSymbol(..., np.uint32) that indicates the number of the current time step
+        offsets: tuple of offsets (constant integers or TypedSymbol(..., np.uint32)) that give the global coordinates
+                 of the local origin
+    """
+    counter = 0
+    while True:
+        keys = (counter, seed) + (0,) * (rng_node._num_keys - 2)
+        node = rng_node(dim, keys=keys, time_step=time_step, offsets=offsets)
+        inserted = False
+        for symbol in node.result_symbols:
+            if not inserted:
+                assignment_list.insert(0, node)
+                inserted = True
+            yield symbol
+        counter += 1

src/pystencils/runhelper/__init__.py

+from pystencils.runhelper.db import Database
+from pystencils.runhelper.parameterstudy import ParameterStudy
+
+__all__ = ['Database', 'ParameterStudy']

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:
+    """NoSQL database for storing simulation results.
+
+    Two backends are supported:
+        * `blitzdb`: simple file-based solution similar to sqlite for SQL databases, stores json files
+                     no server setup required, but slow for larger collections
+        * `mongodb`: mongodb backend via `pymongo`
+
+    A simulation result is stored as an object consisting of
+        * parameters: dict with simulation parameters
+        * results: dict with results
+        * environment: information about the machine, compiler configuration and time
+
+    Args:
+        file: database identifier, for blitzdb pass a directory name here. Database folder is created if it doesn't
+              exist yet. For larger collections use mongodb. In this case pass a pymongo connection string
+              e.g. "mongo://server:9131"
+
+    Example:
+        >>> from tempfile import TemporaryDirectory
+        >>> with TemporaryDirectory() as tmp_dir:
+        ...     db = Database(tmp_dir)  # create database in temporary folder
+        ...     params = {'method': 'finite_diff', 'dx': 1.5}  # some hypothetical simulation parameters
+        ...     db.save(params, result={'error': 1e-6})  # save simulation parameters together with hypothetical results
+        ...     assert db.was_already_simulated(params)  # search for parameters in database
+        ...     assert next(db.filter_params(params))['params'] == params # get data set, keys are 'params', 'results'
+        ...                                                               # and 'env'
+        ...     # get a pandas object with all results matching a query
+        ...     df = db.to_pandas({'dx': 1.5}, remove_prefix=True)
+        ...     # order columns alphabetically (just for doctest output)
+        ...     df.reindex(sorted(df.columns), axis=1)  # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
+                                 dx     error       method
+                pk
+                ...             1.5  0.000001  finite_diff
+    """
+
+    class SimulationResult(blitzdb.Document):
+        pass
+
+    def __init__(self, file: str, serializer_info: tuple = None) -> None:
+        if file.startswith("mongo://"):
+            from pymongo import MongoClient
+            db_name = file[len("mongo://"):]
+            c = MongoClient()
+            self.backend = blitzdb.MongoBackend(c[db_name])
+        else:
+            self.backend = blitzdb.FileBackend(file)
+
+        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.
+
+        Args:
+            params: dict of simulation parameters
+            result: dict of simulation results
+            env: optional environment - if None a default environment with compiler configuration, machine info and time
+                 is used
+            **kwargs: the final object is updated with the keyword arguments
+
+        """
+        document_dict = {
+            'params': params,
+            'result': result,
+            'env': env if env else self.get_environment(),
+        }
+        document_dict.update(kwargs)
+        document = Database.SimulationResult(document_dict, backend=self.backend)
+        document.save()
+        self.backend.commit()
+
+    def filter_params(self, parameter_query: Dict, *args, **kwargs) -> Iterator['SimulationResult']:
+        """Query using simulation parameters.
+
+        See blitzdb documentation for filter
+
+        Args:
+            parameter_query: blitzdb filter dict using only simulation parameters
+            *args: arguments passed to blitzdb filter
+            **kwargs: arguments passed to blitzdb filter
+
+        Returns:
+            generator of SimulationResult, which is a dict-like object with keys 'params', 'result' and 'env'
+        """
+        query = {'params.' + k: v for k, v in parameter_query.items()}
+        return self.filter(query, *args, **kwargs)
+
+    def filter(self, *args, **kwargs):
+        """blitzdb filter on SimulationResult, not only simulation parameters.
+
+        Can be used to filter for results or environment options.
+        The filter dictionary has to have prefixes "params." , "env." or "result."
+        """
+        return self.backend.filter(Database.SimulationResult, *args, **kwargs)
+
+    def was_already_simulated(self, parameters):
+        """Checks if there is at least one simulation result matching the passed parameters."""
+        return len(self.filter({'params': parameters})) > 0
+
+    # Columns with these prefixes are not included in pandas result
+    pandas_columns_to_ignore = ['changedParams.', 'env.']
+
+    def to_pandas(self, parameter_query, remove_prefix=True, drop_constant_columns=False):
+        """Queries for simulations with given parameters and returns them in a pandas data frame.
+
+        Args:
+            parameter_query: see filter method
+            remove_prefix: if True the name of the pandas columns are not prefixed with "params." or "results."
+            drop_constant_columns: if True, all columns are dropped that have the same value is all rows
+
+        Returns:
+            pandas data frame
+        """
+        from pandas import json_normalize
+
+        query_result = self.filter_params(parameter_query)
+        attributes = [e.attributes for e in query_result]
+        if not attributes:
+            return
+        df = json_normalize(attributes)
+        df.set_index('pk', inplace=True)
+
+        if self.pandas_columns_to_ignore:
+            remove_columns_by_prefix(df, self.pandas_columns_to_ignore, inplace=True)
+        if remove_prefix:
+            remove_prefix_in_column_name(df, inplace=True)
+        if drop_constant_columns:
+            df, _ = remove_constant_columns(df)
+
+        return df
+
+    @staticmethod
+    def get_environment():
+        result = {
+            'timestamp': time.mktime(time.gmtime()),
+            'hostname': socket.gethostname(),
+            'cpuCompilerConfig': get_compiler_config(),
+        }
+        try:
+            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 InvalidGitRepositoryError:
+            pass
+
+        return result
+
+# ----------------------------------------- Helper Functions -----------------------------------------------------------
+
+
+def remove_constant_columns(df):
+    """Removes all columns of a pandas data frame that have the same value in all rows."""
+    import pandas as pd
+    remaining_df = df.loc[:, df.apply(pd.Series.nunique) > 1]
+    constants = df.loc[:, df.apply(pd.Series.nunique) <= 1].iloc[0]
+    return remaining_df, constants
+
+
+def remove_columns_by_prefix(df, prefixes: Sequence[str], inplace: bool = False):
+    """Remove all columns from a pandas data frame whose name starts with one of the given prefixes."""
+    if not inplace:
+        df = df.copy()
+
+    for column_name in df.columns:
+        for prefix in prefixes:
+            if column_name.startswith(prefix):
+                del df[column_name]
+    return df
+
+
+def remove_prefix_in_column_name(df, inplace: bool = False):
+    """Removes dotted prefixes from pandas column names.
+
+    A column named 'result.finite_diff.dx' is renamed to 'finite_diff.dx', everything before the first dot is removed.
+    If the column name does not contain a dot, the column name is not changed.
+    """
+    if not inplace:
+        df = df.copy()
+
+    new_column_names = []
+    for column_name in df.columns:
+        if '.' in column_name:
+            new_column_names.append(column_name[column_name.index('.') + 1:])
+        else:
+            new_column_names.append(column_name)
+    df.columns = new_column_names
+    return df

src/pystencils/runhelper/parameterstudy.py

+import datetime
+import itertools
+import json
+import os
+import socket
+from collections import namedtuple
+from copy import deepcopy
+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]
+WeightFunction = Callable[[Dict], int]
+FilterFunction = Callable[[ParameterDict], Optional[ParameterDict]]
+
+
+class ParameterStudy:
+    """Manages and runs multiple configurations locally or distributed and stores results in NoSQL database.
+
+    To run a parameter study, define a run function that takes all parameters as keyword arguments and returns the
+    results as a (possibly nested) dictionary. Then, define the parameter sets that this function should be run with.
+
+    Examples:
+          >>> import tempfile
+          >>>
+          >>> def dummy_run_function(p1, p2, p3, p4):
+          ...     print("Run called with", p1, p2, p3, p4)
+          ...     return { 'result1': p1 * p2, 'result2': p3 + p4 }
+          >>>
+          >>> with tempfile.TemporaryDirectory() as tmp_dir:
+          ...     ps = ParameterStudy(dummy_run_function, database_connector=tmp_dir)
+          ...     ps.add_run({'p1': 5, 'p2': 42, 'p3': 'abc', 'p4': 'def'})
+          ...     ps.add_combinations( [('p1', [1, 2]),
+          ...                           ('p3', ['x', 'y'])], constant_parameters={'p2': 5, 'p4': 'z' })
+          ...     ps.run()
+          ...     ps.run_scenarios_not_in_database()
+          ...     ps.run_from_command_line(argv=['local'])  # alternative to run - exposes a command line interface if
+          ...                                               # no argv is passed. Does not run anything here, because
+          ...                                               # configuration already in database are skipped
+          Run called with 2 5 y z
+          Run called with 2 5 x z
+          Run called with 1 5 y z
+          Run called with 1 5 x z
+          Run called with 5 42 abc def
+
+    Above example runs all parameter combinations locally and stores the returned result in the NoSQL database.
+    It is also possible to distribute the runs to multiple processes, by starting a server on one machine and multiple
+    executing runners on other machines. The server distributes configurations to the runners, collects their results
+    to stores the results in the database.
+    """
+
+    Run = namedtuple("Run", ['parameter_dict', 'weight'])
+
+    def __init__(self, run_function: Callable[..., Dict], runs: Sequence = (),
+                 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, serializer_info)
+
+    def add_run(self, parameter_dict: ParameterDict, weight: int = 1) -> None:
+        """Schedule a dictionary of parameters to run in this parameter study.
+
+        Args:
+            parameter_dict: used as keyword arguments to the run function.
+            weight: weight of the run configuration which should be  proportional to runtime of this case,
+                    used for progress display and distribution to processes.
+        """
+        self.runs.append(self.Run(parameter_dict, weight))
+
+    def add_combinations(self, degrees_of_freedom: Sequence[Tuple[str, Sequence[Any]]],
+                         constant_parameters: Optional[ParameterDict] = None,
+                         filter_function: Optional[FilterFunction] = None,
+                         runtime_weight_function: Optional[WeightFunction] = None) -> None:
+        """Add all possible combinations of given parameters as runs.
+
+        This is a convenience function to simulate all possible parameter combinations of a scenario.
+        Configurations can be filtered and weighted by passing filter- and weighting functions.
+
+        Args:
+            degrees_of_freedom: defines for each parameter the possible values it can take on
+            constant_parameters: parameter dict, for parameters that should not be changed
+            filter_function: optional function that receives a parameter dict and returns the potentially modified dict
+                             or None if this combination should not be added.
+            runtime_weight_function: function mapping a parameter dict to the runtime weight (see weight at add_runs)
+
+        Examples:
+             degrees_of_freedom = [('p1', [1,2]),
+                                   ('p2', ['a', 'b'])]
+             is equivalent to calling add_run four times, with all possible parameter combinations.
+        """
+        parameter_names = [e[0] for e in degrees_of_freedom]
+        parameter_values = [e[1] for e in degrees_of_freedom]
+
+        default_params_dict = {} if constant_parameters is None else constant_parameters
+        for value_tuple in itertools.product(*parameter_values):
+            params_dict = deepcopy(default_params_dict)
+            params_dict.update({name: value for name, value in zip(parameter_names, value_tuple)})
+            params = DotDict(params_dict)
+            if filter_function:
+                params = filter_function(params)
+                if params is None:
+                    continue
+            weight = 1 if not runtime_weight_function else runtime_weight_function(params)
+            self.add_run(params, weight)
+
+    def run(self, process: int = 0, num_processes: int = 1, parameter_update: Optional[ParameterDict] = None) -> None:
+        """Runs all added configurations.
+
+        Args:
+            process: configurations are split into num_processes chunks according to weights and only the
+                     process'th chunk is run. To run all, use process=0 and num_processes=1
+            num_processes: see above
+            parameter_update: Extend/override all configurations with this dictionary.
+        """
+        parameter_update = {} if parameter_update is None else parameter_update
+        own_runs = self._distribute_runs(self.runs, process, num_processes)
+        for run in own_runs:
+            parameter_dict = run.parameter_dict.copy()
+            parameter_dict.update(parameter_update)
+            result = self.run_function(**parameter_dict)
+
+            self.db.save(run.parameter_dict, result, None, changed_params=parameter_update)
+
+    def run_scenarios_not_in_database(self, parameter_update: Optional[ParameterDict] = None) -> None:
+        """Same as run method, but runs only configuration for which no result is in the database yet."""
+        parameter_update = {} if parameter_update is None else parameter_update
+        filtered_runs = self._filter_already_simulated(self.runs)
+        for run in filtered_runs:
+            parameter_dict = run.parameter_dict.copy()
+            parameter_dict.update(parameter_update)
+            result = self.run_function(**parameter_dict)
+
+            self.db.save(run.parameter_dict, result, changed_params=parameter_update)
+
+    def run_server(self, ip: str = "0.0.0.0", port: int = 8342):
+        """Runs server to supply runner clients with scenarios to simulate and collect results from them.
+        Skips scenarios that are already in the database."""
+        from http.server import BaseHTTPRequestHandler, HTTPServer
+        filtered_runs = self._filter_already_simulated(self.runs)
+
+        if not filtered_runs:
+            print("No Scenarios to simulate")
+            return
+
+        class ParameterStudyServer(BaseHTTPRequestHandler):
+            parameterStudy = self
+            all_runs = filtered_runs
+            runs = filtered_runs.copy()
+            currently_running = {}
+            finished_runs = []
+
+            def next_scenario(self, received_json_data):
+                client_name = received_json_data['client_name']
+                if len(self.runs) > 0:
+                    run_status = "%d/%d" % (len(self.finished_runs), len(self.all_runs))
+                    work_status = "%d/%d" % (sum(r.weight for r in self.finished_runs),
+                                             sum(r.weight for r in self.all_runs))
+                    format_args = {
+                        'remaining': len(self.runs),
+                        'time': datetime.datetime.now().strftime("%H:%M:%S"),
+                        'client_name': client_name,
+                        'run_status': run_status,
+                        'work_status': work_status,
+                    }
+
+                    scenario = self.runs.pop(0)
+                    print(" {time} {client_name} fetched scenario. Scenarios: {run_status}, Work: {work_status}"
+                          .format(**format_args))
+                    self.currently_running[client_name] = scenario
+                    return {'status': 'ok', 'params': scenario.parameter_dict}
+                else:
+                    return {'status': 'finished'}
+
+            def result(self, received_json_data):
+                client_name = received_json_data['client_name']
+                run = self.currently_running[client_name]
+                self.finished_runs.append(run)
+                del self.currently_running[client_name]
+                d = received_json_data
+
+                def hash_dict(dictionary):
+                    import hashlib
+                    return hashlib.sha1(json.dumps(dictionary, sort_keys=True).encode()).hexdigest()
+
+                assert hash_dict(d['params']) == hash_dict(run.parameter_dict), \
+                    str(d['params']) + "is not equal to " + str(run.parameter_dict)
+                self.parameterStudy.db.save(run.parameter_dict,
+                                            result=d['result'], env=d['env'], changed_params=d['changed_params'])
+                return {}
+
+            # noinspection PyPep8Naming
+            def do_POST(self) -> None:
+                mapping = {'/next_scenario': self.next_scenario,
+                           '/result': self.result}
+                if self.path in mapping.keys():
+                    data = self._read_contents()
+                    self.send_response(200)
+                    self.send_header("Content-type", "application/json")
+                    self.end_headers()
+                    json_data = json.loads(data)
+                    response = mapping[self.path](json_data)
+                    self.wfile.write(json.dumps(response).encode())
+                else:
+                    self.send_response(400)
+
+            # noinspection PyPep8Naming
+            def do_GET(self):
+                return self.do_POST()
+
+            def _read_contents(self):
+                return self.rfile.read(int(self.headers['Content-Length'])).decode()
+
+            def log_message(self, fmt, *args):
+                return
+
+        print(f"Listening to connections on {ip}:{port}. Scenarios to simulate: {len(filtered_runs)}")
+        server = HTTPServer((ip, port), ParameterStudyServer)
+        while len(ParameterStudyServer.currently_running) > 0 or len(ParameterStudyServer.runs) > 0:
+            server.handle_request()
+        server.handle_request()
+
+    def run_client(self, client_name: str = "{hostname}_{pid}", server: str = 'localhost', port: int = 8342,
+                   parameter_update: Optional[ParameterDict] = None, max_time=None) -> None:
+        """Start runner client that retrieves configuration from server, runs it and reports results back to server.
+
+        Args:
+            client_name: name of the client. Has to be unique for each client.
+                         Placeholders {hostname} and {pid} can be used to generate unique name.
+            server: url to server
+            port: port as specified in run_server
+            parameter_update: Used to override/extend parameters received from the server.
+                              Typical use cases is to set optimization or GPU parameters for some clients to make
+                              some clients simulate on CPU, others on GPU
+            max_time: maximum runtime in seconds: the client runs scenario after scenario, but starts only a new
+                      scenario if not more than max_time seconds have passed since this function was called.
+                      So the time given here should be the total maximum runtime minus a typical runtime for one setup
+        """
+        from urllib.request import urlopen
+        from urllib.error import URLError
+        import time
+        parameter_update = {} if parameter_update is None else parameter_update
+        url = f"http://{server}:{port}"
+        client_name = client_name.format(hostname=socket.gethostname(), pid=os.getpid())
+        start_time = time.time()
+        while True:
+            try:
+                if max_time is not None and (time.time() - start_time) > max_time:
+                    print("Stopping client - maximum time reached")
+                    break
+                http_response = urlopen(url + "/next_scenario",
+                                        data=json.dumps({'client_name': client_name}).encode())
+                scenario = json.loads(http_response.read().decode())
+                if scenario['status'] != 'ok':
+                    break
+                original_params = scenario['params'].copy()
+                scenario['params'].update(parameter_update)
+                result = self.run_function(**scenario['params'])
+
+                answer = {'params': original_params,
+                          'changed_params': parameter_update,
+                          'result': result,
+                          'env': Database.get_environment(),
+                          'client_name': client_name}
+                urlopen(url + '/result', data=json.dumps(answer).encode())
+            except URLError:
+                print(f"Cannot connect to server {url}  retrying in 5 seconds...")
+                sleep(5)
+
+    def run_from_command_line(self, argv: Optional[Sequence[str]] = None) -> None:
+        """Exposes interface to command line with possibility to run directly or distributed via server/client."""
+        from argparse import ArgumentParser
+
+        def server(a):
+            if a.database:
+                self.db = Database(a.database)
+            self.run_server(a.host, a.port)
+
+        def client(a):
+            self.run_client(a.client_name, a.host, a.port, json.loads(a.parameter_override), a.max_time)
+
+        def local(a):
+            if a.database:
+                self.db = Database(a.database)
+            self.run_scenarios_not_in_database(json.loads(a.parameter_override))
+
+        parser = ArgumentParser()
+        subparsers = parser.add_subparsers()
+
+        local_parser = subparsers.add_parser('local', aliases=['l'],
+                                             help="Run scenarios locally which are not yet in database", )
+        local_parser.add_argument("-d", "--database", type=str, default="")
+        local_parser.add_argument("-P", "--parameter_override", type=str, default="{}",
+                                  help="JSON: the parameter dictionary is updated with these parameters. Use this to "
+                                       "set host specific options like GPU call parameters. Enclose in \" ")
+        local_parser.set_defaults(func=local)
+
+        server_parser = subparsers.add_parser('server', aliases=['s'],
+                                              help="Runs server to distribute different scenarios to workers", )
+        server_parser.add_argument("-p", "--port", type=int, default=8342, help="Port to listen on")
+        server_parser.add_argument("-H", "--host", type=str, default="0.0.0.0", help="IP/Hostname to listen on")
+        server_parser.add_argument("-d", "--database", type=str, default="")
+        server_parser.set_defaults(func=server)
+
+        client_parser = subparsers.add_parser('client', aliases=['c'],
+                                              help="Runs a worker client connection to scenario distribution server")
+        client_parser.add_argument("-p", "--port", type=int, default=8342, help="Port to connect to")
+        client_parser.add_argument("-H", "--host", type=str, default="localhost", help="Host or IP to connect to")
+        client_parser.add_argument("-n", "--client_name", type=str, default="{hostname}_{pid}",
+                                   help="Unique client name, you can use {hostname} and {pid} as placeholder")
+        client_parser.add_argument("-P", "--parameter_override", type=str, default="{}",
+                                   help="JSON: the parameter dictionary is updated with these parameters. Use this to "
+                                        "set host specific options like GPU call parameters. Enclose in \" ")
+        client_parser.add_argument("-t", "--max_time", type=int, default=None,
+                                   help="If more than this time in seconds has passed, "
+                                        "the client stops running scenarios.")
+        client_parser.set_defaults(func=client)
+
+        args = parser.parse_args(argv)
+        if not len(vars(args)):
+            parser.print_help()
+        else:
+            args.func(args)
+
+    def _filter_already_simulated(self, all_runs):
+        """Removes all runs from the given list, that are already in the database"""
+        already_simulated = {json.dumps(e.params) for e in self.db.filter({})}
+        return [r for r in all_runs if json.dumps(r.parameter_dict) not in already_simulated]
+
+    @staticmethod
+    def _distribute_runs(all_runs, process, num_processes):
+        """Partitions runs by their weights into num_processes chunks and returns the process's chunk."""
+        sorted_runs = sorted(all_runs, key=lambda e: e.weight, reverse=True)
+        result = sorted_runs[process::num_processes]
+        result.reverse()  # start with faster scenarios
+        return result

src/pystencils/session.py

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

src/pystencils/simp/__init__.py

+from .assignment_collection import AssignmentCollection
+from .simplifications import (
+    add_subexpressions_for_constants,
+    add_subexpressions_for_divisions, add_subexpressions_for_field_reads,
+    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',
+           'insert_aliases', 'insert_zeros', 'insert_constants',
+           'insert_constant_additions', 'insert_constant_multiples',
+           'insert_squares', 'insert_symbol_times_minus_one']

src/pystencils/simp/assignment_collection.py

+import itertools
+from copy import copy
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Set, Union
+
+import sympy as sp
+
+import pystencils
+from pystencils.assignment import Assignment
+from pystencils.simp.simplifications import (sort_assignments_topologically, transform_lhs_and_rhs, transform_rhs)
+from pystencils.sympyextensions import count_operations, fast_subs
+
+
+class AssignmentCollection:
+    """
+    A collection of equations with subexpression definitions, also represented as assignments,
+    that are used in the main equations. AssignmentCollection can be passed to simplification methods.
+    These simplification methods can change the subexpressions, but the number and
+    left hand side of the main equations themselves is not altered.
+    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.
+
+    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
+                                        used to get new symbols that are unique for this 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]] = 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()]
+        if isinstance(subexpressions, Dict):
+            subexpressions = [Assignment(k, v)
+                              for k, v in subexpressions.items()]
+
+        main_assignments = list(itertools.chain.from_iterable(
+            [(a if isinstance(a, Iterable) else [a]) for a in main_assignments]))
+        subexpressions = list(itertools.chain.from_iterable(
+            [(a if isinstance(a, Iterable) else [a]) for a in subexpressions]))
+
+        self.main_assignments = main_assignments
+        self.subexpressions = subexpressions
+
+        if simplification_hints is None:
+            simplification_hints = {}
+
+        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(ctr=max_ctr)
+        else:
+            self.subexpression_symbol_generator = subexpression_symbol_generator
+
+    def add_simplification_hint(self, key: str, value: Any) -> None:
+        """Adds an entry to the simplification_hints dictionary and checks that is does not exist yet."""
+        assert key not in self.simplification_hints, "This hint already exists"
+        self.simplification_hints[key] = value
+
+    def add_subexpression(self, rhs: sp.Expr, lhs: Optional[sp.Symbol] = None, topological_sort=True) -> sp.Symbol:
+        """Adds a subexpression to current collection.
+
+        Args:
+            rhs: right hand side of new subexpression
+            lhs: optional left hand side of new subexpression. If None a new unique symbol is generated.
+            topological_sort: sort the subexpressions topologically after insertion, to make sure that
+                              definition of a symbol comes before its usage. If False, subexpression is appended.
+
+        Returns:
+            left hand side symbol (which could have been generated)
+        """
+        if lhs is None:
+            lhs = next(self.subexpression_symbol_generator)
+        eq = Assignment(lhs, rhs)
+        self.subexpressions.append(eq)
+        if topological_sort:
+            self.topological_sort(sort_subexpressions=True,
+                                  sort_main_assignments=False)
+        return lhs
+
+    def topological_sort(self, sort_subexpressions: bool = True, sort_main_assignments: bool = True) -> None:
+        """Sorts subexpressions and/or main_equations topologically to make sure symbol usage comes after definition."""
+        if sort_subexpressions:
+            self.subexpressions = sort_assignments_topologically(self.subexpressions)
+        if sort_main_assignments:
+            self.main_assignments = sort_assignments_topologically(self.main_assignments)
+
+    # ---------------------------------------------- Properties  -------------------------------------------------------
+
+    @property
+    def all_assignments(self) -> List[Assignment]:
+        """Subexpression and main equations as a single list."""
+        return self.subexpressions + self.main_assignments
+
+    @property
+    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):
+                rhs_symbols.update(eq.rhs.atoms(sp.Symbol))
+            elif isinstance(eq, pystencils.astnodes.Node):
+                rhs_symbols.update(eq.undefined_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]:
+        """All symbols which occur on the left hand side of a main assignment or a subexpression."""
+        bound_symbols_set = set(
+            [assignment.lhs for assignment in self.all_assignments if isinstance(assignment, Assignment)]
+        )
+
+        assert len(bound_symbols_set) == len(list(a for a in self.all_assignments if isinstance(a, Assignment))), \
+            "Not in SSA form - same symbol assigned multiple times"
+
+        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."""
+        return {s.field for s in self.free_symbols if hasattr(s, 'field')}
+
+    @property
+    def bound_fields(self):
+        """All field accessed on the left hand side of a main assignment or a subexpression."""
+        return {s.field for s in self.bound_symbols if hasattr(s, 'field')}
+
+    @property
+    def defined_symbols(self) -> Set[sp.Symbol]:
+        """All symbols which occur as left-hand-sides of one of the main equations"""
+        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):
+        """See :func:`count_operations` """
+        return count_operations(self.all_assignments, only_type=None)
+
+    def atoms(self, *args):
+        return set().union(*[a.atoms(*args) for a in self.all_assignments])
+
+    def dependent_symbols(self, symbols: Iterable[sp.Symbol]) -> Set[sp.Symbol]:
+        """Returns all symbols that depend on one of the passed symbols.
+
+        A symbol 'a' depends on a symbol 'b', if there is an assignment 'a <- some_expression(b)' i.e. when
+        'b' is required to compute 'a'.
+        """
+
+        queue = list(symbols)
+
+        def add_symbols_from_expr(expr):
+            dependent_symbols = expr.atoms(sp.Symbol)
+            for ds in dependent_symbols:
+                queue.append(ds)
+
+        handled_symbols = set()
+        assignment_dict = {e.lhs: e.rhs for e in self.all_assignments}
+
+        while len(queue) > 0:
+            e = queue.pop(0)
+            if e in handled_symbols:
+                continue
+            if e in assignment_dict:
+                add_symbols_from_expr(assignment_dict[e])
+            handled_symbols.add(e)
+
+        return handled_symbols
+
+    def lambdify(self, symbols: Sequence[sp.Symbol], fixed_symbols: Optional[Dict[sp.Symbol, Any]] = None, module=None):
+        """Returns a python function to evaluate this equation collection.
+
+        Args:
+            symbols: symbol(s) which are the parameter for the created function
+            fixed_symbols: dictionary with substitutions, that are applied before sympy's lambdify
+            module: same as sympy.lambdify parameter. Defines which module to use e.g. 'numpy'
+
+        Examples:
+              >>> a, b, c, d = sp.symbols("a b c d")
+              >>> ac = AssignmentCollection([Assignment(c, a + b), Assignment(d, a**2 + b)],
+              ...                           subexpressions=[Assignment(b, a + b / 2)])
+              >>> python_function = ac.lambdify([a], fixed_symbols={b: 2})
+              >>> python_function(4)
+              {c: 6, d: 18}
+        """
+        assignments = self.new_with_substitutions(fixed_symbols, substitute_on_lhs=False) if fixed_symbols else self
+        assignments = assignments.new_without_subexpressions().main_assignments
+        lambdas = {assignment.lhs: sp.lambdify(symbols, assignment.rhs, module) for assignment in assignments}
+
+        def f(*args, **kwargs):
+            return {s: func(*args, **kwargs) for s, func in lambdas.items()}
+
+        return f
+
+    # ---------------------------- Creating new modified collections ---------------------------------------------------
+
+    def copy(self,
+             main_assignments: Optional[List[Assignment]] = None,
+             subexpressions: Optional[List[Assignment]] = None) -> 'AssignmentCollection':
+        """Returns a copy with optionally replaced main_assignments and/or subexpressions."""
+
+        res = copy(self)
+        res.simplification_hints = self.simplification_hints.copy()
+        res.subexpression_symbol_generator = copy(self.subexpression_symbol_generator)
+
+        if main_assignments is not None:
+            res.main_assignments = main_assignments
+        else:
+            res.main_assignments = self.main_assignments.copy()
+
+        if subexpressions is not None:
+            res.subexpressions = subexpressions
+        else:
+            res.subexpressions = self.subexpressions.copy()
+
+        return res
+
+    def new_with_substitutions(self, substitutions: Dict, add_substitutions_as_subexpressions: bool = False,
+                               substitute_on_lhs: bool = True,
+                               sort_topologically: bool = True) -> 'AssignmentCollection':
+        """Returns new object, where terms are substituted according to the passed substitution dict.
+
+        Args:
+            substitutions: dict that is passed to sympy subs, substitutions are done main assignments and subexpressions
+            add_substitutions_as_subexpressions: if True, the substitutions are added as assignments to subexpressions
+            substitute_on_lhs: if False, the substitutions are done only on the right hand side of assignments
+            sort_topologically: if subexpressions are added as substitutions and this parameters is true,
+                                the subexpressions are sorted topologically after insertion
+        Returns:
+            New AssignmentCollection where substitutions have been applied, self is not altered.
+        """
+        transform = transform_lhs_and_rhs if substitute_on_lhs else transform_rhs
+        transformed_subexpressions = transform(self.subexpressions, fast_subs, substitutions)
+        transformed_assignments = transform(self.main_assignments, fast_subs, substitutions)
+
+        if add_substitutions_as_subexpressions:
+            transformed_subexpressions = [Assignment(b, a) for a, b in
+                                          substitutions.items()] + transformed_subexpressions
+            if sort_topologically:
+                transformed_subexpressions = sort_assignments_topologically(transformed_subexpressions)
+        return self.copy(transformed_assignments, transformed_subexpressions)
+
+    def new_merged(self, other: 'AssignmentCollection') -> 'AssignmentCollection':
+        """Returns a new collection which contains self and other. Subexpressions are renamed if they clash."""
+        own_definitions = set([e.lhs for e in self.main_assignments])
+        other_definitions = set([e.lhs for e in other.main_assignments])
+        assert len(own_definitions.intersection(other_definitions)) == 0, \
+            "Cannot merge collections, since both define the same symbols"
+
+        own_subexpression_symbols = {e.lhs: e.rhs for e in self.subexpressions}
+        substitution_dict = {}
+
+        processed_other_subexpression_equations = []
+        for other_subexpression_eq in other.subexpressions:
+            if other_subexpression_eq.lhs in own_subexpression_symbols:
+                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, new_rhs)
+                    processed_other_subexpression_equations.append(new_eq)
+                    substitution_dict[other_subexpression_eq.lhs] = new_lhs
+            else:
+                processed_other_subexpression_equations.append(fast_subs(other_subexpression_eq, substitution_dict))
+
+        processed_other_main_assignments = [fast_subs(eq, substitution_dict) for eq in other.main_assignments]
+        return self.copy(self.main_assignments + processed_other_main_assignments,
+                         self.subexpressions + processed_other_subexpression_equations)
+
+    def new_filtered(self, symbols_to_extract: Iterable[sp.Symbol]) -> 'AssignmentCollection':
+        """Extracts equations that have symbols_to_extract as left hand side, together with necessary subexpressions.
+
+        Returns:
+            new AssignmentCollection, self is not altered
+        """
+        symbols_to_extract = set(symbols_to_extract)
+        dependent_symbols = self.dependent_symbols(symbols_to_extract)
+        new_assignments = []
+        for eq in self.all_assignments:
+            if eq.lhs in symbols_to_extract:
+                new_assignments.append(eq)
+
+        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)
+
+    def new_without_unused_subexpressions(self) -> 'AssignmentCollection':
+        """Returns new collection that only contains subexpressions required to compute the main assignments."""
+        all_lhs = [eq.lhs for eq in self.main_assignments]
+        return self.new_filtered(all_lhs)
+
+    def new_with_inserted_subexpression(self, symbol: sp.Symbol) -> 'AssignmentCollection':
+        """Eliminates the subexpression with the given symbol on its left hand side, by substituting it everywhere."""
+        new_subexpressions = []
+        subs_dict = None
+        for se in self.subexpressions:
+            if se.lhs == symbol:
+                subs_dict = {se.lhs: se.rhs}
+            else:
+                new_subexpressions.append(se)
+        if subs_dict is None:
+            return self
+
+        new_subexpressions = [Assignment(eq.lhs, fast_subs(eq.rhs, subs_dict)) for eq in new_subexpressions]
+        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=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()
+
+        subexpressions_to_keep = set(subexpressions_to_keep)
+
+        kept_subexpressions = []
+        if self.subexpressions[0].lhs in subexpressions_to_keep:
+            substitution_dict = {}
+            kept_subexpressions.append(self.subexpressions[0])
+        else:
+            substitution_dict = {self.subexpressions[0].lhs: self.subexpressions[0].rhs}
+
+        subexpression = [e for e in self.subexpressions]
+        for i in range(1, len(subexpression)):
+            subexpression[i] = fast_subs(subexpression[i], substitution_dict)
+            if subexpression[i].lhs in subexpressions_to_keep:
+                kept_subexpressions.append(subexpression[i])
+            else:
+                substitution_dict[subexpression[i].lhs] = subexpression[i].rhs
+
+        new_assignment = [fast_subs(eq, substitution_dict) for eq in self.main_assignments]
+        return self.copy(new_assignment, kept_subexpressions)
+
+    # ----------------------------------------- Display and Printing   -------------------------------------------------
+
+    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%; ">'
+            line = '<tr {nb}> <td {nb}>$${eq}$$</td>  </tr> '
+            for eq in equations:
+                format_dict = {'eq': sp.latex(eq),
+                               'nb': no_border, }
+                html_table += line.format(**format_dict)
+            html_table += "</table>"
+            return html_table
+
+        result = ""
+        if len(self.subexpressions) > 0:
+            result += "<div>Subexpressions:</div>"
+            result += make_html_equation_table(self.subexpressions)
+        result += "<div>Main Assignments:</div>"
+        result += make_html_equation_table(self.main_assignments)
+        return result
+
+    def __repr__(self):
+        return f"AssignmentCollection: {str(tuple(self.defined_symbols))[1:-1]} <- f{tuple(self.free_symbols)}"
+
+    def __str__(self):
+        result = "Subexpressions:\n"
+        for eq in self.subexpressions:
+            result += f"\t{eq}\n"
+        result += "Main Assignments:\n"
+        for eq in self.main_assignments:
+            result += f"\t{eq}\n"
+        return result
+
+    def __iter__(self):
+        return self.all_assignments.__iter__()
+
+    @property
+    def main_assignments_dict(self):
+        return {a.lhs: a.rhs for a in self.main_assignments}
+
+    @property
+    def subexpressions_dict(self):
+        return {a.lhs: a.rhs for a in self.subexpressions}
+
+    def set_main_assignments_from_dict(self, main_assignments_dict):
+        self.main_assignments = [Assignment(k, v)
+                                 for k, v in main_assignments_dict.items()]
+
+    def set_sub_expressions_from_dict(self, sub_expressions_dict):
+        self.subexpressions = [Assignment(k, v)
+                               for k, v in sub_expressions_dict.items()]
+
+    def find(self, *args, **kwargs):
+        return set.union(
+            *[a.find(*args, **kwargs) for a in self.all_assignments]
+        )
+
+    def match(self, *args, **kwargs):
+        rtn = {}
+        for a in self.all_assignments:
+            partial_result = a.match(*args, **kwargs)
+            if partial_result:
+                rtn.update(partial_result)
+        return rtn
+
+    def subs(self, *args, **kwargs):
+        return AssignmentCollection(
+            main_assignments=[a.subs(*args, **kwargs) for a in self.main_assignments],
+            subexpressions=[a.subs(*args, **kwargs) for a in self.subexpressions]
+        )
+
+    def replace(self, *args, **kwargs):
+        return AssignmentCollection(
+            main_assignments=[a.replace(*args, **kwargs) for a in self.main_assignments],
+            subexpressions=[a.replace(*args, **kwargs) for a in self.subexpressions]
+        )
+
+    def __eq__(self, other):
+        return set(self.all_assignments) == set(other.all_assignments)
+
+    def __bool__(self):
+        return bool(self.all_assignments)
+
+
+class SymbolGen:
+    """Default symbol generator producing number symbols ζ_0, ζ_1, ..."""
+
+    def __init__(self, symbol="xi", dtype=None, ctr=0):
+        self._ctr = ctr
+        self._symbol = symbol
+        self._dtype = dtype
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        name = f"{self._symbol}_{self._ctr}"
+        self._ctr += 1
+        if self._dtype is not None:
+            return pystencils.TypedSymbol(name, self._dtype)
+        return sp.Symbol(name)

src/pystencils/simp/simplifications.py

+from itertools import chain
+from typing import Callable, List, Sequence, Union
+from collections import defaultdict
+
+import sympy as sp
+
+from pystencils.assignment import Assignment
+from pystencils.astnodes import Node
+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]]:
+    """Sorts assignments in topological order, such that symbols used on rhs occur first on a lhs"""
+    edges = []
+    for c1, e1 in enumerate(assignments):
+        if hasattr(e1, 'lhs') and hasattr(e1, 'rhs'):
+            symbols = [e1.lhs]
+        elif isinstance(e1, Node):
+            symbols = e1.symbols_defined
+        else:
+            raise NotImplementedError(f"Cannot sort topologically. Object of type {type(e1)} cannot be handled.")
+
+        for lhs in symbols:
+            for c2, e2 in enumerate(assignments):
+                if isinstance(e2, Assignment) and lhs in e2.rhs.free_symbols:
+                    edges.append((c1, c2))
+                elif isinstance(e2, Node) and lhs in e2.undefined_symbols:
+                    edges.append((c1, c2))
+    return [assignments[i] for i in sp.topological_sort((range(len(assignments)), edges))]
+
+
+def sympy_cse(ac, **kwargs):
+    """Searches for common subexpressions inside the assignment collection.
+
+    Searches is done in both the existing subexpressions as well as the assignments themselves.
+    It uses the sympy subexpression detection to do this. Return a new assignment collection
+    with the additional subexpressions found
+    """
+    symbol_gen = ac.subexpression_symbol_generator
+
+    all_assignments = [e for e in chain(ac.subexpressions, ac.main_assignments) if isinstance(e, Assignment)]
+    other_objects = [e for e in chain(ac.subexpressions, ac.main_assignments) if not isinstance(e, Assignment)]
+    replacements, new_eq = sp.cse(all_assignments, symbols=symbol_gen, **kwargs)
+
+    replacement_eqs = [Assignment(*r) for r in replacements]
+
+    modified_subexpressions = new_eq[:len(ac.subexpressions)]
+    modified_update_equations = new_eq[len(ac.subexpressions):]
+
+    new_subexpressions = sort_assignments_topologically(other_objects + replacement_eqs + modified_subexpressions)
+    return ac.copy(modified_update_equations, new_subexpressions)
+
+
+def sympy_cse_on_assignment_list(assignments: List[Assignment]) -> List[Assignment]:
+    """Extracts common subexpressions from a list of assignments."""
+    from pystencils.simp.assignment_collection import AssignmentCollection
+    ec = AssignmentCollection([], assignments)
+    return sympy_cse(ec).all_assignments
+
+
+def subexpression_substitution_in_existing_subexpressions(ac):
+    """Goes through the subexpressions list and replaces the term in the following subexpressions."""
+    result = []
+    for outer_ctr, s in enumerate(ac.subexpressions):
+        new_rhs = s.rhs
+        for inner_ctr in range(outer_ctr):
+            sub_expr = ac.subexpressions[inner_ctr]
+            new_rhs = subs_additive(new_rhs, sub_expr.lhs, sub_expr.rhs, required_match_replacement=1.0)
+            new_rhs = new_rhs.subs(sub_expr.rhs, sub_expr.lhs)
+        result.append(Assignment(s.lhs, new_rhs))
+
+    return ac.copy(ac.main_assignments, result)
+
+
+def subexpression_substitution_in_main_assignments(ac):
+    """Replaces already existing subexpressions in the equations of the assignment_collection."""
+    result = []
+    for s in ac.main_assignments:
+        new_rhs = s.rhs
+        for sub_expr in ac.subexpressions:
+            new_rhs = subs_additive(new_rhs, sub_expr.lhs, sub_expr.rhs, required_match_replacement=1.0)
+        result.append(Assignment(s.lhs, new_rhs))
+    return ac.copy(result)
+
+
+def add_subexpressions_for_constants(ac):
+    """Extracts constant factors to subexpressions in the given assignment collection.
+
+    SymPy will exclude common factors from a sum only if they are symbols. This simplification
+    can be applied to exclude common numeric constants from multiple terms of a sum. As a consequence,
+    the number of multiplications is reduced and in some cases, more common subexpressions can be found.
+    """
+    constants_to_subexp_dict = defaultdict(lambda: next(ac.subexpression_symbol_generator))
+
+    def visit(expr):
+        args = list(expr.args)
+        if len(args) == 0:
+            return expr
+        if isinstance(expr, sp.Add) or isinstance(expr, sp.Mul):
+            for i, arg in enumerate(args):
+                if is_constant(arg) and abs(arg) != 1:
+                    if arg < 0:
+                        args[i] = - constants_to_subexp_dict[- arg]
+                    else:
+                        args[i] = constants_to_subexp_dict[arg]
+        return expr.func(*(visit(a) for a in args))
+    main_assignments = [Assignment(a.lhs, visit(a.rhs)) for a in ac.main_assignments]
+    subexpressions = [Assignment(a.lhs, visit(a.rhs)) for a in ac.subexpressions]
+
+    symbols_to_collect = set(constants_to_subexp_dict.values())
+
+    main_assignments = [Assignment(a.lhs, recursive_collect(a.rhs, symbols_to_collect, True)) for a in main_assignments]
+    subexpressions = [Assignment(a.lhs, recursive_collect(a.rhs, symbols_to_collect, True)) for a in subexpressions]
+
+    subexpressions = [Assignment(symb, c) for c, symb in constants_to_subexp_dict.items()] + subexpressions
+    return ac.copy(main_assignments=main_assignments, subexpressions=subexpressions)
+
+
+def add_subexpressions_for_divisions(ac):
+    r"""Introduces subexpressions for all divisions which have no constant in the denominator.
+
+    For example :math:`\frac{1}{x}` is replaced while :math:`\frac{1}{3}` is not replaced.
+    """
+    divisors = set()
+
+    def search_divisors(term):
+        if term.func == sp.Pow:
+            if term.exp.is_integer and term.exp.is_number and term.exp < 0:
+                divisors.add(term)
+        else:
+            for a in term.args:
+                search_divisors(a)
+
+    for eq in ac.all_assignments:
+        search_divisors(eq.rhs)
+
+    divisors = sorted(list(divisors), key=lambda x: str(x))
+    new_symbol_gen = ac.subexpression_symbol_generator
+    substitutions = {divisor: new_symbol for new_symbol, divisor in zip(new_symbol_gen, divisors)}
+    return ac.new_with_substitutions(substitutions, add_substitutions_as_subexpressions=True, substitute_on_lhs=False)
+
+
+def add_subexpressions_for_sums(ac):
+    r"""Introduces subexpressions for all sums - i.e. splits addends into subexpressions."""
+    addends = []
+
+    def contains_sum(term):
+        if term.func == sp.Add:
+            return True
+        if term.is_Atom:
+            return False
+        return any([contains_sum(a) for a in term.args])
+
+    def search_addends(term):
+        if term.func == sp.Add:
+            if all([not contains_sum(a) for a in term.args]):
+                addends.extend(term.args)
+        for a in term.args:
+            search_addends(a)
+
+    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, 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, 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 = []
+    if subexpressions:
+        to_iterate = chain(to_iterate, ac.subexpressions)
+    if main_assignments:
+        to_iterate = chain(to_iterate, ac.main_assignments)
+
+    for assignment in to_iterate:
+        if hasattr(assignment, 'lhs') and hasattr(assignment, 'rhs'):
+            field_reads.update(assignment.rhs.atoms(Field.Access))
+
+    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)
+
+
+def transform_rhs(assignment_list, transformation, *args, **kwargs):
+    """Applies a transformation function on the rhs of each element of the passed assignment list
+    If the list also contains other object, like AST nodes, these are ignored.
+    Additional parameters are passed to the transformation function"""
+    return [Assignment(a.lhs, transformation(a.rhs, *args, **kwargs)) if hasattr(a, 'lhs') and hasattr(a, 'rhs') else a
+            for a in assignment_list]
+
+
+def transform_lhs_and_rhs(assignment_list, transformation, *args, **kwargs):
+    return [Assignment(transformation(a.lhs, *args, **kwargs),
+                       transformation(a.rhs, *args, **kwargs))
+            if hasattr(a, 'lhs') and hasattr(a, 'rhs') else a
+            for a in assignment_list]
+
+
+def apply_to_all_assignments(operation: Callable[[sp.Expr], sp.Expr]):
+    """Applies a given operation to all equations in collection."""
+
+    def f(ac):
+        return ac.copy(transform_rhs(ac.main_assignments, operation))
+
+    f.__name__ = operation.__name__
+    return f
+
+
+def apply_on_all_subexpressions(operation: Callable[[sp.Expr], sp.Expr]):
+    """Applies the given operation on all subexpressions of the AC."""
+
+    def f(ac):
+        return ac.copy(ac.main_assignments, transform_rhs(ac.subexpressions, operation))
+
+    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)