Advanced Streaming Indexing WIP

lbmpy/advanced_streaming/__init__.py

-from .flexible_boundaries import AdvancedStreamingBoundaryOffsetInfo, \
-    FlexibleNoSlip, create_flexible_lbm_boundary_kernel
+from .indexing import AdvancedStreamingIndexing
 
-__all__ = [ 'AdvancedStreamingBoundaryOffsetInfo',  
-            'FlexibleNoSlip', 
-            'create_flexible_lbm_boundary_kernel' ]
\ No newline at end of file
+__all__ = [ 'AdvancedStreamingIndexing' ]
\ No newline at end of file

lbmpy/advanced_streaming/indexing.py

+import numpy as np
+import sympy as sp
+import pystencils as ps
+
+from pystencils.data_types import TypedSymbol, create_type
+
+from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
+    StreamPushTwoFieldsAccessor, \
+    AAEvenTimeStepAccessor, \
+    AAOddTimeStepAccessor, \
+    EsoTwistEvenTimeStepAccessor, \
+    EsoTwistOddTimeStepAccessor
+
+class AdvancedStreamingIndexing:
+
+    #   =======================================
+    #       Symbols for usage in boundaries
+    #   =======================================
+
+    def proxy_fields(self):
+        q = len(self.stencil)
+        d = len(self.stencil[0])
+        return ps.fields('f_out({q}), f_in({q}): [{d}}]'.format(q=q, d=d))
+
+    def dir_symbol(self):
+        return TypedSymbol('dir', create_type(np.int64))
+
+    def inverse_dir_symbol(self):
+        return sp.IndexedBase('invdir')
+
+    #   =============================
+    #       Constructor and State
+    #   =============================
+
+    def __init__(self, pdf_field, stencil, between_timesteps = 'both', kernel_type='pull'):
+        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 kernel_type in ['aa', 'esotwist']:
+            raise ValueError('Cannot create an offset info for both kinds of timesteps for kernel type ' + kernel_type)
+
+        odd_to_even = (between_timesteps == 'odd_to_even')
+
+        even_accessors = {
+            'pull': StreamPullTwoFieldsAccessor,
+            'push': StreamPushTwoFieldsAccessor,
+            'aa': AAEvenTimeStepAccessor,
+            'esotwist': EsoTwistEvenTimeStepAccessor
+        }
+
+        odd_accessors = {
+            'pull': StreamPullTwoFieldsAccessor,
+            'push': StreamPushTwoFieldsAccessor,
+            'aa': AAOddTimeStepAccessor,
+            'esotwist': EsoTwistOddTimeStepAccessor
+        }
+
+        try:
+            even_accessor = even_accessors[kernel_type]
+            odd_accessor = odd_accessors[kernel_type]
+        except KeyError:
+            raise ValueError(
+                "Invalid value of parameter 'kernel_type'.", kernel_type)
+
+        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)
+
+        self.outward_accesses = outward_accesses
+        self.inward_accesses = inward_accesses
+
+        self.pdf_field = pdf_field
+        self.stencil = stencil
+        self.directions = ['x', 'y', 'z'][:len(stencil[0])]
+
+        #   Collection of translation arrays required in generated code
+        self.required_arrays = set()
+
+    def _index_array_symbol(self, f_dir, inverse, dir_symbol):
+        assert f_dir in ['in', 'out']
+        name = "f_{d}{inv}_dir_idx".format(d=f_dir, inv='_inv' if inverse else '')
+        self.required_arrays.add(name)
+        return sp.IndexedBase(name)[dir_symbol]
+
+    def _offset_array_symbols(self, f_dir, inverse, dir_symbol):
+        assert f_dir in ['in', 'out']
+        name_base = "f_{d}{inv}_offsets_".format(d=f_dir, inv='_inv' if inverse else '')
+        names = [ name_base + d for d in self.directions ]
+        self.required_arrays |= set(names)
+        return [ sp.IndexedBase(n)[dir_symbol] for n in names ]
+
+
+    #   =================================
+    #       Proxy fields substitution
+    #   =================================
+
+    def substitute_proxies(self, assignments):
+        # Get the accessor lists for the streaming pattern (here AA)
+        outward_accesses = self.outward_accesses
+        inward_accesses = self.inward_accesses
+        
+        f_out, f_in = self.proxy_fields()
+        dir_symbol = self.dir_symbol()
+        inv_dir = self.inverse_dir_symbol()
+
+        # Substitute all proxy field accesses
+
+        if not isinstance(assignments, ps.AssignmentCollection):
+            assignments = ps.AssignmentCollection([assignments])
+        
+        accessor_subs = dict()
+
+        for fa in assignments.atoms(ps.Field.Access):
+            if fa.field == f_out:
+                if fa.offsets == (0,0):
+                    if isinstance(fa.index[0], int):
+                        accessor_subs[fa] = outward_accesses[fa.index[0]]
+                    elif fa.index[0] == dir_symbol:
+                        accessor_subs[fa] = self.pdf_field[ 
+                            self._offset_array_symbols('out', False, dir_symbol)
+                            ](self._index_array_symbol('out', False, dir_symbol))
+                    else:
+                        # other cases like inverse direction, etc.
+                        pass
+                else:
+                    # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
+                    pass
+
+            elif fa.field == f_in:
+                if fa.offsets == (0,0):
+                    if isinstance(fa.index[0], int):
+                        accessor_subs[fa] = inward_accesses[fa.index[0]]
+                    elif fa.index[0] == inv_dir[dir_symbol]:
+                        accessor_subs[fa] = self.pdf_field[ 
+                            self._offset_array_symbols('in', True, dir_symbol)
+                            ](self._index_array_symbol('in', True, dir_symbol))
+                    else:
+                        # other cases
+                        pass
+                else:
+                    # non-trivial neighbour accesses -> add neighbour offset to streaming pattern offsets
+                    pass
+
+            else:
+                pass
+
+        return assignments.new_with_substitutions(accessor_subs)
+
+#   end class AdvancedStreamingIndexing
+
+    
\ No newline at end of file