boundaryhandling.py 20.4 KB
Newer Older
1
2
import numpy as np
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
3
4

from pystencils import create_indexed_kernel
5
from pystencils.assignment import Assignment
6
from pystencils.backends.cbackend import CustomCodeNode
Martin Bauer's avatar
Martin Bauer committed
7
8
from pystencils.boundaries.createindexlist import (
    create_boundary_index_array, numpy_data_type_for_boundary_object)
9
from pystencils.cache import memorycache
Martin Bauer's avatar
Martin Bauer committed
10
from pystencils.data_types import TypedSymbol, create_type
11
from pystencils.datahandling.pycuda import PyCudaArrayHandler
Martin Bauer's avatar
Martin Bauer committed
12
from pystencils.field import Field
13
from pystencils.kernelparameters import FieldPointerSymbol
14

15
16
17
18
19
try:
    # noinspection PyPep8Naming
    import waLBerla as wlb
    if wlb.cpp_available:
        from pystencils.datahandling.parallel_datahandling import ParallelDataHandling
20
21
    else:
        ParallelDataHandling = None
22
23
24
except ImportError:
    ParallelDataHandling = None

25
DEFAULT_FLAG_TYPE = np.uint32
26

Martin Bauer's avatar
Martin Bauer committed
27

28
29
30
31
class FlagInterface:
    """Manages the reservation of bits (i.e. flags) in an array of unsigned integers.

    Examples:
32
33
        >>> from pystencils import create_data_handling
        >>> dh = create_data_handling((4, 5))
34
35
36
37
38
39
40
41
42
43
44
        >>> fi = FlagInterface(dh, 'flag_field', np.uint8)
        >>> assert dh.has_data('flag_field')
        >>> fi.reserve_next_flag()
        2
        >>> fi.reserve_flag(4)
        4
        >>> fi.reserve_next_flag()
        8
    """

    def __init__(self, data_handling, flag_field_name, dtype=DEFAULT_FLAG_TYPE):
Martin Bauer's avatar
Martin Bauer committed
45
        self.flag_field_name = flag_field_name
46
47
        self.domain_flag = dtype(1 << 0)
        self._used_flags = {self.domain_flag}
Martin Bauer's avatar
Martin Bauer committed
48
        self.data_handling = data_handling
49
50
        self.dtype = dtype
        self.max_bits = self.dtype().itemsize * 8
Martin Bauer's avatar
Martin Bauer committed
51
52

        # Add flag field to data handling if it does not yet exist
Martin Bauer's avatar
Martin Bauer committed
53
        if data_handling.has_data(self.flag_field_name):
Martin Bauer's avatar
Martin Bauer committed
54
            raise ValueError("There is already a boundary handling registered at the data handling."
Martin Bauer's avatar
Martin Bauer committed
55
                             "If you want to add multiple handling objects, choose a different name.")
Martin Bauer's avatar
Martin Bauer committed
56

57
        self.flag_field = data_handling.add_array(self.flag_field_name, dtype=self.dtype, cpu=True, gpu=False)
Martin Bauer's avatar
Martin Bauer committed
58
59
60
        ff_ghost_layers = data_handling.ghost_layers_of_field(self.flag_field_name)
        for b in data_handling.iterate(ghost_layers=ff_ghost_layers):
            b[self.flag_field_name].fill(self.domain_flag)
Martin Bauer's avatar
Martin Bauer committed
61

62
63
64
65
66
67
68
    def reserve_next_flag(self):
        for i in range(1, self.max_bits):
            flag = self.dtype(1 << i)
            if flag not in self._used_flags:
                self._used_flags.add(flag)
                assert self._is_power_of_2(flag)
                return flag
69
        raise ValueError(f"All available {self.max_bits} flags are reserved")
70
71
72
73
74

    def reserve_flag(self, flag):
        assert self._is_power_of_2(flag)
        flag = self.dtype(flag)
        if flag in self._used_flags:
75
            raise ValueError(f"The flag {flag} is already reserved")
76
77
78
79
80
81
        self._used_flags.add(flag)
        return flag

    @staticmethod
    def _is_power_of_2(num):
        return num != 0 and ((num & (num - 1)) == 0)
Martin Bauer's avatar
Martin Bauer committed
82
83


84
85
class BoundaryHandling:

Martin Bauer's avatar
Martin Bauer committed
86
87
88
    def __init__(self, data_handling, field_name, stencil, name="boundary_handling", flag_interface=None,
                 target='cpu', openmp=True):
        assert data_handling.has_data(field_name)
89
        assert data_handling.dim == len(stencil[0]), "Dimension of stencil and data handling do not match"
Martin Bauer's avatar
Martin Bauer committed
90
91
92
        self._data_handling = data_handling
        self._field_name = field_name
        self._index_array_name = name + "IndexArrays"
93
        self._target = target
Martin Bauer's avatar
Martin Bauer committed
94
95
        self._openmp = openmp
        self._boundary_object_to_boundary_info = {}
96
97
        self.stencil = stencil
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
98
99
        fi = flag_interface
        self.flag_interface = fi if fi is not None else FlagInterface(data_handling, name + "Flags")
100

101
102
103
104
105
        if ParallelDataHandling and isinstance(self.data_handling, ParallelDataHandling):
            array_handler = PyCudaArrayHandler()
        else:
            array_handler = self.data_handling.array_handler

106
107
108
109
        def to_cpu(gpu_version, cpu_version):
            gpu_version = gpu_version.boundary_object_to_index_list
            cpu_version = cpu_version.boundary_object_to_index_list
            for obj, cpu_arr in cpu_version.items():
110
                array_handler.download(gpu_version[obj], cpu_arr)
111
112
113
114

        def to_gpu(gpu_version, cpu_version):
            gpu_version = gpu_version.boundary_object_to_index_list
            cpu_version = cpu_version.boundary_object_to_index_list
115

116
117
            for obj, cpu_arr in cpu_version.items():
                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
124
        class_.to_cpu = to_cpu
        class_.to_gpu = to_gpu
125
126
        gpu = self._target in data_handling._GPU_LIKE_TARGETS
        data_handling.add_custom_class(self._index_array_name, class_, cpu=True, gpu=gpu)
127
128

    @property
Martin Bauer's avatar
Martin Bauer committed
129
130
    def data_handling(self):
        return self._data_handling
131

Martin Bauer's avatar
Martin Bauer committed
132
133
    def get_flag(self, boundary_obj):
        return self._boundary_object_to_boundary_info[boundary_obj].flag
Martin Bauer's avatar
Martin Bauer committed
134

135
136
    @property
    def shape(self):
Martin Bauer's avatar
Martin Bauer committed
137
        return self._data_handling.shape
138
139
140

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
141
        return self._data_handling.dim
142
143

    @property
Martin Bauer's avatar
Martin Bauer committed
144
    def boundary_objects(self):
145
        return tuple(self._boundary_object_to_boundary_info.keys())
146
147

    @property
Martin Bauer's avatar
Martin Bauer committed
148
149
    def flag_array_name(self):
        return self.flag_interface.flag_field_name
150

Martin Bauer's avatar
Martin Bauer committed
151
152
153
    def get_mask(self, slice_obj, boundary_obj, inverse=False):
        if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain':
            flag = self.flag_interface.domain_flag
154
        else:
Martin Bauer's avatar
Martin Bauer committed
155
            flag = self._boundary_object_to_boundary_info[boundary_obj].flag
156

Martin Bauer's avatar
Martin Bauer committed
157
        arr = self.data_handling.gather_array(self.flag_array_name, slice_obj)
158
159
160
161
162
163
164
165
        if arr is None:
            return None
        else:
            result = np.bitwise_and(arr, flag)
            if inverse:
                result = np.logical_not(result)
            return result

Martin Bauer's avatar
Martin Bauer committed
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
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
        """Sets boundary using either a rectangular slice, a boolean mask or a combination of both.

        Args:
            boundary_obj: instance of a boundary object that should be set
            slice_obj: a slice object (can be created with make_slice[]) that selects a part of the domain where
                       the boundary should be set. If none, the complete domain is selected which makes only sense
                       if a mask_callback is passed. The slice can have ':' placeholders, which are interpreted
                       depending on the 'inner_ghost_layers' parameter i.e. if it is True, the slice extends
                       into the ghost layers
            mask_callback: callback function getting x,y (z) parameters of the cell midpoints and returning a
                          boolean mask with True entries where boundary cells should be set.
                          The x, y, z arrays have 2D/3D shape such that they can be used directly
                          to create the boolean return array. i.e return x < 10 sets boundaries in cells with
                          midpoint x coordinate smaller than 10.
            ghost_layers: see DataHandling.iterate()
            inner_ghost_layers: see DataHandling.iterate()
            replace: by default all other flags are erased in the cells where the boundary is set. To add a
                     boundary condition, set this replace flag to False
186
187
            force_flag_value: flag that should be reserved for this boundary. Has to be an integer that is a power of 2
                              and was not reserved before for another boundary.
188
        """
Martin Bauer's avatar
Martin Bauer committed
189
190
        if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain':
            flag = self.flag_interface.domain_flag
191
        else:
192
193
194
            if force_flag_value:
                self.flag_interface.reserve_flag(force_flag_value)
            flag = self._add_boundary(boundary_obj, force_flag_value)
195

Martin Bauer's avatar
Martin Bauer committed
196
197
        for b in self._data_handling.iterate(slice_obj, ghost_layers=ghost_layers,
                                             inner_ghost_layers=inner_ghost_layers):
Martin Bauer's avatar
Martin Bauer committed
198
199
200
            flag_arr = b[self.flag_interface.flag_field_name]
            if mask_callback is not None:
                mask = mask_callback(*b.midpoint_arrays)
Martin Bauer's avatar
Martin Bauer committed
201
                if replace:
Martin Bauer's avatar
Martin Bauer committed
202
                    flag_arr[mask] = flag
Martin Bauer's avatar
Martin Bauer committed
203
                else:
Martin Bauer's avatar
Martin Bauer committed
204
205
                    np.bitwise_or(flag_arr, flag, where=mask, out=flag_arr)
                    np.bitwise_and(flag_arr, ~self.flag_interface.domain_flag, where=mask, out=flag_arr)
206
            else:
Martin Bauer's avatar
Martin Bauer committed
207
                if replace:
Martin Bauer's avatar
Martin Bauer committed
208
                    flag_arr.fill(flag)
Martin Bauer's avatar
Martin Bauer committed
209
                else:
Martin Bauer's avatar
Martin Bauer committed
210
211
                    np.bitwise_or(flag_arr, flag, out=flag_arr)
                    np.bitwise_and(flag_arr, ~self.flag_interface.domain_flag, out=flag_arr)
Martin Bauer's avatar
Martin Bauer committed
212
213

        self._dirty = True
214

Martin Bauer's avatar
Martin Bauer committed
215
216
        return flag

Martin Bauer's avatar
Martin Bauer committed
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."""
Martin Bauer's avatar
Martin Bauer committed
219
        self._add_boundary(boundary_obj, flag)
220
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
221
        return flag
222
223
224
225

    def prepare(self):
        if not self._dirty:
            return
Martin Bauer's avatar
Martin Bauer committed
226
        self._create_index_fields()
227
228
        self._dirty = False

Martin Bauer's avatar
Martin Bauer committed
229
    def trigger_reinitialization_of_boundary_data(self, **kwargs):
230
231
232
        if self._dirty:
            self.prepare()
        else:
Martin Bauer's avatar
Martin Bauer committed
233
234
            ff_ghost_layers = self._data_handling.ghost_layers_of_field(self.flag_interface.flag_field_name)
            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():
Martin Bauer's avatar
Martin Bauer committed
236
                    self._boundary_data_initialization(b_obj, setter, **kwargs)
237
238
239
240
241

    def __call__(self, **kwargs):
        if self._dirty:
            self.prepare()

242
        for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS):
Martin Bauer's avatar
Martin Bauer committed
243
            for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items():
Martin Bauer's avatar
Martin Bauer committed
244
                kwargs[self._field_name] = b[self._field_name]
Martin Bauer's avatar
Martin Bauer committed
245
                kwargs['indexField'] = idx_arr
246
                data_used_in_kernel = (p.fields[0].name
Martin Bauer's avatar
Martin Bauer committed
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)
Martin Bauer's avatar
Martin Bauer committed
249
250
251
                kwargs.update({name: b[name] for name in data_used_in_kernel})

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

253
254
255
256
    def add_fixed_steps(self, fixed_loop, **kwargs):
        if self._dirty:
            self.prepare()

257
        for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS):
258
259
260
261
            for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items():
                arguments = kwargs.copy()
                arguments[self._field_name] = b[self._field_name]
                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
Martin Bauer's avatar
Martin Bauer committed
264
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.field_name not in arguments)
265
266
267
268
269
                arguments.update({name: b[name] for name in data_used_in_kernel if name not in arguments})

                kernel = self._boundary_object_to_boundary_info[b_obj].kernel
                fixed_loop.add_call(kernel, arguments)

Martin Bauer's avatar
Martin Bauer committed
270
    def geometry_to_vtk(self, file_name='geometry', boundaries='all', ghost_layers=False):
271
272
273
        """
        Writes a VTK field where each cell with the given boundary is marked with 1, other cells are 0
        This can be used to display the simulation geometry in Paraview
274
275
276
277
278
279
280

        Params:
            file_name: vtk filename
            boundaries: boundary object, or special string 'domain' for domain cells or special string 'all' for all
                      boundary conditions.
                      can also  be a sequence, to write multiple boundaries to VTK file
            ghost_layers: number of ghost layers to write, or True for all, False for none
281
282
        """
        if boundaries == 'all':
Martin Bauer's avatar
Martin Bauer committed
283
            boundaries = list(self._boundary_object_to_boundary_info.keys()) + ['domain']
284
285
286
        elif not hasattr(boundaries, "__len__"):
            boundaries = [boundaries]

Martin Bauer's avatar
Martin Bauer committed
287
        masks_to_name = {}
288
        for b in boundaries:
Martin Bauer's avatar
Martin Bauer committed
289
            if b == 'domain':
Martin Bauer's avatar
Martin Bauer committed
290
                masks_to_name[self.flag_interface.domain_flag] = 'domain'
291
            else:
292
293
                flag = self._boundary_object_to_boundary_info[b].flag
                masks_to_name[flag] = b.name
294

Martin Bauer's avatar
Martin Bauer committed
295
296
        writer = self.data_handling.create_vtk_writer_for_flag_array(file_name, self.flag_interface.flag_field_name,
                                                                     masks_to_name, ghost_layers=ghost_layers)
297
298
299
300
        writer(1)

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

Martin Bauer's avatar
Martin Bauer committed
301
302
    def _add_boundary(self, boundary_obj, flag=None):
        if boundary_obj not in self._boundary_object_to_boundary_info:
Martin Bauer's avatar
Martin Bauer committed
303
304
            sym_index_field = Field.create_generic('indexField', spatial_dimensions=1,
                                                   dtype=numpy_data_type_for_boundary_object(boundary_obj, self.dim))
Martin Bauer's avatar
Martin Bauer committed
305
            ast = self._create_boundary_kernel(self._data_handling.fields[self._field_name],
Martin Bauer's avatar
Martin Bauer committed
306
                                               sym_index_field, boundary_obj)
Martin Bauer's avatar
Martin Bauer committed
307
            if flag is None:
308
                flag = self.flag_interface.reserve_next_flag()
Martin Bauer's avatar
Martin Bauer committed
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
            boundary_info = self.BoundaryInfo(boundary_obj, flag=flag, kernel=ast.compile())
            self._boundary_object_to_boundary_info[boundary_obj] = boundary_info
        return self._boundary_object_to_boundary_info[boundary_obj].flag

    def _create_boundary_kernel(self, symbolic_field, symbolic_index_field, boundary_obj):
        return create_boundary_kernel(symbolic_field, symbolic_index_field, self.stencil, boundary_obj,
                                      target=self._target, openmp=self._openmp)

    def _create_index_fields(self):
        dh = self._data_handling
        ff_ghost_layers = dh.ghost_layers_of_field(self.flag_interface.flag_field_name)
        for b in dh.iterate(ghost_layers=ff_ghost_layers):
            flag_arr = b[self.flag_interface.flag_field_name]
            pdf_arr = b[self._field_name]
            index_array_bd = b[self._index_array_name]
            index_array_bd.clear()
Martin Bauer's avatar
Martin Bauer committed
325
            for b_info in self._boundary_object_to_boundary_info.values():
Martin Bauer's avatar
Martin Bauer committed
326
                boundary_obj = b_info.boundary_object
Martin Bauer's avatar
Martin Bauer committed
327
                idx_arr = create_boundary_index_array(flag_arr, self.stencil, b_info.flag,
Martin Bauer's avatar
Martin Bauer committed
328
329
330
                                                      self.flag_interface.domain_flag, boundary_obj,
                                                      ff_ghost_layers, boundary_obj.inner_or_boundary,
                                                      boundary_obj.single_link)
Martin Bauer's avatar
Martin Bauer committed
331
                if idx_arr.size == 0:
332
333
                    continue

Martin Bauer's avatar
Martin Bauer committed
334
335
                boundary_data_setter = BoundaryDataSetter(idx_arr, b.offset, self.stencil, ff_ghost_layers, pdf_arr)
                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
Martin Bauer's avatar
Martin Bauer committed
337
                self._boundary_data_initialization(b_info.boundary_object, boundary_data_setter)
338

Martin Bauer's avatar
Martin Bauer committed
339
340
341
    def _boundary_data_initialization(self, boundary_obj, boundary_data_setter, **kwargs):
        if boundary_obj.additional_data_init_callback:
            boundary_obj.additional_data_init_callback(boundary_data_setter, **kwargs)
342
        if self._target in self._data_handling._GPU_LIKE_TARGETS:
Martin Bauer's avatar
Martin Bauer committed
343
            self._data_handling.to_gpu(self._index_array_name)
344
345

    class BoundaryInfo(object):
Martin Bauer's avatar
Martin Bauer committed
346
        def __init__(self, boundary_obj, flag, kernel):
Martin Bauer's avatar
Martin Bauer committed
347
            self.boundary_object = boundary_obj
348
349
350
351
            self.flag = flag
            self.kernel = kernel

    class IndexFieldBlockData:
352
        def __init__(self, *args, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
353
            self.boundary_object_to_index_list = {}
354
            self.boundary_object_to_data_setter = {}
355
356

        def clear(self):
Martin Bauer's avatar
Martin Bauer committed
357
            self.boundary_object_to_index_list.clear()
358
            self.boundary_object_to_data_setter.clear()
359
360
361
362


class BoundaryDataSetter:

Martin Bauer's avatar
Martin Bauer committed
363
    def __init__(self, index_array, offset, stencil, ghost_layers, pdf_array):
Martin Bauer's avatar
Martin Bauer committed
364
        self.index_array = index_array
365
366
        self.offset = offset
        self.stencil = np.array(stencil)
Martin Bauer's avatar
Martin Bauer committed
367
368
        self.pdf_array = pdf_array.view()
        self.pdf_array.flags.writeable = False
369

Martin Bauer's avatar
Martin Bauer committed
370
371
372
        arr_field_names = index_array.dtype.names
        self.dim = 3 if 'z' in arr_field_names else 2
        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'}
Martin Bauer's avatar
Martin Bauer committed
374
375
        self.coord_map = {0: 'x', 1: 'y', 2: 'z'}
        self.ghost_layers = ghost_layers
376

Martin Bauer's avatar
Martin Bauer committed
377
    def non_boundary_cell_positions(self, coord):
378
        assert coord < self.dim
Martin Bauer's avatar
Martin Bauer committed
379
        return self.index_array[self.coord_map[coord]] + self.offset[coord] - self.ghost_layers + 0.5
380
381

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
382
    def link_offsets(self):
Martin Bauer's avatar
Martin Bauer committed
383
        return self.stencil[self.index_array['dir']]
384
385

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
386
387
    def link_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + 0.5 * self.link_offsets()[:, coord]
388
389

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
390
391
    def boundary_cell_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + self.link_offsets()[:, coord]
392
393

    def __setitem__(self, key, value):
Martin Bauer's avatar
Martin Bauer committed
394
        if key not in self.boundary_data_names:
395
            raise KeyError(f"Invalid boundary data name {key}. Allowed are {self.boundary_data_names}")
Martin Bauer's avatar
Martin Bauer committed
396
        self.index_array[key] = value
397
398

    def __getitem__(self, item):
Martin Bauer's avatar
Martin Bauer committed
399
        if item not in self.boundary_data_names:
400
            raise KeyError(f"Invalid boundary data name {item}. Allowed are {self.boundary_data_names}")
Martin Bauer's avatar
Martin Bauer committed
401
        return self.index_array[item]
402
403


404
class BoundaryOffsetInfo(CustomCodeNode):
405
406
407
408

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

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
409
410
411
    def offset_from_dir(dir_idx, dim):
        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
                      for symbol in BoundaryOffsetInfo._offset_symbols(dim)])
412
413

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
414
415
    def inv_dir(dir_idx):
        return sp.IndexedBase(BoundaryOffsetInfo.INV_DIR_SYMBOL, shape=(1,))[dir_idx]
416
417
418
419
420
421

    # ---------------------------------- Internal ---------------------------------------------

    def __init__(self, stencil):
        dim = len(stencil[0])

Martin Bauer's avatar
Martin Bauer committed
422
        offset_sym = BoundaryOffsetInfo._offset_symbols(dim)
423
424
        code = "\n"
        for i in range(dim):
Martin Bauer's avatar
Martin Bauer committed
425
426
            offset_str = ", ".join([str(d[i]) for d in stencil])
            code += "const int64_t %s [] = { %s };\n" % (offset_sym[i].name, offset_str)
427

Martin Bauer's avatar
Martin Bauer committed
428
        inv_dirs = []
429
        for direction in stencil:
Martin Bauer's avatar
Martin Bauer committed
430
431
            inverse_dir = tuple([-i for i in direction])
            inv_dirs.append(str(stencil.index(inverse_dir)))
432

Martin Bauer's avatar
Martin Bauer committed
433
434
        code += "const int %s [] = { %s };\n" % (self.INV_DIR_SYMBOL.name, ", ".join(inv_dirs))
        offset_symbols = BoundaryOffsetInfo._offset_symbols(dim)
Martin Bauer's avatar
Martin Bauer committed
435
        super(BoundaryOffsetInfo, self).__init__(code, symbols_read=set(),
Martin Bauer's avatar
Martin Bauer committed
436
                                                 symbols_defined=set(offset_symbols + [self.INV_DIR_SYMBOL]))
437
438

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
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


Martin Bauer's avatar
Martin Bauer committed
445
def create_boundary_kernel(field, index_field, stencil, boundary_functor, target='cpu', openmp=True):
446
    elements = [BoundaryOffsetInfo(stencil)]
Martin Bauer's avatar
Martin Bauer committed
447
448
449
    index_arr_dtype = index_field.dtype.numpy_dtype
    dir_symbol = TypedSymbol("dir", index_arr_dtype.fields['dir'][0])
    elements += [Assignment(dir_symbol, index_field[0]('dir'))]
Martin Bauer's avatar
Martin Bauer committed
450
    elements += boundary_functor(field, direction_symbol=dir_symbol, index_field=index_field)
Martin Bauer's avatar
Martin Bauer committed
451
    return create_indexed_kernel(elements, [index_field], target=target, cpu_openmp=openmp)