boundaryhandling.py 20.2 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
12
from pystencils.datahandling import ParallelDataHandling
from pystencils.datahandling.pycuda import PyCudaArrayHandler
Martin Bauer's avatar
Martin Bauer committed
13
from pystencils.field import Field
14
from pystencils.kernelparameters import FieldPointerSymbol
15

16
DEFAULT_FLAG_TYPE = np.uint32
17

Martin Bauer's avatar
Martin Bauer committed
18

19
20
21
22
class FlagInterface:
    """Manages the reservation of bits (i.e. flags) in an array of unsigned integers.

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

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

48
        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
49
50
51
        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
52

53
54
55
56
57
58
59
    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
60
        raise ValueError("All available {} flags are reserved".format(self.max_bits))
61
62
63
64
65

    def reserve_flag(self, flag):
        assert self._is_power_of_2(flag)
        flag = self.dtype(flag)
        if flag in self._used_flags:
66
            raise ValueError("The flag {flag} is already reserved".format(flag=flag))
67
68
69
70
71
72
        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
73
74


75
76
class BoundaryHandling:

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

92
93
94
95
96
97
98
99
100
        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():
                gpu_version[obj].get(cpu_arr)

        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
101
102
103
104
105
106

            if isinstance(self.data_handling, ParallelDataHandling):
                array_handler = PyCudaArrayHandler()
            else:
                array_handler = self.data_handling.array_handler

107
108
            for obj, cpu_arr in cpu_version.items():
                if obj not in gpu_version or gpu_version[obj].shape != cpu_arr.shape:
109
                    gpu_version[obj] = array_handler.to_gpu(cpu_arr)
110
                else:
111
                    array_handler.upload(gpu_version[obj], cpu_arr)
112

113
        class_ = self.IndexFieldBlockData
114
115
        class_.to_cpu = to_cpu
        class_.to_gpu = to_gpu
116
117
        gpu = self._target in data_handling._GPU_LIKE_TARGETS
        data_handling.add_custom_class(self._index_array_name, class_, cpu=True, gpu=gpu)
118
119

    @property
Martin Bauer's avatar
Martin Bauer committed
120
121
    def data_handling(self):
        return self._data_handling
122

Martin Bauer's avatar
Martin Bauer committed
123
124
    def get_flag(self, boundary_obj):
        return self._boundary_object_to_boundary_info[boundary_obj].flag
Martin Bauer's avatar
Martin Bauer committed
125

126
127
    @property
    def shape(self):
Martin Bauer's avatar
Martin Bauer committed
128
        return self._data_handling.shape
129
130
131

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
132
        return self._data_handling.dim
133
134

    @property
Martin Bauer's avatar
Martin Bauer committed
135
    def boundary_objects(self):
136
        return tuple(self._boundary_object_to_boundary_info.keys())
137
138

    @property
Martin Bauer's avatar
Martin Bauer committed
139
140
    def flag_array_name(self):
        return self.flag_interface.flag_field_name
141

Martin Bauer's avatar
Martin Bauer committed
142
143
144
    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
145
        else:
Martin Bauer's avatar
Martin Bauer committed
146
            flag = self._boundary_object_to_boundary_info[boundary_obj].flag
147

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

Martin Bauer's avatar
Martin Bauer committed
187
188
        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
189
190
191
            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
192
                if replace:
Martin Bauer's avatar
Martin Bauer committed
193
                    flag_arr[mask] = flag
Martin Bauer's avatar
Martin Bauer committed
194
                else:
Martin Bauer's avatar
Martin Bauer committed
195
196
                    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)
197
            else:
Martin Bauer's avatar
Martin Bauer committed
198
                if replace:
Martin Bauer's avatar
Martin Bauer committed
199
                    flag_arr.fill(flag)
Martin Bauer's avatar
Martin Bauer committed
200
                else:
Martin Bauer's avatar
Martin Bauer committed
201
202
                    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
203
204

        self._dirty = True
205

Martin Bauer's avatar
Martin Bauer committed
206
207
        return flag

Martin Bauer's avatar
Martin Bauer committed
208
    def set_boundary_where_flag_is_set(self, boundary_obj, flag):
209
        """Adds an (additional) boundary to all cells that have been previously marked with the passed flag."""
Martin Bauer's avatar
Martin Bauer committed
210
        self._add_boundary(boundary_obj, flag)
211
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
212
        return flag
213
214
215
216

    def prepare(self):
        if not self._dirty:
            return
Martin Bauer's avatar
Martin Bauer committed
217
        self._create_index_fields()
218
219
        self._dirty = False

Martin Bauer's avatar
Martin Bauer committed
220
    def trigger_reinitialization_of_boundary_data(self, **kwargs):
221
222
223
        if self._dirty:
            self.prepare()
        else:
Martin Bauer's avatar
Martin Bauer committed
224
225
            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):
226
                for b_obj, setter in b[self._index_array_name].boundary_object_to_data_setter.items():
Martin Bauer's avatar
Martin Bauer committed
227
                    self._boundary_data_initialization(b_obj, setter, **kwargs)
228
229
230
231
232

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

233
        for b in self._data_handling.iterate(gpu=self._target in self._data_handling._GPU_LIKE_TARGETS):
Martin Bauer's avatar
Martin Bauer committed
234
            for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items():
Martin Bauer's avatar
Martin Bauer committed
235
                kwargs[self._field_name] = b[self._field_name]
Martin Bauer's avatar
Martin Bauer committed
236
                kwargs['indexField'] = idx_arr
237
                data_used_in_kernel = (p.fields[0].name
Martin Bauer's avatar
Martin Bauer committed
238
                                       for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters
239
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.fields[0].name not in kwargs)
Martin Bauer's avatar
Martin Bauer committed
240
241
242
                kwargs.update({name: b[name] for name in data_used_in_kernel})

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

244
245
246
247
    def add_fixed_steps(self, fixed_loop, **kwargs):
        if self._dirty:
            self.prepare()

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

        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
272
273
        """
        if boundaries == 'all':
Martin Bauer's avatar
Martin Bauer committed
274
            boundaries = list(self._boundary_object_to_boundary_info.keys()) + ['domain']
275
276
277
        elif not hasattr(boundaries, "__len__"):
            boundaries = [boundaries]

Martin Bauer's avatar
Martin Bauer committed
278
        masks_to_name = {}
279
        for b in boundaries:
Martin Bauer's avatar
Martin Bauer committed
280
            if b == 'domain':
Martin Bauer's avatar
Martin Bauer committed
281
                masks_to_name[self.flag_interface.domain_flag] = 'domain'
282
            else:
283
284
                flag = self._boundary_object_to_boundary_info[b].flag
                masks_to_name[flag] = b.name
285

Martin Bauer's avatar
Martin Bauer committed
286
287
        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)
288
289
290
291
        writer(1)

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

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

Martin Bauer's avatar
Martin Bauer committed
325
326
                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
327
                index_array_bd.boundary_object_to_data_setter[b_info.boundary_object] = boundary_data_setter
Martin Bauer's avatar
Martin Bauer committed
328
                self._boundary_data_initialization(b_info.boundary_object, boundary_data_setter)
329

Martin Bauer's avatar
Martin Bauer committed
330
331
332
    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)
333
        if self._target in self._data_handling._GPU_LIKE_TARGETS:
Martin Bauer's avatar
Martin Bauer committed
334
            self._data_handling.to_gpu(self._index_array_name)
335
336

    class BoundaryInfo(object):
Martin Bauer's avatar
Martin Bauer committed
337
        def __init__(self, boundary_obj, flag, kernel):
Martin Bauer's avatar
Martin Bauer committed
338
            self.boundary_object = boundary_obj
339
340
341
342
            self.flag = flag
            self.kernel = kernel

    class IndexFieldBlockData:
343
        def __init__(self, *args, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
344
            self.boundary_object_to_index_list = {}
345
            self.boundary_object_to_data_setter = {}
346
347

        def clear(self):
Martin Bauer's avatar
Martin Bauer committed
348
            self.boundary_object_to_index_list.clear()
349
            self.boundary_object_to_data_setter.clear()
350
351
352
353


class BoundaryDataSetter:

Martin Bauer's avatar
Martin Bauer committed
354
    def __init__(self, index_array, offset, stencil, ghost_layers, pdf_array):
Martin Bauer's avatar
Martin Bauer committed
355
        self.index_array = index_array
356
357
        self.offset = offset
        self.stencil = np.array(stencil)
Martin Bauer's avatar
Martin Bauer committed
358
359
        self.pdf_array = pdf_array.view()
        self.pdf_array.flags.writeable = False
360

Martin Bauer's avatar
Martin Bauer committed
361
362
363
        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)
364
        self.boundary_data_names = set(self.index_array.dtype.names) - {'x', 'y', 'z', 'dir'}
Martin Bauer's avatar
Martin Bauer committed
365
366
        self.coord_map = {0: 'x', 1: 'y', 2: 'z'}
        self.ghost_layers = ghost_layers
367

Martin Bauer's avatar
Martin Bauer committed
368
    def non_boundary_cell_positions(self, coord):
369
        assert coord < self.dim
Martin Bauer's avatar
Martin Bauer committed
370
        return self.index_array[self.coord_map[coord]] + self.offset[coord] - self.ghost_layers + 0.5
371
372

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
373
    def link_offsets(self):
Martin Bauer's avatar
Martin Bauer committed
374
        return self.stencil[self.index_array['dir']]
375
376

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
377
378
    def link_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + 0.5 * self.link_offsets()[:, coord]
379
380

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
381
382
    def boundary_cell_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + self.link_offsets()[:, coord]
383
384

    def __setitem__(self, key, value):
Martin Bauer's avatar
Martin Bauer committed
385
386
        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
387
        self.index_array[key] = value
388
389

    def __getitem__(self, item):
Martin Bauer's avatar
Martin Bauer committed
390
391
        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
392
        return self.index_array[item]
393
394


395
class BoundaryOffsetInfo(CustomCodeNode):
396
397
398
399

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

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
400
401
402
    def offset_from_dir(dir_idx, dim):
        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
                      for symbol in BoundaryOffsetInfo._offset_symbols(dim)])
403
404

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
405
406
    def inv_dir(dir_idx):
        return sp.IndexedBase(BoundaryOffsetInfo.INV_DIR_SYMBOL, shape=(1,))[dir_idx]
407
408
409
410
411
412

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

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

Martin Bauer's avatar
Martin Bauer committed
413
        offset_sym = BoundaryOffsetInfo._offset_symbols(dim)
414
415
        code = "\n"
        for i in range(dim):
Martin Bauer's avatar
Martin Bauer committed
416
417
            offset_str = ", ".join([str(d[i]) for d in stencil])
            code += "const int64_t %s [] = { %s };\n" % (offset_sym[i].name, offset_str)
418

Martin Bauer's avatar
Martin Bauer committed
419
        inv_dirs = []
420
        for direction in stencil:
Martin Bauer's avatar
Martin Bauer committed
421
422
            inverse_dir = tuple([-i for i in direction])
            inv_dirs.append(str(stencil.index(inverse_dir)))
423

Martin Bauer's avatar
Martin Bauer committed
424
425
        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
426
        super(BoundaryOffsetInfo, self).__init__(code, symbols_read=set(),
Martin Bauer's avatar
Martin Bauer committed
427
                                                 symbols_defined=set(offset_symbols + [self.INV_DIR_SYMBOL]))
428
429

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

433
    INV_DIR_SYMBOL = TypedSymbol("invdir", "int")
434
435


Martin Bauer's avatar
Martin Bauer committed
436
def create_boundary_kernel(field, index_field, stencil, boundary_functor, target='cpu', openmp=True):
437
    elements = [BoundaryOffsetInfo(stencil)]
Martin Bauer's avatar
Martin Bauer committed
438
439
440
    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
441
    elements += boundary_functor(field, direction_symbol=dir_symbol, index_field=index_field)
Martin Bauer's avatar
Martin Bauer committed
442
    return create_indexed_kernel(elements, [index_field], target=target, cpu_openmp=openmp)