Updated boundaries_in_kernel

lbmpy/advanced_streaming/indexing.py

@@ -5,6 +5,7 @@ import pystencils as ps
 from pystencils.data_types import TypedSymbol, create_type
 from pystencils.backends.cbackend import CustomCodeNode
 
+from lbmpy.stencils import get_stencil
 from lbmpy.advanced_streaming.utility import get_accessor, inverse_dir_index, is_inplace, Timestep
 
 from itertools import product
@@ -30,6 +31,8 @@ class BetweenTimestepsIndexing:
 
     @property
     def inverse_dir_symbol(self):
+        """Symbol denoting the inversion of a PDF field index. 
+        Use only at top-level of index to f_out or f_in, otherwise it can't be correctly replaced."""
         return sp.IndexedBase('invdir')
 
     #   =============================
@@ -42,6 +45,9 @@ class BetweenTimestepsIndexing:
             raise ValueError('Cannot create index arrays for both kinds of timesteps for inplace streaming pattern '
                              + streaming_pattern)
 
+        if isinstance(stencil, str):
+            stencil = get_stencil(stencil)
+
         prev_accessor = get_accessor(streaming_pattern, prev_timestep)
         next_accessor = get_accessor(streaming_pattern, prev_timestep.next())
 
@@ -193,9 +199,12 @@ class BetweenTimestepsIndexing:
 class NeighbourOffsetArrays(CustomCodeNode):
 
     @staticmethod
-    def symbolic_neighbour_offset(dir_idx, dim):
-        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
-                      for symbol in NeighbourOffsetArrays._offset_symbols(dim)])
+    def neighbour_offset(dir_idx, stencil):
+        if isinstance(sp.sympify(dir_idx), sp.Integer):
+            return stencil[dir_idx]
+        else:
+            return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
+                         for symbol in NeighbourOffsetArrays._offset_symbols(len(stencil[0]))])
 
     @staticmethod
     def _offset_symbols(dim):

lbmpy/boundaries/boundaries_in_kernel.py

 import sympy as sp
 
 from lbmpy.boundaries.boundaryhandling import LbmWeightInfo
+from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing
+from lbmpy.advanced_streaming.utility import Timestep, get_accessor
 from pystencils.boundaries.boundaryhandling import BoundaryOffsetInfo
 from pystencils.assignment import Assignment
 from pystencils.astnodes import Block, Conditional, LoopOverCoordinate, SympyAssignment
@@ -56,6 +58,8 @@ def border_conditions(direction, field, ghost_layers=1):
         result = sp.And(border_condition, *other_min, *other_max)
         return type_all_numbers(result, loop_ctr.dtype)
 
+#   TODO: Function is called nowhere... remove it?
+
 
 def transformed_boundary_rule(boundary, accessor_func, field, direction_symbol, lb_method, **kwargs):
     tmp_field = field.new_field_with_different_name("t")
@@ -82,27 +86,24 @@ def transformed_boundary_rule(boundary, accessor_func, field, direction_symbol,
     return ac.main_assignments[0].rhs
 
 
-def boundary_conditional(boundary, direction, read_of_next_accessor, lb_method, output_field, cse=False):
+def boundary_conditional(boundary, direction, streaming_pattern, prev_timestep, lb_method, output_field, cse=False):
     stencil = lb_method.stencil
-    tmp_field = output_field.new_field_with_different_name("t")
 
     dir_indices = direction_indices_in_direction(direction, stencil)
+    indexing = BetweenTimestepsIndexing(output_field, lb_method.stencil, prev_timestep, streaming_pattern)
+    f_out, f_in = indexing.proxy_fields
+    inv_dir = indexing.inverse_dir_symbol
 
     assignments = []
     for direction_idx in dir_indices:
-        rule = boundary(tmp_field, direction_idx, lb_method, index_field=None)
-        boundary_subs = boundary_substitutions(lb_method)
-        rule = [a.subs(boundary_subs) for a in rule]
-
-        rhs_substitutions = {tmp_field(i): sym for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
-        offset = stencil[direction_idx]
-        inv_offset = inverse_direction(offset)
-        inv_idx = stencil.index(inv_offset)
-
-        lhs_substitutions = {
-            tmp_field[offset](inv_idx): read_of_next_accessor(output_field, stencil)[inv_idx]}
-        rule = [Assignment(a.lhs.subs(lhs_substitutions), a.rhs.subs(rhs_substitutions)) for a in rule]
-        ac = AssignmentCollection([rule[-1]], rule[:-1]).new_without_subexpressions()
+        rule = boundary(f_out, f_in, direction_idx, inv_dir, lb_method, index_field=None)
+
+        #   rhs: replace f_out by post collision symbols.
+        rhs_substitutions = {f_out(i): sym for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
+        rule = AssignmentCollection([rule]).new_with_substitutions(rhs_substitutions)
+        rule = indexing.substitute_proxies(rule)
+
+        ac = rule.new_without_subexpressions()
         assignments += ac.main_assignments
 
     border_cond = border_conditions(direction, output_field, ghost_layers=1)
@@ -112,8 +113,10 @@ def boundary_conditional(boundary, direction, read_of_next_accessor, lb_method,
     return Conditional(border_cond, Block(assignments))
 
 
-def update_rule_with_push_boundaries(collision_rule, field, boundary_spec, accessor, read_of_next_accessor):
+def update_rule_with_push_boundaries(collision_rule, field, boundary_spec, 
+                                     streaming_pattern='pull', timestep=Timestep.BOTH):
     method = collision_rule.method
+    accessor = get_accessor(streaming_pattern, timestep)
     loads = [Assignment(a, b) for a, b in zip(method.pre_collision_pdf_symbols, accessor.read(field, method.stencil))]
     stores = [Assignment(a, b) for a, b in
               zip(accessor.write(field, method.stencil), method.post_collision_pdf_symbols)]
@@ -122,7 +125,7 @@ def update_rule_with_push_boundaries(collision_rule, field, boundary_spec, acces
     result.subexpressions = loads + result.subexpressions
     result.main_assignments += stores
     for direction, boundary in boundary_spec.items():
-        cond = boundary_conditional(boundary, direction, read_of_next_accessor, method, field)
+        cond = boundary_conditional(boundary, direction, streaming_pattern, timestep, method, field)
         result.main_assignments.append(cond)
 
     if 'split_groups' in result.simplification_hints:

lbmpy/boundaries/boundaryconditions.py

@@ -137,7 +137,7 @@ class UBB(Boundary):
         assert self.dim == lb_method.dim, \
                f"Dimension of UBB ({self.dim}) does not match dimension of method ({lb_method.dim})"
 
-        neighbor_offset = NeighbourOffsetArrays.symbolic_neighbour_offset(direction, lb_method.dim)
+        neighbor_offset = NeighbourOffsetArrays.neighbour_offset(direction, lb_method.stencil)
 
         velocity = tuple(v_i.get_shifted(*neighbor_offset)
                          if isinstance(v_i, Field.Access) and not vel_from_idx_field
@@ -153,7 +153,7 @@ class UBB(Boundary):
         weight_of_direction = LbmWeightInfo.weight_of_direction
         vel_term = 2 / c_s_sq \
             * sum([d_i * v_i for d_i, v_i in zip(neighbor_offset, velocity)]) \
-            * weight_of_direction(direction)
+            * weight_of_direction(direction, lb_method)
 
         # Better alternative: in conserved value computation
         # rename what is currently called density to "virtual_density"
@@ -227,7 +227,7 @@ class NeumannByCopy(Boundary):
         return [NeighbourOffsetArrays(lb_method.stencil)]
 
     def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
-        neighbour_offset = NeighbourOffsetArrays.symbolic_neighbour_offset(dir_symbol, lb_method.dim)
+        neighbour_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
         return [Assignment(f_in(inv_dir[dir_symbol]), f_out(inv_dir[dir_symbol])),
                 Assignment(f_out[neighbour_offset](dir_symbol), f_out(dir_symbol))]
 
@@ -249,7 +249,7 @@ class StreamInConstant(Boundary):
         return [NeighbourOffsetArrays(lb_method.stencil)]
 
     def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
-        neighbour_offset = NeighbourOffsetArrays.symbolic_neighbour_offset(dir_symbol, lb_method.dim)
+        neighbour_offset = NeighbourOffsetArrays.neighbour_offset(dir_symbol, lb_method.stencil)
         return [Assignment(f_in(inv_dir[dir_symbol]), self._constant),
                 Assignment(f_out[neighbour_offset](dir_symbol), self._constant)]
 

lbmpy/boundaries/boundaryhandling.py

@@ -155,8 +155,11 @@ class LbmWeightInfo(CustomCodeNode):
     # --------------------------- Functions to be used by boundaries --------------------------
 
     @staticmethod
-    def weight_of_direction(dir_idx):
-        return sp.IndexedBase(LbmWeightInfo.WEIGHTS_SYMBOL, shape=(1,))[dir_idx]
+    def weight_of_direction(dir_idx, lb_method=None):
+        if isinstance(sp.sympify(dir_idx), sp.Integer):
+            return lb_method.weights[dir_idx].evalf()
+        else:
+            return sp.IndexedBase(LbmWeightInfo.WEIGHTS_SYMBOL, shape=(1,))[dir_idx]
 
     # ---------------------------------- Internal ---------------------------------------------
 

lbmpy/session.py

@@ -3,6 +3,7 @@ import sympy as sp
 
 import lbmpy.plot as plt
 import pystencils as ps
+from lbmpy.advanced_streaming import *
 from lbmpy.boundaries import *
 from lbmpy.creationfunctions import *
 from lbmpy.geometry import *