More flexible boundary handling

boundary conditions can specify how the index list should be built:
- list coordinates of domain or boundary cells (previous always inner)
- list all links or only the first link

boundaries/boundaryconditions.py

@@ -7,6 +7,9 @@ from pystencils.data_types import create_type
 class Boundary:
     """Base class all boundaries should derive from"""
 
+    inner_or_boundary = True
+    single_link = False
+
     def __init__(self, name=None):
         self._name = name
 
@@ -51,17 +54,21 @@ class Boundary:
 
 
 class Neumann(Boundary):
+
+    inner_or_boundary = False
+    single_link = True
+
     def __call__(self, field, direction_symbol, **kwargs):
 
         neighbor = BoundaryOffsetInfo.offset_from_dir(direction_symbol, field.spatial_dimensions)
         if field.index_dimensions == 0:
-            return [Assignment(field[neighbor], field.center)]
+            return [Assignment(field.center, field[neighbor])]
         else:
             from itertools import product
             if not field.has_fixed_index_shape:
                 raise NotImplementedError("Neumann boundary works only for fields with fixed index shape")
             index_iter = product(*(range(i) for i in field.index_shape))
-            return [Assignment(field[neighbor](*idx), field(*idx)) for idx in index_iter]
+            return [Assignment(field(*idx), field[neighbor](*idx)) for idx in index_iter]
 
     def __hash__(self):
         # All boundaries of these class behave equal -> should also be equal
@@ -72,6 +79,10 @@ class Neumann(Boundary):
 
 
 class Dirichlet(Boundary):
+
+    inner_or_boundary = False
+    single_link = True
+
     def __init__(self, value, name="Dirchlet"):
         super().__init__(name)
         self._value = value
@@ -89,7 +100,13 @@ class Dirichlet(Boundary):
             return self._value
 
     def __call__(self, field, direction_symbol, index_field, **kwargs):
-        if self.additional_data:
-            return [Assignment(field.center, index_field("value"))]
+
         if field.index_dimensions == 0:
-            return [Assignment(field.center, self._value)]
+            return [Assignment(field, index_field("value") if self.additional_data else self._value)]
+        elif field.index_dimensions == 1:
+            assert not self.additional_data
+            if not field.has_fixed_index_shape:
+                raise NotImplementedError("Field needs fixed index shape")
+            assert len(self._value) == field.index_shape[0], "Dirichlet value does not match index shape of field"
+            return [Assignment(field(i), self._value[i]) for i in range(field.index_shape[0])]
+        raise NotImplementedError("Dirichlet boundary not implemented for fields with more than one index dimension")

boundaries/boundaryhandling.py

@@ -283,9 +283,11 @@ class BoundaryHandling:
             index_array_bd = b[self._index_array_name]
             index_array_bd.clear()
             for b_info in self._boundary_object_to_boundary_info.values():
+                boundary_obj = b_info.boundary_object
                 idx_arr = create_boundary_index_array(flag_arr, self.stencil, b_info.flag,
-                                                      self.flag_interface.domain_flag, b_info.boundary_object,
-                                                      ff_ghost_layers)
+                                                      self.flag_interface.domain_flag, boundary_obj,
+                                                      ff_ghost_layers, boundary_obj.inner_or_boundary,
+                                                      boundary_obj.single_link)
                 if idx_arr.size == 0:
                     continue
 

boundaries/createindexlist.py

@@ -6,7 +6,8 @@ try:
     import pyximport
 
     pyximport.install(language_level=3)
-    from pystencils.boundaries.createindexlistcython import create_boundary_index_list_2d, create_boundary_index_list_3d
+    from pystencils.boundaries.createindexlistcython import create_boundary_neighbor_index_list_2d, \
+        create_boundary_neighbor_index_list_3d, create_boundary_cell_index_list_2d, create_boundary_cell_index_list_3d
 
     cython_funcs_available = True
 except Exception:
@@ -25,7 +26,8 @@ def numpy_data_type_for_boundary_object(boundary_object, dim):
                     [(i[0], i[1].numpy_dtype) for i in boundary_object.additional_data], align=True)
 
 
-def _create_boundary_index_list_python(flag_field_arr, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil):
+def _create_boundary_neighbor_index_list_python(flag_field_arr, nr_of_ghost_layers, boundary_mask,
+                                                fluid_mask, stencil, single_link):
     coordinate_names = boundary_index_array_coordinate_names[:len(flag_field_arr.shape)]
     index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)])
 
@@ -39,32 +41,88 @@ def _create_boundary_index_list_python(flag_field_arr, nr_of_ghost_layers, bound
             neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)])
             if flag_field_arr[neighbor_cell] & boundary_mask:
                 result.append(cell + (dir_idx,))
+                if single_link:
+                    continue
 
     return np.array(result, dtype=index_arr_dtype)
 
 
-def create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask, nr_of_ghost_layers=1):
+def _create_boundary_cell_index_list_python(flag_field_arr, nr_of_ghost_layers, boundary_mask,
+                                            fluid_mask, stencil, single_link):
+    coordinate_names = boundary_index_array_coordinate_names[:len(flag_field_arr.shape)]
+    index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)])
+
+    result = []
+    gl = nr_of_ghost_layers
+    for cell in itertools.product(*reversed([range(gl, i - gl) for i in flag_field_arr.shape])):
+        cell = cell[::-1]
+        if not flag_field_arr[cell] & boundary_mask:
+            continue
+        for dir_idx, direction in enumerate(stencil):
+            neighbor_cell = tuple([cell_i + dir_i for cell_i, dir_i in zip(cell, direction)])
+            neighbor_is_fluid = False
+            try:
+                neighbor_is_fluid = flag_field_arr[neighbor_cell] & fluid_mask
+            except IndexError:
+                pass
+            if neighbor_is_fluid:
+                result.append(cell + (dir_idx,))
+                if single_link:
+                    continue
+
+    return np.array(result, dtype=index_arr_dtype)
+
+
+def create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask,
+                               nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False):
+    """Creates a numpy array storing links (connections) between domain cells and boundary cells.
+
+    Args:
+        flag_field: flag integer array where boundary and domain cells are marked (interpreted as bit vector)
+        stencil: list of directions, for possible links. When single_link is set to true the order matters, because
+                 then only the first link is added to the list
+        boundary_mask: cells where (cell & mask) is true are considered boundary cells
+        fluid_mask: cells where (cell & mask) is true are considered fluid/inner cells cells
+        nr_of_ghost_layers: only relevant if neighbors is True
+        inner_or_boundary: if true, the result contains the cell coordinates of the domain cells -
+                    if false the boundary cells are listed
+        single_link: if true only the first link is reported from this cell
+
+    """
     dim = len(flag_field.shape)
     coordinate_names = boundary_index_array_coordinate_names[:dim]
     index_arr_dtype = np.dtype([(name, np.int32) for name in coordinate_names] + [(direction_member_name, np.int32)])
 
+    stencil = np.array(stencil, dtype=np.int32)
+    args = (flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil, single_link)
+
     if cython_funcs_available:
-        stencil = np.array(stencil, dtype=np.int32)
         if dim == 2:
-            idx_list = create_boundary_index_list_2d(flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil)
+            if inner_or_boundary:
+                idx_list = create_boundary_neighbor_index_list_2d(*args)
+            else:
+                idx_list = create_boundary_cell_index_list_2d(*args)
         elif dim == 3:
-            idx_list = create_boundary_index_list_3d(flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil)
+            if inner_or_boundary:
+                idx_list = create_boundary_neighbor_index_list_3d(*args)
+            else:
+                idx_list = create_boundary_cell_index_list_3d(*args)
         else:
             raise ValueError("Flag field has to be a 2 or 3 dimensional numpy array")
         return np.array(idx_list, dtype=index_arr_dtype)
     else:
         if flag_field.size > 1e6:
             warnings.warn("Boundary setup may take very long! Consider installing cython to speed it up")
-        return _create_boundary_index_list_python(flag_field, nr_of_ghost_layers, boundary_mask, fluid_mask, stencil)
+        if inner_or_boundary:
+            return _create_boundary_neighbor_index_list_python(*args)
+        else:
+            return _create_boundary_cell_index_list_python(*args)
 
 
-def create_boundary_index_array(flag_field, stencil, boundary_mask, fluid_mask, boundary_object, nr_of_ghost_layers=1):
-    idx_array = create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask, nr_of_ghost_layers)
+def create_boundary_index_array(flag_field, stencil, boundary_mask, fluid_mask, boundary_object,
+                                nr_of_ghost_layers=1, inner_or_boundary=True, single_link=False):
+    idx_array = create_boundary_index_list(flag_field, stencil, boundary_mask, fluid_mask,
+                                           nr_of_ghost_layers, inner_or_boundary, single_link)
     dim = len(flag_field.shape)
 
     if boundary_object.additional_data:

boundaries/createindexlistcython.pyx

@@ -15,9 +15,9 @@ ctypedef fused IntegerType:
 
 @cython.boundscheck(False) # turn off bounds-checking for entire function
 @cython.wraparound(False)  # turn off negative index wrapping for entire function
-def create_boundary_index_list_2d(object[IntegerType, ndim=2] flag_field,
-                                  int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
-                                  object[int, ndim=2] stencil):
+def create_boundary_neighbor_index_list_2d(object[IntegerType, ndim=2] flag_field,
+                                           int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
+                                           object[int, ndim=2] stencil, int single_link):
     cdef int xs, ys, x, y
     cdef int dirIdx, num_directions, dx, dy
 
@@ -33,14 +33,16 @@ def create_boundary_index_list_2d(object[IntegerType, ndim=2] flag_field,
                     dy = stencil[dirIdx,1]
                     if flag_field[x + dx, y + dy] & boundary_mask:
                         boundary_index_list.append((x,y, dirIdx))
+                        if single_link:
+                            break
     return boundary_index_list
 
 
 @cython.boundscheck(False) # turn off bounds-checking for entire function
 @cython.wraparound(False)  # turn off negative index wrapping for entire function
-def create_boundary_index_list_3d(object[IntegerType, ndim=3] flag_field,
-                                  int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
-                                  object[int, ndim=2] stencil):
+def create_boundary_neighbor_index_list_3d(object[IntegerType, ndim=3] flag_field,
+                                           int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
+                                           object[int, ndim=2] stencil, int single_link):
     cdef int xs, ys, zs, x, y, z
     cdef int dirIdx, num_directions, dx, dy, dz
 
@@ -58,6 +60,61 @@ def create_boundary_index_list_3d(object[IntegerType, ndim=3] flag_field,
                         dz = stencil[dirIdx,2]
                         if flag_field[x + dx, y + dy, z + dz] & boundary_mask:
                             boundary_index_list.append((x,y,z, dirIdx))
+                            if single_link:
+                                break
     return boundary_index_list
 
 
+
+@cython.boundscheck(False) # turn off bounds-checking for entire function
+@cython.wraparound(False)  # turn off negative index wrapping for entire function
+def create_boundary_cell_index_list_2d(object[IntegerType, ndim=2] flag_field,
+                                       int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
+                                       object[int, ndim=2] stencil, int single_link):
+    cdef int xs, ys, x, y
+    cdef int dirIdx, num_directions, dx, dy
+
+    xs, ys = flag_field.shape
+    boundary_index_list = []
+    num_directions = stencil.shape[0]
+
+    for y in range(0, ys):
+        for x in range(0, xs):
+            if flag_field[x, y] & boundary_mask:
+                for dirIdx in range(num_directions):
+                    dx = stencil[dirIdx,0]
+                    dy = stencil[dirIdx,1]
+                    if 0 <= x + dx < xs and 0 <= y + dy < ys:
+                        if flag_field[x + dx, y + dy] & fluid_mask:
+                            boundary_index_list.append((x,y, dirIdx))
+                            if single_link:
+                                break
+    return boundary_index_list
+
+
+@cython.boundscheck(False) # turn off bounds-checking for entire function
+@cython.wraparound(False)  # turn off negative index wrapping for entire function
+def create_boundary_cell_index_list_3d(object[IntegerType, ndim=3] flag_field,
+                                       int nr_of_ghost_layers, IntegerType boundary_mask, IntegerType fluid_mask,
+                                       object[int, ndim=2] stencil, int single_link):
+    cdef int xs, ys, zs, x, y, z
+    cdef int dirIdx, num_directions, dx, dy, dz
+
+    xs, ys, zs = flag_field.shape
+    boundary_index_list = []
+    num_directions = stencil.shape[0]
+
+    for z in range(0, zs):
+        for y in range(0, ys):
+            for x in range(0, xs):
+                if flag_field[x, y, z] & boundary_mask:
+                    for dirIdx in range(num_directions):
+                        dx = stencil[dirIdx,0]
+                        dy = stencil[dirIdx,1]
+                        dz = stencil[dirIdx,2]
+                        if 0 <= x + dx < xs and 0 <= y + dy < ys and 0 <= z + dz < zs:
+                            if flag_field[x + dx, y + dy, z + dz] & fluid_mask:
+                                boundary_index_list.append((x,y,z, dirIdx))
+                                if single_link:
+                                    break
+    return boundary_index_list
\ No newline at end of file

field.py

@@ -419,7 +419,7 @@ class Field:
 
         def __new_stage2__(self, field, offsets=(0, 0, 0), idx=None, is_absolute_access=False):
             field_name = field.name
-            offsets_and_index = chain(offsets, idx) if idx is not None else offsets
+            offsets_and_index = (*offsets, *idx) if idx is not None else offsets
             constant_offsets = not any([isinstance(o, sp.Basic) and not o.is_Integer for o in offsets_and_index])
 
             if not idx: