boundaryhandling.py 19.4 KB
Newer Older
1
2
import numpy as np
import sympy as sp
3
from pystencils.assignment import Assignment
Martin Bauer's avatar
Martin Bauer committed
4
from pystencils import Field, TypedSymbol, create_indexed_kernel
5
from pystencils.backends.cbackend import CustomCppCode
Martin Bauer's avatar
Martin Bauer committed
6
from pystencils.boundaries.createindexlist import numpy_data_type_for_boundary_object, create_boundary_index_array
7
from pystencils.cache import memorycache
Martin Bauer's avatar
Martin Bauer committed
8
from pystencils.data_types import create_type
9
from pystencils.kernelparameters import FieldPointerSymbol
10

11
DEFAULT_FLAG_TYPE = np.uint32
12

Martin Bauer's avatar
Martin Bauer committed
13

14
15
16
17
class FlagInterface:
    """Manages the reservation of bits (i.e. flags) in an array of unsigned integers.

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

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

43
        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
44
45
46
        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
47

48
49
50
51
52
53
54
    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
55
        raise ValueError("All available {} flags are reserved".format(self.max_bits))
56
57
58
59
60

    def reserve_flag(self, flag):
        assert self._is_power_of_2(flag)
        flag = self.dtype(flag)
        if flag in self._used_flags:
61
            raise ValueError("The flag {flag} is already reserved".format(flag=flag))
62
63
64
65
66
67
        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
68
69


70
71
class BoundaryHandling:

Martin Bauer's avatar
Martin Bauer committed
72
73
74
    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)
75

Martin Bauer's avatar
Martin Bauer committed
76
77
78
        self._data_handling = data_handling
        self._field_name = field_name
        self._index_array_name = name + "IndexArrays"
79
        self._target = target
Martin Bauer's avatar
Martin Bauer committed
80
81
        self._openmp = openmp
        self._boundary_object_to_boundary_info = {}
82
83
        self.stencil = stencil
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
84
85
        fi = flag_interface
        self.flag_interface = fi if fi is not None else FlagInterface(data_handling, name + "Flags")
86
87

        gpu = self._target == 'gpu'
Martin Bauer's avatar
Martin Bauer committed
88
        data_handling.add_custom_class(self._index_array_name, self.IndexFieldBlockData, cpu=True, gpu=gpu)
89
90

    @property
Martin Bauer's avatar
Martin Bauer committed
91
92
    def data_handling(self):
        return self._data_handling
93

Martin Bauer's avatar
Martin Bauer committed
94
95
    def get_flag(self, boundary_obj):
        return self._boundary_object_to_boundary_info[boundary_obj].flag
Martin Bauer's avatar
Martin Bauer committed
96

97
98
    @property
    def shape(self):
Martin Bauer's avatar
Martin Bauer committed
99
        return self._data_handling.shape
100
101
102

    @property
    def dim(self):
Martin Bauer's avatar
Martin Bauer committed
103
        return self._data_handling.dim
104
105

    @property
Martin Bauer's avatar
Martin Bauer committed
106
    def boundary_objects(self):
107
        return tuple(self._boundary_object_to_boundary_info.keys())
108
109

    @property
Martin Bauer's avatar
Martin Bauer committed
110
111
    def flag_array_name(self):
        return self.flag_interface.flag_field_name
112

Martin Bauer's avatar
Martin Bauer committed
113
114
115
    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
116
        else:
Martin Bauer's avatar
Martin Bauer committed
117
            flag = self._boundary_object_to_boundary_info[boundary_obj].flag
118

Martin Bauer's avatar
Martin Bauer committed
119
        arr = self.data_handling.gather_array(self.flag_array_name, slice_obj)
120
121
122
123
124
125
126
127
        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
128
    def set_boundary(self, boundary_obj, slice_obj=None, mask_callback=None,
129
                     ghost_layers=True, inner_ghost_layers=True, replace=True, force_flag_value=None):
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        """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
148
149
            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.
150
        """
Martin Bauer's avatar
Martin Bauer committed
151
152
        if isinstance(boundary_obj, str) and boundary_obj.lower() == 'domain':
            flag = self.flag_interface.domain_flag
153
        else:
154
155
156
            if force_flag_value:
                self.flag_interface.reserve_flag(force_flag_value)
            flag = self._add_boundary(boundary_obj, force_flag_value)
157

Martin Bauer's avatar
Martin Bauer committed
158
159
        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
160
161
162
            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
163
                if replace:
Martin Bauer's avatar
Martin Bauer committed
164
                    flag_arr[mask] = flag
Martin Bauer's avatar
Martin Bauer committed
165
                else:
Martin Bauer's avatar
Martin Bauer committed
166
167
                    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)
168
            else:
Martin Bauer's avatar
Martin Bauer committed
169
                if replace:
Martin Bauer's avatar
Martin Bauer committed
170
                    flag_arr.fill(flag)
Martin Bauer's avatar
Martin Bauer committed
171
                else:
Martin Bauer's avatar
Martin Bauer committed
172
173
                    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
174
175

        self._dirty = True
176

Martin Bauer's avatar
Martin Bauer committed
177
178
        return flag

Martin Bauer's avatar
Martin Bauer committed
179
    def set_boundary_where_flag_is_set(self, boundary_obj, flag):
180
        """Adds an (additional) boundary to all cells that have been previously marked with the passed flag."""
Martin Bauer's avatar
Martin Bauer committed
181
        self._add_boundary(boundary_obj, flag)
182
        self._dirty = True
Martin Bauer's avatar
Martin Bauer committed
183
        return flag
184
185
186
187

    def prepare(self):
        if not self._dirty:
            return
Martin Bauer's avatar
Martin Bauer committed
188
        self._create_index_fields()
189
190
        self._dirty = False

Martin Bauer's avatar
Martin Bauer committed
191
    def trigger_reinitialization_of_boundary_data(self, **kwargs):
192
193
194
        if self._dirty:
            self.prepare()
        else:
Martin Bauer's avatar
Martin Bauer committed
195
196
            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):
197
                for b_obj, setter in b[self._index_array_name].boundary_object_to_data_setter.items():
Martin Bauer's avatar
Martin Bauer committed
198
                    self._boundary_data_initialization(b_obj, setter, **kwargs)
199
200
201
202
203

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

Martin Bauer's avatar
Martin Bauer committed
204
        for b in self._data_handling.iterate(gpu=self._target == 'gpu'):
Martin Bauer's avatar
Martin Bauer committed
205
            for b_obj, idx_arr in b[self._index_array_name].boundary_object_to_index_list.items():
Martin Bauer's avatar
Martin Bauer committed
206
                kwargs[self._field_name] = b[self._field_name]
Martin Bauer's avatar
Martin Bauer committed
207
                kwargs['indexField'] = idx_arr
208
                data_used_in_kernel = (p.fields[0].name
Martin Bauer's avatar
Martin Bauer committed
209
                                       for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters
210
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.fields[0].name not in kwargs)
Martin Bauer's avatar
Martin Bauer committed
211
212
213
                kwargs.update({name: b[name] for name in data_used_in_kernel})

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

215
216
217
218
219
220
221
222
223
    def add_fixed_steps(self, fixed_loop, **kwargs):
        if self._dirty:
            self.prepare()

        for b in self._data_handling.iterate(gpu=self._target == 'gpu'):
            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
224
                data_used_in_kernel = (p.fields[0].name
225
                                       for p in self._boundary_object_to_boundary_info[b_obj].kernel.parameters
226
                                       if isinstance(p.symbol, FieldPointerSymbol) and p.fields[0].name not in arguments)
227
228
229
230
231
                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
232
    def geometry_to_vtk(self, file_name='geometry', boundaries='all', ghost_layers=False):
233
234
235
        """
        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
Martin Bauer's avatar
Martin Bauer committed
236
        :param file_name: vtk filename
Martin Bauer's avatar
Martin Bauer committed
237
        :param boundaries: boundary object, or special string 'domain' for domain cells or special string 'all' for all
238
239
                         boundary conditions.
                         can also  be a sequence, to write multiple boundaries to VTK file
Martin Bauer's avatar
Martin Bauer committed
240
        :param ghost_layers: number of ghost layers to write, or True for all, False for none
241
242
        """
        if boundaries == 'all':
Martin Bauer's avatar
Martin Bauer committed
243
            boundaries = list(self._boundary_object_to_boundary_info.keys()) + ['domain']
244
245
246
        elif not hasattr(boundaries, "__len__"):
            boundaries = [boundaries]

Martin Bauer's avatar
Martin Bauer committed
247
        masks_to_name = {}
248
        for b in boundaries:
Martin Bauer's avatar
Martin Bauer committed
249
            if b == 'domain':
Martin Bauer's avatar
Martin Bauer committed
250
                masks_to_name[self.flag_interface.domain_flag] = 'domain'
251
            else:
252
253
                flag = self._boundary_object_to_boundary_info[b].flag
                masks_to_name[flag] = b.name
254

Martin Bauer's avatar
Martin Bauer committed
255
256
        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)
257
258
259
260
        writer(1)

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

Martin Bauer's avatar
Martin Bauer committed
261
262
    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
263
264
            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
265
            ast = self._create_boundary_kernel(self._data_handling.fields[self._field_name],
Martin Bauer's avatar
Martin Bauer committed
266
                                               sym_index_field, boundary_obj)
Martin Bauer's avatar
Martin Bauer committed
267
            if flag is None:
268
                flag = self.flag_interface.reserve_next_flag()
Martin Bauer's avatar
Martin Bauer committed
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
            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
285
286
287
288
289
            for b_info in self._boundary_object_to_boundary_info.values():
                idx_arr = create_boundary_index_array(flag_arr, self.stencil, b_info.flag,
                                                      self.flag_interface.domain_flag, b_info.boundary_object,
                                                      ff_ghost_layers)
                if idx_arr.size == 0:
290
291
                    continue

Martin Bauer's avatar
Martin Bauer committed
292
293
                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
294
                index_array_bd.boundary_object_to_data_setter[b_info.boundary_object] = boundary_data_setter
Martin Bauer's avatar
Martin Bauer committed
295
                self._boundary_data_initialization(b_info.boundary_object, boundary_data_setter)
296

Martin Bauer's avatar
Martin Bauer committed
297
298
299
    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)
300
        if self._target == 'gpu':
Martin Bauer's avatar
Martin Bauer committed
301
            self._data_handling.to_gpu(self._index_array_name)
302
303

    class BoundaryInfo(object):
Martin Bauer's avatar
Martin Bauer committed
304
        def __init__(self, boundary_obj, flag, kernel):
Martin Bauer's avatar
Martin Bauer committed
305
            self.boundary_object = boundary_obj
306
307
308
309
            self.flag = flag
            self.kernel = kernel

    class IndexFieldBlockData:
Martin Bauer's avatar
Martin Bauer committed
310
        def __init__(self, *_1, **_2):
Martin Bauer's avatar
Martin Bauer committed
311
            self.boundary_object_to_index_list = {}
312
            self.boundary_object_to_data_setter = {}
313
314

        def clear(self):
Martin Bauer's avatar
Martin Bauer committed
315
            self.boundary_object_to_index_list.clear()
316
            self.boundary_object_to_data_setter.clear()
317
318

        @staticmethod
Martin Bauer's avatar
Martin Bauer committed
319
320
321
        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
322
            for obj, cpu_arr in cpu_version.items():
Martin Bauer's avatar
Martin Bauer committed
323
                gpu_version[obj].get(cpu_arr)
324
325

        @staticmethod
Martin Bauer's avatar
Martin Bauer committed
326
        def to_gpu(gpu_version, cpu_version):
327
            from pycuda import gpuarray
Martin Bauer's avatar
Martin Bauer committed
328
329
            gpu_version = gpu_version.boundary_object_to_index_list
            cpu_version = cpu_version.boundary_object_to_index_list
Martin Bauer's avatar
Martin Bauer committed
330
            for obj, cpu_arr in cpu_version.items():
331
                if obj not in gpu_version or gpu_version[obj].shape != cpu_arr.shape:
Martin Bauer's avatar
Martin Bauer committed
332
                    gpu_version[obj] = gpuarray.to_gpu(cpu_arr)
333
                else:
Martin Bauer's avatar
Martin Bauer committed
334
                    gpu_version[obj].set(cpu_arr)
335
336
337
338


class BoundaryDataSetter:

Martin Bauer's avatar
Martin Bauer committed
339
    def __init__(self, index_array, offset, stencil, ghost_layers, pdf_array):
Martin Bauer's avatar
Martin Bauer committed
340
        self.index_array = index_array
341
342
        self.offset = offset
        self.stencil = np.array(stencil)
Martin Bauer's avatar
Martin Bauer committed
343
344
        self.pdf_array = pdf_array.view()
        self.pdf_array.flags.writeable = False
345

Martin Bauer's avatar
Martin Bauer committed
346
347
348
        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)
349
        self.boundary_data_names = set(self.index_array.dtype.names) - {'x', 'y', 'z', 'dir'}
Martin Bauer's avatar
Martin Bauer committed
350
351
        self.coord_map = {0: 'x', 1: 'y', 2: 'z'}
        self.ghost_layers = ghost_layers
352

Martin Bauer's avatar
Martin Bauer committed
353
    def non_boundary_cell_positions(self, coord):
354
        assert coord < self.dim
Martin Bauer's avatar
Martin Bauer committed
355
        return self.index_array[self.coord_map[coord]] + self.offset[coord] - self.ghost_layers + 0.5
356
357

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
358
    def link_offsets(self):
Martin Bauer's avatar
Martin Bauer committed
359
        return self.stencil[self.index_array['dir']]
360
361

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
362
363
    def link_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + 0.5 * self.link_offsets()[:, coord]
364
365

    @memorycache()
Martin Bauer's avatar
Martin Bauer committed
366
367
    def boundary_cell_positions(self, coord):
        return self.non_boundary_cell_positions(coord) + self.link_offsets()[:, coord]
368
369

    def __setitem__(self, key, value):
Martin Bauer's avatar
Martin Bauer committed
370
371
        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
372
        self.index_array[key] = value
373
374

    def __getitem__(self, item):
Martin Bauer's avatar
Martin Bauer committed
375
376
        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
377
        return self.index_array[item]
378
379
380
381
382
383
384


class BoundaryOffsetInfo(CustomCppCode):

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

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
385
386
387
    def offset_from_dir(dir_idx, dim):
        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
                      for symbol in BoundaryOffsetInfo._offset_symbols(dim)])
388
389

    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
390
391
    def inv_dir(dir_idx):
        return sp.IndexedBase(BoundaryOffsetInfo.INV_DIR_SYMBOL, shape=(1,))[dir_idx]
392
393
394
395
396
397

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

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

Martin Bauer's avatar
Martin Bauer committed
398
        offset_sym = BoundaryOffsetInfo._offset_symbols(dim)
399
400
        code = "\n"
        for i in range(dim):
Martin Bauer's avatar
Martin Bauer committed
401
402
            offset_str = ", ".join([str(d[i]) for d in stencil])
            code += "const int64_t %s [] = { %s };\n" % (offset_sym[i].name, offset_str)
403

Martin Bauer's avatar
Martin Bauer committed
404
        inv_dirs = []
405
        for direction in stencil:
Martin Bauer's avatar
Martin Bauer committed
406
407
            inverse_dir = tuple([-i for i in direction])
            inv_dirs.append(str(stencil.index(inverse_dir)))
408

Martin Bauer's avatar
Martin Bauer committed
409
410
        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
411
        super(BoundaryOffsetInfo, self).__init__(code, symbols_read=set(),
Martin Bauer's avatar
Martin Bauer committed
412
                                                 symbols_defined=set(offset_symbols + [self.INV_DIR_SYMBOL]))
413
414

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

418
    INV_DIR_SYMBOL = TypedSymbol("invdir", "int")
419
420


Martin Bauer's avatar
Martin Bauer committed
421
def create_boundary_kernel(field, index_field, stencil, boundary_functor, target='cpu', openmp=True):
422
    elements = [BoundaryOffsetInfo(stencil)]
Martin Bauer's avatar
Martin Bauer committed
423
424
425
    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
426
    elements += boundary_functor(field, direction_symbol=dir_symbol, index_field=index_field)
Martin Bauer's avatar
Martin Bauer committed
427
    return create_indexed_kernel(elements, [index_field], target=target, cpu_openmp=openmp)