Completed advanced periodicity handling and added fully periodic test scenario

lbmpy/advanced_streaming/__init__.py

@@ -2,7 +2,9 @@ from .boundaryindexing import AdvancedStreamingBoundaryIndexing
 from .boundaryconditions import FlexibleBoundary, FlexibleNoSlip
 from .boundaryhandling import FlexibleLBMBoundaryHandling, \
     create_advanced_streaming_boundary_kernel
+from .communication import get_communication_slices, PeriodicityHandling
 
 __all__ = ['AdvancedStreamingBoundaryIndexing', 'FlexibleBoundary',
            'FlexibleNoSlip', 'create_advanced_streaming_boundary_kernel',
-           'FlexibleLBMBoundaryHandling']
+           'FlexibleLBMBoundaryHandling',
+           'get_communication_slices', 'PeriodicityHandling']

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 get_accessor, numeric_index, numeric_offsets
+from lbmpy.advanced_streaming.utility import is_inplace, get_accessor, numeric_index, numeric_offsets
+from pystencils.datahandling import SerialDataHandling
+from itertools import chain
 
 
-def cut_slice_in_direction(slices, direction):
+def trim_slice_in_direction(slices, direction):
     assert len(slices) == len(direction)
 
     result = []
     for s, d in zip(slices, direction):
+        if isinstance(s, int):
+            result.append(s)
+            continue
         start = s.start + 1 if d == -1 else s.start
         stop = s.stop - 1 if d == 1 else s.stop
         result.append(slice(start, stop, s.step))
@@ -32,37 +37,43 @@ def _get_neighbour_transform(direction, ghost_layers):
 
 
 def _fix_length_one_slices(slices):
-    if isinstance(slices, slice):
-        if slices.stop is None:
-            return slices
+    """Slices of length one are replaced by their start value for correct periodic shifting"""
+    if isinstance(slices, int):
+        return slices
+    elif isinstance(slices, slice):
+        if slices.stop is not None and abs(slices.start - slices.stop) == 1:
+            return slices.start
+        elif slices.stop is None and slices.start == -1:
+            return -1   #   [-1:] also has length one
         else:
-            return slice(slices.start, None if abs(slices.start - slices.stop) == 1 else slices.stop)
-    return tuple(_fix_length_one_slices(s) for s in slices)
+            return slices
+    else:
+        return tuple(_fix_length_one_slices(s) for s in slices)
 
 
-def get_communication_slices(stencil, streaming_pattern='pull', after_timestep='both', ghost_layers=1):
+def get_communication_slices(stencil, comm_stencil=None, streaming_pattern='pull', after_timestep='both', ghost_layers=1):
     """
-    Return the source and destination slices for periodicity handling or communication between
-    blocks. This function returns a dictionary which contains a list of tuples 
-    (idx, src_slice, dst_slice) for each stencil direction. Here, idx is an index to the pdf field,
-    src_slice the source region of the current block, and dst_slice the destination region in the
-    adjacent block.
+    Return the source and destination slices for periodicity handling or communication between blocks.
 
-    :param lb_method: The lattice boltzmann method to be used.
+    :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 ghost_layers: Number of ghost layers in each direction.
 
     """
+    if comm_stencil is None:
+        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)
-    all_slices = dict()
+    slices_per_comm_direction = dict()
 
-    for comm_dir in stencil:
+    for comm_dir in comm_stencil:
         if all(d == 0 for d in comm_dir):
             continue
 
-        slices_for_dir = dict()
+        slices_for_dir = []
 
         for streaming_dir in set(extend_dir(comm_dir)) & set(stencil):
             d = stencil.index(streaming_dir)
@@ -72,18 +83,74 @@ def get_communication_slices(stencil, streaming_pattern='pull', after_timestep='
             tangential_dir = tuple(s - c for s, c in zip(streaming_dir, comm_dir))
             origin_slice = get_slice_before_ghost_layer(comm_dir, ghost_layers=ghost_layers, thickness=1)
             origin_slice = _fix_length_one_slices(origin_slice)
-            src_slice = shift_slice(cut_slice_in_direction(origin_slice, tangential_dir), write_offsets)
+            src_slice = shift_slice(trim_slice_in_direction(origin_slice, tangential_dir), write_offsets)
 
             neighbour_transform = _get_neighbour_transform(comm_dir, ghost_layers)
             dst_slice = shift_slice(src_slice, neighbour_transform)
 
-            slices_for_dir[write_index] = ((src_slice, dst_slice))
+            src_slice = src_slice + (write_index, )
+            dst_slice = dst_slice + (write_index, )
+
+            slices_for_dir.append((src_slice, dst_slice))
 
-        all_slices[comm_dir] = slices_for_dir
-    return all_slices
+        slices_per_comm_direction[comm_dir] = slices_for_dir
+    return slices_per_comm_direction
 
 
 class PeriodicityHandling:
 
-    def __init__(self, lb_method, streaming_pattern, data_handling, pdf_field_name, periodicities):
-        pass
+    def __init__(self, stencil, data_handling, pdf_field_name,
+                 streaming_pattern='pull', zeroth_timestep='both',
+                 ghost_layers=1, gpu=False):
+        if not isinstance(data_handling, SerialDataHandling):
+            raise ValueError('Only single node data handling is supported!')
+
+        self.stencil = stencil
+        self.dh = data_handling
+        self.pdf_field_name = pdf_field_name
+        periodicity = data_handling.periodicity
+        self.inplace_pattern = is_inplace(streaming_pattern)
+        self.gpu = gpu
+
+        def is_copy_direction(direction):
+            for d, p in zip(direction, periodicity):
+                if d != 0 and not p:
+                    return False
+
+            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:
+            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()))
+
+    def __call__(self, timestep_modulus='both'):
+        if self.gpu:
+            self._periodicity_handling_gpu(timestep_modulus)
+        else:
+            self._periodicity_handling_cpu(timestep_modulus)
+
+    def _periodicity_handling_cpu(self, 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]
+        for src, dst in comm_slices:
+            arr[dst] = arr[src]
+
+    def _periodicity_handling_gpu(self, timestep):
+        raise NotImplementedError()
\ No newline at end of file

lbmpy/advanced_streaming/utility.py

@@ -39,6 +39,11 @@ def get_accessor(streaming_pattern, timestep) -> PdfFieldAccessor:
     else:
         return odd_accessors[streaming_pattern]
 
+def is_inplace(streaming_pattern):
+    if streaming_pattern not in streaming_patterns:
+        raise ValueError('Invalid streaming pattern', streaming_pattern)
+
+    return streaming_pattern in ['aa', 'esotwist'] 
 
 def numeric_offsets(field_access: ps.Field.Access):
     return tuple(int(o) for o in field_access.offsets)

lbmpy_tests/advanced_streaming/test_communication.py

@@ -11,24 +11,48 @@ import pytest
 @pytest.mark.parametrize('timestep', ['even', 'odd'])
 def test_slices_not_empty(stencil, streaming_pattern, timestep):
     stencil = get_stencil(stencil)
-    arr = np.zeros( (4,) * len(stencil[0]) )
-    slices = get_communication_slices(stencil, streaming_pattern, timestep, 1)
-    for _, slices_dict in slices.items():
-        for _, comm_slices in slices_dict.items():
-            assert all(s != 0 for s in arr[comm_slices[0]].shape)
-            assert all(s != 0 for s in arr[comm_slices[1]].shape)
+    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)
+    for _, slices_list in slices.items():
+        for src, dst in slices_list:
+            assert all(s != 0 for s in arr[src].shape)
+            assert all(s != 0 for s in arr[dst].shape)
+
+
+@pytest.mark.parametrize('stencil', ['D2Q9', 'D3Q19', 'D3Q27'])
+@pytest.mark.parametrize('streaming_pattern', streaming_patterns)
+@pytest.mark.parametrize('timestep', ['even', '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)
+    for _, slices_list in slices.items():
+        for src, dst in slices_list:
+            src_shape = arr[src].shape
+            dst_shape = arr[dst].shape
+            assert src_shape == dst_shape
 
 
 @pytest.mark.parametrize('stencil', ['D2Q9', 'D3Q19', 'D3Q27'])
 def test_pull_communication_slices(stencil):
     stencil = get_stencil(stencil)
 
-    slices = get_communication_slices(stencil, 'pull', 'both', 1)
+    slices = get_communication_slices(
+        stencil, streaming_pattern='pull', after_timestep='both', ghost_layers=1)
     for i, d in enumerate(stencil):
         if i == 0:
             continue
 
-        src, dst = slices[d][i]
+        for s in slices[d]:
+            if s[0][-1] == i:
+                src = s[0][:-1]
+                dst = s[1][:-1]
+                break
+        
         inner_slice = _fix_length_one_slices(get_slice_before_ghost_layer(d, ghost_layers=1))
         inv_dir = (-e for e in d)
         gl_slice = _fix_length_one_slices(get_ghost_region_slice(inv_dir, ghost_layers=1))

lbmpy_tests/advanced_streaming/test_fully_periodic_flow.py

+import numpy as np
+import sympy as sp
+
+from pystencils.datahandling import create_data_handling
+from pystencils import create_kernel
+from pystencils.plot import scalar_field, vector_field, vector_field_magnitude
+
+from lbmpy.creationfunctions import create_lb_collision_rule, create_lb_function
+from lbmpy.macroscopic_value_kernels import flexible_macroscopic_values_getter, flexible_macroscopic_values_setter
+from lbmpy.stencils import get_stencil
+
+from lbmpy.advanced_streaming import PeriodicityHandling
+from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns
+
+import pytest
+from numpy.testing import assert_allclose, assert_array_equal
+
+all_results = dict()
+
+
+@pytest.mark.parametrize('stencil', ['D2Q9', 'D3Q19', 'D3Q27'])
+@pytest.mark.parametrize('streaming_pattern', streaming_patterns)
+def test_fully_periodic_flow(stencil, streaming_pattern):
+
+    #   Stencil
+    stencil = get_stencil(stencil)
+    q = len(stencil)
+    dim = len(stencil[0])
+
+    #   Streaming
+    inplace = is_inplace(streaming_pattern)
+    timesteps = ['even', 'odd'] if inplace else ['both']
+    zeroth_timestep = timesteps[0]
+
+    #   Data Handling and PDF fields
+    domain_size = (30,) * dim
+    periodicity = (True,) * dim
+
+    dh = create_data_handling(domain_size=domain_size, periodicity=periodicity)
+
+    pdfs = dh.add_array('pdfs', q)
+    if not inplace:
+        pdfs_tmp = dh.add_array_like('pdfs_tmp', pdfs.name)
+
+    #   LBM Streaming and Collision
+    method_params = {
+        'stencil': stencil,
+        'method': 'srt',
+        'relaxation_rate': 1.0
+    }
+
+    optimization = {
+        'symbolic_field': pdfs,
+        'target': 'cpu'
+    }
+
+    if not inplace:
+        optimization['symbolic_temporary_field'] = pdfs_tmp
+
+    lb_collision = create_lb_collision_rule(optimization=optimization, **method_params)
+    lb_method = lb_collision.method
+
+    lb_kernels = []
+    if inplace:
+        lb_kernels.append(create_lb_function(collision_rule=lb_collision,
+                                             optimization=optimization,
+                                             kernel_type=streaming_pattern + '_even',
+                                             **method_params))
+        lb_kernels.append(create_lb_function(collision_rule=lb_collision,
+                                             optimization=optimization,
+                                             kernel_type=streaming_pattern + '_odd',
+                                             **method_params))
+    else:
+        if streaming_pattern == 'pull':
+            kernel_type = 'stream_pull_collide'
+        elif streaming_pattern == 'push':
+            kernel_type = 'collide_stream_push'
+        lb_kernels.append(create_lb_function(collision_rule=lb_collision,
+                                             optimization=optimization,
+                                             kernel_type=kernel_type,
+                                             **method_params))
+
+    #   Macroscopic Values
+    density = 1.0
+    density_field = dh.add_array('rho', 1)
+    u_x = 0.01
+    velocity = (u_x,) * dim
+    velocity_field = dh.add_array('u', dim)
+
+    u_ref = np.full(domain_size + (dim,), u_x)
+
+    setter = flexible_macroscopic_values_setter(
+        lb_method, density, velocity, pdfs,
+        streaming_pattern=streaming_pattern, previous_timestep=zeroth_timestep)
+    setter_kernel = create_kernel(setter, ghost_layers=1).compile()
+
+    getter_kernels = []
+    for t in timesteps:
+        getter = flexible_macroscopic_values_getter(
+            lb_method, density_field, velocity_field, pdfs,
+            streaming_pattern=streaming_pattern, previous_timestep=t)
+        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)
+
+    # Initialization and Timestep
+    t_modulus = 0
+
+    def init():
+        global t_modulus
+        t_modulus = 0
+        dh.run_kernel(setter_kernel)
+
+    def one_step():
+        global t_modulus
+
+        # Periodicty
+        periodicity_handler(timesteps[t_modulus])
+
+        # Macroscopic Values
+        dh.run_kernel(getter_kernels[t_modulus])
+
+        # Here, the next time step begins
+        t_modulus = (t_modulus + 1) % len(timesteps)
+
+        # LBM Step
+        dh.run_kernel(lb_kernels[t_modulus])
+
+        # Field Swaps
+        if not inplace:
+            dh.swap(pdfs.name, pdfs_tmp.name)
+
+    #   Run the simulation
+    init()
+
+    for _ in range(100):
+        one_step()
+
+    #   Evaluation
+
+    u = dh.gather_array(velocity_field.name)
+
+    #   Equal to the steady-state velocity field up to numerical errors
+    assert_allclose(u, u_ref)
+
+    #   Exactly equal to other streaming patterns!
+    global all_results
+    for prev_pattern, prev_u in all_results.items():
+        assert_array_equal(
+            u, prev_u, err_msg=f'Velocity field for {streaming_pattern} differed from {prev_pattern}!')