macroscopic value kernels with extended test cases

lbmpy/macroscopic_value_kernels.py

@@ -4,6 +4,8 @@ from copy import deepcopy
 from lbmpy.simplificationfactory import create_simplification_strategy
 from pystencils.field import Field, get_layout_of_array
 
+from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor
+
 
 def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
     """Assignments to initialize the pdf field with equilibrium"""
@@ -28,8 +30,27 @@ def macroscopic_values_getter(lb_method, density, velocity, pdfs):
 
 macroscopic_values_setter = pdf_initialization_assignments
 
+#   Extensions for any streaming patterns
+
+
+def flexible_macroscopic_values_setter(lb_method, density, velocity,
+                                       pdf_field,
+                                       previous_step_accessor=StreamPullTwoFieldsAccessor):
+    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,
+                                       previous_step_accessor=StreamPullTwoFieldsAccessor):
+    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
+    return macroscopic_values_getter(lb_method, density, velocity, write_accesses)
 
-def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None, field_layout='numpy', target='cpu'):
+
+def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
+                                      ghost_layers=1, iteration_slice=None,
+                                      field_layout='numpy', target='cpu',
+                                      previous_step_accessor=StreamPullTwoFieldsAccessor):
     """
     Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
 
@@ -37,8 +58,13 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
         lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
         output_quantities: sequence of quantities to compute e.g. ['density', 'velocity']
         pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
+        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                      that should be excluded from the iteration. If None, the number of ghost layers 
+                      is determined automatically and assumed to be equal for all dimensions.        
+        iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout for output field, also used for pdf field if pdf_arr is not given
         target: 'cpu' or 'gpu'
+        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
         a function to compute macroscopic values:
@@ -83,15 +109,21 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
             output_mapping[output_quantity] = output_mapping[output_quantity][0]
 
     stencil = lb_method.stencil
-    pdf_symbols = [pdf_field(i) for i in range(len(stencil))]
+    # pdf_symbols = [pdf_field(i) for i in range(len(stencil))]
+
+    #   Extension for any kind of streaming rule
+    pdf_symbols = previous_step_accessor.write(pdf_field, stencil)
+
     eqs = cqc.output_equations_from_pdfs(pdf_symbols, output_mapping).all_assignments
 
     if target == 'cpu':
         import pystencils.cpu as cpu
-        kernel = cpu.make_python_function(cpu.create_kernel(eqs))
+        kernel = cpu.make_python_function(cpu.create_kernel(
+            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
     elif target == 'gpu':
         import pystencils.gpucuda as gpu
-        kernel = gpu.make_python_function(gpu.create_cuda_kernel(eqs))
+        kernel = gpu.make_python_function(gpu.create_cuda_kernel(
+            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
     else:
         raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
 
@@ -107,7 +139,10 @@ def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None
     return getter
 
 
-def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None, field_layout='numpy', target='cpu'):
+def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None, 
+                                      ghost_layers=1, iteration_slice=None,
+                                      field_layout='numpy', target='cpu', 
+                                      previous_step_accessor=StreamPullTwoFieldsAccessor):
     """
     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
@@ -116,8 +151,13 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
         lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
         quantities_to_set: map from conserved quantity name to fixed value or numpy array
         pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
+        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
+                      that should be excluded from the iteration. If None, the number of ghost layers 
+                      is determined automatically and assumed to be equal for all dimensions.        
+        iteration_slice: if not None, iteration is done only over this slice of the field
         field_layout: layout of the pdf field if pdf_arr was not given
         target: 'cpu' or 'gpu'
+        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
 
     Returns:
         function taking pdf array as single argument and which sets the field to the given values
@@ -155,7 +195,10 @@ def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None
     else:
         eq = eq.new_without_subexpressions()
 
-    substitutions = {sym: pdf_field(i) for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
+    #   Extension for any streaming rule
+    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
+
+    substitutions = {sym: write_accesses[i] for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
     eq = eq.new_with_substitutions(substitutions).all_assignments
 
     if target == 'cpu':

lbmpy_tests/test_macroscopic_value_kernels.py

@@ -3,30 +3,59 @@ import numpy as np
 from lbmpy.creationfunctions import create_lb_method
 from lbmpy.macroscopic_value_kernels import (
     compile_macroscopic_values_getter, compile_macroscopic_values_setter)
+from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
+    StreamPushTwoFieldsAccessor, \
+    AAEvenTimeStepAccessor, \
+    AAOddTimeStepAccessor, \
+    EsoTwistEvenTimeStepAccessor, \
+    EsoTwistOddTimeStepAccessor
 
+import pytest
+from itertools import product
 
-def test_set_get_density_velocity_with_fields():
-    for stencil in ['D2Q9', 'D3Q19']:
-        for force_model in ['guo', 'luo', 'none']:
-            for compressible in [True, False]:
-                force = (0.1, 0.12, -0.17)
-                method = create_lb_method(stencil=stencil, force_model=force_model, method='trt',
-                                          compressible=compressible, force=force)
-                size = (3, 7, 4)[:method.dim]
-                pdf_arr = np.zeros(size + (len(method.stencil),))
-                density_input_field = 1 + 0.2 * (np.random.random_sample(size) - 0.5)
-                velocity_input_field = 0.1 * (np.random.random_sample(size + (method.dim, )) - 0.5)
-                setter = compile_macroscopic_values_setter(method, pdf_arr=pdf_arr,
-                                                           quantities_to_set={'density': density_input_field,
-                                                                              'velocity': velocity_input_field}, )
-                setter(pdf_arr)
+accessors = [StreamPullTwoFieldsAccessor,
+             StreamPushTwoFieldsAccessor,
+             AAEvenTimeStepAccessor,
+             AAOddTimeStepAccessor,
+             EsoTwistEvenTimeStepAccessor,
+             EsoTwistOddTimeStepAccessor]
+
+stencils = ['D2Q9', 'D3Q19']
+force_models = ['guo', 'luo', 'none']
+compressibilities = [True, False]
+
+
+@pytest.mark.parametrize('stencil,force_model,compressible,accessor', product(stencils, force_models, compressibilities, accessors))
+def test_set_get_density_velocity_with_fields(stencil, force_model, compressible, accessor):
+    force = (0.1, 0.12, -0.17)
+    method = create_lb_method(stencil=stencil, force_model=force_model, method='trt',
+                              compressible=compressible, force=force)
+    ghost_layers = 1
+    inner_size = (3, 7, 4)[:method.dim]
+    total_size = tuple(s + 2 * ghost_layers for s in inner_size)
+    pdf_arr = np.zeros(total_size + (len(method.stencil),))
+
+    inner_slice = (slice(ghost_layers, -ghost_layers, 1), ) * method.dim
+    density_input_field = np.zeros(total_size)
+    velocity_input_field = np.zeros(total_size + (method.dim, ))
+    density_input_field[inner_slice] = 1 + 0.2 * (np.random.random_sample(inner_size) - 0.5)
+    velocity_input_field[inner_slice] = 0.1 * \
+        (np.random.random_sample(inner_size + (method.dim, )) - 0.5)
+
+    quantities_to_set = {'density': density_input_field, 'velocity': velocity_input_field}
+
+    setter = compile_macroscopic_values_setter(method, pdf_arr=pdf_arr, quantities_to_set=quantities_to_set,
+                                               ghost_layers=ghost_layers, previous_step_accessor=accessor)
+    setter(pdf_arr)
+
+    getter = compile_macroscopic_values_getter(method, ['density', 'velocity'], pdf_arr=pdf_arr, 
+                                               ghost_layers=ghost_layers, previous_step_accessor=accessor)
 
-                getter = compile_macroscopic_values_getter(method, ['density', 'velocity'], pdf_arr=pdf_arr)
-                density_output_field = np.empty_like(density_input_field)
-                velocity_output_field = np.empty_like(velocity_input_field)
-                getter(pdfs=pdf_arr, density=density_output_field, velocity=velocity_output_field)
-                np.testing.assert_almost_equal(density_input_field, density_output_field)
-                np.testing.assert_almost_equal(velocity_input_field, velocity_output_field)
+    density_output_field = np.zeros_like(density_input_field)
+    velocity_output_field = np.zeros_like(velocity_input_field)
+    getter(pdfs=pdf_arr, density=density_output_field, velocity=velocity_output_field)
+    np.testing.assert_almost_equal(density_input_field, density_output_field)
+    np.testing.assert_almost_equal(velocity_input_field, velocity_output_field)
 
 
 def test_set_get_constant_velocity():