started with list lbm implementation

- extended pystencils
- Visualization for list LBM
- not yet working

listlbm.py

+import numpy as np
+
+
+def get_fluid_cell_array(flag_arr, fluid_flag, ghost_layers):
+    inner_slice = [slice(ghost_layers, -ghost_layers)] * len(flag_arr.shape)
+    flag_arr_no_gl = flag_arr[inner_slice]
+    return np.argwhere(np.bitwise_and(flag_arr_no_gl, fluid_flag)) + ghost_layers
+
+
+def get_cell_index(fluid_cell_arr, cell):
+    assert len(fluid_cell_arr.shape) == len(cell)
+    first_coord = [slice(None, None)] + [0] * (len(fluid_cell_arr.shape) - 1)
+    begin = np.searchsorted(fluid_cell_arr[first_coord], cell[0], 'left')
+    end = np.searchsorted(fluid_cell_arr[first_coord], cell[0], 'right')
+    if begin == end:
+        raise ValueError("Element not found")
+
+    if len(fluid_cell_arr.shape) == 1:
+        return begin
+    else:
+        sub_array = fluid_cell_arr[begin:end, 1]
+        return begin + get_cell_index(sub_array, cell[1:])
+
+
+def get_index_array(fluid_cell_arr, flag_arr, ghost_layers, stencil, fluid_flag, noslip_flag):
+
+    def pdf_index(cell_index, direction_index):
+        return cell_index + direction_index * len(fluid_cell_arr)
+
+    def inverse_idx(idx):
+        return stencil.index(tuple(-d_i for d_i in stencil[idx]))
+
+    result = []
+    ctr = 0
+    for direction_idx, direction in enumerate(stencil):
+        for own_cell_idx, cell in enumerate(fluid_cell_arr):
+            inv_neighbor_cell = np.array([cell_i - dir_i for cell_i, dir_i in zip(cell, direction)])
+
+            if flag_arr[tuple(inv_neighbor_cell)] & fluid_flag:
+                neighbor_cell_idx = get_cell_index(fluid_cell_arr, inv_neighbor_cell)
+                result.append(pdf_index(neighbor_cell_idx, direction_idx))
+            elif flag_arr[tuple(inv_neighbor_cell)] & noslip_flag:  # no-slip before periodicity!
+                result.append(pdf_index(own_cell_idx, inverse_idx(direction_idx)))
+            else:
+                # periodicity handling
+                # print(inv_neighbor_cell, end="")
+                at_border = False
+                for i, x_i in enumerate(inv_neighbor_cell):
+                    if x_i == (ghost_layers - 1):
+                        inv_neighbor_cell[i] += flag_arr.shape[i] - (2 * ghost_layers)
+                        at_border = True
+                    elif x_i == flag_arr.shape[i] - ghost_layers:
+                        inv_neighbor_cell[i] -= flag_arr.shape[i] - (2 * ghost_layers)
+                        at_border = True
+                if at_border:
+                    assert flag_arr[tuple(inv_neighbor_cell)] & fluid_flag
+                    neighbor_cell_idx = get_cell_index(fluid_cell_arr, inv_neighbor_cell)
+                    result.append(pdf_index(neighbor_cell_idx, direction_idx))
+                else:
+                    raise ValueError("Could not find neighbor for {} direction {}".format(cell, direction))
+
+            ctr += 1 # TODO
+    return np.array(result, dtype=np.uint32)
+
+
+def plot_index_array(fluid_cell_arr, stencil, ghost_layers=1, index_arr=None, **kwargs):
+    """Visualizes index array.
+
+    Args:
+        fluid_cell_arr: array of fluid cells
+        stencil: stencil
+        ghost_layers: number of ghost layers
+        index_arr: index array, or None to show pdf index
+        **kwargs: passed to LbGrid.add_arrow
+
+    Returns:
+        LbGrid
+    """
+    from lbmpy.plot2d import LbGrid
+    x_max, y_max = np.max(fluid_cell_arr, axis=0)
+    grid = LbGrid(x_max + 1 - ghost_layers, y_max + 1 - ghost_layers)
+    index = 0
+    for direction in stencil:
+        for fluid_cell in fluid_cell_arr:
+            annotation = index_arr[index] if index_arr is not None else index
+            grid.add_arrow((fluid_cell[0] - ghost_layers, fluid_cell[1] - ghost_layers),
+                           direction, direction, str(annotation), **kwargs)
+            index += 1
+    return grid
+
+
+if __name__ == '__main__':
+    x = np.arange(9 * 7).reshape(9, 7)
+    ind = np.argwhere(x > 50)
+    print(ind)
+    res = get_cell_index(ind, (7, 2))
+    print("result index", res)
+    print("element ", ind[res])
+
+

plot2d.py

 from itertools import cycle
 import matplotlib.patches as patches
+from matplotlib.text import Text
 
 from pystencils import make_slice
 from pystencils.plot2d import *
@@ -102,7 +103,8 @@ class LbGrid:
                 self._patches.append(patches.Rectangle((x, y), 1.0, 1.0, fill=False, linewidth=3, color='#bbbbbb'))
 
         self._cellBoundaries = dict()  # mapping cell to rectangle patch
-        self._arrows = dict()  # mapping (cell, direction) tuples to arrow patches
+        self.arrows = dict()  # mapping (cell, direction) tuples to arrow patches
+        self.annotations = dict()
 
     def add_cell_boundary(self, cell, **kwargs):
         """Draws a rectangle around a single cell. Keyword arguments are passed to the matplotlib Rectangle patch"""
@@ -117,14 +119,16 @@ class LbGrid:
             for y in range(self._yCells):
                 self.add_cell_boundary((x, y), **kwargs)
 
-    def add_arrow(self, cell, arrow_position, arrow_direction, **kwargs):
+    def add_arrow(self, cell, arrow_position, arrow_direction, annotation='', **kwargs):
         """
         Draws an arrow in a cell. If an arrow exists already at this position, it is replaced.
 
-        :param cell: cell coordinate as tuple (x,y)
-        :param arrow_position: each cell has 9 possible positions specified as tuple e.g. upper left (-1, 1)
-        :param arrow_direction: direction of the arrow as (x,y) tuple
-        :param kwargs: arguments passed directly to the FancyArrow patch of matplotlib
+        Args:
+            cell: cell coordinate as tuple (x,y)
+            arrow_position: each cell has 9 possible positions specified as tuple e.g. upper left (-1, 1)
+            arrow_direction: direction of the arrow as (x,y) tuple
+            annotation: text to display at end of arrow
+            kwargs: arguments passed directly to the FancyArrow patch of matplotlib
         """
         cell_midpoint = (0.5 + cell[0], 0.5 + cell[1])
 
@@ -133,19 +137,22 @@ class LbGrid:
 
         if arrow_position == (0, 0):
             del kwargs['width']
-            self._arrows[(cell, arrow_position)] = patches.Circle(cell_midpoint, radius=0.03, **kwargs)
+            self.arrows[(cell, arrow_position)] = patches.Circle(cell_midpoint, radius=0.03, **kwargs)
+            if annotation:
+                self.annotations[(cell, arrow_position)] = Text(*cell_midpoint, annotation)
         else:
             arrow_midpoint = (cell_midpoint[0] + arrow_position[0] * 0.25,
                               cell_midpoint[1] + arrow_position[1] * 0.25)
             length = 0.75
             arrow_start = (arrow_midpoint[0] - arrow_direction[0] * 0.25 * length,
                            arrow_midpoint[1] - arrow_direction[1] * 0.25 * length)
-
-            patch = patches.FancyArrow(arrow_start[0], arrow_start[1],
-                                       0.25 * length * arrow_direction[0],
-                                       0.25 * length * arrow_direction[1],
-                                       **kwargs)
-            self._arrows[(cell, arrow_position)] = patch
+            arrow_direction = (0.25 * length * arrow_direction[0],
+                               0.25 * length * arrow_direction[1])
+            arrow_end = tuple(a + b for a, b in zip(arrow_start, arrow_direction))
+            patch = patches.FancyArrow(*arrow_start, *arrow_direction, **kwargs)
+            self.arrows[(cell, arrow_position)] = patch
+            if annotation:
+                self.annotations[(cell, arrow_position)] = Text(*arrow_end, annotation)
 
     def fill_with_default_arrows(self, **kwargs):
         """Fills the complete grid with the default pdf arrows"""
@@ -164,9 +171,12 @@ class LbGrid:
         for p in self._patches:
             ax.add_patch(p)
 
-        for arrow_patch in self._arrows.values():
+        for arrow_patch in self.arrows.values():
             ax.add_patch(arrow_patch)
 
+        for text_obj in self.annotations.values():
+            ax.add_artist(text_obj)
+
         offset = 0.1
         ax.set_xlim(-offset, self._xCells+offset)
         ax.set_xlim(-offset, self._xCells + offset)