pystencils/include/opencl_stdint.h

-#ifndef OPENCL_STDINT
-#define OPENCL_STDINT
-
-typedef unsigned int      uint_t;
-
-typedef signed char       int8_t;
-typedef signed short      int16_t;
-typedef signed int        int32_t;
-typedef signed long int   int64_t;
-typedef unsigned char     uint8_t;
-typedef unsigned short    uint16_t;
-typedef unsigned int      uint32_t;
-typedef unsigned long int uint64_t;
-
-#endif

pystencils/node_collection.py

-from typing import List, Union
-
-import sympy
-import sympy as sp
-from sympy.codegen import Assignment
-from sympy.codegen.rewriting import ReplaceOptim, optimize
-
-from pystencils.astnodes import Block, Node, SympyAssignment
-from pystencils.backends.cbackend import CustomCodeNode
-from pystencils.functions import DivFunc
-from pystencils.simp import AssignmentCollection
-
-
-class NodeCollection:
-    def __init__(self, assignments: List[Union[Node, Assignment]]):
-        self.all_assignments = assignments
-
-        if all((isinstance(a, Assignment) for a in assignments)):
-            self.is_Nodes = False
-            self.is_Assignments = True
-        elif all((isinstance(n, Node) for n in assignments)):
-            self.is_Nodes = True
-            self.is_Assignments = False
-        else:
-            raise ValueError(f'The list "{assignments}" is mixed. Pass either a list of "pystencils.Assignments" '
-                             f'or a list of "pystencils.astnodes.Node')
-
-        self.simplification_hints = {}
-
-    @staticmethod
-    def from_assignment_collection(assignment_collection: AssignmentCollection):
-        nodes = list()
-        for assignemt in assignment_collection.all_assignments:
-            if isinstance(assignemt, Assignment):
-                nodes.append(SympyAssignment(assignemt.lhs, assignemt.rhs))
-            elif isinstance(assignemt, Node):
-                nodes.append(assignemt)
-            else:
-                raise ValueError(f"Unknown node in the AssignmentCollection: {assignemt}")
-
-        return NodeCollection(nodes)
-
-    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), sp.Mul(*([p.base] * -p.exp), evaluate=False))
-            ))
-        sympy_optimisations = [evaluate_constant_terms, evaluate_pow]
-
-        if self.is_Nodes:
-            def visitor(node):
-                if isinstance(node, CustomCodeNode):
-                    return node
-                elif isinstance(node, Block):
-                    return node.func([visitor(child) for child in node.args])
-                elif isinstance(node, 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]
-        else:
-            self.all_assignments = [Assignment(a.lhs, optimize(a.rhs, sympy_optimisations))
-                                    if hasattr(a, 'lhs')
-                                    else a for a in self.all_assignments]

pystencils_tests/test_create_kernel_backwards_compability.py

-import pytest
-
-import pystencils as ps
-import numpy as np
-
-
-# This test aims to trigger deprication warnings. Thus the warnings should not be displayed in the warning summary.
-import pystencils.config
-
-
-def test_create_kernel_backwards_compatibility():
-    size = (30, 20)
-
-    src_field_string = np.random.rand(*size)
-    src_field_enum = np.copy(src_field_string)
-    src_field_config = np.copy(src_field_string)
-    dst_field_string = np.zeros(size)
-    dst_field_enum = np.zeros(size)
-    dst_field_config = np.zeros(size)
-
-    f = ps.Field.create_from_numpy_array("f", src_field_enum)
-    d = ps.Field.create_from_numpy_array("d", dst_field_enum)
-
-    jacobi = ps.Assignment(d[0, 0], (f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
-    ast_enum = ps.create_kernel(jacobi, target=ps.Target.CPU).compile()
-    with pytest.warns(DeprecationWarning):
-        ast_string = ps.create_kernel(jacobi, target='cpu').compile()
-    # noinspection PyTypeChecker
-    with pytest.warns(DeprecationWarning):
-        ast_config = ps.create_kernel(jacobi, config=pystencils.config.CreateKernelConfig(target='cpu')).compile()
-    ast_enum(f=src_field_enum, d=dst_field_enum)
-    ast_string(f=src_field_string, d=dst_field_string)
-    ast_config(f=src_field_config, d=dst_field_config)
-
-    error = np.sum(np.abs(dst_field_enum - dst_field_string))
-    np.testing.assert_almost_equal(error, 0.0)
-    error = np.sum(np.abs(dst_field_enum - dst_field_config))
-    np.testing.assert_almost_equal(error, 0.0)

pystencils_tests/test_indexed_kernels.py

-import numpy as np
-import pystencils as ps
-from pystencils import Assignment, Field, CreateKernelConfig, create_kernel, Target
-
-
-def test_indexed_kernel():
-    arr = np.zeros((3, 4))
-    dtype = np.dtype([('x', int), ('y', int), ('value', arr.dtype)])
-    index_arr = np.zeros((3,), dtype=dtype)
-    index_arr[0] = (0, 2, 3.0)
-    index_arr[1] = (1, 3, 42.0)
-    index_arr[2] = (2, 1, 5.0)
-
-    indexed_field = Field.create_from_numpy_array('index', index_arr)
-    normal_field = Field.create_from_numpy_array('f', arr)
-    update_rule = Assignment(normal_field[0, 0], indexed_field('value'))
-
-    config = CreateKernelConfig(index_fields=[indexed_field])
-    ast = create_kernel([update_rule], config=config)
-    kernel = ast.compile()
-    kernel(f=arr, index=index_arr)
-    code = ps.get_code_str(kernel)
-    for i in range(index_arr.shape[0]):
-        np.testing.assert_allclose(arr[index_arr[i]['x'], index_arr[i]['y']], index_arr[i]['value'], atol=1e-13)
-
-
-def test_indexed_cuda_kernel():
-    try:
-        import pycuda
-    except ImportError:
-        pycuda = None
-
-    if pycuda:
-        import pycuda.gpuarray as gpuarray
-
-        arr = np.zeros((3, 4))
-        dtype = np.dtype([('x', int), ('y', int), ('value', arr.dtype)])
-        index_arr = np.zeros((3,), dtype=dtype)
-        index_arr[0] = (0, 2, 3.0)
-        index_arr[1] = (1, 3, 42.0)
-        index_arr[2] = (2, 1, 5.0)
-
-        indexed_field = Field.create_from_numpy_array('index', index_arr)
-        normal_field = Field.create_from_numpy_array('f', arr)
-        update_rule = Assignment(normal_field[0, 0], indexed_field('value'))
-
-        config = CreateKernelConfig(target=Target.GPU, index_fields=[indexed_field])
-        ast = create_kernel([update_rule], config=config)
-        kernel = ast.compile()
-
-        gpu_arr = gpuarray.to_gpu(arr)
-        gpu_index_arr = gpuarray.to_gpu(index_arr)
-        kernel(f=gpu_arr, index=gpu_index_arr)
-        gpu_arr.get(arr)
-        for i in range(index_arr.shape[0]):
-            np.testing.assert_allclose(arr[index_arr[i]['x'], index_arr[i]['y']], index_arr[i]['value'], atol=1e-13)
-    else:
-        print("Did not run test on GPU since no pycuda is available")

pystencils_tests/test_print_infinity.py

-import pytest
-
-import pystencils
-from sympy import oo
-
-
-@pytest.mark.parametrize('type', ('float32', 'float64', 'int64'))
-@pytest.mark.parametrize('negative', (False, 'Negative'))
-@pytest.mark.parametrize('target', (pystencils.Target.CPU, pystencils.Target.GPU))
-def test_print_infinity(type, negative, target):
-
-    x = pystencils.fields(f'x:  {type}[1d]')
-
-    if negative:
-        assignment = pystencils.Assignment(x.center, -oo)
-    else:
-        assignment = pystencils.Assignment(x.center, oo)
-    ast = pystencils.create_kernel(assignment, data_type=type, target=target)
-
-    if target == pystencils.Target.GPU:
-        pytest.importorskip('pycuda')
-
-    ast.compile()
-
-    print(ast.compile().code)

pystencils_tests/test_print_unsupported_node.py

-# -*- coding: utf-8 -*-
-#
-# Copyright © 2019 Stephan Seitz <stephan.seitz@fau.de>
-#
-# Distributed under terms of the GPLv3 license.
-
-"""
-
-"""
-import pytest
-
-import pystencils
-from pystencils.backends.cbackend import CBackend
-
-
-class UnsupportedNode(pystencils.astnodes.Node):
-
-    def __init__(self):
-        super().__init__()
-
-
-def test_print_unsupported_node():
-    with pytest.raises(NotImplementedError, match='CBackend does not support node of type UnsupportedNode'):
-        CBackend()(UnsupportedNode())

pystencils_tests/test_sliced_iteration.py

-import numpy as np
-import sympy as sp
-
-from pystencils import Assignment, Field, TypedSymbol, create_kernel, make_slice
-from pystencils.simp import sympy_cse_on_assignment_list
-
-
-def test_sliced_iteration():
-    size = (4, 4)
-    src_arr = np.ones(size)
-    dst_arr = np.zeros_like(src_arr)
-    src_field = Field.create_from_numpy_array('src', src_arr)
-    dst_field = Field.create_from_numpy_array('dst', dst_arr)
-
-    a, b = sp.symbols("a b")
-    update_rule = Assignment(dst_field[0, 0],
-                             (a * src_field[0, 1] + a * src_field[0, -1] +
-                              b * src_field[1, 0] + b * src_field[-1, 0]) / 4)
-
-    x_end = TypedSymbol("x_end", "int")
-    s = make_slice[1:x_end, 1]
-    x_end_value = size[1] - 1
-    kernel = create_kernel(sympy_cse_on_assignment_list([update_rule]), iteration_slice=s).compile()
-
-    kernel(src=src_arr, dst=dst_arr, a=1.0, b=1.0, x_end=x_end_value)
-
-    expected_result = np.zeros(size)
-    expected_result[1:x_end_value, 1] = 1
-    np.testing.assert_almost_equal(expected_result, dst_arr)

pystencils_tests/test_transformations.py

-import pystencils as ps
-from pystencils import TypedSymbol
-from pystencils.astnodes import LoopOverCoordinate, SympyAssignment
-from pystencils.typing import create_type
-from pystencils.transformations import filtered_tree_iteration, get_loop_hierarchy, get_loop_counter_symbol_hierarchy
-
-
-def test_loop_information():
-    f, g = ps.fields("f, g: double[2D]")
-    update_rule = ps.Assignment(g[0, 0], f[0, 0])
-
-    ast = ps.create_kernel(update_rule)
-    inner_loops = [l for l in filtered_tree_iteration(ast, LoopOverCoordinate, stop_type=SympyAssignment)
-                   if l.is_innermost_loop]
-
-    loop_order = []
-    for i in get_loop_hierarchy(inner_loops[0].args[0]):
-        loop_order.append(i)
-
-    assert loop_order == [0, 1]
-
-    loop_symbols = get_loop_counter_symbol_hierarchy(inner_loops[0].args[0])
-
-    assert loop_symbols == [TypedSymbol("ctr_1", create_type("int"), nonnegative=True),
-                            TypedSymbol("ctr_0", create_type("int"), nonnegative=True)]

pytest.ini

 [pytest]
+testpaths = src tests doc/notebooks
+pythonpath = src
 python_files = test_*.py *_test.py scenario_*.py
 norecursedirs = *.egg-info .git .cache .ipynb_checkpoints htmlcov
 addopts = --doctest-modules --durations=20  --cov-config pytest.ini
@@ -17,20 +19,21 @@ filterwarnings =
 
 [run]
 branch = True
-source = pystencils
-         pystencils_tests
+source = src/pystencils
+         tests
 
 omit = doc/*
-       pystencils_tests/*
+       tests/*
        setup.py
+       quicktest.py
        conftest.py
        versioneer.py
-       pystencils/jupytersetup.py
-       pystencils/cpu/msvc_detection.py
-       pystencils/sympy_gmpy_bug_workaround.py
-       pystencils/cache.py
-       pystencils/pacxx/benchmark.py
-       pystencils/_version.py
+       src/pystencils/jupytersetup.py
+       src/pystencils/cpu/msvc_detection.py
+       src/pystencils/sympy_gmpy_bug_workaround.py
+       src/pystencils/cache.py
+       src/pystencils/pacxx/benchmark.py
+       src/pystencils/_version.py
        venv/
 
 [report]

quicktest.py

+#!/usr/bin/env python3
+
+from contextlib import redirect_stdout
+import io
+from tests.test_quicktests import (
+    test_basic_kernel,
+    test_basic_blocking_staggered,
+    test_basic_vectorization,
+)
+
+quick_tests = [
+    test_basic_kernel,
+    test_basic_blocking_staggered,
+    test_basic_vectorization,
+]
+
+if __name__ == "__main__":
+    print("Running pystencils quicktests")
+    for qt in quick_tests:
+        print(f"   -> {qt.__name__}")
+        with redirect_stdout(io.StringIO()):
+            qt()

setup.cfg

-# See the docstring in versioneer.py for instructions. Note that you must
-# re-run 'versioneer.py setup' after changing this section, and commit the
-# resulting files.
-
-[versioneer]
-VCS = git
-style = pep440
-versionfile_source = pystencils/_version.py
-versionfile_build = pystencils/_version.py
-tag_prefix = release/
-parentdir_prefix = pystencils-

setup.py

-import distutils
-import io
-import os
-from contextlib import redirect_stdout
-from importlib import import_module
+from setuptools import setup, __version__ as setuptools_version
 
-import setuptools
+if int(setuptools_version.split('.')[0]) < 61:
+    raise Exception(
+        "[ERROR] pystencils requires at least setuptools version 61 to install.\n"
+        "If this error occurs during an installation via pip, it is likely that there is a conflict between "
+        "versions of setuptools installed by pip and the system package manager. "
+        "In this case, it is recommended to install pystencils into a virtual environment instead."
+    )
 
 import versioneer
 
-try:
-    import cython  # noqa
-
-    USE_CYTHON = True
-except ImportError:
-    USE_CYTHON = False
-
-quick_tests = [
-    'test_quicktests.test_basic_kernel',
-    'test_quicktests.test_basic_blocking_staggered',
-    'test_quicktests.test_basic_vectorization',
-]
-
-
-class SimpleTestRunner(distutils.cmd.Command):
-    """A custom command to run selected tests"""
-
-    description = 'run some quick tests'
-    user_options = []
-
-    @staticmethod
-    def _run_tests_in_module(test):
-        """Short test runner function - to work also if py.test is not installed."""
-        test = f'pystencils_tests.{test}'
-        mod, function_name = test.rsplit('.', 1)
-        if isinstance(mod, str):
-            mod = import_module(mod)
-
-        func = getattr(mod, function_name)
-        print(f"   -> {function_name} in {mod.__name__}")
-        with redirect_stdout(io.StringIO()):
-            func()
-
-    def initialize_options(self):
-        pass
-
-    def finalize_options(self):
-        pass
-
-    def run(self):
-        """Run command."""
-        for test in quick_tests:
-            self._run_tests_in_module(test)
-
-
-def readme():
-    with open('README.md') as f:
-        return f.read()
-
-
-def cython_extensions(*extensions):
-    from distutils.extension import Extension
-    if USE_CYTHON:
-        ext = '.pyx'
-        result = [Extension(e, [os.path.join(*e.split(".")) + ext]) for e in extensions]
-        from Cython.Build import cythonize
-        result = cythonize(result, language_level=3)
-        return result
-    elif all([os.path.exists(os.path.join(*e.split(".")) + '.c') for e in extensions]):
-        ext = '.c'
-        result = [Extension(e, [os.path.join(*e.split(".")) + ext]) for e in extensions]
-        return result
-    else:
-        return None
-
 
 def get_cmdclass():
-    cmdclass = {"quicktest": SimpleTestRunner}
-    cmdclass.update(versioneer.get_cmdclass())
-    return cmdclass
-
+    return versioneer.get_cmdclass()
 
-setuptools.setup(name='pystencils',
-                 description='Speeding up stencil computations on CPUs and GPUs',
-                 version=versioneer.get_version(),
-                 long_description=readme(),
-                 long_description_content_type="text/markdown",
-                 author='Martin Bauer, Jan Hönig, Markus Holzer',
-                 license='AGPLv3',
-                 author_email='cs10-codegen@fau.de',
-                 url='https://i10git.cs.fau.de/pycodegen/pystencils/',
-                 packages=['pystencils'] + ['pystencils.' + s for s in setuptools.find_packages('pystencils')],
-                 install_requires=['sympy>=1.6,<=1.10', 'numpy>=1.8.0', 'appdirs', 'joblib'],
-                 package_data={'pystencils': ['include/*.h',
-                                              'backends/cuda_known_functions.txt',
-                                              'backends/opencl1.1_known_functions.txt',
-                                              'boundaries/createindexlistcython.c',
-                                              'boundaries/createindexlistcython.pyx']},
-                 ext_modules=cython_extensions("pystencils.boundaries.createindexlistcython"),
-                 classifiers=[
-                     'Development Status :: 4 - Beta',
-                     'Framework :: Jupyter',
-                     'Topic :: Software Development :: Code Generators',
-                     'Topic :: Scientific/Engineering :: Physics',
-                     'Intended Audience :: Developers',
-                     'Intended Audience :: Science/Research',
-                     'License :: OSI Approved :: GNU Affero General Public License v3 or later (AGPLv3+)',
-                 ],
-                 project_urls={
-                     "Bug Tracker": "https://i10git.cs.fau.de/pycodegen/pystencils/issues",
-                     "Documentation": "http://pycodegen.pages.walberla.net/pystencils/",
-                     "Source Code": "https://i10git.cs.fau.de/pycodegen/pystencils",
-                 },
-                 extras_require={
-                     'gpu': ['pycuda'],
-                     'alltrafos': ['islpy', 'py-cpuinfo'],
-                     'bench_db': ['blitzdb', 'pymongo', 'pandas'],
-                     'interactive': ['matplotlib', 'ipy_table', 'imageio', 'jupyter', 'pyevtk', 'rich', 'graphviz'],
-                     'doc': ['sphinx', 'sphinx_rtd_theme', 'nbsphinx',
-                             'sphinxcontrib-bibtex', 'sphinx_autodoc_typehints', 'pandoc'],
-                     'use_cython': ['Cython']
-                 },
-                 tests_require=['pytest',
-                                'pytest-cov',
-                                'pytest-html',
-                                'ansi2html',
-                                'pytest-xdist',
-                                'flake8',
-                                'nbformat',
-                                'nbconvert',
-                                'ipython',
-                                'randomgen>=1.18'],
 
-                 python_requires=">=3.8",
-                 cmdclass=get_cmdclass()
-                 )
+setup(
+    version=versioneer.get_version(),
+    cmdclass=get_cmdclass(),
+)

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

@@ -2,18 +2,19 @@
 from .enums import Backend, Target
 from . import fd
 from . import stencil as stencil
-from .assignment import Assignment, assignment_from_stencil
-from pystencils.typing.typed_sympy import TypedSymbol
-from .datahandling import create_data_handling
+from .assignment import Assignment, AddAugmentedAssignment, assignment_from_stencil
+from .typing.typed_sympy import TypedSymbol
 from .display_utils import get_code_obj, get_code_str, show_code, to_dot
 from .field import Field, FieldType, fields
 from .config import CreateKernelConfig
+from .cache import clear_cache
 from .kernel_decorator import kernel, kernel_config
 from .kernelcreation import create_kernel, create_staggered_kernel
 from .simp import AssignmentCollection
 from .slicing import make_slice
 from .spatial_coordinates import x_, x_staggered, x_staggered_vector, x_vector, y_, y_staggered, z_, z_staggered
 from .sympyextensions import SymbolCreator
+from .datahandling import create_data_handling
 
 __all__ = ['Field', 'FieldType', 'fields',
            'TypedSymbol',
@@ -23,10 +24,11 @@ __all__ = ['Field', 'FieldType', 'fields',
            'Target', 'Backend',
            'show_code', 'to_dot', 'get_code_obj', 'get_code_str',
            'AssignmentCollection',
-           'Assignment',
+           'Assignment', 'AddAugmentedAssignment',
            'assignment_from_stencil',
            'SymbolCreator',
            'create_data_handling',
+           'clear_cache',
            'kernel', 'kernel_config',
            'x_', 'y_', 'z_',
            'x_staggered', 'y_staggered', 'z_staggered',
@@ -34,7 +36,5 @@ __all__ = ['Field', 'FieldType', 'fields',
            'fd',
            'stencil']
 
-from ._version import get_versions
-
-__version__ = get_versions()['version']
-del get_versions
+from . import _version
+__version__ = _version.get_versions()['version']

pystencils/alignedarray.py → src/pystencils/alignedarray.py

 import numpy as np
-from pystencils.typing import numpy_name_to_c
 
 
 def aligned_empty(shape, byte_alignment=True, dtype=np.float64, byte_offset=0, order='C', align_inner_coordinate=True):
@@ -21,26 +20,26 @@ def aligned_empty(shape, byte_alignment=True, dtype=np.float64, byte_offset=0, o
         from pystencils.backends.simd_instruction_sets import (get_supported_instruction_sets, get_cacheline_size,
                                                                get_vector_instruction_set)
 
-        type_name = numpy_name_to_c(np.dtype(dtype).name)
         instruction_sets = get_supported_instruction_sets()
         if instruction_sets is None:
             byte_alignment = 64
         elif byte_alignment == 'cacheline':
             cacheline_sizes = [get_cacheline_size(is_name) for is_name in instruction_sets]
-            if all([s is None for s in cacheline_sizes]):
-                widths = [get_vector_instruction_set(type_name, is_name)['width'] * np.dtype(dtype).itemsize
+            if all([s is None for s in cacheline_sizes]) or \
+                    max([s for s in cacheline_sizes if s is not None]) > 0x100000:
+                widths = [get_vector_instruction_set(dtype, is_name)['width'] * np.dtype(dtype).itemsize
                           for is_name in instruction_sets
-                          if type(get_vector_instruction_set(type_name, is_name)['width']) is int]
+                          if type(get_vector_instruction_set(dtype, is_name)['width']) is int]
                 byte_alignment = 64 if all([s is None for s in widths]) else max(widths)
             else:
                 byte_alignment = max([s for s in cacheline_sizes if s is not None])
-        elif not any([type(get_vector_instruction_set(type_name, is_name)['width']) is int
+        elif not any([type(get_vector_instruction_set(dtype, is_name)['width']) is int
                       for is_name in instruction_sets]):
             byte_alignment = 64
         else:
-            byte_alignment = max([get_vector_instruction_set(type_name, is_name)['width'] * np.dtype(dtype).itemsize
+            byte_alignment = max([get_vector_instruction_set(dtype, is_name)['width'] * np.dtype(dtype).itemsize
                                   for is_name in instruction_sets
-                                  if type(get_vector_instruction_set(type_name, is_name)['width']) is int])
+                                  if type(get_vector_instruction_set(dtype, is_name)['width']) is int])
     if (not align_inner_coordinate) or (not hasattr(shape, '__len__')):
         size = np.prod(shape)
         d = np.dtype(dtype)
@@ -78,7 +77,7 @@ def aligned_empty(shape, byte_alignment=True, dtype=np.float64, byte_offset=0, o
         return tmp
 
 
-def aligned_zeros(shape, byte_alignment=True, dtype=float, byte_offset=0, order='C', align_inner_coordinate=True):
+def aligned_zeros(shape, byte_alignment=True, dtype=np.float64, byte_offset=0, order='C', align_inner_coordinate=True):
     arr = aligned_empty(shape, dtype=dtype, byte_offset=byte_offset,
                         order=order, byte_alignment=byte_alignment, align_inner_coordinate=align_inner_coordinate)
     x = np.zeros((), arr.dtype)
@@ -86,7 +85,7 @@ def aligned_zeros(shape, byte_alignment=True, dtype=float, byte_offset=0, order=
     return arr
 
 
-def aligned_ones(shape, byte_alignment=True, dtype=float, byte_offset=0, order='C', align_inner_coordinate=True):
+def aligned_ones(shape, byte_alignment=True, dtype=np.float64, byte_offset=0, order='C', align_inner_coordinate=True):
     arr = aligned_empty(shape, dtype=dtype, byte_offset=byte_offset,
                         order=order, byte_alignment=byte_alignment, align_inner_coordinate=align_inner_coordinate)
     x = np.ones((), arr.dtype)

pystencils/assignment.py → src/pystencils/assignment.py

 import numpy as np
 import sympy as sp
-from sympy.codegen.ast import Assignment
+from sympy.codegen.ast import Assignment, AugmentedAssignment, AddAugmentedAssignment
 from sympy.printing.latex import LatexPrinter
 
-__all__ = ['Assignment', 'assignment_from_stencil']
+__all__ = ['Assignment', 'AugmentedAssignment', 'AddAugmentedAssignment', 'assignment_from_stencil']
 
 
 def print_assignment_latex(printer, expr):
+    binop = f"{expr.binop}=" if isinstance(expr, AugmentedAssignment) else ''
     """sympy cannot print Assignments as Latex. Thus, this function is added to the sympy Latex printer"""
     printed_lhs = printer.doprint(expr.lhs)
     printed_rhs = printer.doprint(expr.rhs)
-    return fr"{printed_lhs} \leftarrow {printed_rhs}"
+    return fr"{printed_lhs} \leftarrow_{{{binop}}} {printed_rhs}"
 
 
 def assignment_str(assignment):
-    return fr"{assignment.lhs} ← {assignment.rhs}"
+    op = f"{assignment.binop}=" if isinstance(assignment, AugmentedAssignment) else '←'
+    return fr"{assignment.lhs} {op} {assignment.rhs}"
 
 
 _old_new = sp.codegen.ast.Assignment.__new__
@@ -32,6 +34,9 @@ Assignment.__str__ = assignment_str
 Assignment.__new__ = _Assignment__new__
 LatexPrinter._print_Assignment = print_assignment_latex
 
+AugmentedAssignment.__str__ = assignment_str
+LatexPrinter._print_AugmentedAssignment = print_assignment_latex
+
 sp.MutableDenseMatrix.__hash__ = lambda self: hash(tuple(self))