Finished periodicity, started parallel.

lbmpy/sparse/__init__.py

-from .mapping import SparseLbBoundaryMapper, SparseLbMapper, SparseLbPeriodicityMapper
+from .mapping import SparseLbBoundaryMapper, SparseLbMapper, SparseLbPeriodicityMapper, SparseLbCommunicationMapper
 from .update_rule_sparse import (
     create_lb_update_rule_sparse, create_macroscopic_value_getter_sparse,
     create_macroscopic_value_setter_sparse, create_symbolic_list)
 
-__all__ = ['SparseLbBoundaryMapper', 'SparseLbPeriodicityMapper', 'SparseLbMapper', 'create_lb_update_rule_sparse',
+__all__ = ['SparseLbBoundaryMapper', 'SparseLbPeriodicityMapper', 'SparseLbMapper', 'SparseLbCommunicationMapper', 'create_lb_update_rule_sparse',
            'create_macroscopic_value_setter_sparse', 'create_macroscopic_value_getter_sparse',
            'create_symbolic_list']

lbmpy/sparse/mapping.py

@@ -144,6 +144,7 @@ class SparseLbMapper:
                 at_border = [n_cell_i == 0 or n_cell_i == ds_i+1 for n_cell_i, ds_i in zip(cell, self.domain_size)]
                 #print("write:", cell, "read:", inv_neighbor_cell)
                 if (True in at_border):
+                    #print("at border")
                     result.append(pdf_index(own_cell_idx, direction_idx, len(self)))
                 elif flag_arr[tuple(inv_neighbor_cell)] & fluid_boundary_mask:
                     neighbor_cell_idx = self.cell_idx(tuple(inv_neighbor_cell))
@@ -154,7 +155,6 @@ class SparseLbMapper:
                     result.append(pdf_index(own_cell_idx, inverse_idx(stencil, direction_idx), len(self)))
                     #print("write:", cell, "read:", cell, "direction:", inverse_idx(stencil, direction_idx))
                 else:
-                    print("HEllooo")
                     #This is where we end up if inv_neighbor_cell is not fluid/no_slip/ghost cell;
                     #legacy periodicity handler
                     #print("N is not normal")
@@ -175,6 +175,7 @@ class SparseLbMapper:
 
         index_array = np.array(result, dtype=np.uint32)
         index_arr = index_array.reshape([len(stencil) - 1, len(self.fluid_coordinates)])
+        #index_arr = index_array.reshape([len(stencil) - 1, len(self.fluid_coordinates)-self._num_ghost_cells])
         index_arr = index_arr.swapaxes(0, 1) 
         return index_arr
     
@@ -219,13 +220,14 @@ class SparseLbPeriodicityMapper:
     def get_kernel(self):
         return self._kernel
     
-    def direction_index(self, direction):
-        direction = tuple(direction)
-        for direction_idx, dir_coordinates in enumerate(self.mapping.stencil):
-            if dir_coordinates == direction:
-                return direction_idx
-        raise ValueError("Could not find index for direction {}".format(direction))
-        
+    def at_border(self, cell):
+        return True in [cell_i == 0 or cell_i  == ds_i-1 for cell_i, ds_i in zip(cell, self.domain_size)]
+    
+    def ghost_cells(self):
+        border_bool = [self.at_border(cell) for cell in self.mapping._coordinate_arr[:self.mapping.num_fluid_cells]]
+        fluid_border_coord = self.mapping._coordinate_arr[:self.mapping.num_fluid_cells][border_bool]
+        return fluid_border_coord
+    
     def get_pdf_read_idx(self, read, write, direction, direction_idx):  
         if self.flag_arr[tuple(read)] & self.no_slip_flag: # Read cell is no-slip!
             # Change read cell to that cell that later wants to pull from the solid cell, get inverse PDF from it
@@ -241,21 +243,13 @@ class SparseLbPeriodicityMapper:
     
     def get_assignment(self, direction, cell, bool_slice):
         cell_idx = self.mapping.cell_idx(cell)
-        direction_idx = self.direction_index(direction)
+        direction_idx = self.mapping.stencil.index(direction)
         inv_neighbor_cell = [(cell_i - 2*dir_i*int(bs_i))%ds_i for cell_i, dir_i, ds_i, bs_i in zip(cell, direction, self.domain_size, bool_slice)]
         #print("write:", cell, "read:", tuple(inv_neighbor_cell))
         pdf_write_idx = pdf_index(cell_idx, direction_idx, len(self.mapping))
         pdf_read_idx = self.get_pdf_read_idx(inv_neighbor_cell, cell, direction, direction_idx)
         return [pdf_write_idx, pdf_read_idx]
     
-    def ghost_cells(self):
-        fluid_border_coord = []
-        for cell in self.mapping._coordinate_arr[:self.mapping.num_fluid_cells]:
-            bool_border = [cell_i == 0 or cell_i == ds_i-1 for cell_i, ds_i in zip(cell, self.domain_size)]
-            if True in bool_border:
-                fluid_border_coord.append(cell)
-        return fluid_border_coord
-    
     def create_periodic_index_array(self):
         stencil = self.mapping.stencil
         fluid_border_coord = self.ghost_cells()
@@ -272,9 +266,10 @@ class SparseLbPeriodicityMapper:
                 #print(cell, ", range", periodic_coord_range, "bool", bool_range)
                 if all(bool_range): # This ghost cell needs this direction value
                     # block_direction: where to put this part of index_array (which block to send to...)
-                    block_direction = self.direction_index([int(bp_i)*dir_i for dir_i, bp_i in zip(direction, bool_slice)]) 
-                    result[direction_idx-1][block_direction-1].append(self.get_assignment(direction, cell, bool_slice))
-                    #print("Goes into", direction_idx-1, block_direction-1)
+                    block_direction = [int(bp_i)*dir_i for dir_i, bp_i in zip(direction, bool_slice)]
+                    block_index = self.mapping.stencil.index(tuple(block_direction))
+                    result[direction_idx-1][block_index-1].append(self.get_assignment(direction, cell, bool_slice))
+                    #print("Goes into", direction_idx-1, block_index-1)
         # Flatten result array:
         flattened_result = []
         for block in result:
@@ -296,6 +291,97 @@ class SparseLbPeriodicityMapper:
         _ast = kernel
         self._kernel = _ast.compile()
         return kernel
+
+class SparseLbCommunicationMapper:
+    
+    def __init__(self, mapping: SparseLbMapper, dh, pdf_field_sparse):
+        index_field = Field.create_generic("idx", spatial_dimensions=1, index_dimensions=1, dtype=np.int64)
+        self._coordinates_all = None
+        self.pdf_field_sparse = pdf_field_sparse
+        self.index_field = index_field
+        self.mapping = mapping
+        self.flag_arr = mapping._flag_arr
+        self.domain_size = self.flag_arr.shape
+        self.fluid_flag = mapping.fluid_flag
+        self.no_slip_flag = mapping.no_slip_flag
+        self._index_arrays = None
+        self._kernel = None
+        self._dirty = True
+    
+    def at_border_not_ghost(self, cell):
+        if True in [cell_i == 0 or cell_i  == ds_i-1 for cell_i, ds_i in zip(cell, self.domain_size)]:
+            return False
+        return True in [cell_i == 1 or cell_i  == ds_i-2 for cell_i, ds_i in zip(cell, self.domain_size)]
+    
+    def border_cells(self):
+        border_bool = [self.at_border_not_ghost(cell) for cell in self._coordinates_all]
+        fluid_border_coord = self._coordinates_all[border_bool]
+        return fluid_border_coord
+    
+    #def get_pdf_read_idx(self, read, write, direction, direction_idx):  
+    #    if self.flag_arr[tuple(read)] & self.no_slip_flag: # Read cell is no-slip!
+    #        # Change read cell to that cell that later wants to pull from the solid cell, get inverse PDF from it
+    #        read = [cell_i + dir_i for cell_i, dir_i in zip(write, direction)]
+    #        if self.flag_arr[tuple(read)] & self.no_slip_flag: 
+    #            # New read cell is no-slip as well! This can be a problem if pdf_index ends up being out-of-bounds, so the only operation is to pull read from the same place as write, changing nothing
+    #            read = write
+    #        else:
+    #            # Direction has to be flipped
+    #            direction_idx = inverse_idx(self.mapping.stencil, direction_idx) 
+    #        #print("Neighbor is solid, new read is", tuple(read), "at dir", direction_idx)
+    #    return pdf_index(self.mapping.cell_idx(tuple(read)), direction_idx, len(self.mapping))
+    
+    def get_assignment(self, direction, cell, bool_slice):
+        cell = tuple(cell)
+        if self.flag_arr[cell] & self.no_slip_flag:
+            #print("pack:", cell, "is solid")
+            return -1
+        cell_idx = self.mapping.cell_idx(cell)
+        direction_idx = self.mapping.stencil.index(direction)
+        #Maybe check n_ghost_cell later for solid?? Dunno???
+        neighbor_ghost_cell = [(cell_i + dir_i*int(bs_i)) for cell_i, dir_i, bs_i in zip(cell, direction, bool_slice)]
+        #print("pack:", cell, "ghost:", tuple(neighbor_ghost_cell))
+        pdf_cell_idx = pdf_index(cell_idx, direction_idx, len(self.mapping))
+        return pdf_cell_idx
+    
+    def create_packages(self):
+        stencil = self.mapping.stencil
+        self._coordinates_all = np.argwhere(np.bitwise_or(self.mapping._flag_arr, 1)).astype(np.uint32)
+        result = [[[] for j in range(0, len(stencil)-1)] for i in range(0, len(stencil)-1)]
+        for direction_idx, direction in enumerate(stencil):
+            if all(d_i == 0 for d_i in direction): # Direction (0,0) irrelevant
+                continue
+            #print("\n New direction:", direction, ", ", direction_idx)
+            periodic_slice_coord = [[int((ds_i-1)*(1+dir_i)/2-dir_i)] if dir_i != 0 else [] for i, (dir_i, ds_i) in enumerate(zip(direction, self.domain_size))]   
+            for cell in self.border_cells():
+                bool_slice = [cell_i in slice_i for (cell_i, slice_i) in zip(cell, periodic_slice_coord)]
+                #periodic_coord_range = [[1 - dir_i, ds_i - 2 - dir_i] for dir_i, ds_i in zip(direction, self.domain_size)]
+                #bool_range = [pcr_i[0] <= cell_i <= pcr_i[1] for pcr_i, cell_i in zip(periodic_coord_range, cell)]
+                #print(cell, ", range", periodic_coord_range, "bool", bool_range)
+                #print(cell, "per slice:", periodic_slice_coord, bool_slice)
+                if True in bool_slice: # This ghost cell needs this direction value
+                    # block_direction: where to put this part of index_array (which block to send to...)
+                    block_direction = [int(bp_i)*dir_i for dir_i, bp_i in zip(direction, bool_slice)]
+                    block_index = self.mapping.stencil.index(tuple(block_direction))
+                    result[block_index-1][direction_idx-1].append(self.get_assignment(direction, cell, bool_slice))
+                    #print("Goes into", block_index-1, direction_idx-1)
+        # Flatten result array:
+        return result
+        flattened_result = []
+        for block in result:
+            for i_a in block:
+                if (i_a):
+                    flattened_result.append(np.array(i_a))
+        flattened_result = np.array(flattened_result)
+        self._index_arrays = flattened_result
+        self._dirty = False
+        return flattened_result
+        
+        
+        
+    
+    
+    
     
 class SparseLbBoundaryMapper:
     NEIGHBOR_IDX_NAME = 'nidx{}'