1import sympy as sp 

2 

3from pystencils.field import create_numpy_array_with_layout, get_layout_of_array 

4 

5 

6class SliceMaker(object): 

7 def __getitem__(self, item): 

8 return item 

9 

10 

11make_slice = SliceMaker() 

12 

13 

14class SlicedGetter(object): 

15 def __init__(self, function_returning_array): 

16 self._functionReturningArray = function_returning_array 

17 

18 def __getitem__(self, item): 

19 return self._functionReturningArray(item) 

20 

21 

22class SlicedGetterDataHandling: 

23 def __init__(self, data_handling, name): 

24 self.dh = data_handling 

25 self.name = name 

26 

27 def __getitem__(self, slice_obj): 

28 if slice_obj is None: 

29 slice_obj = make_slice[:, :] if self.data_handling.dim == 2 else make_slice[:, :, 0.5] 

30 return self.dh.gather_array(self.name, slice_obj).squeeze() 

31 

32 

33def normalize_slice(slices, sizes): 

34 """Converts slices with floating point and/or negative entries to integer slices""" 

35 

36 if len(slices) != len(sizes): 36 ↛ 37line 36 didn't jump to line 37, because the condition on line 36 was never true

37 raise ValueError("Slice dimension does not match sizes") 

38 

39 result = [] 

40 

41 for s, size in zip(slices, sizes): 

42 if type(s) is int: 42 ↛ 43line 42 didn't jump to line 43, because the condition on line 42 was never true

43 if s < 0: 

44 s = size + s 

45 result.append(s) 

46 continue 

47 if type(s) is float: 47 ↛ 48line 47 didn't jump to line 48, because the condition on line 47 was never true

48 result.append(int(s * size)) 

49 continue 

50 

51 assert (type(s) is slice) 

52 

53 if s.start is None: 53 ↛ 55line 53 didn't jump to line 55, because the condition on line 53 was never false

54 new_start = 0 

55 elif type(s.start) is float: 

56 new_start = int(s.start * size) 

57 elif not isinstance(s.start, sp.Basic) and s.start < 0: 

58 new_start = size + s.start 

59 else: 

60 new_start = s.start 

61 

62 if s.stop is None: 62 ↛ 64line 62 didn't jump to line 64, because the condition on line 62 was never false

63 new_stop = size 

64 elif type(s.stop) is float: 

65 new_stop = int(s.stop * size) 

66 elif not isinstance(s.stop, sp.Basic) and s.stop < 0: 

67 new_stop = size + s.stop 

68 else: 

69 new_stop = s.stop 

70 

71 result.append(slice(new_start, new_stop, s.step if s.step is not None else 1)) 

72 

73 return tuple(result) 

74 

75 

76def shift_slice(slices, offset): 

77 def shift_slice_component(slice_comp, shift_offset): 

78 if slice_comp is None: 

79 return None 

80 elif isinstance(slice_comp, int): 

81 return slice_comp + shift_offset 

82 elif isinstance(slice_comp, float): 

83 return slice_comp # relative entries are not shifted 

84 elif isinstance(slice_comp, slice): 

85 return slice(shift_slice_component(slice_comp.start, shift_offset), 

86 shift_slice_component(slice_comp.stop, shift_offset), 

87 slice_comp.step) 

88 else: 

89 raise ValueError() 

90 

91 if hasattr(offset, '__len__'): 

92 return tuple(shift_slice_component(k, off) for k, off in zip(slices, offset)) 

93 else: 

94 if isinstance(slices, slice) or isinstance(slices, int) or isinstance(slices, float): 

95 return shift_slice_component(slices, offset) 

96 else: 

97 return tuple(shift_slice_component(k, offset) for k in slices) 

98 

99 

100def slice_from_direction(direction_name, dim, normal_offset=0, tangential_offset=0): 

101 """ 

102 Create a slice from a direction named by compass scheme: 

103 i.e. 'N' for north returns same as make_slice[:, -1] 

104 the naming is: 

105 - x: W, E (west, east) 

106 - y: S, N (south, north) 

107 - z: B, T (bottom, top) 

108 Also combinations are allowed like north-east 'NE' 

109 

110 :param direction_name: name of direction as explained above 

111 :param dim: dimension of the returned slice (should be 2 or 3) 

112 :param normal_offset: the offset in 'normal' direction: e.g. slice_from_direction('N',2, normal_offset=2) 

113 would return make_slice[:, -3] 

114 :param tangential_offset: offset in the other directions: e.g. slice_from_direction('N',2, tangential_offset=2) 

115 would return make_slice[2:-2, -1] 

116 """ 

117 if tangential_offset == 0: 

118 result = [slice(None, None, None)] * dim 

119 else: 

120 result = [slice(tangential_offset, -tangential_offset, None)] * dim 

121 

122 normal_slice_high, normal_slice_low = -1 - normal_offset, normal_offset 

123 

124 for dim_idx, (low_name, high_name) in enumerate([('W', 'E'), ('S', 'N'), ('B', 'T')]): 

125 if low_name in direction_name: 

126 assert high_name not in direction_name, "Invalid direction name" 

127 result[dim_idx] = normal_slice_low 

128 if high_name in direction_name: 

129 assert low_name not in direction_name, "Invalid direction name" 

130 result[dim_idx] = normal_slice_high 

131 return tuple(result) 

132 

133 

134def remove_ghost_layers(arr, index_dimensions=0, ghost_layers=1): 

135 if ghost_layers <= 0: 135 ↛ 137line 135 didn't jump to line 137, because the condition on line 135 was never false

136 return arr 

137 dimensions = len(arr.shape) 

138 spatial_dimensions = dimensions - index_dimensions 

139 indexing = [slice(ghost_layers, -ghost_layers, None), ] * spatial_dimensions 

140 indexing += [slice(None, None, None)] * index_dimensions 

141 return arr[tuple(indexing)] 

142 

143 

144def add_ghost_layers(arr, index_dimensions=0, ghost_layers=1, layout=None): 

145 dimensions = len(arr.shape) 

146 spatial_dimensions = dimensions - index_dimensions 

147 new_shape = [e + 2 * ghost_layers for e in arr.shape[:spatial_dimensions]] + list(arr.shape[spatial_dimensions:]) 

148 if layout is None: 

149 layout = get_layout_of_array(arr) 

150 result = create_numpy_array_with_layout(new_shape, layout) 

151 result.fill(0.0) 

152 indexing = [slice(ghost_layers, -ghost_layers, None), ] * spatial_dimensions 

153 indexing += [slice(None, None, None)] * index_dimensions 

154 result[tuple(indexing)] = arr 

155 return result 

156 

157 

158def get_slice_before_ghost_layer(direction, ghost_layers=1, thickness=None, full_slice=False): 

159 """ 

160 Returns slicing expression for region before ghost layer 

161 :param direction: tuple specifying direction of slice 

162 :param ghost_layers: number of ghost layers 

163 :param thickness: thickness of the slice, defaults to number of ghost layers 

164 :param full_slice: if true also the ghost cells in directions orthogonal to direction are contained in the 

165 returned slice. Example (d=W ): if full_slice then also the ghost layer in N-S and T-B 

166 are included, otherwise only inner cells are returned 

167 """ 

168 if not thickness: 

169 thickness = ghost_layers 

170 full_slice_inc = ghost_layers if not full_slice else 0 

171 slices = [] 

172 for dir_component in direction: 

173 if dir_component == -1: 

174 s = slice(ghost_layers, thickness + ghost_layers) 

175 elif dir_component == 0: 

176 end = -full_slice_inc 

177 s = slice(full_slice_inc, end if end != 0 else None) 

178 elif dir_component == 1: 

179 start = -thickness - ghost_layers 

180 end = -ghost_layers 

181 s = slice(start if start != 0 else None, end if end != 0 else None) 

182 else: 

183 raise ValueError("Invalid direction: only -1, 0, 1 components are allowed") 

184 slices.append(s) 

185 return tuple(slices) 

186 

187 

188def get_ghost_region_slice(direction, ghost_layers=1, thickness=None, full_slice=False): 

189 """ 

190 Returns slice of ghost region. For parameters see :func:`get_slice_before_ghost_layer` 

191 """ 

192 if not thickness: 

193 thickness = ghost_layers 

194 assert thickness > 0 

195 assert thickness <= ghost_layers 

196 full_slice_inc = ghost_layers if not full_slice else 0 

197 slices = [] 

198 for dir_component in direction: 

199 if dir_component == -1: 

200 s = slice(ghost_layers - thickness, ghost_layers) 

201 elif dir_component == 0: 

202 end = -full_slice_inc 

203 s = slice(full_slice_inc, end if end != 0 else None) 

204 elif dir_component == 1: 

205 start = -ghost_layers 

206 end = - ghost_layers + thickness 

207 s = slice(start if start != 0 else None, end if end != 0 else None) 

208 else: 

209 raise ValueError("Invalid direction: only -1, 0, 1 components are allowed") 

210 slices.append(s) 

211 return tuple(slices) 

212 

213 

214def get_periodic_boundary_src_dst_slices(stencil, ghost_layers=1, thickness=None): 

215 src_dst_slice_tuples = [] 

216 

217 for d in stencil: 

218 if sum([abs(e) for e in d]) == 0: 

219 continue 

220 inv_dir = (-e for e in d) 

221 src = get_slice_before_ghost_layer(inv_dir, ghost_layers, thickness=thickness, full_slice=False) 

222 dst = get_ghost_region_slice(d, ghost_layers, thickness=thickness, full_slice=False) 

223 src_dst_slice_tuples.append((src, dst)) 

224 return src_dst_slice_tuples 

225 

226 

227def get_periodic_boundary_functor(stencil, ghost_layers=1, thickness=None): 

228 """ 

229 Returns a function that applies periodic boundary conditions 

230 :param stencil: sequence of directions e.g. ( [0,1], [0,-1] ) for y periodicity 

231 :param ghost_layers: how many ghost layers the array has 

232 :param thickness: how many of the ghost layers to copy, None means 'all' 

233 :return: function that takes a single array and applies the periodic copy operation 

234 """ 

235 src_dst_slice_tuples = get_periodic_boundary_src_dst_slices(stencil, ghost_layers, thickness) 

236 

237 def functor(pdfs, **_): 

238 for src_slice, dst_slice in src_dst_slice_tuples: 

239 pdfs[dst_slice] = pdfs[src_slice] 

240 

241 return functor 

242 

243 

244def slice_intersection(slice1, slice2): 

245 slice1 = [s if not isinstance(s, int) else slice(s, s + 1, None) for s in slice1] 

246 slice2 = [s if not isinstance(s, int) else slice(s, s + 1, None) for s in slice2] 

247 

248 new_min = [max(s1.start, s2.start) for s1, s2 in zip(slice1, slice2)] 

249 new_max = [min(s1.stop, s2.stop) for s1, s2 in zip(slice1, slice2)] 

250 if any(max_p - min_p < 0 for min_p, max_p in zip(new_min, new_max)): 

251 return None 

252 

253 return [slice(min_p, max_p, None) for min_p, max_p in zip(new_min, new_max)]