Added different pdf storage schemes (AA pattern, esotwist) to high level interface

lbmpy/creationfunctions.py

@@ -184,7 +184,9 @@ from pystencils.simp import add_subexpressions_for_field_reads
 from pystencils.stencil import have_same_entries
 import lbmpy.forcemodels as forcemodels
 from lbmpy.simplificationfactory import create_simplification_strategy
-from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, PeriodicTwoFieldsAccessor, CollideOnlyInplaceAccessor
+from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, PeriodicTwoFieldsAccessor, CollideOnlyInplaceAccessor, \
+    EsoTwistEvenTimeStepAccessor, EsoTwistOddTimeStepAccessor, AAEvenTimeStepAccessor, AAOddTimeStepAccessor, \
+    StreamPushTwoFieldsAccessor, PdfFieldAccessor
 from lbmpy.updatekernels import create_lbm_kernel, create_stream_pull_with_output_kernel
 
 
@@ -256,18 +258,33 @@ def create_lb_update_rule(collision_rule=None, optimization={}, **kwargs):
     else:
         dst_field = src_field.new_field_with_different_name(params['temporary_field_name'])
 
-    if params['kernel_type'] == 'stream_pull_collide':
+    kernel_type = params['kernel_type']
+    if isinstance(kernel_type, PdfFieldAccessor):
+        accessor = kernel_type
+        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
+    elif params['kernel_type'] == 'stream_pull_collide':
         accessor = StreamPullTwoFieldsAccessor
         if any(opt_params['builtin_periodicity']):
             accessor = PeriodicTwoFieldsAccessor(opt_params['builtin_periodicity'], ghost_layers=1)
         return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
     elif params['kernel_type'] == 'collide_only':
-        result = create_lbm_kernel(collision_rule, src_field, src_field, CollideOnlyInplaceAccessor)
+        result = create_lbm_kernel(collision_rule, src_field, src_field, CollideOnlyInplaceAccessor())
         return add_subexpressions_for_field_reads(result, subexpressions=False, main_assignments=True)
     elif params['kernel_type'] == 'stream_pull_only':
         return create_stream_pull_with_output_kernel(lb_method, src_field, dst_field, params['output'])
     else:
-        raise ValueError("Invalid value of parameter 'kernel_type'", params['kernel_type'])
+        kernel_type_to_accessor = {
+            'collide_stream_push': StreamPushTwoFieldsAccessor,
+            'esotwist_even': EsoTwistEvenTimeStepAccessor,
+            'esotwist_odd': EsoTwistOddTimeStepAccessor,
+            'aa_even': AAEvenTimeStepAccessor,
+            'aa_odd': AAOddTimeStepAccessor,
+        }
+        try:
+            accessor = kernel_type_to_accessor[kernel_type]()
+        except KeyError:
+            raise ValueError("Invalid value of parameter 'kernel_type'", params['kernel_type'])
+        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
 
 
 @disk_cache_no_fallback

lbmpy/fieldaccess.py

@@ -14,7 +14,7 @@ __all__ = ['PdfFieldAccessor', 'CollideOnlyInplaceAccessor', 'StreamPullTwoField
 from pystencils.astnodes import LoopOverCoordinate
 
 
-class PdfFieldAccessor(object):
+class PdfFieldAccessor:
     """
     Defines how data is read and written in an LBM time step.
 
@@ -34,11 +34,20 @@ class PdfFieldAccessor(object):
         """Returns sequence  of field accesses for all stencil values where pdfs are written to"""
         pass
 
+    @property
+    @abc.abstractmethod
+    def is_inplace(self):
+        """True if accessor writes the same entries that are read. In this case all values are first read into
+        temporaries."""
+        pass
+
 
 # ----------------------------------------------- Implementation -------------------------------------------------------
 
 
 class CollideOnlyInplaceAccessor(PdfFieldAccessor):
+    is_inplace = True
+
     @staticmethod
     def read(field, stencil):
         return [field(i) for i in range(len(stencil))]
@@ -49,6 +58,8 @@ class CollideOnlyInplaceAccessor(PdfFieldAccessor):
 
 
 class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
+    is_inplace = False
+
     @staticmethod
     def read(field, stencil):
         return [field[inverse_direction(d)](i) for i, d in enumerate(stencil)]
@@ -58,6 +69,18 @@ class StreamPullTwoFieldsAccessor(PdfFieldAccessor):
         return [field(i) for i in range(len(stencil))]
 
 
+class StreamPushTwoFieldsAccessor(PdfFieldAccessor):
+    is_inplace = False
+
+    @staticmethod
+    def read(field, stencil):
+        return [field(i) for i in range(len(stencil))]
+
+    @staticmethod
+    def write(field, stencil):
+        return [field[d](i) for i, d in enumerate(stencil)]
+
+
 class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
     """Access scheme that builds periodicity into the kernel.
 
@@ -66,6 +89,8 @@ class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
     of assumed ghost layers of the field. For the periodic kernel itself no ghost layers are required,
     however other kernels might need them.
     """
+    is_inplace = False
+
     def __init__(self, periodicity, ghost_layers=0):
         self._periodicity = periodicity
         self._ghostLayers = ghost_layers
@@ -103,6 +128,8 @@ class PeriodicTwoFieldsAccessor(PdfFieldAccessor):
 
 
 class AAEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
     @staticmethod
     def read(field, stencil):
         return [field(i) for i in range(len(stencil))]
@@ -113,6 +140,8 @@ class AAEvenTimeStepAccessor(PdfFieldAccessor):
 
 
 class AAOddTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
     @staticmethod
     def read(field, stencil):
         res = []
@@ -128,6 +157,8 @@ class AAOddTimeStepAccessor(PdfFieldAccessor):
 
 
 class EsoTwistOddTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
     @staticmethod
     def read(field, stencil):
         result = []
@@ -147,6 +178,8 @@ class EsoTwistOddTimeStepAccessor(PdfFieldAccessor):
 
 
 class EsoTwistEvenTimeStepAccessor(PdfFieldAccessor):
+    is_inplace = True
+
     @staticmethod
     def read(field, stencil):
         result = []

lbmpy/macroscopic_value_kernels.py

@@ -14,6 +14,20 @@ def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
     return setter_eqs
 
 
+def macroscopic_values_getter(lb_method, density, velocity, pdfs):
+    cqc = lb_method.conserved_quantity_computation
+    assert not (velocity is None and density is None)
+    output_spec = {}
+    if velocity is not None:
+        output_spec['velocity'] = velocity
+    if density is not None:
+        output_spec['density'] = density
+    return cqc.output_equations_from_pdfs(pdfs, output_spec)
+
+
+macroscopic_values_setter = pdf_initialization_assignments
+
+
 def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None, field_layout='numpy', target='cpu'):
     """
     Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)

lbmpy/updatekernels.py

@@ -3,6 +3,7 @@ import sympy as sp
 
 from pystencils import Field, Assignment, AssignmentCollection
 from pystencils.field import create_numpy_array_with_layout, layout_string_to_tuple
+from pystencils.simp import add_subexpressions_for_field_reads
 from pystencils.sympyextensions import fast_subs
 from lbmpy.methods.abstractlbmethod import LbmCollisionRule
 from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor
@@ -17,13 +18,17 @@ def create_lbm_kernel(collision_rule, input_field, output_field, accessor):
     Args:
         collision_rule:  instance of LbmCollisionRule, defining the collision step
         input_field: field used for reading pdf values
-        output_field: field used for writing pdf values (may be the same as input field for certain accessors)
+        output_field: field used for writing pdf values
+                      if accessor.is_inplace this parameter is ignored
         accessor: instance of PdfFieldAccessor, defining where to read and write values
                   to create e.g. a fused stream-collide kernel
 
     Returns:
         LbmCollisionRule where pre- and post collision symbols have been replaced
     """
+    if accessor.is_inplace:
+        output_field = input_field
+
     method = collision_rule.method
     pre_collision_symbols = method.pre_collision_pdf_symbols
     post_collision_symbols = method.post_collision_pdf_symbols
@@ -44,6 +49,9 @@ def create_lbm_kernel(collision_rule, input_field, output_field, accessor):
             new_split_groups.append([fast_subs(e, substitutions) for e in split_group])
         result.simplification_hints['split_groups'] = new_split_groups
 
+    if accessor.is_inplace:
+        result = add_subexpressions_for_field_reads(result, subexpressions=False, main_assignments=True)
+
     return result