1""" 

2This module contains function that simplify the iteration over walberla's distributed data structure. 

3These function simplify the iteration over rectangular slices, managing the mapping between block local coordinates and 

4global coordinates. 

5""" 

6import numpy as np 

7 

8from pystencils.datahandling.datahandling_interface import Block 

9from pystencils.slicing import normalize_slice 

10 

11try: 

12 # noinspection PyPep8Naming 

13 import waLBerla as wlb 

14except ImportError: 

15 wlb = None 

16 

17 

18def block_iteration(blocks, ghost_layers, dim=3, access_prefix=''): 

19 """Simple iteration over parallel walberla domain. 

20 

21 Iterator that simplifies the access to field data by automatically converting from walberla fields to 

22 numpy arrays 

23 

24 Args: 

25 blocks: walberla block data structure 

26 ghost_layers: how many ghost layers to include (outer and inner) 

27 dim: walberla's block data structure is 3D - 2D domains can be done by setting z_size=1 

28 if dim=2 is set here, the third coordinate of the returned fields is accessed at z=0 automatically 

29 access_prefix: see documentation of sliced_block_iteration 

30 """ 

31 for block in blocks: 

32 cell_interval = blocks.getBlockCellBB(block) 

33 cell_interval.expand(ghost_layers) 

34 local_slice = [slice(0, w, None) for w in cell_interval.size] 

35 if dim == 2: 

36 local_slice[2] = ghost_layers 

37 yield ParallelBlock(block, cell_interval.min[:dim], tuple(local_slice), ghost_layers, access_prefix) 

38 

39 

40def sliced_block_iteration(blocks, slice_obj=None, inner_ghost_layers=1, outer_ghost_layers=1, dim=3, access_prefix=''): 

41 """Iterates of all blocks that have an intersection with the given slice object. 

42 

43 For intersection blocks a Block object is yielded 

44 

45 Args: 

46 blocks: walberla block data structure 

47 slice_obj: a slice (i.e. rectangular sub-region), can be created with make_slice[] 

48 inner_ghost_layers: how many ghost layers are included in the local slice and the optional index arrays 

49 outer_ghost_layers: slices can have relative coordinates e.g. make_slice[0.2, :, :] 

50 when computing absolute values, the domain size is needed. This parameter 

51 specifies how many ghost layers are taken into account for this operation. 

52 dim: set to 2 for pseudo 2D simulation (i.e. where z coordinate of blocks has extent 1) 

53 the arrays returned when indexing the block 

54 access_prefix: when accessing block data, this prefix is prepended to the access name 

55 mostly used to switch between CPU and GPU field access (gpu fields are added with a 

56 certain prefix 'gpu_') 

57 

58 Example: 

59 assume no slice is given, then slice_normalization_ghost_layers effectively sets how much ghost layers at the 

60 border of the domain are included. The inner_ghost_layers parameter specifies how many inner ghost layers are 

61 included 

62 """ 

63 if slice_obj is None: 

64 slice_obj = tuple([slice(None, None, None)] * dim) 

65 if dim == 2: 

66 slice_obj += (inner_ghost_layers,) 

67 

68 domain_cell_bb = blocks.getDomainCellBB() 

69 domain_extent = [s + 2 * outer_ghost_layers for s in domain_cell_bb.size] 

70 slice_obj = normalize_slice(slice_obj, domain_extent) 

71 target_cell_bb = wlb.CellInterval.fromSlice(slice_obj) 

72 target_cell_bb.shift(*[a - outer_ghost_layers for a in domain_cell_bb.min]) 

73 

74 for block in blocks: 

75 intersection = blocks.getBlockCellBB(block).getExpanded(inner_ghost_layers) 

76 intersection.intersect(target_cell_bb) 

77 if intersection.empty(): 

78 continue 

79 

80 local_target_bb = blocks.transformGlobalToLocal(block, intersection) 

81 local_target_bb.shift(inner_ghost_layers, inner_ghost_layers, inner_ghost_layers) 

82 local_slice = local_target_bb.toSlice(False) 

83 if dim == 2: 

84 local_slice = (local_slice[0], local_slice[1], inner_ghost_layers) 

85 yield ParallelBlock(block, intersection.min[:dim], local_slice, inner_ghost_layers, access_prefix) 

86 

87 

88# ----------------------------- Implementation details ----------------------------------------------------------------- 

89 

90 

91class SerialBlock(Block): 

92 """Simple mock-up block that is used for SerialDataHandling.""" 

93 def __init__(self, field_dict, offset, local_slice): 

94 super(SerialBlock, self).__init__(offset, local_slice) 

95 self._fieldDict = field_dict 

96 

97 def __getitem__(self, data_name): 

98 result = self._fieldDict[data_name] 

99 if isinstance(result, np.ndarray): 99 ↛ 101line 99 didn't jump to line 101, because the condition on line 99 was never false

100 result = result[self._localSlice] 

101 return result 

102 

103 

104class ParallelBlock(Block): 

105 def __init__(self, block, offset, local_slice, inner_ghost_layers, name_prefix): 

106 super(ParallelBlock, self).__init__(offset, local_slice) 

107 self._block = block 

108 self._gls = inner_ghost_layers 

109 self._name_prefix = name_prefix 

110 

111 def __getitem__(self, data_name): 

112 result = self._block[self._name_prefix + data_name] 

113 type_name = type(result).__name__ 

114 if 'GhostLayerField' in type_name: 

115 result = wlb.field.toArray(result, with_ghost_layers=self._gls) 

116 result = self._normalize_array_shape(result) 

117 elif 'GpuField' in type_name: 

118 result = wlb.cuda.toGpuArray(result, with_ghost_layers=self._gls) 

119 result = self._normalize_array_shape(result) 

120 return result 

121 

122 def _normalize_array_shape(self, arr): 

123 if arr.shape[-1] == 1 and len(arr.shape) == 4: 

124 arr = arr[..., 0] 

125 return arr[self._localSlice]