Refactored time step as Enum

lbmpy/advanced_streaming/__init__.py

@@ -3,8 +3,10 @@ from .boundaryconditions import FlexibleBoundary, FlexibleNoSlip
 from .boundaryhandling import FlexibleLBMBoundaryHandling, \
     create_advanced_streaming_boundary_kernel
 from .communication import get_communication_slices, PeriodicityHandling
+from .utility import Timestep, get_accessor
 
 __all__ = ['AdvancedStreamingBoundaryIndexing', 'FlexibleBoundary',
            'FlexibleNoSlip', 'create_advanced_streaming_boundary_kernel',
            'FlexibleLBMBoundaryHandling',
-           'get_communication_slices', 'PeriodicityHandling']
+           'get_communication_slices', 'PeriodicityHandling',
+           'Timestep', 'get_accessor']

lbmpy/advanced_streaming/boundaryhandling.py

-from lbmpy.advanced_streaming import AdvancedStreamingBoundaryIndexing
+from lbmpy.advanced_streaming.boundaryindexing import AdvancedStreamingBoundaryIndexing
+from lbmpy.advanced_streaming.utility import Timestep
 from pystencils import Field, Assignment, create_indexed_kernel
 from pystencils.boundaries import BoundaryHandling
 from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object
@@ -16,7 +17,7 @@ class FlexibleLBMBoundaryHandling(BoundaryHandling):
         self._lb_method = lb_method
         self._streaming_pattern = streaming_pattern
         self._two_fields_kernel = streaming_pattern in ['pull', 'push']
-        self._between_timesteps = None
+        self._prev_timestep = None
         super(FlexibleLBMBoundaryHandling, self).__init__(data_handling, pdf_field_name, lb_method.stencil,
                                                           name, flag_interface, target, openmp)
 
@@ -25,16 +26,13 @@ class FlexibleLBMBoundaryHandling(BoundaryHandling):
     #   ------------------------- Overridden methods of pystencils.BoundaryHandling -------------------------
 
     @property
-    def between_timesteps(self):
-        return self._between_timesteps
+    def prev_timestep(self):
+        return self._prev_timestep
 
-    def __call__(self, between_timesteps='both', **kwargs):
-        if between_timesteps not in ['both', 'odd_to_even', 'even_to_odd']:
-            raise ValueError("Invalid value of between_timesteps argument:", between_timesteps)
-
-        self._between_timesteps = between_timesteps
+    def __call__(self, prev_timestep=Timestep.BOTH, **kwargs):
+        self._prev_timestep = prev_timestep
         super(FlexibleLBMBoundaryHandling, self).__call__(**kwargs)
-        self._between_timesteps = None
+        self._prev_timestep = None
 
     def add_fixed_steps(self, fixed_loop, **kwargs):
         raise NotImplementedError("Adding to fixed loop is not supported by FlexibleLBMBoundaryHandling")
@@ -49,56 +47,49 @@ class FlexibleLBMBoundaryHandling(BoundaryHandling):
         if boundary_obj not in self._boundary_object_to_boundary_info:
             sym_index_field = Field.create_generic('indexField', spatial_dimensions=1,
                                                    dtype=numpy_data_type_for_boundary_object(boundary_obj, self.dim))
-            even_to_odd_kernel = self._create_boundary_kernel(self._data_handling.fields[self._field_name],
-                                                              sym_index_field, boundary_obj, 'even_to_odd').compile()
-            odd_to_even_kernel = self._create_boundary_kernel(self._data_handling.fields[self._field_name],
-                                                              sym_index_field, boundary_obj, 'odd_to_even').compile()
+            kernels = [ self._create_boundary_kernel(
+                                self._data_handling.fields[self._field_name], sym_index_field, boundary_obj, Timestep.EVEN).compile(),
+                        self._create_boundary_kernel(
+                                self._data_handling.fields[self._field_name], sym_index_field, boundary_obj, Timestep.ODD).compile() ]
             if flag is None:
                 flag = self.flag_interface.reserve_next_flag()
-            boundary_info = self.FlexibleBoundaryInfo(self, boundary_obj, flag, even_to_odd_kernel, odd_to_even_kernel)
+            boundary_info = self.FlexibleBoundaryInfo(self, boundary_obj, flag, kernels)
             self._boundary_object_to_boundary_info[boundary_obj] = boundary_info
         return self._boundary_object_to_boundary_info[boundary_obj].flag
 
-    def _create_boundary_kernel(self, symbolic_field, symbolic_index_field, boundary_obj, between_timesteps='both'):
+    def _create_boundary_kernel(self, symbolic_field, symbolic_index_field, boundary_obj, prev_timestep=Timestep.BOTH):
         return create_advanced_streaming_boundary_kernel(
             symbolic_field, symbolic_index_field, self._lb_method, boundary_obj,
-            between_timesteps=between_timesteps, streaming_pattern=self._streaming_pattern,
+            prev_timestep=prev_timestep, streaming_pattern=self._streaming_pattern,
             target=self._target, openmp=self._openmp)
 
     class FlexibleBoundaryInfo(object):
 
         @property
         def kernel(self):
-            between_timesteps = self._boundary_handling.between_timesteps
-            if between_timesteps is None:
+            prev_timestep = self._boundary_handling.prev_timestep
+            if prev_timestep is None:
                 raise Exception(
                     "The flexible boundary kernel property was accessed while " 
                     + "there was no boundary handling in progress.")
-            if between_timesteps == 'both':
-                raise ValueError("No boundary kernel can be selected for both kinds of time steps!")
-            assert between_timesteps in ['odd_to_even', 'even_to_odd']
-            if between_timesteps == 'even_to_odd':
-                return self._even_to_odd_kernel
-            else:
-                return self._odd_to_even_kernel
-
-        def __init__(self, boundary_handling, boundary_obj, flag, even_to_odd_kernel, odd_to_even_kernel):
+            return self._kernels[prev_timestep]
+
+        def __init__(self, boundary_handling, boundary_obj, flag, kernels):
             self._boundary_handling = boundary_handling
             self.boundary_object = boundary_obj
             self.flag = flag
-            self._odd_to_even_kernel = odd_to_even_kernel
-            self._even_to_odd_kernel = even_to_odd_kernel
+            self._kernels = kernels
 
 # end class FlexibleLBMBoundaryHandling
 
 
 def create_advanced_streaming_boundary_kernel(pdf_field, index_field, lb_method, boundary_functor,
-                                              between_timesteps='both', streaming_pattern='pull',
+                                              prev_timestep=Timestep.BOTH, streaming_pattern='pull',
                                               target='cpu', openmp=True):
     index_dtype = index_field.dtype.numpy_dtype.fields['dir'][0]
     offsets_dtype = index_field.dtype.numpy_dtype.fields['x'][0]
     indexing = AdvancedStreamingBoundaryIndexing(
-        pdf_field, lb_method.stencil, between_timesteps, streaming_pattern, index_dtype, offsets_dtype)
+        pdf_field, lb_method.stencil, prev_timestep, streaming_pattern, index_dtype, offsets_dtype)
 
     f_out, f_in = indexing.proxy_fields
     dir_symbol = indexing.dir_symbol

lbmpy/advanced_streaming/boundaryindexing.py

@@ -5,7 +5,7 @@ import pystencils as ps
 from pystencils.data_types import TypedSymbol, create_type
 from pystencils.backends.cbackend import CustomCodeNode
 
-from lbmpy.advanced_streaming.utility import get_accessor, inverse_dir_index
+from lbmpy.advanced_streaming.utility import get_accessor, inverse_dir_index, is_inplace, Timestep
 
 from itertools import product
 
@@ -32,30 +32,17 @@ class AdvancedStreamingBoundaryIndexing:
     #       Constructor and State
     #   =============================
 
-    def __init__(self, pdf_field, stencil, between_timesteps='both', streaming_pattern='pull',
+    def __init__(self, pdf_field, stencil, prev_timestep=Timestep.BOTH, streaming_pattern='pull',
                  index_dtype=np.int32, offsets_dtype=np.int32):
-        if between_timesteps not in ['both', 'odd_to_even', 'even_to_odd']:
-            raise ValueError(
-                "Invalid value of parameter 'between_timesteps'."
-                + " Must be one of 'both', 'odd_to_even' or 'even_to_odd'.",
-                between_timesteps)
-
-        if between_timesteps == 'both' and streaming_pattern in ['aa', 'esotwist']:
-            raise ValueError('Cannot create index arrays for both kinds of timesteps for kernel type '
+        if prev_timestep == Timestep.BOTH and is_inplace(streaming_pattern):
+            raise ValueError('Cannot create index arrays for both kinds of timesteps for inplace streaming pattern '
                              + streaming_pattern)
 
-        odd_to_even = (between_timesteps == 'odd_to_even')
+        prev_accessor = get_accessor(streaming_pattern, prev_timestep)
+        next_accessor = get_accessor(streaming_pattern, prev_timestep.next())
 
-        even_accessor = get_accessor(streaming_pattern, 'even')
-        odd_accessor = get_accessor(streaming_pattern, 'odd')
-
-        if between_timesteps == 'both':
-            assert even_accessor == odd_accessor
-
-        outward_accesses = (
-            odd_accessor if odd_to_even else even_accessor).write(pdf_field, stencil)
-        inward_accesses = (
-            even_accessor if odd_to_even else odd_accessor).read(pdf_field, stencil)
+        outward_accesses = prev_accessor.write(pdf_field, stencil)
+        inward_accesses = next_accessor.read(pdf_field, stencil)
 
         self._accesses = {'out': outward_accesses, 'in': inward_accesses}
 

lbmpy/advanced_streaming/communication.py

 from pystencils import Field
 from pystencils.slicing import shift_slice, get_slice_before_ghost_layer
-from lbmpy.advanced_streaming.utility import is_inplace, get_accessor, numeric_index, numeric_offsets
+from lbmpy.advanced_streaming.utility import is_inplace, get_accessor, numeric_index, numeric_offsets, Timestep, get_timesteps
 from pystencils.datahandling import SerialDataHandling
 from itertools import chain
 
@@ -52,14 +52,14 @@ def _fix_length_one_slices(slices):
 
 
 def get_communication_slices(
-        stencil, comm_stencil=None, streaming_pattern='pull', after_timestep='both', ghost_layers=1):
+        stencil, comm_stencil=None, streaming_pattern='pull', prev_timestep=Timestep.BOTH, ghost_layers=1):
     """
     Return the source and destination slices for periodicity handling or communication between blocks.
 
     :param stencil: The stencil used by the LB method.
     :param comm_stencil: The stencil defining the communication directions. If None, it will be set to stencil.
     :param streaming_pattern: The streaming pattern.
-    :param after_timestep: Timestep after which communication is run; either 'even', 'odd' or 'both'.
+    :param prev_timestep: Timestep after which communication is run.
     :param ghost_layers: Number of ghost layers in each direction.
 
     """
@@ -67,7 +67,7 @@ def get_communication_slices(
         comm_stencil = stencil
 
     pdfs = Field.create_generic('pdfs', spatial_dimensions=len(stencil[0]), index_shape=(len(stencil),))
-    write_accesses = get_accessor(streaming_pattern, after_timestep).write(pdfs, stencil)
+    write_accesses = get_accessor(streaming_pattern, prev_timestep).write(pdfs, stencil)
     slices_per_comm_direction = dict()
 
     for comm_dir in comm_stencil:
@@ -101,8 +101,7 @@ def get_communication_slices(
 class PeriodicityHandling:
 
     def __init__(self, stencil, data_handling, pdf_field_name,
-                 streaming_pattern='pull', zeroth_timestep='both',
-                 ghost_layers=1, gpu=False):
+                 streaming_pattern='pull', ghost_layers=1, gpu=False):
         if not isinstance(data_handling, SerialDataHandling):
             raise ValueError('Only single node data handling is supported!')
 
@@ -121,37 +120,46 @@ class PeriodicityHandling:
             return True
 
         copy_directions = tuple(filter(is_copy_direction, stencil[1:]))
-        if self.inplace_pattern:
-            self.comm_slices = dict()
-            for timestep in ['even', 'odd']:
-                slices_per_comm_dir = get_communication_slices(stencil=stencil,
-                                                               comm_stencil=copy_directions,
-                                                               streaming_pattern=streaming_pattern,
-                                                               after_timestep=timestep,
-                                                               ghost_layers=ghost_layers)
-                self.comm_slices[timestep] = list(chain.from_iterable(v for k, v in slices_per_comm_dir.items()))
-        else:
+        # if self.inplace_pattern:
+        #     self.comm_slices = dict()
+        #     for timestep in [Timestep.EVEN, Timestep.ODD]:
+        #         slices_per_comm_dir = get_communication_slices(stencil=stencil,
+        #                                                        comm_stencil=copy_directions,
+        #                                                        streaming_pattern=streaming_pattern,
+        #                                                        prev_timestep=timestep,
+        #                                                        ghost_layers=ghost_layers)
+        #         self.comm_slices[timestep] = list(chain.from_iterable(v for k, v in slices_per_comm_dir.items()))
+        # else:
+        #     slices_per_comm_dir = get_communication_slices(stencil=stencil,
+        #                                                    comm_stencil=copy_directions,
+        #                                                    streaming_pattern=streaming_pattern,
+        #                                                    prev_timestep=Timestep.BOTH,
+        #                                                    ghost_layers=ghost_layers)
+        #     self.comm_slices = list(chain.from_iterable(v for k, v in slices_per_comm_dir.items()))
+
+        self.comm_slices = []
+        for timestep in get_timesteps(streaming_pattern):
             slices_per_comm_dir = get_communication_slices(stencil=stencil,
-                                                           comm_stencil=copy_directions,
-                                                           streaming_pattern=streaming_pattern,
-                                                           after_timestep='both',
-                                                           ghost_layers=ghost_layers)
-            self.comm_slices = list(chain.from_iterable(v for k, v in slices_per_comm_dir.items()))
+                                                            comm_stencil=copy_directions,
+                                                            streaming_pattern=streaming_pattern,
+                                                            prev_timestep=timestep,
+                                                            ghost_layers=ghost_layers)
+            self.comm_slices.append(list(chain.from_iterable(v for k, v in slices_per_comm_dir.items())))
 
-    def __call__(self, timestep_modulus='both'):
+    def __call__(self, prev_timestep=Timestep.BOTH):
         if self.gpu:
-            self._periodicity_handling_gpu(timestep_modulus)
+            self._periodicity_handling_gpu(prev_timestep)
         else:
-            self._periodicity_handling_cpu(timestep_modulus)
+            self._periodicity_handling_cpu(prev_timestep)
 
-    def _periodicity_handling_cpu(self, timestep):
+    def _periodicity_handling_cpu(self, prev_timestep):
         arr = self.dh.cpu_arrays[self.pdf_field_name]
-        if timestep == 'both':
-            comm_slices = self.comm_slices
-        else:
-            comm_slices = self.comm_slices[timestep]
+        # if prev_timestep == 'both':
+        #     comm_slices = self.comm_slices
+        # else:
+        comm_slices = self.comm_slices[prev_timestep.idx]
         for src, dst in comm_slices:
             arr[dst] = arr[src]
 
-    def _periodicity_handling_gpu(self, timestep):
+    def _periodicity_handling_gpu(self, prev_timestep):
         raise NotImplementedError()

lbmpy/advanced_streaming/utility.py

@@ -7,9 +7,22 @@ from lbmpy.fieldaccess import PdfFieldAccessor, \
     EsoTwistOddTimeStepAccessor
 
 import pystencils as ps
+from enum import IntEnum
+
+class Timestep(IntEnum):
+    EVEN = 0
+    ODD = 1
+    BOTH = 2
+
+    def next(self):
+        return self if self == Timestep.BOTH else Timestep((self + 1) % 2)
+
+    @property
+    def idx(self):
+        """To use this timestep as an array index"""
+        return self % 2
 
 streaming_patterns = ['push', 'pull', 'aa', 'esotwist']
-timesteps = ['even', 'odd', 'both']
 
 even_accessors = {
     'pull': StreamPullTwoFieldsAccessor,
@@ -26,15 +39,12 @@ odd_accessors = {
 }
 
 
-def get_accessor(streaming_pattern, timestep) -> PdfFieldAccessor:
+def get_accessor(streaming_pattern : str, timestep : Timestep) -> PdfFieldAccessor:
     if streaming_pattern not in streaming_patterns:
         raise ValueError(
             "Invalid value of parameter 'streaming_pattern'.", streaming_pattern)
 
-    if timestep not in timesteps:
-        raise ValueError("Invalid time step:", timestep)
-
-    if timestep == 'even':
+    if timestep == Timestep.EVEN:
         return even_accessors[streaming_pattern]
     else:
         return odd_accessors[streaming_pattern]
@@ -46,6 +56,8 @@ def is_inplace(streaming_pattern):
 
     return streaming_pattern in ['aa', 'esotwist']
 
+def get_timesteps(streaming_pattern):
+    return (Timestep.EVEN, Timestep.ODD) if is_inplace(streaming_pattern) else (Timestep.BOTH, )
 
 def numeric_offsets(field_access: ps.Field.Access):
     return tuple(int(o) for o in field_access.offsets)
@@ -63,7 +75,7 @@ class AccessPdfValues:
     """Allows to access values from a PDF array correctly depending on 
     the streaming pattern."""
 
-    def __init__(self, pdf_field, stencil, streaming_pattern='pull', timestep='both', accessor=None):
+    def __init__(self, pdf_field, stencil, streaming_pattern='pull', timestep=Timestep.BOTH, accessor=None):
         if accessor is None:
             accessor = get_accessor(streaming_pattern, timestep)
         self.read_accs = accessor.read(pdf_field, stencil)

lbmpy/macroscopic_value_kernels.py

@@ -4,7 +4,7 @@ from copy import deepcopy
 from lbmpy.simplificationfactory import create_simplification_strategy
 from pystencils.field import Field, get_layout_of_array
 
-from lbmpy.advanced_streaming.utility import get_accessor
+from lbmpy.advanced_streaming.utility import get_accessor, Timestep
 
 
 def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
@@ -34,14 +34,14 @@ macroscopic_values_setter = pdf_initialization_assignments
 
 
 def flexible_macroscopic_values_setter(lb_method, density, velocity,
-                                       pdf_field, streaming_pattern='pull', previous_timestep='both'):
+                                       pdf_field, streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
     write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
     return pdf_initialization_assignments(lb_method, density, velocity, write_accesses)
 
 
 def flexible_macroscopic_values_getter(lb_method, density, velocity,
-                                       pdf_field, streaming_pattern='pull', previous_timestep='both'):
+                                       pdf_field, streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
     write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
     return macroscopic_values_getter(lb_method, density, velocity, write_accesses)
@@ -50,7 +50,7 @@ def flexible_macroscopic_values_getter(lb_method, density, velocity,
 def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
                                       ghost_layers=1, iteration_slice=None,
                                       field_layout='numpy', target='cpu',
-                                      streaming_pattern='pull', previous_timestep='both'):
+                                      streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
 
@@ -140,7 +140,7 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
 def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None, 
                                       ghost_layers=1, iteration_slice=None,
                                       field_layout='numpy', target='cpu', 
-                                      streaming_pattern='pull', previous_timestep='both'):
+                                      streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Creates a function that sets a pdf field to specified macroscopic quantities
     The returned function can be called with the pdf field to set as single argument

lbmpy_tests/advanced_streaming/test_flexible_noslip.py → lbmpy_tests/advanced_streaming/test_advanced_streaming_noslip.py

 import numpy as np
 import sympy as sp
 
-from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
-    StreamPushTwoFieldsAccessor, \
-    AAEvenTimeStepAccessor, \
-    AAOddTimeStepAccessor, \
-    EsoTwistEvenTimeStepAccessor, \
-    EsoTwistOddTimeStepAccessor
-
-from lbmpy.advanced_streaming import FlexibleNoSlip, create_advanced_streaming_boundary_kernel
+from lbmpy.advanced_streaming import FlexibleNoSlip, create_advanced_streaming_boundary_kernel, Timestep
 from lbmpy.advanced_streaming.utility import even_accessors, odd_accessors, streaming_patterns, inverse_dir_index, AccessPdfValues
 from lbmpy.creationfunctions import create_lb_method
 from lbmpy.stencils import get_stencil
@@ -25,8 +18,8 @@ import pytest
 
 @pytest.mark.parametrize("stencil", [ 'D2Q9', 'D3Q19', 'D3Q27'])
 @pytest.mark.parametrize("streaming_pattern", streaming_patterns)
-@pytest.mark.parametrize("timestep_type", ['odd_to_even', 'even_to_odd'])
-def test_flexible_noslip_single_cell(stencil, streaming_pattern, timestep_type):
+@pytest.mark.parametrize("prev_timestep", [Timestep.EVEN, Timestep.ODD])
+def test_advanced_streaming_noslip_single_cell(stencil, streaming_pattern, prev_timestep):
     """
     Flexible NoSlip Test
     """
@@ -36,14 +29,8 @@ def test_flexible_noslip_single_cell(stencil, streaming_pattern, timestep_type):
     dim = len(stencil[0])
     pdf_field = ps.fields(f'pdfs({q}): [{dim}D]')
 
-    even_accessor = even_accessors[streaming_pattern]
-    odd_accessor = odd_accessors[streaming_pattern]
-
-    prev_accessor = (odd_accessor if timestep_type == 'odd_to_even' else even_accessor)
-    next_accessor = (even_accessor if timestep_type == 'odd_to_even' else odd_accessor)
-
-    prev_pdf_access = AccessPdfValues(pdf_field, stencil, accessor=prev_accessor)
-    next_pdf_access = AccessPdfValues(pdf_field, stencil, accessor=next_accessor)
+    prev_pdf_access = AccessPdfValues(pdf_field, stencil, streaming_pattern, prev_timestep)
+    next_pdf_access = AccessPdfValues(pdf_field, stencil, streaming_pattern, prev_timestep.next())
 
     pdfs = np.zeros((3,) * dim + (q,))
     pos = (1,) * dim
@@ -51,17 +38,17 @@ def test_flexible_noslip_single_cell(stencil, streaming_pattern, timestep_type):
         prev_pdf_access.write_pdf(pdfs, pos, d, d)
 
     lb_method = create_lb_method(stencil=stencil, method='srt')
-    flex_noslip = FlexibleNoSlip()
+    noslip = FlexibleNoSlip()
 
-    index_struct_dtype = numpy_data_type_for_boundary_object(flex_noslip, dim)
+    index_struct_dtype = numpy_data_type_for_boundary_object(noslip, dim)
 
     index_field = Field('indexVector', FieldType.INDEXED, index_struct_dtype, layout=[0],
                     shape=(TypedSymbol("indexVectorSize", create_type(np.int64)), 1), strides=(1, 1))
     index_vector = np.array([ pos + (d,) for d in range(q) ], dtype=index_struct_dtype)
 
     flex_ast = create_advanced_streaming_boundary_kernel(pdf_field, 
-                    index_field, lb_method, flex_noslip,
-                    between_timesteps=timestep_type,
+                    index_field, lb_method, noslip,
+                    prev_timestep=prev_timestep,
                     streaming_pattern=streaming_pattern)
 
     flex_kernel = flex_ast.compile()

lbmpy_tests/advanced_streaming/test_communication.py

@@ -2,19 +2,19 @@ import numpy as np
 from lbmpy.stencils import get_stencil
 from pystencils.slicing import get_slice_before_ghost_layer, get_ghost_region_slice
 from lbmpy.advanced_streaming.communication import get_communication_slices, _fix_length_one_slices
-from lbmpy.advanced_streaming.utility import streaming_patterns
+from lbmpy.advanced_streaming.utility import streaming_patterns, Timestep
 
 import pytest
 
 @pytest.mark.parametrize('stencil', ['D2Q9', 'D3Q19', 'D3Q27'])
 @pytest.mark.parametrize('streaming_pattern', streaming_patterns)
-@pytest.mark.parametrize('timestep', ['even', 'odd'])
+@pytest.mark.parametrize('timestep', [Timestep.EVEN, Timestep.ODD])
 def test_slices_not_empty(stencil, streaming_pattern, timestep):
     stencil = get_stencil(stencil)
     dim = len(stencil[0])
     q = len(stencil)
     arr = np.zeros( (4,) * dim + (q,) )
-    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, after_timestep=timestep, ghost_layers=1)
+    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, prev_timestep=timestep, ghost_layers=1)
     for _, slices_list in slices.items():
         for src, dst in slices_list:
             assert all(s != 0 for s in arr[src].shape)
@@ -23,13 +23,13 @@ def test_slices_not_empty(stencil, streaming_pattern, timestep):
 
 @pytest.mark.parametrize('stencil', ['D2Q9', 'D3Q19', 'D3Q27'])
 @pytest.mark.parametrize('streaming_pattern', streaming_patterns)
-@pytest.mark.parametrize('timestep', ['even', 'odd'])
+@pytest.mark.parametrize('timestep', [Timestep.EVEN, Timestep.ODD])
 def test_src_dst_same_shape(stencil, streaming_pattern, timestep):
     stencil = get_stencil(stencil)
     dim = len(stencil[0])
     q = len(stencil)
     arr = np.zeros( (4,) * dim + (q,) )
-    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, after_timestep=timestep, ghost_layers=1)
+    slices = get_communication_slices(stencil, streaming_pattern=streaming_pattern, prev_timestep=timestep, ghost_layers=1)
     for _, slices_list in slices.items():
         for src, dst in slices_list:
             src_shape = arr[src].shape
@@ -42,7 +42,7 @@ def test_pull_communication_slices(stencil):
     stencil = get_stencil(stencil)
 
     slices = get_communication_slices(
-        stencil, streaming_pattern='pull', after_timestep='both', ghost_layers=1)
+        stencil, streaming_pattern='pull', prev_timestep=Timestep.BOTH, ghost_layers=1)
     for i, d in enumerate(stencil):
         if i == 0:
             continue

lbmpy_tests/advanced_streaming/test_fully_periodic_flow.py

@@ -10,7 +10,7 @@ from lbmpy.macroscopic_value_kernels import flexible_macroscopic_values_getter,
 from lbmpy.stencils import get_stencil
 
 from lbmpy.advanced_streaming import PeriodicityHandling
-from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns
+from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns, get_timesteps
 
 import pytest
 from numpy.testing import assert_allclose, assert_array_equal
@@ -29,7 +29,7 @@ def test_fully_periodic_flow(stencil, streaming_pattern):
 
     #   Streaming
     inplace = is_inplace(streaming_pattern)
-    timesteps = ['even', 'odd'] if inplace else ['both']
+    timesteps = get_timesteps(streaming_pattern)
     zeroth_timestep = timesteps[0]
 
     #   Data Handling and PDF fields
@@ -102,37 +102,35 @@ def test_fully_periodic_flow(stencil, streaming_pattern):
         getter_kernels.append(create_kernel(getter, ghost_layers=1).compile())
 
     #   Periodicity
-    periodicity_handler = PeriodicityHandling(
-        stencil, dh, pdfs.name,
-        streaming_pattern=streaming_pattern, zeroth_timestep=zeroth_timestep)
+    periodicity_handler = PeriodicityHandling(stencil, dh, pdfs.name, streaming_pattern=streaming_pattern)
 
     # Initialization and Timestep
-    t_modulus = 0
+    current_timestep = zeroth_timestep
 
     def init():
-        global t_modulus
-        t_modulus = 0
+        global current_timestep
+        current_timestep = zeroth_timestep
         dh.run_kernel(setter_kernel)
 
     def one_step():
-        global t_modulus
+        global current_timestep
 
         # Periodicty
-        periodicity_handler(timesteps[t_modulus])
-
-        # Macroscopic Values
-        dh.run_kernel(getter_kernels[t_modulus])
+        periodicity_handler(current_timestep)
 
         # Here, the next time step begins
-        t_modulus = (t_modulus + 1) % len(timesteps)
+        current_timestep = current_timestep.next()
 
         # LBM Step
-        dh.run_kernel(lb_kernels[t_modulus])
+        dh.run_kernel(lb_kernels[current_timestep.idx])
 
         # Field Swaps
         if not inplace:
             dh.swap(pdfs.name, pdfs_tmp.name)
 
+        # Macroscopic Values
+        dh.run_kernel(getter_kernels[current_timestep.idx])
+
     #   Run the simulation
     init()
 

lbmpy_tests/advanced_streaming/test_periodic_pipe_with_force.py

@@ -10,7 +10,7 @@ from lbmpy.macroscopic_value_kernels import flexible_macroscopic_values_getter,
 from lbmpy.stencils import get_stencil
 
 from lbmpy.advanced_streaming import PeriodicityHandling, FlexibleNoSlip, FlexibleLBMBoundaryHandling
-from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns
+from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns, Timestep, get_timesteps
 
 import pytest
 from numpy.testing import assert_allclose
@@ -29,8 +29,7 @@ class PeriodicPipeFlow: