diff --git a/listlbm.py b/listlbm.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd798463bd6273a95ed505deafd0b2ce58aa25b3
--- /dev/null
+++ b/listlbm.py
@@ -0,0 +1,100 @@
+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])
+
+
diff --git a/plot2d.py b/plot2d.py
index d25231590804bb8f79bfa06f04e874156e645056..6c8bf45798bd75ceca7fd529f305997c1086e3ad 100644
--- a/plot2d.py
+++ b/plot2d.py
@@ -1,5 +1,6 @@
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)