Completed Force on Boundary, and UBB

lbmpy/advanced_streaming/boundaryconditions.py

-import pystencils as ps
+import sympy as sp
+from pystencils import Assignment, Field
+from lbmpy.boundaries.boundaryhandling import LbmWeightInfo
+from pystencils.data_types import create_type
+from pystencils.sympyextensions import get_symmetric_part
+from lbmpy.simplificationfactory import create_simplification_strategy
+from lbmpy.advanced_streaming.indexing import NeighbourOffsetArraysForStencil
 
 
 class FlexibleBoundary:
@@ -45,8 +51,7 @@ class FlexibleBoundary:
         data-name to data for each element that should be initialized"""
         return None
 
-    @property
-    def additional_code_nodes(self):
+    def get_additional_code_nodes(self, lb_method):
         """Return a list of code nodes that will be added in the generated code before the index field loop."""
         return []
 
@@ -76,7 +81,7 @@ class FlexibleNoSlip(FlexibleBoundary):
     """
 
     def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
-        return ps.Assignment(f_in(inv_dir[dir_symbol]), f_out(dir_symbol))
+        return Assignment(f_in(inv_dir[dir_symbol]), f_out(dir_symbol))
 
     def __hash__(self):
         return hash(self.name)
@@ -87,3 +92,84 @@ class FlexibleNoSlip(FlexibleBoundary):
         return self.name == other.name
 
 # end class FlexibleNoSlip
+
+
+class FlexibleUBB(FlexibleBoundary):
+    """Velocity bounce back boundary condition, enforcing specified velocity at obstacle"""
+
+    def __init__(self, velocity, adapt_velocity_to_force=False, dim=None, name=None):
+        """
+        Args:
+            velocity: can either be a constant, an access into a field, or a callback function.
+                      The callback functions gets a numpy record array with members, 'x','y','z', 'dir' (direction)
+                      and 'velocity' which has to be set to the desired velocity of the corresponding link
+            adapt_velocity_to_force:
+        """
+        super(FlexibleUBB, self).__init__(name)
+        self._velocity = velocity
+        self._adaptVelocityToForce = adapt_velocity_to_force
+        if callable(self._velocity) and not dim:
+            raise ValueError("When using a velocity callback the dimension has to be specified with the dim parameter")
+        elif not callable(self._velocity):
+            dim = len(velocity)
+        self.dim = dim
+
+    @property
+    def additional_data(self):
+        if callable(self._velocity):
+            return [('vel_%d' % (i,), create_type("double")) for i in range(self.dim)]
+        else:
+            return []
+
+    @property
+    def additional_data_init_callback(self):
+        if callable(self._velocity):
+            return self._velocity
+
+    @property
+    def get_additional_code_nodes(self, lb_method):
+        return [LbmWeightInfo(lb_method), NeighbourOffsetArraysForStencil(lb_method.stencil)]
+    
+
+    def __call__(self, f_out, f_in, dir_symbol, inv_dir, lb_method, index_field):
+        vel_from_idx_field = callable(self._velocity)
+        vel = [index_field(f'vel_{i}') for i in range(self.dim)] if vel_from_idx_field else self._velocity
+        direction = dir_symbol
+
+        assert self.dim == lb_method.dim, f"Dimension of UBB ({self.dim}) does not match dimension of method ({lb_method.dim})"
+
+        neighbor_offset = NeighbourOffsetArraysForStencil.symbolic_neighbour_offset_from_dir(
+            direction, lb_method.dim)
+
+        velocity = tuple(v_i.get_shifted(*neighbor_offset)
+                         if isinstance(v_i, Field.Access) and not vel_from_idx_field
+                         else v_i
+                         for v_i in vel)
+
+        if self._adaptVelocityToForce:
+            cqc = lb_method.conserved_quantity_computation
+            shifted_vel_eqs = cqc.equilibrium_input_equations_from_init_values(velocity=velocity)
+            velocity = [eq.rhs for eq in shifted_vel_eqs.new_filtered(cqc.first_order_moment_symbols).main_assignments]
+
+        c_s_sq = sp.Rational(1, 3)
+        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)
+
+        # Better alternative: in conserved value computation
+        # rename what is currently called density to "virtual_density"
+        # provide a new quantity density, which is constant in case of incompressible models
+        if not lb_method.conserved_quantity_computation.zero_centered_pdfs:
+            cqc = lb_method.conserved_quantity_computation
+            density_symbol = sp.Symbol("rho")
+            pdf_field_accesses = [f_out(i) for i in range(len(lb_method.stencil))]
+            density_equations = cqc.output_equations_from_pdfs(pdf_field_accesses, {'density': density_symbol})
+            density_symbol = lb_method.conserved_quantity_computation.defined_symbols()['density']
+            result = density_equations.all_assignments
+            result += [Assignment(f_in(inv_dir[direction]),
+                                  f_out(direction) - vel_term * density_symbol)]
+            return result
+        else:
+            return [Assignment(f_in(inv_dir[direction]),
+                               f_out(direction) - vel_term)]

lbmpy/advanced_streaming/boundaryhandling.py

@@ -2,6 +2,7 @@ import numpy as np
 from lbmpy.advanced_streaming.indexing import BetweenTimestepsIndexing
 from lbmpy.advanced_streaming.utility import is_inplace, Timestep, AccessPdfValues
 from pystencils import Field, Assignment, create_indexed_kernel
+from pystencils.stencil import inverse_direction
 from pystencils.boundaries import BoundaryHandling
 from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object
 
@@ -112,7 +113,32 @@ class FlexibleLBMBoundaryHandling(BoundaryHandling):
         return dh.reduce_float_sequence(list(result), 'sum')
 
     def _force_on_boundary(self, boundary_obj, prev_timestep):
-        raise NotImplementedError()
+        dh = self._data_handling
+        ff_ghost_layers = dh.ghost_layers_of_field(self.flag_interface.flag_field_name)
+        method = self._lb_method
+        stencil = np.array(method.stencil)
+        result = np.zeros(self.dim)
+
+        for b in dh.iterate(ghost_layers=ff_ghost_layers):
+            obj_to_ind_list = b[self._index_array_name].boundary_object_to_index_list
+            pdf_array = b[self._field_name]
+            if boundary_obj in obj_to_ind_list:
+                ind_arr = obj_to_ind_list[boundary_obj]
+                acc_out = AccessPdfValues(dh.fields[self._field_name], self._lb_method.stencil,
+                                          streaming_pattern=self._streaming_pattern, timestep=prev_timestep,
+                                          streaming_dir='out')
+                acc_in = AccessPdfValues(dh.fields[self._field_name], self._lb_method.stencil,
+                                         streaming_pattern=self._streaming_pattern, timestep=prev_timestep.next(),
+                                         streaming_dir='in')
+                acc_fluid = acc_out if boundary_obj.inner_or_boundary else acc_in
+                acc_boundary = acc_in if boundary_obj.inner_or_boundary else acc_out
+                fluid_values = acc_fluid.collect_from_index_list(pdf_array, ind_arr)
+                boundary_values = acc_boundary.collect_from_index_list(pdf_array, ind_arr)
+                values = fluid_values + boundary_values
+                forces = stencil[ind_arr['dir']] * values[:, np.newaxis]
+                result += forces.sum(axis=0)
+
+        return dh.reduce_float_sequence(list(result), 'sum')
 
 # end class FlexibleLBMBoundaryHandling
 
@@ -140,7 +166,7 @@ def create_advanced_streaming_boundary_kernel(pdf_field, index_field, lb_method,
 
     #   Code Elements ahead of the loop
     index_arrs_node = indexing.create_code_node()
-    for node in boundary_functor.additional_code_nodes[::-1]:
+    for node in boundary_functor.get_additional_code_nodes(lb_method)[::-1]:
         kernel.body.insert_front(node)
     kernel.body.insert_front(index_arrs_node)
     return kernel

lbmpy/advanced_streaming/indexing.py

@@ -9,6 +9,8 @@ from lbmpy.advanced_streaming.utility import get_accessor, inverse_dir_index, is
 
 from itertools import product
 
+def _array_pattern(dtype, name, content):
+    return f"const {str(dtype)} {name} [] = {{ {','.join(str(c) for c in content)} }}; \n"
 
 class BetweenTimestepsIndexing:
 
@@ -167,21 +169,16 @@ class BetweenTimestepsIndexing:
 
                 index_array_symbol = indexing._index_array_symbol(f_dir, inv)
                 symbols_defined.add(index_array_symbol)
-                acc_indices = ", ".join([str(i) for i in indices])
-                code += self._array_pattern(indexing._index_dtype, index_array_symbol.name, acc_indices)
+                code += _array_pattern(indexing._index_dtype, index_array_symbol.name, indices)
 
                 offset_array_symbols = indexing._offset_array_symbols(f_dir, inv)
                 symbols_defined |= set(offset_array_symbols)
                 for d, arrsymb in enumerate(offset_array_symbols):
-                    acc_offsets = ", ".join([str(o) for o in offsets[d]])
-                    code += self._array_pattern(indexing._offsets_dtype, arrsymb.name, acc_offsets)
+                    code += _array_pattern(indexing._offsets_dtype, arrsymb.name, offsets[d])
 
             super(BetweenTimestepsIndexing.TranslationArraysNode, self).__init__(
                 code, symbols_read=set(), symbols_defined=symbols_defined)
 
-        def _array_pattern(self, dtype, name, content):
-            return f"const {str(dtype)} {name} [] = {{ {content} }}; \n"
-
         def __str__(self):
             return "Variable PDF Access Translation Arrays"
 
@@ -189,3 +186,27 @@ class BetweenTimestepsIndexing:
             return "Variable PDF Access Translation Arrays"
 
 #   end class AdvancedStreamingIndexing
+
+class NeighbourOffsetArraysForStencil(CustomCodeNode):
+
+    @staticmethod
+    def symbolic_neighbour_offset_from_dir(dir_idx, dim):
+        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
+                      for symbol in NeighbourOffsetArraysForStencil._offset_symbols(dim)])
+
+    @staticmethod
+    def _offset_symbols(dim):
+        return [TypedSymbol(f"neighbour_offset_{d}", create_type(np.int64)) for d in ['x', 'y', 'z'][:dim]]
+
+    def __init__(self, stencil, offsets_dtype=np.int32):
+        offsets_dtype = create_type(offsets_dtype)
+        dim = len(stencil[0])
+
+        array_symbols = NeighbourOffsetArraysForStencil._offset_symbols(dim)
+        code = "\n"
+        for i, arrsymb in enumerate(array_symbols):
+            code += _array_pattern(offsets_dtype, arrsymb.name, (d[i] for d in stencil))
+
+        offset_symbols = NeighbourOffsetArraysForStencil._offset_symbols(dim)
+        super(NeighbourOffsetArraysForStencil, self).__init__(code, symbols_read=set(),
+                                                 symbols_defined=set(offset_symbols))
\ No newline at end of file

lbmpy/advanced_streaming/utility.py

@@ -104,9 +104,15 @@ class AccessPdfValues:
         i = numeric_index(self.accs[d])[0]
         return pdf_arr[pos + (i,)]
 
+    def read_multiple(self, pdf_arr, indices):
+        """Returns PDF values for a list of index tuples (x, y, [z,] dir)"""
+        return np.array([self.read_pdf(pdf_arr, idx[:-1], idx[-1]) for idx in indices])
+
     def collect_from_index_list(self, pdf_arr, index_list):
-        def to_coordinate_tuple(idx):
-            return tuple(idx[v] for v in ('x', 'y', 'z')[:len(idx) - 1])
+        """To collect PDF values according to an pystencils boundary handling index list"""
+        def to_index_tuple(idx):
+            return tuple(idx[v] for v in ('x', 'y', 'z')[:len(idx) - 1] + ('dir',))
+
+        return self.read_multiple(pdf_arr, (to_index_tuple(idx) for idx in index_list))
 
-        gen = [self.read_pdf(pdf_arr, to_coordinate_tuple(idx), idx['dir']) for idx in index_list]
-        return np.array(gen)
+