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
20
21
22
try:
    # noinspection PyPep8Naming
    import waLBerla as wlb
    if wlb.cpp_available:
        from pystencils.datahandling.parallel_datahandling import ParallelDataHandling
except ImportError:
    ParallelDataHandling = None

23
DEFAULT_FLAG_TYPE = np.uint32
24

Martin Bauer's avatar
Martin Bauer committed
25

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

    Examples:
30
31
        >>> from pystencils import create_data_handling
        >>> dh = create_data_handling((4, 5))
32
33
34
35
36
37
38
39
40
41
42
        >>> 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
43
        self.flag_field_name = flag_field_name
44
45
        self.domain_flag = dtype(1 << 0)
        self._used_flags = {self.domain_flag}
Martin Bauer's avatar
Martin Bauer committed
46
        self.data_handling = data_handling
47
48
        self.dtype = dtype
        self.max_bits = self.dtype().itemsize * 8
Martin Bauer's avatar
Martin Bauer committed
49
50

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

55
        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
56
57
58
        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
59

60
61
62
63
64
65
66
    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
67
        raise ValueError("All available {} flags are reserved".format(self.max_bits))
68
69
70
71
72

    def reserve_flag(self, flag):
        assert self._is_power_of_2(flag)
        flag = self.dtype(flag)
        if flag in self._used_flags:
73
            raise ValueError("The flag {flag} is already reserved".format(flag=flag))
74
75
76
77
78
79
        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
80
81


82
83
class BoundaryHandling:

Martin Bauer's avatar
Martin Bauer committed
84
85
86
    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)
87
        assert data_handling.dim == len(stencil[0]), "Dimension of stencil and data handling do not match"
Martin Bauer's avatar
Martin Bauer committed
88
89
90
        self._data_handling = data_handling
        self._field_name = field_name
        self._index_array_name = name + "IndexArrays"
91
        self._target = target
Martin Bauer's avatar
Martin Bauer committed
92
93
        self._openmp = openmp
        self._boundary_object_to_boundary_info = {}
94
95
        self.stencil = stencil
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
96
97
        fi = flag_interface
        self.flag_interface = fi if fi is not None else FlagInterface(data_handling, name + "Flags")
98

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

104
105
106
107
        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():
108
                array_handler.download(gpu_version[obj], cpu_arr)
109
110
111
112

        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
113

114
115
            for obj, cpu_arr in cpu_version.items():
                if obj not in gpu_version or gpu_version[obj].shape != cpu_arr.shape:
116
                    gpu_version[obj] = array_handler.to_gpu(cpu_arr)
117
                else:
118
                    array_handler.upload(gpu_version[obj], cpu_arr)
119

120
        class_ = self.IndexFieldBlockData
121
122
        class_.to_cpu = to_cpu
        class_.to_gpu = to_gpu
123
124
        gpu = self._target in data_handling._GPU_LIKE_TARGETS
        data_handling.add_custom_class(self._index_array_name, class_, cpu=True, gpu=gpu)
125
126

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

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

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

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

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

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

Martin Bauer's avatar
Martin Bauer committed
149
150
151
    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
152
        else:
Martin Bauer's avatar
Martin Bauer committed
153
            flag = self._boundary_object_to_boundary_info[boundary_obj].flag
154

Martin Bauer's avatar
Martin Bauer committed
155
        arr = self.data_handling.gather_array(self.flag_array_name, slice_obj)
156
157
158
159
160
161
162
163
        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
164
    def set_boundary(self, boundary_obj, slice_obj=None, mask_callback=None,
165
                     ghost_layers=True, inner_ghost_layers=True, replace=True, force_flag_value=None):
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
        """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
184
185
            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.
186
        """
Martin Bauer's avatar
Martin Bauer committed
187
188
        if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain':
            flag = self.flag_interface.domain_flag
189
        else:
190
191
192
            if force_flag_value:
                self.flag_interface.reserve_flag(force_flag_value)
            flag = self._add_boundary(boundary_obj, force_flag_value)
193

Martin Bauer's avatar
Martin Bauer committed
194
195
        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
196
197
198
            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
199
                if replace:
Martin Bauer's avatar
Martin Bauer committed
200
                    flag_arr[mask] = flag
Martin Bauer's avatar
Martin Bauer committed
201
                else:
Martin Bauer's avatar
Martin Bauer committed
202
203
                    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)
204
            else:
Martin Bauer's avatar
Martin Bauer committed
205
                if replace:
Martin Bauer's avatar
Martin Bauer committed
206
                    flag_arr.fill(flag)
Martin Bauer's avatar
Martin Bauer committed
207
                else:
Martin Bauer's avatar
Martin Bauer committed
208
209
                    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
210
211

        self._dirty = True
212

Martin Bauer's avatar
Martin Bauer committed
213
214
        return flag

Martin Bauer's avatar
Martin Bauer committed
215
    def set_boundary_where_flag_is_set(self, boundary_obj, flag):
216
        """Adds an (additional) boundary to all cells that have been previously marked with the passed flag."""
Martin Bauer's avatar
Martin Bauer committed
217
        self._add_boundary(boundary_obj, flag)
218
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
219
        return flag
220
221
222
223

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

Martin Bauer's avatar
Martin Bauer committed
227
    def trigger_reinitialization_of_boundary_data(self, **kwargs):
228
229
230
        if self._dirty:
            self.prepare()
        else:
Martin Bauer's avatar
Martin Bauer committed
231
232
            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):
233
                for b_obj, setter in b[self._index_array_name].boundary_object_to_data_setter.items():
Martin Bauer's avatar
Martin Bauer committed
234
                    self._boundary_data_initialization(b_obj, setter, **kwargs)
235
236
237
238
239

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

240
        for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS):
Martin Bauer's avatar
Martin Bauer committed
241
            for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items():
Martin Bauer's avatar
Martin Bauer committed
242
                kwargs[self._field_name] = b[self._field_name]
Martin Bauer's avatar
Martin Bauer committed
243
                kwargs['indexField'] = idx_arr
244
                data_used_in_kernel = (p.fields[0].name
Martin Bauer's avatar
Martin Bauer committed
245
                                       for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters
246
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.fields[0].name not in kwargs)
Martin Bauer's avatar
Martin Bauer committed
247
248
249
                kwargs.update({name: b[name] for name in data_used_in_kernel})

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

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

255
        for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS):
256
257
258
259
            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
260
                data_used_in_kernel = (p.fields[0].name
261
                                       for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters
Martin Bauer's avatar
Martin Bauer committed
262
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.field_name not in arguments)
263
264
265
266
267
                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
268
    def geometry_to_vtk(self, file_name='geometry', boundaries='all', ghost_layers=False):
269
270
271
        """
        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
272
273
274
275
276
277
278

        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
279
280
        """
        if boundaries == 'all':
Martin Bauer's avatar
Martin Bauer committed
281
            boundaries = list(self._boundary_object_to_boundary_info.keys()) + ['domain']
282
283
284
        elif not hasattr(boundaries, "__len__"):
            boundaries = [boundaries]

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

Martin Bauer's avatar
Martin Bauer committed
293
294
        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)
295
296
297
298
        writer(1)

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

Martin Bauer's avatar
Martin Bauer committed
299
300
    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
301
302
            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
303
            ast = self._create_boundary_kernel(self._data_handling.fields[self._field_name],
Martin Bauer's avatar
Martin Bauer committed
304
                                               sym_index_field, boundary_obj)
Martin Bauer's avatar
Martin Bauer committed
305
            if flag is None:
306
                flag = self.flag_interface.reserve_next_flag()
Martin Bauer's avatar
Martin Bauer committed
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
            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
323
            for b_info in self._boundary_object_to_boundary_info.values():
Martin Bauer's avatar
Martin Bauer committed
324
                boundary_obj = b_info.boundary_object
Martin Bauer's avatar
Martin Bauer committed
325
                idx_arr = create_boundary_index_array(flag_arr, self.stencil, b_info.flag,
Martin Bauer's avatar
Martin Bauer committed
326
327
328
                                                      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
329
                if idx_arr.size == 0:
330
331
                    continue

Martin Bauer's avatar
Martin Bauer committed
332
333
                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
334
                index_array_bd.boundary_object_to_data_setter[b_info.boundary_object] = boundary_data_setter
Martin Bauer's avatar
Martin Bauer committed
335
                self._boundary_data_initialization(b_info.boundary_object, boundary_data_setter)
336

Martin Bauer's avatar
Martin Bauer committed
337
338
339
    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)
340
        if self._target in self._data_handling._GPU_LIKE_TARGETS:
Martin Bauer's avatar
Martin Bauer committed
341
            self._data_handling.to_gpu(self._index_array_name)
342
343

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

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

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


class BoundaryDataSetter:

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

Martin Bauer's avatar
Martin Bauer committed
368
369
370
        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)
371
        self.boundary_data_names = set(self.index_array.dtype.names) - {'x', 'y', 'z', 'dir'}
Martin Bauer's avatar
Martin Bauer committed
372
373
        self.coord_map = {0: 'x', 1: 'y', 2: 'z'}
        self.ghost_layers = ghost_layers
374

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

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

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

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

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

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


402
class BoundaryOffsetInfo(CustomCodeNode):
403
404
405
406

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

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

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

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

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

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

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

Martin Bauer's avatar
Martin Bauer committed
431
432
        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
433
        super(BoundaryOffsetInfo, self).__init__(code, symbols_read=set(),
Martin Bauer's avatar
Martin Bauer committed
434
                                                 symbols_defined=set(offset_symbols + [self.INV_DIR_SYMBOL]))
435
436

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
437
    def _offset_symbols(dim):
438
        return [TypedSymbol("c%s" % (d,), create_type(np.int64)) for d in ['x', 'y', 'z'][:dim]]
439

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


Martin Bauer's avatar
Martin Bauer committed
443
def create_boundary_kernel(field, index_field, stencil, boundary_functor, target='cpu', openmp=True):
444
    elements = [BoundaryOffsetInfo(stencil)]
Martin Bauer's avatar
Martin Bauer committed
445
446
447
    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
448
    elements += boundary_functor(field, direction_symbol=dir_symbol, index_field=index_field)
Martin Bauer's avatar
Martin Bauer committed
449
    return create_indexed_kernel(elements, [index_field], target=target, cpu_openmp=openmp)