pystencils_tests/test_field.py → tests/test_field.py

@@ -4,22 +4,34 @@ import sympy as sp
 
 import pystencils as ps
 from pystencils import TypedSymbol
-from pystencils.data_types import create_type
-from pystencils.field import Field, FieldType, layout_string_to_tuple
+from pystencils.typing import create_type
+from pystencils.field import Field, FieldType, layout_string_to_tuple, spatial_layout_string_to_tuple
 
 
 def test_field_basic():
-    f = Field.create_generic('f', spatial_dimensions=2)
+    f = Field.create_generic("f", spatial_dimensions=2)
     assert FieldType.is_generic(f)
-    assert f['E'] == f[1, 0]
-    assert f['N'] == f[0, 1]
-    assert '_' in f.center._latex('dummy')
-
-    assert f.index_to_physical(index_coordinates=sp.Matrix([0, 0]), staggered=False)[0] == 0
-    assert f.index_to_physical(index_coordinates=sp.Matrix([0, 0]), staggered=False)[1] == 0
-
-    assert f.physical_to_index(physical_coordinates=sp.Matrix([0, 0]), staggered=False)[0] == 0
-    assert f.physical_to_index(physical_coordinates=sp.Matrix([0, 0]), staggered=False)[1] == 0
+    assert f["E"] == f[1, 0]
+    assert f["N"] == f[0, 1]
+    assert "_" in f.center._latex("dummy")
+
+    assert (
+        f.index_to_physical(index_coordinates=sp.Matrix([0, 0]), staggered=False)[0]
+        == 0
+    )
+    assert (
+        f.index_to_physical(index_coordinates=sp.Matrix([0, 0]), staggered=False)[1]
+        == 0
+    )
+
+    assert (
+        f.physical_to_index(physical_coordinates=sp.Matrix([0, 0]), staggered=False)[0]
+        == 0
+    )
+    assert (
+        f.physical_to_index(physical_coordinates=sp.Matrix([0, 0]), staggered=False)[1]
+        == 0
+    )
 
     f1 = f.new_field_with_different_name("f1")
     assert f1.ndim == f.ndim
@@ -28,7 +40,7 @@ def test_field_basic():
     fixed = ps.fields("f(5, 5) : double[20, 20]")
     assert fixed.neighbor_vector((1, 1)).shape == (5, 5)
 
-    f = Field.create_fixed_size('f', (10, 10), strides=(80, 8), dtype=np.float64)
+    f = Field.create_fixed_size("f", (10, 10), strides=(80, 8), dtype=np.float64)
     assert f.spatial_strides == (10, 1)
     assert f.index_strides == ()
     assert f.center_vector == sp.Matrix([f.center])
@@ -37,20 +49,21 @@ def test_field_basic():
     assert f1.ndim == f.ndim
     assert f1.values_per_cell() == f.values_per_cell()
 
-    f = Field.create_fixed_size('f', (8, 8, 2, 2), index_dimensions=2)
-    assert f.center_vector == sp.Matrix([[f(0, 0), f(0, 1)],
-                                         [f(1, 0), f(1, 1)]])
+    f = Field.create_fixed_size("f", (8, 8, 2, 2), index_dimensions=2)
+    assert f.center_vector == sp.Matrix([[f(0, 0), f(0, 1)], [f(1, 0), f(1, 1)]])
     field_access = f[1, 1]
     assert field_access.nr_of_coordinates == 2
-    assert field_access.offset_name == 'NE'
+    assert field_access.offset_name == "NE"
     neighbor = field_access.neighbor(coord_id=0, offset=-2)
     assert neighbor.offsets == (-1, 1)
-    assert '_' in neighbor._latex('dummy')
+    assert "_" in neighbor._latex("dummy")
 
-    f = Field.create_fixed_size('f', (8, 8, 2, 2, 2), index_dimensions=3)
-    assert f.center_vector == sp.Array([[[f(i, j, k) for k in range(2)] for j in range(2)] for i in range(2)])
+    f = Field.create_fixed_size("f", (8, 8, 2, 2, 2), index_dimensions=3)
+    assert f.center_vector == sp.Array(
+        [[[f(i, j, k) for k in range(2)] for j in range(2)] for i in range(2)]
+    )
 
-    f = Field.create_generic('f', spatial_dimensions=5, index_dimensions=2)
+    f = Field.create_generic("f", spatial_dimensions=5, index_dimensions=2)
     field_access = f[1, -1, 2, -3, 0](1, 0)
     assert field_access.offsets == (1, -1, 2, -3, 0)
     assert field_access.index == (1, 0)
@@ -60,61 +73,71 @@ def test_error_handling():
     struct_dtype = np.dtype([('a', np.int32), ('b', np.float64), ('c', np.uint32)])
     Field.create_generic('f', spatial_dimensions=2, index_dimensions=0, dtype=struct_dtype)
     with pytest.raises(ValueError) as e:
-        Field.create_generic('f', spatial_dimensions=2, index_dimensions=1, dtype=struct_dtype)
-    assert 'index dimension' in str(e.value)
+        Field.create_generic(
+            "f", spatial_dimensions=2, index_dimensions=1, dtype=struct_dtype
+        )
+    assert "index dimension" in str(e.value)
 
-    arr = np.array([[[(1,)*3, (2,)*3, (3,)*3]]*2], dtype=struct_dtype)
-    Field.create_from_numpy_array('f', arr, index_dimensions=0)
+    arr = np.array([[[(1,) * 3, (2,) * 3, (3,) * 3]] * 2], dtype=struct_dtype)
+    Field.create_from_numpy_array("f", arr, index_dimensions=0)
     with pytest.raises(ValueError) as e:
-        Field.create_from_numpy_array('f', arr, index_dimensions=1)
-    assert 'Structured arrays' in str(e.value)
+        Field.create_from_numpy_array("f", arr, index_dimensions=1)
+    assert "Structured arrays" in str(e.value)
 
     arr = np.zeros([3, 3, 3])
-    Field.create_from_numpy_array('f', arr, index_dimensions=2)
+    Field.create_from_numpy_array("f", arr, index_dimensions=2)
     with pytest.raises(ValueError) as e:
-        Field.create_from_numpy_array('f', arr, index_dimensions=3)
-    assert 'Too many' in str(e.value)
+        Field.create_from_numpy_array("f", arr, index_dimensions=3)
+    assert "Too many" in str(e.value)
 
-    Field.create_fixed_size('f', (3, 2, 4), index_dimensions=0, dtype=struct_dtype, layout='reverse_numpy')
+    Field.create_fixed_size(
+        "f", (3, 2, 4), index_dimensions=0, dtype=struct_dtype, layout="reverse_numpy"
+    )
     with pytest.raises(ValueError) as e:
-        Field.create_fixed_size('f', (3, 2, 4), index_dimensions=1, dtype=struct_dtype, layout='reverse_numpy')
-    assert 'Structured arrays' in str(e.value)
-
-    f = Field.create_fixed_size('f', (10, 10))
+        Field.create_fixed_size(
+            "f",
+            (3, 2, 4),
+            index_dimensions=1,
+            dtype=struct_dtype,
+            layout="reverse_numpy",
+        )
+    assert "Structured arrays" in str(e.value)
+
+    f = Field.create_fixed_size("f", (10, 10))
     with pytest.raises(ValueError) as e:
         f[1]
-    assert 'Wrong number of spatial indices' in str(e.value)
+    assert "Wrong number of spatial indices" in str(e.value)
 
-    f = Field.create_generic('f', spatial_dimensions=2, index_shape=(3,))
+    f = Field.create_generic("f", spatial_dimensions=2, index_shape=(3,))
     with pytest.raises(ValueError) as e:
         f(3)
-    assert 'out of bounds' in str(e.value)
+    assert "out of bounds" in str(e.value)
 
-    f = Field.create_fixed_size('f', (10, 10, 3, 4), index_dimensions=2)
+    f = Field.create_fixed_size("f", (10, 10, 3, 4), index_dimensions=2)
     with pytest.raises(ValueError) as e:
         f(3, 0)
-    assert 'out of bounds' in str(e.value)
+    assert "out of bounds" in str(e.value)
 
     with pytest.raises(ValueError) as e:
         f(1, 0)(1, 0)
-    assert 'Indexing an already indexed' in str(e.value)
+    assert "Indexing an already indexed" in str(e.value)
 
     with pytest.raises(ValueError) as e:
         f(1)
-    assert 'Wrong number of indices' in str(e.value)
+    assert "Wrong number of indices" in str(e.value)
 
     with pytest.raises(ValueError) as e:
-        Field.create_generic('f', spatial_dimensions=2, layout='wrong')
-    assert 'Unknown layout descriptor' in str(e.value)
+        Field.create_generic("f", spatial_dimensions=2, layout="wrong")
+    assert "Unknown layout descriptor" in str(e.value)
 
-    assert layout_string_to_tuple('fzyx', dim=4) == (3, 2, 1, 0)
+    assert layout_string_to_tuple("fzyx", dim=4) == (3, 2, 1, 0)
     with pytest.raises(ValueError) as e:
-        layout_string_to_tuple('wrong', dim=4)
-    assert 'Unknown layout descriptor' in str(e.value)
+        layout_string_to_tuple("wrong", dim=4)
+    assert "Unknown layout descriptor" in str(e.value)
 
 
 def test_decorator_scoping():
-    dst = ps.fields('dst : double[2D]')
+    dst = ps.fields("dst : double[2D]")
 
     def f1():
         a = sp.Symbol("a")
@@ -134,7 +157,7 @@ def test_decorator_scoping():
 
 
 def test_string_creation():
-    x, y, z = ps.fields('  x(4),    y(3,5) z : double[  3,  47]')
+    x, y, z = ps.fields("  x(4),    y(3,5) z : double[  3,  47]")
     assert x.index_shape == (4,)
     assert y.index_shape == (3, 5)
     assert z.spatial_shape == (3, 47)
@@ -142,19 +165,85 @@ def test_string_creation():
 
 def test_itemsize():
 
-    x = ps.fields('x: float32[1d]')
-    y = ps.fields('y:  float64[2d]')
-    i = ps.fields('i:  int16[1d]')
+    x = ps.fields("x: float32[1d]")
+    y = ps.fields("y:  float64[2d]")
+    i = ps.fields("i:  int16[1d]")
 
     assert x.itemsize == 4
     assert y.itemsize == 8
     assert i.itemsize == 2
 
 
+def test_spatial_memory_layout_descriptors():
+    assert (
+        spatial_layout_string_to_tuple("AoS", 3)
+        == spatial_layout_string_to_tuple("aos", 3)
+        == spatial_layout_string_to_tuple("ZYXF", 3)
+        == spatial_layout_string_to_tuple("zyxf", 3)
+        == (2, 1, 0)
+    )
+    assert (
+        spatial_layout_string_to_tuple("SoA", 3)
+        == spatial_layout_string_to_tuple("soa", 3)
+        == spatial_layout_string_to_tuple("FZYX", 3)
+        == spatial_layout_string_to_tuple("fzyx", 3)
+        == spatial_layout_string_to_tuple("f", 3)
+        == spatial_layout_string_to_tuple("F", 3)
+        == (2, 1, 0)
+    )
+    assert (
+        spatial_layout_string_to_tuple("c", 3)
+        == spatial_layout_string_to_tuple("C", 3)
+        == (0, 1, 2)
+    )
+
+    assert spatial_layout_string_to_tuple("C", 5) == (0, 1, 2, 3, 4)
+
+    with pytest.raises(ValueError):
+        spatial_layout_string_to_tuple("aos", -1)
+
+    with pytest.raises(ValueError):
+        spatial_layout_string_to_tuple("aos", 4)
+
+
+def test_memory_layout_descriptors():
+    assert (
+        layout_string_to_tuple("AoS", 4)
+        == layout_string_to_tuple("aos", 4)
+        == layout_string_to_tuple("ZYXF", 4)
+        == layout_string_to_tuple("zyxf", 4)
+        == (2, 1, 0, 3)
+    )
+    assert (
+        layout_string_to_tuple("SoA", 4)
+        == layout_string_to_tuple("soa", 4)
+        == layout_string_to_tuple("FZYX", 4)
+        == layout_string_to_tuple("fzyx", 4)
+        == layout_string_to_tuple("f", 4)
+        == layout_string_to_tuple("F", 4)
+        == (3, 2, 1, 0)
+    )
+    assert (
+        layout_string_to_tuple("c", 4)
+        == layout_string_to_tuple("C", 4)
+        == (0, 1, 2, 3)
+    )
+
+    assert layout_string_to_tuple("C", 5) == (0, 1, 2, 3, 4)
+
+    with pytest.raises(ValueError):
+        layout_string_to_tuple("aos", -1)
+
+    with pytest.raises(ValueError):
+        layout_string_to_tuple("aos", 5)
+
+
 def test_staggered():
 
     # D2Q5
-    j1, j2, j3 = ps.fields('j1(2), j2(2,2), j3(2,2,2) : double[2D]', field_type=FieldType.STAGGERED)
+    j1, j2, j3 = ps.fields(
+        "j1(2), j2(2,2), j3(2,2,2) : double[2D]", field_type=FieldType.STAGGERED
+    )
 
     assert j1[0, 1](1) == j1.staggered_access((0, sp.Rational(1, 2)))
     assert j1[0, 1](1) == j1.staggered_access(np.array((0, sp.Rational(1, 2))))
@@ -163,7 +252,7 @@ def test_staggered():
     assert j1[0, 1](1) == j1.staggered_access("N")
     assert j1[0, 0](1) == j1.staggered_access("S")
     assert j1.staggered_vector_access("N") == sp.Matrix([j1.staggered_access("N")])
-    assert j1.staggered_stencil_name == 'D2Q5'
+    assert j1.staggered_stencil_name == "D2Q5"
 
     assert j1.physical_coordinates[0] == TypedSymbol("ctr_0", create_type("int"), nonnegative=True)
     assert j1.physical_coordinates[1] == TypedSymbol("ctr_1", create_type("int"), nonnegative=True)
@@ -176,28 +265,40 @@ def test_staggered():
 
     assert j2[0, 1](1, 1) == j2.staggered_access((0, sp.Rational(1, 2)), 1)
     assert j2[0, 1](1, 1) == j2.staggered_access("N", 1)
-    assert j2.staggered_vector_access("N") == sp.Matrix([j2.staggered_access("N", 0), j2.staggered_access("N", 1)])
+    assert j2.staggered_vector_access("N") == sp.Matrix(
+        [j2.staggered_access("N", 0), j2.staggered_access("N", 1)]
+    )
 
     assert j3[0, 1](1, 1, 1) == j3.staggered_access((0, sp.Rational(1, 2)), (1, 1))
     assert j3[0, 1](1, 1, 1) == j3.staggered_access("N", (1, 1))
-    assert j3.staggered_vector_access("N") == sp.Matrix([[j3.staggered_access("N", (i, j))
-                                                        for j in range(2)] for i in range(2)])
+    assert j3.staggered_vector_access("N") == sp.Matrix(
+        [[j3.staggered_access("N", (i, j)) for j in range(2)] for i in range(2)]
+    )
 
     # D2Q9
-    k1, k2 = ps.fields('k1(4), k2(2) : double[2D]', field_type=FieldType.STAGGERED)
+    k1, k2 = ps.fields("k1(4), k2(2) : double[2D]", field_type=FieldType.STAGGERED)
 
     assert k1[1, 1](2) == k1.staggered_access("NE")
     assert k1[0, 0](2) == k1.staggered_access("SW")
     assert k1[0, 0](3) == k1.staggered_access("NW")
-    
+
     a = k1.staggered_access("NE")
-    assert a._staggered_offset(a.offsets, a.index[0]) == [sp.Rational(1, 2), sp.Rational(1, 2)]
+    assert a._staggered_offset(a.offsets, a.index[0]) == [
+        sp.Rational(1, 2),
+        sp.Rational(1, 2),
+    ]
     a = k1.staggered_access("SW")
-    assert a._staggered_offset(a.offsets, a.index[0]) == [sp.Rational(-1, 2), sp.Rational(-1, 2)]
+    assert a._staggered_offset(a.offsets, a.index[0]) == [
+        sp.Rational(-1, 2),
+        sp.Rational(-1, 2),
+    ]
     a = k1.staggered_access("NW")
-    assert a._staggered_offset(a.offsets, a.index[0]) == [sp.Rational(-1, 2), sp.Rational(1, 2)]
+    assert a._staggered_offset(a.offsets, a.index[0]) == [
+        sp.Rational(-1, 2),
+        sp.Rational(1, 2),
+    ]
 
     # sign reversed when using as flux field
-    r = ps.fields('r(2) : double[2D]', field_type=FieldType.STAGGERED_FLUX)
+    r = ps.fields("r(2) : double[2D]", field_type=FieldType.STAGGERED_FLUX)
     assert r[0, 0](0) == r.staggered_access("W")
     assert -r[1, 0](0) == r.staggered_access("E")

pystencils_tests/test_field_equality.ipynb → tests/test_field_equality.ipynb

 %% Cell type:code id: tags:
 
 ``` python
 from pystencils.session import *
-from pystencils.data_types import cast_func
 ```
 
 %% Cell type:markdown id: tags:
 
 ## Test field equality behaviour
 
 %% Cell type:markdown id: tags:
 
 Fields create with same parameters are equal
 
 %% Cell type:code id: tags:
 
 ``` python
 f1 = ps.Field.create_generic('f', spatial_dimensions=2, index_dimensions=0)
 f2 = ps.Field.create_generic('f', spatial_dimensions=2, index_dimensions=0)
 
 assert f1 == f2
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 print("Field ids equal in accesses: ", id(f1.center._field) == id(f2.center._field))
 print("Field accesses equal: ", f1.center == f2.center)
 ```
 
 %% Output
 
     Field ids equal in accesses:  True
     Field accesses equal:  True
 
 %% Cell type:code id: tags:
 
 ``` python
 f1 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0)
 f2 = ps.Field.create_generic('f', spatial_dimensions=2, index_dimensions=0)
 assert f1 != f2
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 f1 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0)
 f2 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0, dtype=np.float32)
 assert f1 != f2
 ```
 
 %% Cell type:markdown id: tags:
 
 Properties of fields:
 - `field_type`: enum distinguishing normal from index or buffer fields
 - `_dtype`: data type of field elements
 - `_layout`: tuple indicating the memory linearization order
 - `shape`: size of field for each dimension
 - `strides`: number of elements to jump over to increase coordinate of this dimension by one
 - `latex_name`: optional display name when field is printed as latex
 
 Equality compare of fields:
 - field has `__eq__` and ``__hash__`` overridden
 - all parameter but `latex_name` are considered for equality
 
 %% Cell type:markdown id: tags:
 
 ## Test field access equality behaviour
 
 %% Cell type:code id: tags:
 
 ``` python
 f1 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0)
 f2 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0)
 assert f1.center == f2.center
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 f1 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0)
 f2 = ps.Field.create_generic('f', spatial_dimensions=1, index_dimensions=0, dtype=np.float32)
 assert f1.center != f2.center
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 def print_field_accesses_debug(expr):
     from pystencils import Field
     fas = list(expr.atoms(Field.Access))
     fields = list({e.field for e in fas})
     print("Field Accesses:")
     for fa in fas:
         print(f"   - {fa}, hash {hash(fa)}, offsets {fa._offsets}, index {fa._index}, {fa._hashable_content()}")
     print("")
     for i in range(len(fas)):
         for j in range(len(fas)):
             if not i < j:
                 continue
             print( f"   -> {i},{j}  {fas[i]} == {fas[j]}: {fas[i] == {fas[j]}}")
 
     print("Fields")
     for f in fields:
         print(f"  - {f}, {id(f)}, shape {f.shape}, strides {f.strides}, {f._dtype}, {f.field_type}, layout {f.layout}")
     print("")
     for i in range(len(fields)):
         for j in range(len(fields)):
             if not i < j:
                 continue
             print(f"  - {fields[i]} == {fields[j]}: {fields[i] == fields[j]}, ids equal {id(fields[i])==id(fields[j])}, hash equal {hash(fields[i])==hash(fields[j])}")
 ```
 
 %% Cell type:code id: tags:
 
 ``` python
 print_field_accesses_debug(f1.center * f2.center)
 ```
 
 %% Output
 
     Field Accesses:
-       - f[0], hash -3276894289571194847, offsets (0,), index (), (('f_C', ('commutative', True)), ((0,), (_size_f_0,), (_stride_f_0,), 3146377891102027609, <FieldType.GENERIC: 0>, 'f', None), 0)
-       - f[0], hash -1516451775709390846, offsets (0,), index (), (('f_C', ('commutative', True)), ((0,), (_size_f_0,), (_stride_f_0,), -1421177580377734245, <FieldType.GENERIC: 0>, 'f', None), 0)
+       - f[0], hash -8859424145258271267, offsets (0,), index (), ((('f_C', ('commutative', True), ('complex', True), ('extended_real', True), ('finite', True), ('hermitian', True), ('imaginary', False), ('infinite', False), ('real', True)), 2305067722319023373), ((0,), (_size_f_0,), (_stride_f_0,), <FieldType.GENERIC: 0>, 'f', None, double), 0)
+       - f[0], hash -6454673863007224785, offsets (0,), index (), ((('f_C', ('commutative', True), ('complex', True), ('extended_real', True), ('finite', True), ('hermitian', True), ('imaginary', False), ('infinite', False), ('real', True)), 4093629613697528859), ((0,), (_size_f_0,), (_stride_f_0,), <FieldType.GENERIC: 0>, 'f', None, float), 0)
     
        -> 0,1  f[0] == f[0]: False
     Fields
-      - f, 140548694371968, shape (_size_f_0,), strides (_stride_f_0,), double, FieldType.GENERIC, layout (0,)
-      - f, 140548693963104, shape (_size_f_0,), strides (_stride_f_0,), float, FieldType.GENERIC, layout (0,)
+      - f, 4881406800, shape (_size_f_0,), strides (_stride_f_0,), double, FieldType.GENERIC, layout (0,)
+      - f, 4881445024, shape (_size_f_0,), strides (_stride_f_0,), float, FieldType.GENERIC, layout (0,)
     
       - f == f: False, ids equal False, hash equal False

pystencils_tests/test_floor_ceil_int_optimization.py → tests/test_floor_ceil_int_optimization.py

@@ -11,7 +11,7 @@
 import sympy as sp
 
 import pystencils
-from pystencils.data_types import create_type
+from pystencils.typing import create_type
 
 
 def test_floor_ceil_int_optimization():

pystencils_tests/test_fvm.py → tests/test_fvm.py

@@ -4,6 +4,8 @@ import numpy as np
 import pytest
 from itertools import product
 from pystencils.rng import random_symbol
+from pystencils.astnodes import SympyAssignment
+from pystencils.node_collection import NodeCollection
 
 
 def advection_diffusion(dim: int):
@@ -315,7 +317,6 @@ def diffusion_reaction(fluctuations: bool):
             fluct = sp.sqrt(2 * dens * D) * sp.sqrt(1 / length) * stencil_factor
             # add fluctuations
             fluct *= 2 * (next(rng_symbol_gen) - 0.5) * sp.sqrt(3)
-            
             flux.main_assignments[i] = ps.Assignment(flux.main_assignments[i].lhs, flux.main_assignments[i].rhs + fluct)
         
         # Add the folding to the flux, so that the random numbers persist through the ghostlayers.
@@ -323,26 +324,30 @@ def diffusion_reaction(fluctuations: bool):
                 ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
         flux.subs(fold)
 
-    r_flux = ps.AssignmentCollection([ps.Assignment(j_fields[i].center, 0) for i in range(species)])
+    r_flux = NodeCollection([SympyAssignment(j_fields[i].center, 0) for i in range(species)])
     reaction = r_rate_const
     for i in range(species):
         reaction *= sp.Pow(n_fields[i].center, r_order[i])
-    if(fluctuations):
-        rng_symbol_gen = random_symbol(r_flux.subexpressions, dim=dh.dim)
+    new_assignments = []
+    if fluctuations:
+        rng_symbol_gen = random_symbol(new_assignments, dim=dh.dim)
         reaction_fluctuations = sp.sqrt(sp.Abs(reaction)) * 2 * (next(rng_symbol_gen) - 0.5) * sp.sqrt(3)
         reaction_fluctuations *= sp.Min(1, sp.Abs(reaction**2))
     else:
         reaction_fluctuations = 0.0
     for i in range(species):
-        r_flux.main_assignments[i] = ps.Assignment(
+        r_flux.all_assignments[i] = SympyAssignment(
             r_flux_fields[i].center, (reaction + reaction_fluctuations) * r_coefs[i])
+    [r_flux.all_assignments.insert(0, new) for new in new_assignments]
 
-    continuity_assignments.append(ps.Assignment(n_fields[0].center, n_fields[0].center + r_flux_fields[0].center))
+    continuity_assignments = [SympyAssignment(*assignment.args) for assignment in continuity_assignments]
+    continuity_assignments.append(SympyAssignment(n_fields[0].center, n_fields[0].center + r_flux_fields[0].center))
 
     flux_kernel = ps.create_staggered_kernel(flux).compile()
     reaction_kernel = ps.create_kernel(r_flux).compile()
 
-    pde_kernel = ps.create_kernel(continuity_assignments).compile()
+    config = ps.CreateKernelConfig(allow_double_writes=True)
+    pde_kernel = ps.create_kernel(continuity_assignments, config=config).compile()
 
     sync_conc = dh.synchronization_function([n_fields[0].name, n_fields[1].name])
 
@@ -412,7 +417,7 @@ advection_diffusion_fluctuations.runners = {}
 @pytest.mark.parametrize("density", [27.0, 56.5])
 @pytest.mark.parametrize("fluctuations", [False, True])
 @pytest.mark.longrun
-def test_diffusion_reaction(density, velocity, fluctuations):
+def test_diffusion_reaction(fluctuations, density, velocity):
     diffusion_reaction.runner = diffusion_reaction(fluctuations)
     diffusion_reaction.runner(density, velocity)
 
@@ -617,3 +622,20 @@ def test_source_stencil(stencil):
             assert len(diff.atoms(ps.field.Field.Access)) == 1
         else:
             assert len(diff.atoms(ps.field.Field.Access)) == 2
+
+
+def test_fvm_staggered_simplification():
+    D = sp.Symbol("D")
+    data_type = "float64"
+
+    c = ps.fields(f"c: {data_type}[2D]", layout='fzyx')
+    j = ps.fields(f"j(2): {data_type}[2D]", layout='fzyx', field_type=ps.FieldType.STAGGERED_FLUX)
+
+    grad_c = sp.Matrix([ps.fd.diff(c, i) for i in range(c.spatial_dimensions)])
+
+    ek = ps.fd.FVM1stOrder(c, flux=-D * grad_c)
+
+    ast = ps.create_staggered_kernel(ek.discrete_flux(j))
+    code = ps.get_code_str(ast)
+
+    assert '_size_c_0 - 1 < _size_c_0 - 1' not in code

pystencils_tests/test_global_definitions.py → tests/test_global_definitions.py

@@ -2,7 +2,7 @@ import sympy
 
 import pystencils.astnodes
 from pystencils.backends.cbackend import CBackend
-from pystencils.data_types import TypedSymbol
+from pystencils.typing import TypedSymbol
 
 
 class BogusDeclaration(pystencils.astnodes.Node):
@@ -95,7 +95,7 @@ def test_global_definitions_with_global_symbol():
     z, x, y = pystencils.fields("z, y, x: [2d]")
 
     normal_assignments = pystencils.AssignmentCollection([pystencils.Assignment(
-        z[0, 0], x[0, 0] * sympy.log(x[0, 0] * y[0, 0]))], [])
+        z[0, 0], x[0, 0] * x[0, 0] * y[0, 0])], [])
 
     ast = pystencils.create_kernel(normal_assignments)
     print(pystencils.show_code(ast))
@@ -115,7 +115,7 @@ def test_global_definitions_without_global_symbol():
     z, x, y = pystencils.fields("z, y, x: [2d]")
 
     normal_assignments = pystencils.AssignmentCollection([pystencils.Assignment(
-        z[0, 0], x[0, 0] * sympy.log(x[0, 0] * y[0, 0]))], [])
+        z[0, 0], x[0, 0] * x[0, 0] * y[0, 0])], [])
 
     ast = pystencils.create_kernel(normal_assignments)
     print(pystencils.show_code(ast))

pystencils_tests/test_cudagpu.py → tests/test_gpu.py

+import pytest
+
 import numpy as np
-import pycuda.autoinit
-import pycuda.gpuarray as gpuarray
 import sympy as sp
+import math
 from scipy.ndimage import convolve
 
-from pystencils import Assignment, Field, fields
-from pystencils.gpucuda import BlockIndexing, create_cuda_kernel, make_python_function
-from pystencils.gpucuda.indexing import LineIndexing
+from pystencils import Assignment, Field, fields, CreateKernelConfig, create_kernel, Target, get_code_str
+from pystencils.gpu import BlockIndexing
 from pystencils.simp import sympy_cse_on_assignment_list
-from pystencils.slicing import add_ghost_layers, make_slice, remove_ghost_layers
+from pystencils.slicing import add_ghost_layers, make_slice, remove_ghost_layers, normalize_slice
+
+try:
+    import cupy as cp
+    device_numbers = range(cp.cuda.runtime.getDeviceCount())
+except ImportError:
+    device_numbers = []
+    cp = None
 
 
 def test_averaging_kernel():
+    pytest.importorskip('cupy')
     size = (40, 55)
     src_arr = np.random.rand(*size)
     src_arr = add_ghost_layers(src_arr)
@@ -22,13 +30,14 @@ def test_averaging_kernel():
     update_rule = Assignment(dst_field[0, 0],
                              (src_field[0, 1] + src_field[0, -1] + src_field[1, 0] + src_field[-1, 0]) / 4)
 
-    ast = create_cuda_kernel(sympy_cse_on_assignment_list([update_rule]))
-    kernel = make_python_function(ast)
+    config = CreateKernelConfig(target=Target.GPU)
+    ast = create_kernel(sympy_cse_on_assignment_list([update_rule]), config=config)
+    kernel = ast.compile()
 
-    gpu_src_arr = gpuarray.to_gpu(src_arr)
-    gpu_dst_arr = gpuarray.to_gpu(dst_arr)
+    gpu_src_arr = cp.asarray(src_arr)
+    gpu_dst_arr = cp.asarray(dst_arr)
     kernel(src=gpu_src_arr, dst=gpu_dst_arr)
-    gpu_dst_arr.get(dst_arr)
+    dst_arr = gpu_dst_arr.get()
 
     stencil = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]) / 4.0
     reference = convolve(remove_ghost_layers(src_arr), stencil, mode='constant', cval=0.0)
@@ -37,24 +46,26 @@ def test_averaging_kernel():
 
 
 def test_variable_sized_fields():
+    pytest.importorskip('cupy')
     src_field = Field.create_generic('src', spatial_dimensions=2)
     dst_field = Field.create_generic('dst', spatial_dimensions=2)
 
     update_rule = Assignment(dst_field[0, 0],
                              (src_field[0, 1] + src_field[0, -1] + src_field[1, 0] + src_field[-1, 0]) / 4)
 
-    ast = create_cuda_kernel(sympy_cse_on_assignment_list([update_rule]))
-    kernel = make_python_function(ast)
+    config = CreateKernelConfig(target=Target.GPU)
+    ast = create_kernel(sympy_cse_on_assignment_list([update_rule]), config=config)
+    kernel = ast.compile()
 
     size = (3, 3)
     src_arr = np.random.rand(*size)
     src_arr = add_ghost_layers(src_arr)
     dst_arr = np.zeros_like(src_arr)
 
-    gpu_src_arr = gpuarray.to_gpu(src_arr)
-    gpu_dst_arr = gpuarray.to_gpu(dst_arr)
+    gpu_src_arr = cp.asarray(src_arr)
+    gpu_dst_arr = cp.asarray(dst_arr)
     kernel(src=gpu_src_arr, dst=gpu_dst_arr)
-    gpu_dst_arr.get(dst_arr)
+    dst_arr = gpu_dst_arr.get()
 
     stencil = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 0]]) / 4.0
     reference = convolve(remove_ghost_layers(src_arr), stencil, mode='constant', cval=0.0)
@@ -63,6 +74,7 @@ def test_variable_sized_fields():
 
 
 def test_multiple_index_dimensions():
+    pytest.importorskip('cupy')
     """Sums along the last axis of a numpy array"""
     src_size = (7, 6, 4)
     dst_size = src_size[:2]
@@ -76,13 +88,14 @@ def test_multiple_index_dimensions():
     update_rule = Assignment(dst_field[0, 0],
                              sum([src_field[offset[0], offset[1]](i) for i in range(src_size[-1])]))
 
-    ast = create_cuda_kernel([update_rule])
-    kernel = make_python_function(ast)
+    config = CreateKernelConfig(target=Target.GPU)
+    ast = create_kernel([update_rule], config=config)
+    kernel = ast.compile()
 
-    gpu_src_arr = gpuarray.to_gpu(src_arr)
-    gpu_dst_arr = gpuarray.to_gpu(dst_arr)
+    gpu_src_arr = cp.asarray(src_arr)
+    gpu_dst_arr = cp.asarray(dst_arr)
     kernel(src=gpu_src_arr, dst=gpu_dst_arr)
-    gpu_dst_arr.get(dst_arr)
+    dst_arr = gpu_dst_arr.get()
 
     reference = np.zeros_like(dst_arr)
     gl = np.max(np.abs(np.array(offset, dtype=int)))
@@ -94,6 +107,7 @@ def test_multiple_index_dimensions():
 
 
 def test_ghost_layer():
+    pytest.importorskip('cupy')
     size = (6, 5)
     src_arr = np.ones(size)
     dst_arr = np.zeros_like(src_arr)
@@ -102,13 +116,15 @@ def test_ghost_layer():
 
     update_rule = Assignment(dst_field[0, 0], src_field[0, 0])
     ghost_layers = [(1, 2), (2, 1)]
-    ast = create_cuda_kernel([update_rule], ghost_layers=ghost_layers, indexing_creator=LineIndexing)
-    kernel = make_python_function(ast)
 
-    gpu_src_arr = gpuarray.to_gpu(src_arr)
-    gpu_dst_arr = gpuarray.to_gpu(dst_arr)
+    config = CreateKernelConfig(target=Target.GPU, ghost_layers=ghost_layers, gpu_indexing="line")
+    ast = create_kernel(sympy_cse_on_assignment_list([update_rule]), config=config)
+    kernel = ast.compile()
+
+    gpu_src_arr = cp.asarray(src_arr)
+    gpu_dst_arr = cp.asarray(dst_arr)
     kernel(src=gpu_src_arr, dst=gpu_dst_arr)
-    gpu_dst_arr.get(dst_arr)
+    dst_arr = gpu_dst_arr.get()
 
     reference = np.zeros_like(src_arr)
     reference[ghost_layers[0][0]:-ghost_layers[0][1], ghost_layers[1][0]:-ghost_layers[1][1]] = 1
@@ -116,25 +132,29 @@ def test_ghost_layer():
 
 
 def test_setting_value():
+    pytest.importorskip('cupy')
     arr_cpu = np.arange(25, dtype=np.float64).reshape(5, 5)
-    arr_gpu = gpuarray.to_gpu(arr_cpu)
+    arr_gpu = cp.asarray(arr_cpu)
 
     iteration_slice = make_slice[:, :]
     f = Field.create_generic("f", 2)
     update_rule = [Assignment(f(0), sp.Symbol("value"))]
-    ast = create_cuda_kernel(update_rule, iteration_slice=iteration_slice, indexing_creator=LineIndexing)
 
-    kernel = make_python_function(ast)
+    config = CreateKernelConfig(target=Target.GPU, gpu_indexing="line", iteration_slice=iteration_slice)
+    ast = create_kernel(sympy_cse_on_assignment_list(update_rule), config=config)
+    kernel = ast.compile()
+
     kernel(f=arr_gpu, value=np.float64(42.0))
     np.testing.assert_equal(arr_gpu.get(), np.ones((5, 5)) * 42.0)
 
 
 def test_periodicity():
-    from pystencils.gpucuda.periodicity import get_periodic_boundary_functor as periodic_gpu
+    pytest.importorskip('cupy')
+    from pystencils.gpu.periodicity import get_periodic_boundary_functor as periodic_gpu
     from pystencils.slicing import get_periodic_boundary_functor as periodic_cpu
 
     arr_cpu = np.arange(50, dtype=np.float64).reshape(5, 5, 2)
-    arr_gpu = gpuarray.to_gpu(arr_cpu)
+    arr_gpu = cp.asarray(arr_cpu)
 
     periodicity_stencil = [(1, 0), (-1, 0), (1, 1)]
     periodic_gpu_kernel = periodic_gpu(periodicity_stencil, (5, 5), 1, 2)
@@ -143,23 +163,95 @@ def test_periodicity():
     cpu_result = np.copy(arr_cpu)
     periodic_cpu_kernel(cpu_result)
 
-    gpu_result = np.copy(arr_cpu)
     periodic_gpu_kernel(pdfs=arr_gpu)
-    arr_gpu.get(gpu_result)
+    gpu_result = arr_gpu.get()
     np.testing.assert_equal(cpu_result, gpu_result)
 
 
-def test_block_indexing():
+@pytest.mark.parametrize("device_number", device_numbers)
+def test_block_indexing(device_number):
+    pytest.importorskip('cupy')
     f = fields("f: [3D]")
-    bi = BlockIndexing(f, make_slice[:, :, :], block_size=(16, 8, 2), permute_block_size_dependent_on_layout=False)
+    s = normalize_slice(make_slice[:, :, :], f.spatial_shape)
+    bi = BlockIndexing(s, f.layout, block_size=(16, 8, 2),
+                       permute_block_size_dependent_on_layout=False)
     assert bi.call_parameters((3, 2, 32))['block'] == (3, 2, 32)
     assert bi.call_parameters((32, 2, 32))['block'] == (16, 2, 8)
 
-    bi = BlockIndexing(f, make_slice[:, :, :], block_size=(32, 1, 1), permute_block_size_dependent_on_layout=False)
+    bi = BlockIndexing(s, f.layout, block_size=(32, 1, 1),
+                       permute_block_size_dependent_on_layout=False)
     assert bi.call_parameters((1, 16, 16))['block'] == (1, 16, 2)
 
-    bi = BlockIndexing(f, make_slice[:, :, :], block_size=(16, 8, 2), maximum_block_size="auto")
+    bi = BlockIndexing(s, f.layout, block_size=(16, 8, 2),
+                       maximum_block_size="auto", device_number=device_number)
+
     # This function should be used if number of needed registers is known. Can be determined with func.num_regs
-    blocks = bi.limit_block_size_by_register_restriction([1024, 1024, 1], 1000)
+    registers_per_thread = 1000
+    blocks = bi.limit_block_size_by_register_restriction([1024, 1024, 1], registers_per_thread)
+
+    if cp.cuda.runtime.is_hip:
+        max_registers_per_block = cp.cuda.runtime.deviceGetAttribute(71, device_number)
+    else:
+        device = cp.cuda.Device(device_number)
+        da = device.attributes
+        max_registers_per_block = da.get("MaxRegistersPerBlock")
+
+    assert np.prod(blocks) * registers_per_thread < max_registers_per_block
 
-    assert sum(blocks) < sum([1024, 1024, 1])
+
+@pytest.mark.parametrize('gpu_indexing', ("block", "line"))
+@pytest.mark.parametrize('layout', ("C", "F"))
+@pytest.mark.parametrize('shape', ((5, 5, 5, 5), (3, 17, 387, 4), (23, 44, 21, 11)))
+def test_four_dimensional_kernel(gpu_indexing, layout, shape):
+    pytest.importorskip('cupy')
+    n_elements = np.prod(shape)
+
+    arr_cpu = np.arange(n_elements, dtype=np.float64).reshape(shape, order=layout)
+    arr_gpu = cp.asarray(arr_cpu)
+
+    iteration_slice = make_slice[:, :, :, :]
+    f = Field.create_from_numpy_array("f", arr_cpu)
+    update_rule = [Assignment(f.center, sp.Symbol("value"))]
+
+    config = CreateKernelConfig(target=Target.GPU, gpu_indexing=gpu_indexing, iteration_slice=iteration_slice)
+    ast = create_kernel(update_rule, config=config)
+    kernel = ast.compile()
+
+    kernel(f=arr_gpu, value=np.float64(42.0))
+    np.testing.assert_equal(arr_gpu.get(), np.ones(shape) * 42.0)
+
+
+@pytest.mark.parametrize('start', (1, 5))
+@pytest.mark.parametrize('end', (-1, -2, -3, -4))
+@pytest.mark.parametrize('step', (1, 2, 3, 4))
+@pytest.mark.parametrize('shape', ([55, 60], [77, 101, 80], [44, 64, 66]))
+def test_guards_with_iteration_slices(start, end, step, shape):
+    iter_slice = tuple([slice(start, end, step)] * len(shape))
+
+    kernel_config_gpu = CreateKernelConfig(target=Target.GPU, iteration_slice=iter_slice)
+    field_1 = fields(f"f(1) : double{list(shape)}")
+    assignment = Assignment(field_1.center, 1)
+    ast = create_kernel(assignment, config=kernel_config_gpu)
+    code_str = get_code_str(ast)
+
+    test_strings = list()
+    iteration_ranges = list()
+    for i, s in enumerate(iter_slice):
+        e = ((shape[i] + end) - s.start) / s.step
+        e = math.ceil(e) + s.start
+        test_strings.append(f"{s.start} < {e}")
+
+        a = s.start
+        counter = 0
+        while a < e:
+            a += 1
+            counter += 1
+        iteration_ranges.append(counter)
+
+    # check if the expected if statement is in the GPU code
+    for s in test_strings:
+        assert s in code_str
+
+    # check if these bounds lead to same lengths as the range function would produce
+    for i in range(len(iter_slice)):
+        assert iteration_ranges[i] == len(range(iter_slice[i].start, shape[i] + end, iter_slice[i].step))

tests/test_half_precision.py

+import pytest
+import platform
+
+import numpy as np
+import pystencils as ps
+
+
+@pytest.mark.parametrize('target', (ps.Target.CPU, ps.Target.GPU))
+def test_half_precison(target):
+    if target == ps.Target.CPU:
+        if not platform.machine() in ['arm64', 'aarch64']:
+            pytest.xfail("skipping half precision test on non arm platform")
+
+        if 'clang' not in ps.cpu.cpujit.get_compiler_config()['command']:
+            pytest.xfail("skipping half precision because clang compiler is not used")
+
+    if target == ps.Target.GPU:
+        pytest.importorskip("cupy")
+
+    dh = ps.create_data_handling(domain_size=(10, 10), default_target=target)
+
+    f1 = dh.add_array("f1", values_per_cell=1, dtype=np.float16)
+    dh.fill("f1", 1.0, ghost_layers=True)
+    f2 = dh.add_array("f2", values_per_cell=1, dtype=np.float16)
+    dh.fill("f2", 2.0, ghost_layers=True)
+
+    f3 = dh.add_array("f3", values_per_cell=1, dtype=np.float16)
+    dh.fill("f3", 0.0, ghost_layers=True)
+
+    up = ps.Assignment(f3.center, f1.center + 2.1 * f2.center)
+
+    config = ps.CreateKernelConfig(target=dh.default_target, default_number_float=np.float32)
+    ast = ps.create_kernel(up, config=config)
+
+    kernel = ast.compile()
+
+    dh.run_kernel(kernel)
+    dh.all_to_cpu()
+
+    assert np.all(dh.cpu_arrays[f3.name] == 5.2)
+    assert dh.cpu_arrays[f3.name].dtype == np.float16

tests/test_indexed_kernels.py

+import sympy as sp
+import numpy as np
+import pytest
+
+import pystencils as ps
+from pystencils import Assignment, Field, CreateKernelConfig, create_kernel, Target
+from pystencils.transformations import filtered_tree_iteration
+from pystencils.typing import BasicType, FieldPointerSymbol, PointerType, TypedSymbol
+
+
+@pytest.mark.parametrize('target', [ps.Target.CPU, ps.Target.GPU])
+def test_indexed_kernel(target):
+    if target == Target.GPU:
+        pytest.importorskip("cupy")
+        import cupy as cp
+
+    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, index_fields=[indexed_field])
+    ast = create_kernel([update_rule], config=config)
+    kernel = ast.compile()
+
+    if target == Target.CPU:
+        kernel(f=arr, index=index_arr)
+    else:
+        gpu_arr = cp.asarray(arr)
+        gpu_index_arr = cp.ndarray(index_arr.shape, dtype=index_arr.dtype)
+        gpu_index_arr.set(index_arr)
+        kernel(f=gpu_arr, index=gpu_index_arr)
+        arr = gpu_arr.get()
+    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)
+
+
+@pytest.mark.parametrize('index_size', ("fixed", "variable"))
+@pytest.mark.parametrize('array_size', ("3D", "2D", "10, 12", "13, 17, 19"))
+@pytest.mark.parametrize('target', (ps.Target.CPU, ps.Target.GPU))
+@pytest.mark.parametrize('dtype', ("float64", "float32"))
+def test_indexed_domain_kernel(index_size, array_size, target, dtype):
+    dtype = BasicType(dtype)
+
+    f = ps.fields(f'f(1): {dtype.numpy_dtype.name}[{array_size}]')
+    g = ps.fields(f'g(1): {dtype.numpy_dtype.name}[{array_size}]')
+
+    index = TypedSymbol("index", dtype=BasicType(np.int16))
+    if index_size == "variable":
+        index_src = TypedSymbol("_size_src", dtype=BasicType(np.int16))
+        index_dst = TypedSymbol("_size_dst", dtype=BasicType(np.int16))
+    else:
+        index_src = 16
+        index_dst = 16
+    pointer_type = PointerType(dtype, const=False, restrict=True, double_pointer=True)
+    const_pointer_type = PointerType(dtype, const=True, restrict=True, double_pointer=True)
+
+    src = sp.IndexedBase(TypedSymbol(f"_data_{f.name}", dtype=const_pointer_type), shape=index_src)
+    dst = sp.IndexedBase(TypedSymbol(f"_data_{g.name}", dtype=pointer_type), shape=index_dst)
+
+    update_rule = [ps.Assignment(FieldPointerSymbol("f", dtype, const=True), src[index + 1]),
+                   ps.Assignment(FieldPointerSymbol("g", dtype, const=False), dst[index + 1]),
+                   ps.Assignment(g.center, f.center)]
+
+    ast = ps.create_kernel(update_rule, target=target)
+
+    code = ps.get_code_str(ast)
+    assert f"const {dtype.c_name} * RESTRICT _data_f = (({dtype.c_name} * RESTRICT const)(_data_f[index + 1]));" in code
+    assert f"{dtype.c_name} * RESTRICT  _data_g = (({dtype.c_name} * RESTRICT )(_data_g[index + 1]));" in code
+
+    if target == Target.CPU:
+        assert code.count("for") == f.spatial_dimensions + 1
+

pystencils_tests/test_json_backend.py → tests/test_json_backend.py

@@ -21,13 +21,15 @@ def test_json_backend():
     a = sympy.Symbol('a')
 
     assignments = pystencils.AssignmentCollection({
-        z[0, 0]: x[0, 0] * sympy.log(a * x[0, 0] * y[0, 0])
+        z[0, 0]: x[0, 0] * a * x[0, 0] * y[0, 0]
     })
 
     ast = pystencils.create_kernel(assignments)
 
-    print(print_json(ast))
-    print(print_yaml(ast))
+    pj = print_json(ast)
+    # print(pj)
+    py = print_yaml(ast)
+    # print(py)
 
     temp_dir = tempfile.TemporaryDirectory()
     write_json(temp_dir.name + '/test.json', ast)

tests/test_json_serializer.py

+"""
+Test the pystencils-specific JSON encoder and serializer as used in the Database class.
+"""
+
+import numpy as np
+import tempfile
+
+from pystencils.config import CreateKernelConfig
+from pystencils import Target, Field
+from pystencils.runhelper.db import Database, PystencilsJsonSerializer
+
+
+def test_json_serializer():
+
+    dtype = np.float32
+
+    index_arr = np.zeros((3,), dtype=dtype)
+    indexed_field = Field.create_from_numpy_array('index', index_arr)
+
+    # create pystencils config
+    config = CreateKernelConfig(target=Target.CPU, function_name='dummy_config', data_type=dtype,
+                                index_fields=[indexed_field])
+
+    # create dummy database
+    temp_dir = tempfile.TemporaryDirectory()
+    db = Database(file=temp_dir.name, serializer_info=('pystencils_serializer', PystencilsJsonSerializer))
+
+    db.save(params={'config': config}, result={'test': 'dummy'})

tests/test_jupyter_extensions.ipynb

+%% Cell type:code id: tags:
+
+``` python
+from pystencils.session import *
+```
+
+%% Cell type:code id: tags:
+
+``` python
+dh = ps.create_data_handling(domain_size=(256, 256), periodicity=True)
+c_field = dh.add_array('c')
+dh.fill("c", 0.0, ghost_layers=True)
+```
+
+%% Cell type:code id: tags:
+
+``` python
+for x in range(129):
+    for y in range(258):
+        dh.cpu_arrays['c'][x, y] = 1.0
+```
+
+%% Cell type:code id: tags:
+
+``` python
+plt.scalar_field(dh.cpu_arrays["c"])
+```
+
+%% Output
+
+    <matplotlib.image.AxesImage at 0x117081c10>
+
+
+
+%% Cell type:code id: tags:
+
+``` python
+ur = ps.Assignment(c_field[0, 0], c_field[1, 0])
+config = ps.CreateKernelConfig(target=dh.default_target, cpu_openmp=False, skip_independence_check=True)
+ast = ps.create_kernel(ur, config=config)
+kernel = ast.compile()
+```
+
+%% Cell type:code id: tags:
+
+``` python
+c_sync = dh.synchronization_function_cpu(['c'])
+```
+
+%% Cell type:code id: tags:
+
+``` python
+def timeloop(steps=10):
+    for i in range(steps):
+        c_sync()
+        dh.run_kernel(kernel)
+    return dh.gather_array('c')
+```
+
+%% Cell type:code id: tags:
+
+``` python
+ps.jupyter.set_display_mode('video')
+```
+
+%% Cell type:code id: tags:
+
+``` python
+ani = ps.plot.scalar_field_animation(timeloop, rescale=True, frames=12)
+ps.jupyter.display_animation(ani)
+```
+
+%% Output
+
+    <IPython.core.display.HTML object>
+
+%% Cell type:code id: tags:
+
+``` python
+ps.jupyter.set_display_mode('image_update')
+```
+
+%% Cell type:code id: tags:
+
+``` python
+ani = ps.plot.scalar_field_animation(timeloop, rescale=True, frames=12)
+ps.jupyter.display_animation(ani)
+```
+
+%% Output
+
+
+
+%% Cell type:code id: tags:
+
+``` python
+def grid_update_function(image):
+    for i in range(40):
+        c_sync()
+        dh.run_kernel(kernel)
+    return dh.gather_array('c')
+```
+
+%% Cell type:code id: tags:
+
+``` python
+animation = ps.jupyter.make_imshow_animation(dh.cpu_arrays["c"], grid_update_function, frames=300)
+```
+
+%% Output
+
+
+
+%% Cell type:code id: tags:
+
+``` python
+ps.jupyter.set_display_mode("video")
+ps.jupyter.set_display_mode("window")
+ps.jupyter.set_display_mode("image_update")
+ps.jupyter.activate_ipython()
+```
+
+%% Output
+
+    ---------------------------------------------------------------------------
+    ImportError                               Traceback (most recent call last)
+Cell     In[14], line 2
+          1 ps.jupyter.set_display_mode("video")
+    ----> 2 ps.jupyter.set_display_mode("window")
+          3 ps.jupyter.set_display_mode("image_update")
+          4 ps.jupyter.activate_ipython()
+File     ~/pystencils/pystencils/src/pystencils/jupyter.py:115, in set_display_mode(mode)
+        113     display_animation_func = display_as_html_video
+        114 elif animation_display_mode == 'window':
+    --> 115     ipython.magic("matplotlib qt")
+        116     display_animation_func = display_in_extra_window
+        117 elif animation_display_mode == 'image_update':
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/IPython/core/interactiveshell.py:2539, in InteractiveShell.magic(self, arg_s)
+       2537 magic_name, _, magic_arg_s = arg_s.partition(' ')
+       2538 magic_name = magic_name.lstrip(prefilter.ESC_MAGIC)
+    -> 2539 return self.run_line_magic(magic_name, magic_arg_s, _stack_depth=2)
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/IPython/core/interactiveshell.py:2417, in InteractiveShell.run_line_magic(self, magic_name, line, _stack_depth)
+       2415     kwargs['local_ns'] = self.get_local_scope(stack_depth)
+       2416 with self.builtin_trap:
+    -> 2417     result = fn(*args, **kwargs)
+       2419 # The code below prevents the output from being displayed
+       2420 # when using magics with decodator @output_can_be_silenced
+       2421 # when the last Python token in the expression is a ';'.
+       2422 if getattr(fn, magic.MAGIC_OUTPUT_CAN_BE_SILENCED, False):
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/IPython/core/magics/pylab.py:99, in PylabMagics.matplotlib(self, line)
+         97     print("Available matplotlib backends: %s" % backends_list)
+         98 else:
+    ---> 99     gui, backend = self.shell.enable_matplotlib(args.gui.lower() if isinstance(args.gui, str) else args.gui)
+        100     self._show_matplotlib_backend(args.gui, backend)
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/IPython/core/interactiveshell.py:3603, in InteractiveShell.enable_matplotlib(self, gui)
+       3599         print('Warning: Cannot change to a different GUI toolkit: %s.'
+       3600                 ' Using %s instead.' % (gui, self.pylab_gui_select))
+       3601         gui, backend = pt.find_gui_and_backend(self.pylab_gui_select)
+    -> 3603 pt.activate_matplotlib(backend)
+       3604 configure_inline_support(self, backend)
+       3606 # Now we must activate the gui pylab wants to use, and fix %run to take
+       3607 # plot updates into account
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/IPython/core/pylabtools.py:360, in activate_matplotlib(backend)
+        355 # Due to circular imports, pyplot may be only partially initialised
+        356 # when this function runs.
+        357 # So avoid needing matplotlib attribute-lookup to access pyplot.
+        358 from matplotlib import pyplot as plt
+    --> 360 plt.switch_backend(backend)
+        362 plt.show._needmain = False
+        363 # We need to detect at runtime whether show() is called by the user.
+        364 # For this, we wrap it into a decorator which adds a 'called' flag.
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/matplotlib/pyplot.py:271, in switch_backend(newbackend)
+        268 # have to escape the switch on access logic
+        269 old_backend = dict.__getitem__(rcParams, 'backend')
+    --> 271 backend_mod = importlib.import_module(
+        272     cbook._backend_module_name(newbackend))
+        274 required_framework = _get_required_interactive_framework(backend_mod)
+        275 if required_framework is not None:
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/importlib/__init__.py:126, in import_module(name, package)
+        124             break
+        125         level += 1
+    --> 126 return _bootstrap._gcd_import(name[level:], package, level)
+File     <frozen importlib._bootstrap>:1204, in _gcd_import(name, package, level)
+File     <frozen importlib._bootstrap>:1176, in _find_and_load(name, import_)
+File     <frozen importlib._bootstrap>:1147, in _find_and_load_unlocked(name, import_)
+File     <frozen importlib._bootstrap>:690, in _load_unlocked(spec)
+File     <frozen importlib._bootstrap_external>:940, in exec_module(self, module)
+File     <frozen importlib._bootstrap>:241, in _call_with_frames_removed(f, *args, **kwds)
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/matplotlib/backends/backend_qt5agg.py:7
+          4 from .. import backends
+          6 backends._QT_FORCE_QT5_BINDING = True
+    ----> 7 from .backend_qtagg import (    # noqa: F401, E402 # pylint: disable=W0611
+          8     _BackendQTAgg, FigureCanvasQTAgg, FigureManagerQT, NavigationToolbar2QT,
+          9     FigureCanvasAgg, FigureCanvasQT)
+         12 @_BackendQTAgg.export
+         13 class _BackendQT5Agg(_BackendQTAgg):
+         14     pass
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/matplotlib/backends/backend_qtagg.py:9
+          5 import ctypes
+          7 from matplotlib.transforms import Bbox
+    ----> 9 from .qt_compat import QT_API, _enum
+         10 from .backend_agg import FigureCanvasAgg
+         11 from .backend_qt import QtCore, QtGui, _BackendQT, FigureCanvasQT
+File     /opt/local/Library/Frameworks/Python.framework/Versions/3.11/lib/python3.11/site-packages/matplotlib/backends/qt_compat.py:135
+        133         break
+        134     else:
+    --> 135         raise ImportError(
+        136             "Failed to import any of the following Qt binding modules: {}"
+        137             .format(", ".join(_ETS.values())))
+        138 else:  # We should not get there.
+        139     raise AssertionError(f"Unexpected QT_API: {QT_API}")
+    ImportError: Failed to import any of the following Qt binding modules: PyQt6, PySide6, PyQt5, PySide2

tests/test_logarithm.py

+import pytest
+import numpy as np
+import sympy as sp
+
+import pystencils as ps
+
+
+@pytest.mark.parametrize('dtype', ["float64", "float32"])
+def test_log(dtype):
+    a = sp.Symbol("a")
+    x = ps.fields(f'x: {dtype}[1d]')
+
+    assignments = ps.AssignmentCollection({x.center(): sp.log(a)})
+
+    ast = ps.create_kernel(assignments)
+    code = ps.get_code_str(ast)
+    kernel = ast.compile()
+
+    # ps.show_code(ast)
+
+    if dtype == "float64":
+        assert "float" not in code
+
+    array = np.zeros((10,), dtype=dtype)
+    kernel(x=array, a=100)
+    assert np.allclose(array, 4.60517019)

pystencils_tests/test_loop_cutting.py → tests/test_loop_cutting.py

@@ -50,7 +50,9 @@ def test_staggered_iteration():
                                  sum(f[o] for o in offsets_in_plane(d, -1, dim)))
             cond = sp.And(*[conditions[i] for i in range(dim) if d != i])
             eqs.append(Conditional(cond, eq))
-        func = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)]).compile()
+        # TODO: correct type hint
+        config = ps.CreateKernelConfig(target=ps.Target.CPU, ghost_layers=[(1, 0), (1, 0), (1, 0)])
+        func = ps.create_kernel(eqs, config=config).compile()
 
         # --- Built-in optimized
         expressions = []
@@ -93,7 +95,8 @@ def test_staggered_iteration_manual():
     cond = sp.And(*[conditions2])
     eqs.append(Conditional(cond, eq))
 
-    kernel_ast = create_kernel(eqs, ghost_layers=[(1, 0), (1, 0), (1, 0)])
+    config = ps.CreateKernelConfig(target=ps.Target.CPU, ghost_layers=[(1, 0), (1, 0), (1, 0)])
+    kernel_ast = ps.create_kernel(eqs, config=config)
 
     func = make_python_function(kernel_ast)
     func(f=f_arr, s=s_arr_ref)

pystencils_tests/test_match_subs_for_assignment_collection.py → tests/test_match_subs_for_assignment_collection.py

@@ -11,12 +11,12 @@
 import sympy as sp
 
 import pystencils
-from pystencils.data_types import create_type
+from pystencils.typing import TypedSymbol, BasicType
 
 
 def test_wild_typed_symbol():
     x = pystencils.fields('x:  float32[3d]')
-    typed_symbol = pystencils.data_types.TypedSymbol('a', create_type('float64'))
+    typed_symbol = TypedSymbol('a', BasicType('float64'))
 
     assert x.center().match(sp.Wild('w1'))
     assert typed_symbol.match(sp.Wild('w1'))