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

8from pystencils.datahandling.datahandling_interface import Block

9from pystencils.slicing import normalize_slice

11try:

12 # noinspection PyPep8Naming

13 import waLBerla as wlb

14except ImportError:

15 wlb = None

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

19 """Simple iteration over parallel walberla domain.

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

22 numpy arrays

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)

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.

43 For intersection blocks a Block object is yielded

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_')

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,)

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])

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

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)

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

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

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

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

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

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]