serial_datahandling.py 17.5 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
import itertools
Martin Bauer's avatar
Martin Bauer committed
2
from typing import Sequence, Union
3
import numpy as np
4
import time
5
from pystencils import Field
6
from pystencils.datahandling.datahandling_interface import DataHandling
Martin Bauer's avatar
Martin Bauer committed
7
from pystencils.field import layout_string_to_tuple, spatial_layout_string_to_tuple, create_numpy_array_with_layout
8
from pystencils.datahandling.blockiteration import SerialBlock
Martin Bauer's avatar
Martin Bauer committed
9
from pystencils.slicing import normalize_slice, remove_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
10
11
12
13
from pystencils.utils import DotDict

try:
    import pycuda.gpuarray as gpuarray
Martin Bauer's avatar
Martin Bauer committed
14
    import pycuda.autoinit  # NOQA
Martin Bauer's avatar
Martin Bauer committed
15
16
17
except ImportError:
    gpuarray = None

18

Martin Bauer's avatar
Martin Bauer committed
19
20
class SerialDataHandling(DataHandling):

Martin Bauer's avatar
Martin Bauer committed
21
    def __init__(self, domain_size: Sequence[int], default_ghost_layers: int = 1, default_layout: str = 'SoA',
Martin Bauer's avatar
Martin Bauer committed
22
                 periodicity: Union[bool, Sequence[bool]] = False, default_target: str = 'cpu') -> None:
Martin Bauer's avatar
Martin Bauer committed
23
        """
Martin Bauer's avatar
Martin Bauer committed
24
25
26
27
28
29
30
31
        Creates a data handling for single node simulations.

        Args:
            domain_size: size of the spatial domain as tuple
            default_ghost_layers: default number of ghost layers used, if not overridden in add_array() method
            default_layout: default layout used, if  not overridden in add_array() method
            default_target: either 'cpu' or 'gpu' . If set to 'gpu' for each array also a GPU version is allocated
                            if not overwritten in add_array, and synchronization functions are for the GPU by default
Martin Bauer's avatar
Martin Bauer committed
32
        """
33
        super(SerialDataHandling, self).__init__()
Martin Bauer's avatar
Martin Bauer committed
34
        self._domainSize = tuple(domain_size)
Martin Bauer's avatar
Martin Bauer committed
35
36
        self.default_ghost_layers = default_ghost_layers
        self.default_layout = default_layout
Martin Bauer's avatar
Martin Bauer committed
37
        self._fields = DotDict()
Martin Bauer's avatar
Martin Bauer committed
38
39
40
41
42
        self.cpu_arrays = DotDict()
        self.gpu_arrays = DotDict()
        self.custom_data_cpu = DotDict()
        self.custom_data_gpu = DotDict()
        self._custom_data_transfer_functions = {}
43

44
45
46
47
48
49
        if periodicity is None or periodicity is False:
            periodicity = [False] * self.dim
        if periodicity is True:
            periodicity = [True] * self.dim

        self._periodicity = periodicity
Martin Bauer's avatar
Martin Bauer committed
50
        self._field_information = {}
Martin Bauer's avatar
Martin Bauer committed
51
        self.default_target = default_target
52
        self._start_time = time.perf_counter()
Martin Bauer's avatar
Martin Bauer committed
53
54
55
56
57

    @property
    def dim(self):
        return len(self._domainSize)

Martin Bauer's avatar
Martin Bauer committed
58
59
60
61
62
63
64
65
    @property
    def shape(self):
        return self._domainSize

    @property
    def periodicity(self):
        return self._periodicity

Martin Bauer's avatar
Martin Bauer committed
66
67
68
69
    @property
    def fields(self):
        return self._fields

Martin Bauer's avatar
Martin Bauer committed
70
71
    def ghost_layers_of_field(self, name):
        return self._field_information[name]['ghost_layers']
72

Martin Bauer's avatar
Martin Bauer committed
73
74
    def values_per_cell(self, name):
        return self._field_information[name]['values_per_cell']
75

Martin Bauer's avatar
Martin Bauer committed
76
77
78
    def add_array(self, name, values_per_cell=1, dtype=np.float64, latex_name=None, ghost_layers=None, layout=None,
                  cpu=True, gpu=None, alignment=False):
        if ghost_layers is None:
Martin Bauer's avatar
Martin Bauer committed
79
            ghost_layers = self.default_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
80
        if layout is None:
Martin Bauer's avatar
Martin Bauer committed
81
            layout = self.default_layout
82
        if gpu is None:
Martin Bauer's avatar
Martin Bauer committed
83
            gpu = self.default_target == 'gpu'
Martin Bauer's avatar
Martin Bauer committed
84
85

        kwargs = {
Martin Bauer's avatar
Martin Bauer committed
86
            'shape': tuple(s + 2 * ghost_layers for s in self._domainSize),
Martin Bauer's avatar
Martin Bauer committed
87
88
            'dtype': dtype,
        }
89
90
91
92
93
94

        if not hasattr(values_per_cell, '__len__'):
            values_per_cell = (values_per_cell, )
        if len(values_per_cell) == 1 and values_per_cell[0] == 1:
            values_per_cell = ()

Martin Bauer's avatar
Martin Bauer committed
95
96
97
        self._field_information[name] = {
            'ghost_layers': ghost_layers,
            'values_per_cell': values_per_cell,
Martin Bauer's avatar
Martin Bauer committed
98
99
            'layout': layout,
            'dtype': dtype,
Martin Bauer's avatar
Martin Bauer committed
100
            'alignment': alignment,
Martin Bauer's avatar
Martin Bauer committed
101
102
        }

103
104
105
106
107
        index_dimensions = len(values_per_cell)
        kwargs['shape'] = kwargs['shape'] + values_per_cell

        if index_dimensions > 0:
            layout_tuple = layout_string_to_tuple(layout, self.dim + index_dimensions)
Martin Bauer's avatar
Martin Bauer committed
108
        else:
Martin Bauer's avatar
Martin Bauer committed
109
            layout_tuple = spatial_layout_string_to_tuple(layout, self.dim)
110

Martin Bauer's avatar
Martin Bauer committed
111
        # cpu_arr is always created - since there is no create_pycuda_array_with_layout()
Martin Bauer's avatar
Martin Bauer committed
112
113
114
        byte_offset = ghost_layers * np.dtype(dtype).itemsize
        cpu_arr = create_numpy_array_with_layout(layout=layout_tuple, alignment=alignment,
                                                 byte_offset=byte_offset, **kwargs)
115

Martin Bauer's avatar
Martin Bauer committed
116
117
118
        if alignment and gpu:
            raise NotImplementedError("Alignment for GPU fields not supported")

Martin Bauer's avatar
Martin Bauer committed
119
        if cpu:
Martin Bauer's avatar
Martin Bauer committed
120
            if name in self.cpu_arrays:
Martin Bauer's avatar
Martin Bauer committed
121
                raise ValueError("CPU Field with this name already exists")
Martin Bauer's avatar
Martin Bauer committed
122
            self.cpu_arrays[name] = cpu_arr
Martin Bauer's avatar
Martin Bauer committed
123
        if gpu:
Martin Bauer's avatar
Martin Bauer committed
124
            if name in self.gpu_arrays:
Martin Bauer's avatar
Martin Bauer committed
125
                raise ValueError("GPU Field with this name already exists")
Martin Bauer's avatar
Martin Bauer committed
126
            self.gpu_arrays[name] = gpuarray.to_gpu(cpu_arr)
Martin Bauer's avatar
Martin Bauer committed
127

Martin Bauer's avatar
Martin Bauer committed
128
        assert all(f.name != name for f in self.fields.values()), "Symbolic field with this name already exists"
Martin Bauer's avatar
Martin Bauer committed
129
130
        self.fields[name] = Field.create_from_numpy_array(name, cpu_arr, index_dimensions=index_dimensions)
        self.fields[name].latex_name = latex_name
131
        return self.fields[name]
Martin Bauer's avatar
Martin Bauer committed
132

Martin Bauer's avatar
Martin Bauer committed
133
134
    def add_custom_data(self, name, cpu_creation_function,
                        gpu_creation_function=None, cpu_to_gpu_transfer_func=None, gpu_to_cpu_transfer_func=None):
135

Martin Bauer's avatar
Martin Bauer committed
136
137
        if cpu_creation_function and gpu_creation_function:
            if cpu_to_gpu_transfer_func is None or gpu_to_cpu_transfer_func is None:
138
                raise ValueError("For GPU data, both transfer functions have to be specified")
Martin Bauer's avatar
Martin Bauer committed
139
            self._custom_data_transfer_functions[name] = (cpu_to_gpu_transfer_func, gpu_to_cpu_transfer_func)
140

Martin Bauer's avatar
Martin Bauer committed
141
142
143
144
        assert name not in self.custom_data_cpu
        if cpu_creation_function:
            assert name not in self.cpu_arrays
            self.custom_data_cpu[name] = cpu_creation_function()
145

Martin Bauer's avatar
Martin Bauer committed
146
147
148
        if gpu_creation_function:
            assert name not in self.gpu_arrays
            self.custom_data_gpu[name] = gpu_creation_function()
149

Martin Bauer's avatar
Martin Bauer committed
150
    def has_data(self, name):
Martin Bauer's avatar
Martin Bauer committed
151
        return name in self.fields
152

Martin Bauer's avatar
Martin Bauer committed
153
154
155
156
157
158
    def add_array_like(self, name, name_of_template_field, latex_name=None, cpu=True, gpu=None):
        return self.add_array(name, latex_name=latex_name, cpu=cpu, gpu=gpu,
                              **self._field_information[name_of_template_field])

    def iterate(self, slice_obj=None, gpu=False, ghost_layers=True, inner_ghost_layers=True):
        if ghost_layers is True:
Martin Bauer's avatar
Martin Bauer committed
159
            ghost_layers = self.default_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
160
161
162
163
164
165
166
167
168
169
170
171
172
        elif ghost_layers is False:
            ghost_layers = 0
        elif isinstance(ghost_layers, str):
            ghost_layers = self.ghost_layers_of_field(ghost_layers)

        if slice_obj is None:
            slice_obj = (slice(None, None, None),) * self.dim
        slice_obj = normalize_slice(slice_obj, tuple(s + 2 * ghost_layers for s in self._domainSize))
        slice_obj = tuple(s if type(s) is slice else slice(s, s + 1, None) for s in slice_obj)

        arrays = self.gpu_arrays if gpu else self.cpu_arrays
        custom_data_dict = self.custom_data_gpu if gpu else self.custom_data_cpu
        iter_dict = custom_data_dict.copy()
173
        for name, arr in arrays.items():
Martin Bauer's avatar
Martin Bauer committed
174
175
            field_gls = self._field_information[name]['ghost_layers']
            if field_gls < ghost_layers:
176
                continue
Martin Bauer's avatar
Martin Bauer committed
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
            arr = remove_ghost_layers(arr, index_dimensions=len(arr.shape) - self.dim,
                                      ghost_layers=field_gls - ghost_layers)
            iter_dict[name] = arr

        offset = tuple(s.start - ghost_layers for s in slice_obj)
        yield SerialBlock(iter_dict, offset, slice_obj)

    def gather_array(self, name, slice_obj=None, ghost_layers=False, **kwargs):
        gl_to_remove = self._field_information[name]['ghost_layers']
        if isinstance(ghost_layers, int):
            gl_to_remove -= ghost_layers
        if ghost_layers is True:
            gl_to_remove = 0
        arr = self.cpu_arrays[name]
        ind_dimensions = self.fields[name].index_dimensions
        spatial_dimensions = self.fields[name].spatial_dimensions

        arr = remove_ghost_layers(arr, index_dimensions=ind_dimensions, ghost_layers=gl_to_remove)

        if slice_obj is not None:
            normalized_slice = normalize_slice(slice_obj[:spatial_dimensions], arr.shape[:spatial_dimensions])
            normalized_slice = tuple(s if type(s) is slice else slice(s, s + 1, None) for s in normalized_slice)
            normalized_slice += slice_obj[spatial_dimensions:]
            arr = arr[normalized_slice]
201
202
203
204
        else:
            arr = arr.view()
        arr.flags.writeable = False
        return arr
Martin Bauer's avatar
Martin Bauer committed
205

206
207
208
    def swap(self, name1, name2, gpu=None):
        if gpu is None:
            gpu = self.default_target == "gpu"
209
210
        arr = self.gpu_arrays if gpu else self.cpu_arrays
        arr[name1], arr[name2] = arr[name2], arr[name1]
Martin Bauer's avatar
Martin Bauer committed
211
212
213
214
215
216
217
218
219

    def all_to_cpu(self):
        for name in (self.cpu_arrays.keys() & self.gpu_arrays.keys()) | self._custom_data_transfer_functions.keys():
            self.to_cpu(name)

    def all_to_gpu(self):
        for name in (self.cpu_arrays.keys() & self.gpu_arrays.keys()) | self._custom_data_transfer_functions.keys():
            self.to_gpu(name)

Martin Bauer's avatar
Martin Bauer committed
220
    def run_kernel(self, kernel_function, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
221
        arrays = self.gpu_arrays if kernel_function.ast.backend == 'gpucuda' else self.cpu_arrays
Martin Bauer's avatar
Martin Bauer committed
222
        kernel_function(**arrays, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
223

224
225
226
227
228
229
    def get_kernel_kwargs(self, kernel_function, **kwargs):
        result = {}
        result.update(self.gpu_arrays if kernel_function.ast.backend == 'gpucuda' else self.cpu_arrays)
        result.update(kwargs)
        return [result]

Martin Bauer's avatar
Martin Bauer committed
230
231
232
233
    def to_cpu(self, name):
        if name in self._custom_data_transfer_functions:
            transfer_func = self._custom_data_transfer_functions[name][1]
            transfer_func(self.custom_data_gpu[name], self.custom_data_cpu[name])
234
        else:
Martin Bauer's avatar
Martin Bauer committed
235
            self.gpu_arrays[name].get(self.cpu_arrays[name])
Martin Bauer's avatar
Martin Bauer committed
236

Martin Bauer's avatar
Martin Bauer committed
237
238
239
240
    def to_gpu(self, name):
        if name in self._custom_data_transfer_functions:
            transfer_func = self._custom_data_transfer_functions[name][0]
            transfer_func(self.custom_data_gpu[name], self.custom_data_cpu[name])
241
        else:
Martin Bauer's avatar
Martin Bauer committed
242
            self.gpu_arrays[name].set(self.cpu_arrays[name])
243

Martin Bauer's avatar
Martin Bauer committed
244
245
    def is_on_gpu(self, name):
        return name in self.gpu_arrays
246

Martin Bauer's avatar
Martin Bauer committed
247
248
    def synchronization_function_cpu(self, names, stencil_name=None, **_):
        return self.synchronization_function(names, stencil_name, 'cpu')
249

Martin Bauer's avatar
Martin Bauer committed
250
251
    def synchronization_function_gpu(self, names, stencil_name=None, **_):
        return self.synchronization_function(names, stencil_name, 'gpu')
252

Martin Bauer's avatar
Martin Bauer committed
253
    def synchronization_function(self, names, stencil=None, target=None, **_):
254
        if target is None:
Martin Bauer's avatar
Martin Bauer committed
255
            target = self.default_target
256
257
258
259
        assert target in ('cpu', 'gpu')
        if not hasattr(names, '__len__') or type(names) is str:
            names = [names]

Martin Bauer's avatar
Martin Bauer committed
260
        filtered_stencil = []
Martin Bauer's avatar
Martin Bauer committed
261
262
        neighbors = [-1, 0, 1]

263
        if (stencil is None and self.dim == 2) or (stencil is not None and stencil.startswith('D2')):
Martin Bauer's avatar
Martin Bauer committed
264
            directions = itertools.product(*[neighbors] * 2)
265
        elif (stencil is None and self.dim == 3) or (stencil is not None and stencil.startswith('D3')):
Martin Bauer's avatar
Martin Bauer committed
266
267
268
269
270
            directions = itertools.product(*[neighbors] * 3)
        else:
            raise ValueError("Invalid stencil")

        for direction in directions:
Martin Bauer's avatar
Martin Bauer committed
271
            use_direction = True
272
            if direction == (0, 0) or direction == (0, 0, 0):
Martin Bauer's avatar
Martin Bauer committed
273
                use_direction = False
274
275
            for component, periodicity in zip(direction, self._periodicity):
                if not periodicity and component != 0:
Martin Bauer's avatar
Martin Bauer committed
276
277
278
                    use_direction = False
            if use_direction:
                filtered_stencil.append(direction)
279

Martin Bauer's avatar
Martin Bauer committed
280
        result = []
281
        for name in names:
Martin Bauer's avatar
Martin Bauer committed
282
            gls = self._field_information[name]['ghost_layers']
283
284
285
286
287
288
289
290
            values_per_cell = self._field_information[name]['values_per_cell']
            if values_per_cell == ():
                values_per_cell = (1, )
            if len(values_per_cell) == 1:
                values_per_cell = values_per_cell[0]
            else:
                raise NotImplementedError("Synchronization of this field is not supported: " + name)

Martin Bauer's avatar
Martin Bauer committed
291
            if len(filtered_stencil) > 0:
292
                if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
293
294
                    from pystencils.slicing import get_periodic_boundary_functor
                    result.append(get_periodic_boundary_functor(filtered_stencil, ghost_layers=gls))
295
                else:
Martin Bauer's avatar
Martin Bauer committed
296
297
298
                    from pystencils.gpucuda.periodicity import get_periodic_boundary_functor as boundary_func
                    result.append(boundary_func(filtered_stencil, self._domainSize,
                                                index_dimensions=self.fields[name].index_dimensions,
299
                                                index_dim_shape=values_per_cell,
Martin Bauer's avatar
Martin Bauer committed
300
301
                                                dtype=self.fields[name].dtype.numpy_dtype,
                                                ghost_layers=gls))
302
303

        if target == 'cpu':
Martin Bauer's avatar
Martin Bauer committed
304
305
306
            def result_functor():
                for arr_name, func in zip(names, result):
                    func(pdfs=self.cpu_arrays[arr_name])
307
        else:
Martin Bauer's avatar
Martin Bauer committed
308
309
310
            def result_functor():
                for arr_name, func in zip(names, result):
                    func(pdfs=self.gpu_arrays[arr_name])
311

Martin Bauer's avatar
Martin Bauer committed
312
        return result_functor
Martin Bauer's avatar
Martin Bauer committed
313

314
    @property
Martin Bauer's avatar
Martin Bauer committed
315
    def array_names(self):
316
317
318
        return tuple(self.fields.keys())

    @property
Martin Bauer's avatar
Martin Bauer committed
319
320
    def custom_data_names(self):
        return tuple(self.custom_data_cpu.keys())
321

Martin Bauer's avatar
Martin Bauer committed
322
    def reduce_float_sequence(self, sequence, operation, all_reduce=False) -> np.array:
Martin Bauer's avatar
Martin Bauer committed
323
324
        return np.array(sequence)

Martin Bauer's avatar
Martin Bauer committed
325
    def reduce_int_sequence(self, sequence, operation, all_reduce=False) -> np.array:
Martin Bauer's avatar
Martin Bauer committed
326
327
        return np.array(sequence)

Martin Bauer's avatar
Martin Bauer committed
328
    def create_vtk_writer(self, file_name, data_names, ghost_layers=False):
329
        from pystencils.datahandling.vtk import image_to_vtk
Martin Bauer's avatar
Martin Bauer committed
330
331

        def writer(step):
Martin Bauer's avatar
Martin Bauer committed
332
333
334
335
            full_file_name = "%s_%08d" % (file_name, step,)
            cell_data = {}
            for name in data_names:
                field = self._get_field_with_given_number_of_ghost_layers(name, ghost_layers)
Martin Bauer's avatar
Martin Bauer committed
336
                if self.dim == 2:
337
                    field = field[:, :, np.newaxis]
Martin Bauer's avatar
Martin Bauer committed
338
                if len(field.shape) == 3:
Martin Bauer's avatar
Martin Bauer committed
339
                    cell_data[name] = np.ascontiguousarray(field)
Martin Bauer's avatar
Martin Bauer committed
340
                elif len(field.shape) == 4:
Martin Bauer's avatar
Martin Bauer committed
341
342
343
                    values_per_cell = field.shape[-1]
                    if values_per_cell == self.dim:
                        field = [np.ascontiguousarray(field[..., i]) for i in range(values_per_cell)]
Martin Bauer's avatar
Martin Bauer committed
344
345
                        if len(field) == 2:
                            field.append(np.zeros_like(field[0]))
Martin Bauer's avatar
Martin Bauer committed
346
                        cell_data[name] = tuple(field)
Martin Bauer's avatar
Martin Bauer committed
347
                    else:
Martin Bauer's avatar
Martin Bauer committed
348
349
                        for i in range(values_per_cell):
                            cell_data["%s[%d]" % (name, i)] = np.ascontiguousarray(field[..., i])
Martin Bauer's avatar
Martin Bauer committed
350
                else:
351
352
                    raise NotImplementedError("VTK export for fields with more than one index "
                                              "coordinate not implemented")
Martin Bauer's avatar
Martin Bauer committed
353
            image_to_vtk(full_file_name, cell_data=cell_data)
Martin Bauer's avatar
Martin Bauer committed
354
355
        return writer

Martin Bauer's avatar
Martin Bauer committed
356
    def create_vtk_writer_for_flag_array(self, file_name, data_name, masks_to_name, ghost_layers=False):
357
        from pystencils.datahandling.vtk import image_to_vtk
Martin Bauer's avatar
Martin Bauer committed
358
359

        def writer(step):
Martin Bauer's avatar
Martin Bauer committed
360
361
            full_file_name = "%s_%08d" % (file_name, step,)
            field = self._get_field_with_given_number_of_ghost_layers(data_name, ghost_layers)
Martin Bauer's avatar
Martin Bauer committed
362
363
            if self.dim == 2:
                field = field[:, :, np.newaxis]
364
            cell_data = {name: np.ascontiguousarray(np.bitwise_and(field, field.dtype.type(mask)) > 0, dtype=np.uint8)
Martin Bauer's avatar
Martin Bauer committed
365
366
                         for mask, name in masks_to_name.items()}
            image_to_vtk(full_file_name, cell_data=cell_data)
Martin Bauer's avatar
Martin Bauer committed
367
368
369

        return writer

Martin Bauer's avatar
Martin Bauer committed
370
371
372
373
    def _get_field_with_given_number_of_ghost_layers(self, name, ghost_layers):
        actual_ghost_layers = self.ghost_layers_of_field(name)
        if ghost_layers is True:
            ghost_layers = actual_ghost_layers
374

Martin Bauer's avatar
Martin Bauer committed
375
        gl_to_remove = actual_ghost_layers - ghost_layers
376
        ind_dims = len(self._field_information[name]['values_per_cell'])
Martin Bauer's avatar
Martin Bauer committed
377
        return remove_ghost_layers(self.cpu_arrays[name], ind_dims, gl_to_remove)
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397

    def log(self, *args, level='INFO'):
        level = level.upper()
        message = " ".join(str(e) for e in args)

        time_running = time.perf_counter() - self._start_time
        spacing = 7 - len(str(int(time_running)))
        message = "[{: <8}]{}({:.3f} sec) {} ".format(level, spacing * '-', time_running, message)
        print(message, flush=True)

    def log_on_root(self, *args, level='INFO'):
        self.log(*args, level=level)

    @property
    def is_root(self):
        return True

    @property
    def world_rank(self):
        return 0
Martin Bauer's avatar
Martin Bauer committed
398
399
400

    def save_all(self, file):
        np.savez_compressed(file, **self.cpu_arrays)
401
402
403
404
405
406
407
408
409

    def load_all(self, file):
        file_contents = np.load(file)
        for arr_name, arr_contents in self.cpu_arrays.items():
            if arr_name not in file_contents:
                print("Skipping read data {} because there is no data with this name in data handling".format(arr_name))
                continue
            if file_contents[arr_name].shape != arr_contents.shape:
                print("Skipping read data {} because shapes don't match. "
410
411
                      "Read array shape {}, existing array shape {}".format(arr_name, file_contents[arr_name].shape,
                                                                            arr_contents.shape))
412
413
                continue
            np.copyto(arr_contents, file_contents[arr_name])