tests/test_transformations.py → tests/_old/test_transformations.py

@@ -3,11 +3,11 @@ import numpy as np
 
 import pystencils as ps
 from pystencils import fields, TypedSymbol
-from pystencils.astnodes import LoopOverCoordinate, SympyAssignment
+from pystencils.sympyextensions.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,
-    iterate_loops_by_depth, split_inner_loop, loop_blocking
+    iterate_loops_by_depth, split_inner_loop
 )
 
 from pystencils.cpu import add_pragmas

tests/test_conditional_vec.py → tests/_todo/test_conditional_vec.py

@@ -3,7 +3,7 @@ import sympy as sp
 import pytest
 
 import pystencils as ps
-from pystencils.astnodes import Block, Conditional, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, Conditional, SympyAssignment
 from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets, get_vector_instruction_set
 from pystencils.enums import Target
 from pystencils.cpu.vectorization import vec_all, vec_any

tests/test_dot_printer.py → tests/_todo/test_dot_printer.py

 import pystencils as ps
 
-from pystencils.astnodes import Block, Conditional, SympyAssignment
+from pystencils.sympyextensions.astnodes import Block, Conditional, SympyAssignment
 
 
 def test_dot_print():

tests/test_fvm.py → tests/_todo/test_fvm.py

@@ -4,8 +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
+# from pystencils.sympyextensions.astnodes import SympyAssignment
+# from pystencils.node_collection import NodeCollection
 
 
 def advection_diffusion(dim: int):
@@ -182,8 +182,8 @@ def advection_diffusion_fluctuations(dim: int):
         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.
-    fold = {ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
-            ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
+    fold = {pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
+                pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
     flux.subs(fold)
 
     flux_kernel = ps.create_staggered_kernel(flux).compile()
@@ -320,8 +320,8 @@ def diffusion_reaction(fluctuations: bool):
             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.
-        fold = {ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
-                ps.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
+        fold = {pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i):
+                    pystencils.sympyextensions.astnodes.LoopOverCoordinate.get_loop_counter_symbol(i) % L[i] for i in range(len(L))}
         flux.subs(fold)
 
     r_flux = NodeCollection([SympyAssignment(j_fields[i].center, 0) for i in range(species)])

tests/test_indexed_kernels.py → tests/_todo/test_indexed_kernels.py

@@ -8,39 +8,6 @@ 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))

tests/test_loop_cutting.py → tests/_todo/test_loop_cutting.py

@@ -4,10 +4,10 @@ import sympy as sp
 import pytest
 
 import pystencils as ps
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
 from pystencils.field import Field, FieldType
-from pystencils.astnodes import Conditional, LoopOverCoordinate, SympyAssignment
-from pystencils.cpu import create_kernel, make_python_function
+from pystencils.sympyextensions.astnodes import Conditional, LoopOverCoordinate, SympyAssignment
+from pystencils.cpu import make_python_function
 from pystencils.kernelcreation import create_staggered_kernel
 from pystencils.transformations import (
     cleanup_blocks, cut_loop, move_constants_before_loop, simplify_conditionals)

tests/test_modulo.py → tests/_todo/test_modulo.py

@@ -3,7 +3,7 @@ import pytest
 import numpy as np
 import sympy as sp
 import pystencils as ps
-from pystencils.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment
+from pystencils.sympyextensions.astnodes import LoopOverCoordinate, Conditional, Block, SympyAssignment
 
 SLICE_LIST = [False,
               ps.make_slice[1:-1:2, 1:-1:2],

tests/test_random.py → tests/_todo/test_random.py

@@ -3,7 +3,7 @@ import numpy as np
 import pytest
 
 import pystencils as ps
-from pystencils.astnodes import SympyAssignment
+from pystencils.sympyextensions.astnodes import SympyAssignment
 from pystencils.node_collection import NodeCollection
 from pystencils.rng import PhiloxFourFloats, PhiloxTwoDoubles, AESNIFourFloats, AESNITwoDoubles, random_symbol
 from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets

tests/test_vectorization.py → tests/_todo/test_vectorization.py

@@ -2,11 +2,10 @@ import numpy as np
 
 import pytest
 
-import pystencils.config
 import sympy as sp
 
 import pystencils as ps
-import pystencils.astnodes as ast
+import pystencils.sympyextensions.astnodes as ast
 from pystencils.backends.simd_instruction_sets import get_supported_instruction_sets, get_vector_instruction_set
 from pystencils.cpu.vectorization import vectorize
 from pystencils.enums import Target
@@ -141,7 +140,7 @@ def test_aligned_and_nt_stores(openmp, instruction_set=instruction_set):
            'assume_inner_stride_one': True}
     update_rule = [ps.Assignment(f.center(), 0.25 * (g[-1, 0] + g[1, 0] + g[0, -1] + g[0, 1]))]
     # Without the base pointer spec, the inner store is not aligned
-    config = pystencils.config.CreateKernelConfig(target=dh.default_target, cpu_vectorize_info=opt, cpu_openmp=openmp)
+    config = ps.CreateKernelConfig(target=dh.default_target, cpu_vectorize_info=opt, cpu_openmp=openmp)
     ast = ps.create_kernel(update_rule, config=config)
     if instruction_set in ['sse'] or instruction_set.startswith('avx'):
         assert 'stream' in ast.instruction_set
@@ -166,7 +165,7 @@ def test_nt_stores_symbolic_size(instruction_set=instruction_set):
     update_rule = [ps.Assignment(f.center(), 0.0), ps.Assignment(g.center(), 0.0)]
     opt = {'instruction_set': instruction_set, 'assume_aligned': True, 'nontemporal': True,
            'assume_inner_stride_one': True}
-    config = pystencils.config.CreateKernelConfig(target=Target.CPU, cpu_vectorize_info=opt)
+    config = ps.CreateKernelConfig(target=Target.CPU, cpu_vectorize_info=opt)
     ast = ps.create_kernel(update_rule, config=config)
     # ps.show_code(ast)
     ast.compile()
@@ -187,7 +186,7 @@ def test_inplace_update(instruction_set=instruction_set):
         f1 @= 2 * s.tmp0
         f2 @= 2 * s.tmp0
 
-    config = pystencils.config.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
+    config = ps.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
     ast = ps.create_kernel(update_rule, config=config)
     kernel = ast.compile()
     kernel(f=arr)
@@ -379,7 +378,7 @@ def test_issue40(*_):
     eq = [ps.Assignment(sp.Symbol('rho'), 1.0),
           ps.Assignment(src[0, 0](0), sp.Rational(4, 9) * sp.Symbol('rho'))]
 
-    config = pystencils.config.CreateKernelConfig(target=Target.CPU, cpu_vectorize_info=opt, data_type='float64')
+    config = ps.CreateKernelConfig(target=Target.CPU, cpu_vectorize_info=opt, data_type='float64')
     ast = ps.create_kernel(eq, config=config)
 
     code = ps.get_code_str(ast)

tests/test_vectorization_specific.py → tests/_todo/test_vectorization_specific.py

@@ -2,7 +2,6 @@ import pytest
 
 import numpy as np
 
-import pystencils.config
 import sympy as sp
 
 import pystencils as ps
@@ -30,7 +29,7 @@ def test_vectorisation_varying_arch(instruction_set):
         f1 @= 2 * s.tmp0
         f2 @= 2 * s.tmp0
 
-    config = pystencils.config.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
+    config = ps.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
     ast = ps.create_kernel(update_rule, config=config)
     kernel = ast.compile()
     kernel(f=arr)
@@ -49,7 +48,7 @@ def test_vectorized_abs(instruction_set, dtype):
     f, g = ps.fields(f=arr, g=arr)
     update_rule = [ps.Assignment(g.center(), sp.Abs(f.center()))]
 
-    config = pystencils.config.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
+    config = ps.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set})
     ast = ps.create_kernel(update_rule, config=config)
 
     func = ast.compile()
@@ -66,20 +65,20 @@ def test_strided(instruction_set, dtype):
     if 'storeS' not in get_vector_instruction_set(dtype, instruction_set) \
             and instruction_set not in ['avx512', 'avx512vl', 'rvv'] and not instruction_set.startswith('sve'):
         with pytest.warns(UserWarning) as warn:
-            config = pystencils.config.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set},
+            config = ps.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set},
                                                           default_number_float=dtype)
             ast = ps.create_kernel(update_rule, config=config)
             assert 'Could not vectorize loop' in warn[0].message.args[0]
     else:
         with pytest.warns(None) as warn:
-            config = pystencils.config.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set},
+            config = ps.CreateKernelConfig(cpu_vectorize_info={'instruction_set': instruction_set},
                                                           default_number_float=dtype)
             ast = ps.create_kernel(update_rule, config=config)
             assert len(warn) == 0
 
     # ps.show_code(ast)
     func = ast.compile()
-    ref_config = pystencils.config.CreateKernelConfig(default_number_float=dtype)
+    ref_config = ps.CreateKernelConfig(default_number_float=dtype)
     ref_func = ps.create_kernel(update_rule, config=ref_config).compile()
 
     # For some reason other array creations fail on the emulated ppc pipeline
@@ -115,7 +114,7 @@ def test_alignment_and_correct_ghost_layers(gl_field, gl_kernel, instruction_set
     update_rule = ps.Assignment(dst[0, 0], src[0, 0])
     opt = {'instruction_set': instruction_set, 'assume_aligned': True,
            'nontemporal': True, 'assume_inner_stride_one': True}
-    config = pystencils.config.CreateKernelConfig(target=dh.default_target,
+    config = ps.CreateKernelConfig(target=dh.default_target,
                                                   cpu_vectorize_info=opt, ghost_layers=gl_kernel)
     ast = ps.create_kernel(update_rule, config=config)
     kernel = ast.compile()
@@ -152,7 +151,7 @@ def test_vectorization_other(instruction_set, function):
 @pytest.mark.parametrize('instruction_set', supported_instruction_sets)
 @pytest.mark.parametrize('field_layout', ('fzyx', 'zyxf'))
 def test_square_root(dtype, instruction_set, field_layout):
-    config = pystencils.config.CreateKernelConfig(data_type=dtype,
+    config = ps.CreateKernelConfig(data_type=dtype,
                                                   default_number_float=dtype,
                                                   cpu_vectorize_info={'instruction_set': instruction_set,
                                                                       'assume_inner_stride_one': True,
@@ -195,7 +194,7 @@ def test_square_root_2(dtype, instruction_set, padding):
 @pytest.mark.parametrize('instruction_set', supported_instruction_sets)
 @pytest.mark.parametrize('padding', (True, False))
 def test_pow(dtype, instruction_set, padding):
-    config = pystencils.config.CreateKernelConfig(data_type=dtype,
+    config = ps.CreateKernelConfig(data_type=dtype,
                                                   default_number_float=dtype,
                                                   cpu_vectorize_info={'instruction_set': instruction_set,
                                                                       'assume_inner_stride_one': True,

tests/codegen/test_config.py

+import pytest
+
+from dataclasses import dataclass
+import numpy as np
+from pystencils.codegen.config import (
+    BasicOption,
+    Option,
+    Category,
+    ConfigBase,
+    CreateKernelConfig,
+    CpuOptions
+)
+from pystencils.field import Field, FieldType
+from pystencils.types.quick import Int, UInt, Fp, Ptr
+from pystencils.types import PsVectorType
+
+
+def test_descriptors():
+
+    @dataclass
+    class SampleCategory(ConfigBase):
+        val1: BasicOption[int] = BasicOption(2)
+        val2: Option[bool, str | bool] = Option(False)
+
+        @val2.validate
+        def validate_val2(self, v: str | bool):
+            if isinstance(v, str):
+                if v.lower() in ("off", "false", "no"):
+                    return False
+                elif v.lower() in ("on", "true", "yes"):
+                    return True
+
+                raise ValueError()
+            else:
+                return v
+
+    @dataclass
+    class SampleConfig(ConfigBase):
+        cat: Category[SampleCategory] = Category(SampleCategory())
+        val: BasicOption[str] = BasicOption("fallback")
+
+    cfg = SampleConfig()
+
+    #   Check unset and default values
+    assert cfg.val is None
+    assert cfg.get_option("val") == "fallback"
+
+    #   Check setting
+    cfg.val = "test"
+    assert cfg.val == "test"
+    assert cfg.get_option("val") == "test"
+    assert cfg.is_option_set("val")
+
+    #   Check unsetting
+    cfg.val = None
+    assert not cfg.is_option_set("val")
+    assert cfg.val is None
+
+    #   Check category
+    assert cfg.cat.val1 is None
+    assert cfg.cat.get_option("val1") == 2
+    assert cfg.cat.val2 is None
+    assert cfg.cat.get_option("val2") is False
+
+    #   Check copy on category setting
+    c = SampleCategory(32, "on")
+    cfg.cat = c
+    assert cfg.cat.val1 == 32
+    assert cfg.cat.val2 is True
+
+    assert cfg.cat is not c
+    c.val1 = 13
+    assert cfg.cat.val1 == 32
+
+    #   Check that category objects on two config objects are not the same
+    cfg1 = SampleConfig()
+    cfg2 = SampleConfig()
+
+    assert cfg1.cat is not cfg2.cat
+
+
+def test_category_init():
+    cfg1 = CreateKernelConfig()
+    cfg2 = CreateKernelConfig()
+
+    assert cfg1.cpu is not cfg2.cpu
+    assert cfg1.cpu.openmp is not cfg2.cpu.openmp
+    assert cfg1.cpu.vectorize is not cfg2.cpu.vectorize
+    assert cfg1.gpu is not cfg2.gpu
+
+
+def test_category_copy():
+    cfg = CreateKernelConfig()
+    cpu_repl = CpuOptions()
+    cpu_repl.openmp.num_threads = 42
+
+    cfg.cpu = cpu_repl
+    assert cfg.cpu.openmp.num_threads == 42
+    assert cfg.cpu is not cpu_repl
+    assert cfg.cpu.openmp is not cpu_repl.openmp
+
+
+def test_config_validation():
+    #   Check index dtype validation
+    cfg = CreateKernelConfig(index_dtype="int32")
+    assert cfg.index_dtype == Int(32)
+    cfg.index_dtype = np.uint64
+    assert cfg.index_dtype == UInt(64)
+
+    with pytest.raises(ValueError):
+        _ = CreateKernelConfig(index_dtype=np.float32)
+
+    with pytest.raises(ValueError):
+        cfg.index_dtype = "double"
+
+    #   Check default dtype validation
+    cfg = CreateKernelConfig(default_dtype="float32")
+    assert cfg.default_dtype == Fp(32)
+    cfg.default_dtype = np.int64
+    assert cfg.default_dtype == Int(64)
+
+    with pytest.raises(ValueError):
+        cfg.default_dtype = PsVectorType(Fp(64), 4)
+
+    with pytest.raises(ValueError):
+        _ = CreateKernelConfig(default_dtype=Ptr(Fp(32)))
+
+    #   Check index field validation
+    idx_field = Field.create_generic(
+        "idx", spatial_dimensions=1, field_type=FieldType.INDEXED
+    )
+    cfg.index_field = idx_field
+    assert cfg.index_field == idx_field
+
+    with pytest.raises(ValueError):
+        cfg.index_field = Field.create_generic(
+            "idx", spatial_dimensions=1, field_type=FieldType.GENERIC
+        )
+
+
+def test_override():
+    cfg1 = CreateKernelConfig()
+    cfg1.function_name = "test"
+    cfg1.cpu.openmp.schedule = "dynamic"
+    cfg1.gpu.manual_launch_grid = False
+    cfg1.allow_double_writes = True
+
+    cfg2 = CreateKernelConfig()
+    cfg2.function_name = "func"
+    cfg2.cpu.openmp.schedule = "static(5)"
+    cfg2.cpu.vectorize.lanes = 12
+    cfg2.allow_double_writes = False
+
+    cfg1.override(cfg2)
+
+    assert cfg1.function_name == "func"
+    assert cfg1.cpu.openmp.schedule == "static(5)"
+    assert cfg1.cpu.openmp.enable is None
+    assert cfg1.cpu.vectorize.lanes == 12
+    assert cfg1.cpu.vectorize.assume_aligned is None
+    assert cfg1.allow_double_writes is False

tests/test_finite_differences.py → tests/fd/test_finite_differences.py

@@ -2,7 +2,7 @@ import sympy as sp
 import pytest
 
 import pystencils as ps
-from pystencils.astnodes import LoopOverCoordinate
+from pystencils import DEFAULTS
 from pystencils.fd import diff, diffusion, Discretization2ndOrder
 from pystencils.fd.spatial import discretize_spatial, fd_stencils_isotropic, fd_stencils_standard, \
     fd_stencils_forth_order_isotropic
@@ -71,7 +71,7 @@ def test_staggered_laplacian():
 def test_staggered_combined():
     from pystencils.fd import diff
     f = ps.fields("f : double[2D]")
-    x, y = [LoopOverCoordinate.get_loop_counter_symbol(i) for i in range(2)]
+    x, y = DEFAULTS.spatial_counters[:2]
     dx = sp.symbols("dx")
 
     expr = diff(x * diff(f, 0) + y * diff(f, 1), 0)