1import numpy as np 

2import sympy as sp 

3 

4from pystencils import create_indexed_kernel 

5from pystencils.assignment import Assignment 

6from pystencils.backends.cbackend import CustomCodeNode 

7from pystencils.boundaries.createindexlist import ( 

8 create_boundary_index_array, numpy_data_type_for_boundary_object) 

9from pystencils.cache import memorycache 

10from pystencils.data_types import TypedSymbol, create_type 

11from pystencils.datahandling.pycuda import PyCudaArrayHandler 

12from pystencils.field import Field 

13from pystencils.kernelparameters import FieldPointerSymbol 

14 

15try: 

16 # noinspection PyPep8Naming 

17 import waLBerla as wlb 

18 if wlb.cpp_available: 

19 from pystencils.datahandling.parallel_datahandling import ParallelDataHandling 

20 else: 

21 ParallelDataHandling = None 

22except ImportError: 

23 ParallelDataHandling = None 

24 

25DEFAULT_FLAG_TYPE = np.uint32 

26 

27 

28class FlagInterface: 

29 """Manages the reservation of bits (i.e. flags) in an array of unsigned integers. 

30 

31 Examples: 

32 >>> from pystencils import create_data_handling 

33 >>> dh = create_data_handling((4, 5)) 

34 >>> fi = FlagInterface(dh, 'flag_field', np.uint8) 

35 >>> assert dh.has_data('flag_field') 

36 >>> fi.reserve_next_flag() 

37 2 

38 >>> fi.reserve_flag(4) 

39 4 

40 >>> fi.reserve_next_flag() 

41 8 

42 """ 

43 

44 def __init__(self, data_handling, flag_field_name, dtype=DEFAULT_FLAG_TYPE): 

45 self.flag_field_name = flag_field_name 

46 self.domain_flag = dtype(1 << 0) 

47 self._used_flags = {self.domain_flag} 

48 self.data_handling = data_handling 

49 self.dtype = dtype 

50 self.max_bits = self.dtype().itemsize * 8 

51 

52 # Add flag field to data handling if it does not yet exist 

53 if data_handling.has_data(self.flag_field_name): 

54 raise ValueError("There is already a boundary handling registered at the data handling." 

55 "If you want to add multiple handling objects, choose a different name.") 

56 

57 self.flag_field = data_handling.add_array(self.flag_field_name, dtype=self.dtype, cpu=True, gpu=False) 

58 ff_ghost_layers = data_handling.ghost_layers_of_field(self.flag_field_name) 

59 for b in data_handling.iterate(ghost_layers=ff_ghost_layers): 

60 b[self.flag_field_name].fill(self.domain_flag) 

61 

62 def reserve_next_flag(self): 

63 for i in range(1, self.max_bits): 

64 flag = self.dtype(1 << i) 

65 if flag not in self._used_flags: 

66 self._used_flags.add(flag) 

67 assert self._is_power_of_2(flag) 

68 return flag 

69 raise ValueError(f"All available {self.max_bits} flags are reserved") 

70 

71 def reserve_flag(self, flag): 

72 assert self._is_power_of_2(flag) 

73 flag = self.dtype(flag) 

74 if flag in self._used_flags: 

75 raise ValueError(f"The flag {flag} is already reserved") 

76 self._used_flags.add(flag) 

77 return flag 

78 

79 @staticmethod 

80 def _is_power_of_2(num): 

81 return num != 0 and ((num & (num - 1)) == 0) 

82 

83 

84class BoundaryHandling: 

85 

86 def __init__(self, data_handling, field_name, stencil, name="boundary_handling", flag_interface=None, 

87 target='cpu', openmp=True): 

88 assert data_handling.has_data(field_name) 

89 assert data_handling.dim == len(stencil[0]), "Dimension of stencil and data handling do not match" 

90 self._data_handling = data_handling 

91 self._field_name = field_name 

92 self._index_array_name = name + "IndexArrays" 

93 self._target = target 

94 self._openmp = openmp 

95 self._boundary_object_to_boundary_info = {} 

96 self.stencil = stencil 

97 self._dirty = True 

98 fi = flag_interface 

99 self.flag_interface = fi if fi is not None else FlagInterface(data_handling, name + "Flags") 

100 

101 if ParallelDataHandling and isinstance(self.data_handling, ParallelDataHandling): 

102 array_handler = PyCudaArrayHandler() 

103 else: 

104 array_handler = self.data_handling.array_handler 

105 

106 def to_cpu(gpu_version, cpu_version): 

107 gpu_version = gpu_version.boundary_object_to_index_list 

108 cpu_version = cpu_version.boundary_object_to_index_list 

109 for obj, cpu_arr in cpu_version.items(): 

110 array_handler.download(gpu_version[obj], cpu_arr) 

111 

112 def to_gpu(gpu_version, cpu_version): 

113 gpu_version = gpu_version.boundary_object_to_index_list 

114 cpu_version = cpu_version.boundary_object_to_index_list 

115 

116 for obj, cpu_arr in cpu_version.items(): 

117 if obj not in gpu_version or gpu_version[obj].shape != cpu_arr.shape: 

118 gpu_version[obj] = array_handler.to_gpu(cpu_arr) 

119 else: 

120 array_handler.upload(gpu_version[obj], cpu_arr) 

121 

122 class_ = self.IndexFieldBlockData 

123 class_.to_cpu = to_cpu 

124 class_.to_gpu = to_gpu 

125 gpu = self._target in data_handling._GPU_LIKE_TARGETS 

126 data_handling.add_custom_class(self._index_array_name, class_, cpu=True, gpu=gpu) 

127 

128 @property 

129 def data_handling(self): 

130 return self._data_handling 

131 

132 def get_flag(self, boundary_obj): 

133 return self._boundary_object_to_boundary_info[boundary_obj].flag 

134 

135 @property 

136 def shape(self): 

137 return self._data_handling.shape 

138 

139 @property 

140 def dim(self): 

141 return self._data_handling.dim 

142 

143 @property 

144 def boundary_objects(self): 

145 return tuple(self._boundary_object_to_boundary_info.keys()) 

146 

147 @property 

148 def flag_array_name(self): 

149 return self.flag_interface.flag_field_name 

150 

151 def get_mask(self, slice_obj, boundary_obj, inverse=False): 

152 if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain': 

153 flag = self.flag_interface.domain_flag 

154 else: 

155 flag = self._boundary_object_to_boundary_info[boundary_obj].flag 

156 

157 arr = self.data_handling.gather_array(self.flag_array_name, slice_obj) 

158 if arr is None: 

159 return None 

160 else: 

161 result = np.bitwise_and(arr, flag) 

162 if inverse: 

163 result = np.logical_not(result) 

164 return result 

165 

166 def set_boundary(self, boundary_obj, slice_obj=None, mask_callback=None, 

167 ghost_layers=True, inner_ghost_layers=True, replace=True, force_flag_value=None): 

168 """Sets boundary using either a rectangular slice, a boolean mask or a combination of both. 

169 

170 Args: 

171 boundary_obj: instance of a boundary object that should be set 

172 slice_obj: a slice object (can be created with make_slice[]) that selects a part of the domain where 

173 the boundary should be set. If none, the complete domain is selected which makes only sense 

174 if a mask_callback is passed. The slice can have ':' placeholders, which are interpreted 

175 depending on the 'inner_ghost_layers' parameter i.e. if it is True, the slice extends 

176 into the ghost layers 

177 mask_callback: callback function getting x,y (z) parameters of the cell midpoints and returning a 

178 boolean mask with True entries where boundary cells should be set. 

179 The x, y, z arrays have 2D/3D shape such that they can be used directly 

180 to create the boolean return array. i.e return x < 10 sets boundaries in cells with 

181 midpoint x coordinate smaller than 10. 

182 ghost_layers: see DataHandling.iterate() 

183 inner_ghost_layers: see DataHandling.iterate() 

184 replace: by default all other flags are erased in the cells where the boundary is set. To add a 

185 boundary condition, set this replace flag to False 

186 force_flag_value: flag that should be reserved for this boundary. Has to be an integer that is a power of 2 

187 and was not reserved before for another boundary. 

188 """ 

189 if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain': 

190 flag = self.flag_interface.domain_flag 

191 else: 

192 if force_flag_value: 

193 self.flag_interface.reserve_flag(force_flag_value) 

194 flag = self._add_boundary(boundary_obj, force_flag_value) 

195 

196 for b in self._data_handling.iterate(slice_obj, ghost_layers=ghost_layers, 

197 inner_ghost_layers=inner_ghost_layers): 

198 flag_arr = b[self.flag_interface.flag_field_name] 

199 if mask_callback is not None: 

200 mask = mask_callback(*b.midpoint_arrays) 

201 if replace: 

202 flag_arr[mask] = flag 

203 else: 

204 np.bitwise_or(flag_arr, flag, where=mask, out=flag_arr) 

205 np.bitwise_and(flag_arr, ~self.flag_interface.domain_flag, where=mask, out=flag_arr) 

206 else: 

207 if replace: 

208 flag_arr.fill(flag) 

209 else: 

210 np.bitwise_or(flag_arr, flag, out=flag_arr) 

211 np.bitwise_and(flag_arr, ~self.flag_interface.domain_flag, out=flag_arr) 

212 

213 self._dirty = True 

214 

215 return flag 

216 

217 def set_boundary_where_flag_is_set(self, boundary_obj, flag): 

218 """Adds an (additional) boundary to all cells that have been previously marked with the passed flag.""" 

219 self._add_boundary(boundary_obj, flag) 

220 self._dirty = True 

221 return flag 

222 

223 def prepare(self): 

224 if not self._dirty: 

225 return 

226 self._create_index_fields() 

227 self._dirty = False 

228 

229 def trigger_reinitialization_of_boundary_data(self, **kwargs): 

230 if self._dirty: 

231 self.prepare() 

232 else: 

233 ff_ghost_layers = self._data_handling.ghost_layers_of_field(self.flag_interface.flag_field_name) 

234 for b in self._data_handling.iterate(ghost_layers=ff_ghost_layers): 

235 for b_obj, setter in b[self._index_array_name].boundary_object_to_data_setter.items(): 

236 self._boundary_data_initialization(b_obj, setter, **kwargs) 

237 

238 def __call__(self, **kwargs): 

239 if self._dirty: 

240 self.prepare() 

241 

242 for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS): 

243 for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items(): 

244 kwargs[self._field_name] = b[self._field_name] 

245 kwargs['indexField'] = idx_arr 

246 data_used_in_kernel = (p.fields[0].name 

247 for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters 

248 if isinstance(p.symbol, FieldPointerSymbol) and p.fields[0].name not in kwargs) 

249 kwargs.update({name: b[name] for name in data_used_in_kernel}) 

250 

251 self._boundary_object_to_boundary_info[b_obj].kernel(**kwargs) 

252 

253 def add_fixed_steps(self, fixed_loop, **kwargs): 

254 if self._dirty: 

255 self.prepare() 

256 

257 for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS): 

258 for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items(): 

259 arguments = kwargs.copy() 

260 arguments[self._field_name] = b[self._field_name] 

261 arguments['indexField'] = idx_arr 

262 data_used_in_kernel = (p.fields[0].name 

263 for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters 

264 if isinstance(p.symbol, FieldPointerSymbol) and p.field_name not in arguments) 

265 arguments.update({name: b[name] for name in data_used_in_kernel if name not in arguments}) 

266 

267 kernel = self._boundary_object_to_boundary_info[b_obj].kernel 

268 fixed_loop.add_call(kernel, arguments) 

269 

270 def geometry_to_vtk(self, file_name='geometry', boundaries='all', ghost_layers=False): 

271 """ 

272 Writes a VTK field where each cell with the given boundary is marked with 1, other cells are 0 

273 This can be used to display the simulation geometry in Paraview 

274 

275 Params: 

276 file_name: vtk filename 

277 boundaries: boundary object, or special string 'domain' for domain cells or special string 'all' for all 

278 boundary conditions. 

279 can also be a sequence, to write multiple boundaries to VTK file 

280 ghost_layers: number of ghost layers to write, or True for all, False for none 

281 """ 

282 if boundaries == 'all': 

283 boundaries = list(self._boundary_object_to_boundary_info.keys()) + ['domain'] 

284 elif not hasattr(boundaries, "__len__"): 

285 boundaries = [boundaries] 

286 

287 masks_to_name = {} 

288 for b in boundaries: 

289 if b == 'domain': 

290 masks_to_name[self.flag_interface.domain_flag] = 'domain' 

291 else: 

292 flag = self._boundary_object_to_boundary_info[b].flag 

293 masks_to_name[flag] = b.name 

294 

295 writer = self.data_handling.create_vtk_writer_for_flag_array(file_name, self.flag_interface.flag_field_name, 

296 masks_to_name, ghost_layers=ghost_layers) 

297 writer(1) 

298 

299 # ------------------------------ Implementation Details ------------------------------------------------------------ 

300 

301 def _add_boundary(self, boundary_obj, flag=None): 

302 if boundary_obj not in self._boundary_object_to_boundary_info: 

303 sym_index_field = Field.create_generic('indexField', spatial_dimensions=1, 

304 dtype=numpy_data_type_for_boundary_object(boundary_obj, self.dim)) 

305 ast = self._create_boundary_kernel(self._data_handling.fields[self._field_name], 

306 sym_index_field, boundary_obj) 

307 if flag is None: 

308 flag = self.flag_interface.reserve_next_flag() 

309 boundary_info = self.BoundaryInfo(boundary_obj, flag=flag, kernel=ast.compile()) 

310 self._boundary_object_to_boundary_info[boundary_obj] = boundary_info 

311 return self._boundary_object_to_boundary_info[boundary_obj].flag 

312 

313 def _create_boundary_kernel(self, symbolic_field, symbolic_index_field, boundary_obj): 

314 return create_boundary_kernel(symbolic_field, symbolic_index_field, self.stencil, boundary_obj, 

315 target=self._target, cpu_openmp=self._openmp) 

316 

317 def _create_index_fields(self): 

318 dh = self._data_handling 

319 ff_ghost_layers = dh.ghost_layers_of_field(self.flag_interface.flag_field_name) 

320 for b in dh.iterate(ghost_layers=ff_ghost_layers): 

321 flag_arr = b[self.flag_interface.flag_field_name] 

322 pdf_arr = b[self._field_name] 

323 index_array_bd = b[self._index_array_name] 

324 index_array_bd.clear() 

325 for b_info in self._boundary_object_to_boundary_info.values(): 

326 boundary_obj = b_info.boundary_object 

327 idx_arr = create_boundary_index_array(flag_arr, self.stencil, b_info.flag, 

328 self.flag_interface.domain_flag, boundary_obj, 

329 ff_ghost_layers, boundary_obj.inner_or_boundary, 

330 boundary_obj.single_link) 

331 if idx_arr.size == 0: 

332 continue 

333 

334 boundary_data_setter = BoundaryDataSetter(idx_arr, b.offset, self.stencil, ff_ghost_layers, pdf_arr) 

335 index_array_bd.boundary_object_to_index_list[b_info.boundary_object] = idx_arr 

336 index_array_bd.boundary_object_to_data_setter[b_info.boundary_object] = boundary_data_setter 

337 self._boundary_data_initialization(b_info.boundary_object, boundary_data_setter) 

338 

339 def _boundary_data_initialization(self, boundary_obj, boundary_data_setter, **kwargs): 

340 if boundary_obj.additional_data_init_callback: 

341 boundary_obj.additional_data_init_callback(boundary_data_setter, **kwargs) 

342 if self._target in self._data_handling._GPU_LIKE_TARGETS: 

343 self._data_handling.to_gpu(self._index_array_name) 

344 

345 class BoundaryInfo(object): 

346 def __init__(self, boundary_obj, flag, kernel): 

347 self.boundary_object = boundary_obj 

348 self.flag = flag 

349 self.kernel = kernel 

350 

351 class IndexFieldBlockData: 

352 def __init__(self, *args, **kwargs): 

353 self.boundary_object_to_index_list = {} 

354 self.boundary_object_to_data_setter = {} 

355 

356 def clear(self): 

357 self.boundary_object_to_index_list.clear() 

358 self.boundary_object_to_data_setter.clear() 

359 

360 

361class BoundaryDataSetter: 

362 

363 def __init__(self, index_array, offset, stencil, ghost_layers, pdf_array): 

364 self.index_array = index_array 

365 self.offset = offset 

366 self.stencil = np.array(stencil) 

367 self.pdf_array = pdf_array.view() 

368 self.pdf_array.flags.writeable = False 

369 

370 arr_field_names = index_array.dtype.names 

371 self.dim = 3 if 'z' in arr_field_names else 2 

372 assert 'x' in arr_field_names and 'y' in arr_field_names and 'dir' in arr_field_names, str(arr_field_names) 

373 self.boundary_data_names = set(self.index_array.dtype.names) - {'x', 'y', 'z', 'dir'} 

374 self.coord_map = {0: 'x', 1: 'y', 2: 'z'} 

375 self.ghost_layers = ghost_layers 

376 

377 def non_boundary_cell_positions(self, coord): 

378 assert coord < self.dim 

379 return self.index_array[self.coord_map[coord]] + self.offset[coord] - self.ghost_layers + 0.5 

380 

381 @memorycache() 

382 def link_offsets(self): 

383 return self.stencil[self.index_array['dir']] 

384 

385 @memorycache() 

386 def link_positions(self, coord): 

387 return self.non_boundary_cell_positions(coord) + 0.5 * self.link_offsets()[:, coord] 

388 

389 @memorycache() 

390 def boundary_cell_positions(self, coord): 

391 return self.non_boundary_cell_positions(coord) + self.link_offsets()[:, coord] 

392 

393 def __setitem__(self, key, value): 

394 if key not in self.boundary_data_names: 

395 raise KeyError(f"Invalid boundary data name {key}. Allowed are {self.boundary_data_names}") 

396 self.index_array[key] = value 

397 

398 def __getitem__(self, item): 

399 if item not in self.boundary_data_names: 

400 raise KeyError(f"Invalid boundary data name {item}. Allowed are {self.boundary_data_names}") 

401 return self.index_array[item] 

402 

403 

404class BoundaryOffsetInfo(CustomCodeNode): 

405 

406 # --------------------------- Functions to be used by boundaries -------------------------- 

407 

408 @staticmethod 

409 def offset_from_dir(dir_idx, dim): 

410 return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx] 

411 for symbol in BoundaryOffsetInfo._offset_symbols(dim)]) 

412 

413 @staticmethod 

414 def inv_dir(dir_idx): 

415 return sp.IndexedBase(BoundaryOffsetInfo.INV_DIR_SYMBOL, shape=(1,))[dir_idx] 

416 

417 # ---------------------------------- Internal --------------------------------------------- 

418 

419 def __init__(self, stencil): 

420 dim = len(stencil[0]) 

421 

422 offset_sym = BoundaryOffsetInfo._offset_symbols(dim) 

423 code = "\n" 

424 for i in range(dim): 

425 offset_str = ", ".join([str(d[i]) for d in stencil]) 

426 code += "const int64_t %s [] = { %s };\n" % (offset_sym[i].name, offset_str) 

427 

428 inv_dirs = [] 

429 for direction in stencil: 

430 inverse_dir = tuple([-i for i in direction]) 

431 inv_dirs.append(str(stencil.index(inverse_dir))) 

432 

433 code += "const int %s [] = { %s };\n" % (self.INV_DIR_SYMBOL.name, ", ".join(inv_dirs)) 

434 offset_symbols = BoundaryOffsetInfo._offset_symbols(dim) 

435 super(BoundaryOffsetInfo, self).__init__(code, symbols_read=set(), 

436 symbols_defined=set(offset_symbols + [self.INV_DIR_SYMBOL])) 

437 

438 @staticmethod 

439 def _offset_symbols(dim): 

440 return [TypedSymbol(f"c{d}", create_type(np.int64)) for d in ['x', 'y', 'z'][:dim]] 

441 

442 INV_DIR_SYMBOL = TypedSymbol("invdir", "int") 

443 

444 

445def create_boundary_kernel(field, index_field, stencil, boundary_functor, target='cpu', **kernel_creation_args): 

446 elements = [BoundaryOffsetInfo(stencil)] 

447 index_arr_dtype = index_field.dtype.numpy_dtype 

448 dir_symbol = TypedSymbol("dir", index_arr_dtype.fields['dir'][0]) 

449 elements += [Assignment(dir_symbol, index_field[0]('dir'))] 

450 elements += boundary_functor(field, direction_symbol=dir_symbol, index_field=index_field) 

451 return create_indexed_kernel(elements, [index_field], target=target, **kernel_creation_args)