merged macroscopic value kernel implementations

lbmpy/advanced_streaming/__init__.py

 from .indexing import BetweenTimestepsIndexing, NeighbourOffsetArrays
-from .communication import get_communication_slices, PeriodicityHandling
+from .communication import get_communication_slices, LBMPeriodicityHandling
 from .utility import Timestep, get_accessor, is_inplace, get_timesteps, \
     numeric_index, numeric_offsets, inverse_dir_index, AccessPdfValues
 
 __all__ = ['BetweenTimestepsIndexing', 'NeighbourOffsetArrays',
-           'get_communication_slices', 'PeriodicityHandling',
+           'get_communication_slices', 'LBMPeriodicityHandling',
            'Timestep', 'get_accessor', 'is_inplace', 'get_timesteps',
            'numeric_index', 'numeric_offsets', 'inverse_dir_index', 'AccessPdfValues']

lbmpy/advanced_streaming/communication.py

@@ -139,7 +139,7 @@ def periodic_pdf_copy_kernel(pdf_field, src_slice, dst_slice,
         raise ValueError('Invalid target:', target)
 
 
-class PeriodicityHandling:
+class LBMPeriodicityHandling:
 
     def __init__(self, stencil, data_handling, pdf_field_name,
                  streaming_pattern='pull', ghost_layers=1, target='cpu',

lbmpy/macroscopic_value_kernels.py

 import functools
 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, Timestep
 
 
-def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
+def pdf_initialization_assignments(lb_method, density, velocity, pdfs,
+                                   streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """Assignments to initialize the pdf field with equilibrium"""
+    if isinstance(pdfs, Field):
+        previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
+        field_accesses = previous_step_accessor.write(pdfs, lb_method.stencil)
+    elif streaming_pattern == 'pull':
+        field_accesses = pdfs
+    else:
+        raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+                         + f"initialization assignments for streaming pattern {streaming_pattern}.")
+
     cqc = lb_method.conserved_quantity_computation
     inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity)
     setter_eqs = lb_method.get_equilibrium(conserved_quantity_equations=inp_eqs)
-    setter_eqs = setter_eqs.new_with_substitutions({sym: pdfs[i]
+    setter_eqs = setter_eqs.new_with_substitutions({sym: field_accesses[i]
                                                     for i, sym in enumerate(lb_method.post_collision_pdf_symbols)})
     return setter_eqs
 
 
-def macroscopic_values_getter(lb_method, density, velocity, pdfs):
+def macroscopic_values_getter(lb_method, density, velocity, pdfs,
+                              streaming_pattern='pull', previous_timestep=Timestep.BOTH):
+    if isinstance(pdfs, Field):
+        previous_step_accessor = get_accessor(streaming_pattern, previous_timestep)
+        field_accesses = previous_step_accessor.write(pdfs, lb_method.stencil)
+    elif streaming_pattern == 'pull':
+        field_accesses = pdfs
+    else:
+        raise ValueError("Invalid value of pdfs: A PDF field reference is required to derive "
+                         + f"getter assignments for streaming pattern {streaming_pattern}.")
     cqc = lb_method.conserved_quantity_computation
     assert not (velocity is None and density is None)
     output_spec = {}
@@ -25,27 +43,11 @@ def macroscopic_values_getter(lb_method, density, velocity, pdfs):
         output_spec['velocity'] = velocity
     if density is not None:
         output_spec['density'] = density
-    return cqc.output_equations_from_pdfs(pdfs, output_spec)
+    return cqc.output_equations_from_pdfs(field_accesses, output_spec)
 
 
 macroscopic_values_setter = pdf_initialization_assignments
 
-#   Extensions for any streaming patterns
-
-
-def flexible_macroscopic_values_setter(lb_method, density, velocity,
-                                       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=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)
-
 
 def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
                                       ghost_layers=1, iteration_slice=None,
@@ -137,9 +139,9 @@ 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, 
+def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None,
                                       ghost_layers=1, iteration_slice=None,
-                                      field_layout='numpy', target='cpu', 
+                                      field_layout='numpy', target='cpu',
                                       streaming_pattern='pull', previous_timestep=Timestep.BOTH):
     """
     Creates a function that sets a pdf field to specified macroscopic quantities

lbmpy_tests/advanced_streaming/test_fully_periodic_flow.py

@@ -6,10 +6,10 @@ 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.macroscopic_value_kernels import macroscopic_values_getter, macroscopic_values_setter
 from lbmpy.stencils import get_stencil
 
-from lbmpy.advanced_streaming import PeriodicityHandling
+from lbmpy.advanced_streaming import LBMPeriodicityHandling
 from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns, get_timesteps
 
 import pytest
@@ -114,20 +114,20 @@ def test_fully_periodic_flow(target, stencil, streaming_pattern):
 
     u_ref = np.full(domain_size + (dim,), u_x)
 
-    setter = flexible_macroscopic_values_setter(
+    setter = macroscopic_values_setter(
         lb_method, density, velocity, pdfs,
         streaming_pattern=streaming_pattern, previous_timestep=zeroth_timestep)
     setter_kernel = create_kernel(setter, ghost_layers=1, target=target).compile()
 
     getter_kernels = []
     for t in timesteps:
-        getter = flexible_macroscopic_values_getter(
+        getter = 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, target=target).compile())
 
     #   Periodicity
-    periodicity_handler = PeriodicityHandling(stencil, dh, pdfs.name, streaming_pattern=streaming_pattern, target=target)
+    periodicity_handler = LBMPeriodicityHandling(stencil, dh, pdfs.name, streaming_pattern=streaming_pattern, target=target)
 
     # Initialization and Timestep
     current_timestep = zeroth_timestep

lbmpy_tests/advanced_streaming/test_noslip_ubb.py

-import numpy as np
-import sympy as sp
-import matplotlib.pyplot as plt
-
-from pystencils.datahandling import create_data_handling
-from pystencils import create_kernel
-from pystencils.slicing import slice_from_direction, make_slice
-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.plot import boundary_handling as plot_boundaries
-
-from lbmpy.advanced_streaming import PeriodicityHandling
-from lbmpy.boundaries import LatticeBoltzmannBoundaryHandling, NoSlip, UBB
-from lbmpy.advanced_streaming.utility import is_inplace
-
-from numpy.testing import assert_allclose
-from test_periodic_pipe_with_force import PeriodicPipeFlow
-
-stencil = get_stencil('D2Q9')
-streaming_pattern = 'pull'
-noslip_flow = PeriodicPipeFlow(stencil, streaming_pattern, NoSlip())
-ubb_flow = PeriodicPipeFlow(stencil, streaming_pattern, UBB((0.0, 0.0)))
-
-noslip_flow.init()
-noslip_flow.run(100)
-noslip_u = noslip_flow.get_trimmed_velocity_array()
-
-ubb_flow.init()
-ubb_flow.run(100)
-ubb_u = ubb_flow.get_trimmed_velocity_array()
\ No newline at end of file

lbmpy_tests/advanced_streaming/test_periodic_pipe_with_force.py

@@ -6,10 +6,10 @@ from pystencils import create_kernel
 from pystencils.slicing import make_slice
 
 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.macroscopic_value_kernels import macroscopic_values_getter, macroscopic_values_setter
 from lbmpy.stencils import get_stencil
 
-from lbmpy.advanced_streaming import PeriodicityHandling
+from lbmpy.advanced_streaming import LBMPeriodicityHandling
 from lbmpy.boundaries import NoSlip, LatticeBoltzmannBoundaryHandling
 from lbmpy.advanced_streaming.utility import is_inplace, streaming_patterns, Timestep, get_timesteps
 
@@ -96,20 +96,20 @@ class PeriodicPipeFlow:
         self.velocity = (u_x,) * self.dim
         self.velocity_field = self.dh.add_array('u', self.dim)
 
-        setter = flexible_macroscopic_values_setter(
+        setter = macroscopic_values_setter(
             self.lb_method, self.density, self.velocity, self.pdfs,
             streaming_pattern=self.streaming_pattern, previous_timestep=self.zeroth_timestep)
         self.init_kernel = create_kernel(setter, ghost_layers=1).compile()
 
         self.getter_kernels = []
         for t in self.timesteps:
-            getter = flexible_macroscopic_values_getter(
+            getter = macroscopic_values_getter(
                 self.lb_method, self.density_field, self.velocity_field, self.pdfs,
                 streaming_pattern=self.streaming_pattern, previous_timestep=t)
             self.getter_kernels.append(create_kernel(getter, ghost_layers=1).compile())
 
         #   Periodicity
-        self.periodicity_handler = PeriodicityHandling(self.stencil, self.dh, self.pdfs.name, streaming_pattern=self.streaming_pattern)
+        self.periodicity_handler = LBMPeriodicityHandling(self.stencil, self.dh, self.pdfs.name, streaming_pattern=self.streaming_pattern)
 
         #   Boundary Handling
         self.wall = wall_boundary