serial_datahandling.py 19.1 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
import itertools
Martin Bauer's avatar
Martin Bauer committed
2
import time
Martin Bauer's avatar
Martin Bauer committed
3
from typing import Sequence, Union
Martin Bauer's avatar
Martin Bauer committed
4

5
import numpy as np
Martin Bauer's avatar
Martin Bauer committed
6

7
from pystencils.datahandling.blockiteration import SerialBlock
Martin Bauer's avatar
Martin Bauer committed
8
from pystencils.datahandling.datahandling_interface import DataHandling
9
from pystencils.datahandling.pycuda import PyCudaArrayHandler, PyCudaNotAvailableHandler
10
from pystencils.datahandling.pyopencl import PyOpenClArrayHandler
Jan Hönig's avatar
Jan Hönig committed
11
from pystencils.enums import Target
Martin Bauer's avatar
Martin Bauer committed
12
from pystencils.field import (
13
14
    Field, FieldType, create_numpy_array_with_layout, layout_string_to_tuple,
    spatial_layout_string_to_tuple)
Martin Bauer's avatar
Martin Bauer committed
15
from pystencils.slicing import normalize_slice, remove_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
16
17
from pystencils.utils import DotDict

18

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

21
22
23
24
25
    def __init__(self,
                 domain_size: Sequence[int],
                 default_ghost_layers: int = 1,
                 default_layout: str = 'SoA',
                 periodicity: Union[bool, Sequence[bool]] = False,
Jan Hönig's avatar
Jan Hönig committed
26
                 default_target: Target = Target.CPU,
27
                 opencl_queue=None,
28
                 opencl_ctx=None,
29
                 array_handler=None) -> None:
Martin Bauer's avatar
Martin Bauer committed
30
        """
Martin Bauer's avatar
Martin Bauer committed
31
32
33
34
35
36
        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
Jan Hönig's avatar
Jan Hönig committed
37
38
39
            default_target: `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
40
        """
41
        super(SerialDataHandling, self).__init__()
Martin Bauer's avatar
Martin Bauer committed
42
        self._domainSize = tuple(domain_size)
Martin Bauer's avatar
Martin Bauer committed
43
44
        self.default_ghost_layers = default_ghost_layers
        self.default_layout = default_layout
Martin Bauer's avatar
Martin Bauer committed
45
        self._fields = DotDict()
Martin Bauer's avatar
Martin Bauer committed
46
47
48
49
50
        self.cpu_arrays = DotDict()
        self.gpu_arrays = DotDict()
        self.custom_data_cpu = DotDict()
        self.custom_data_gpu = DotDict()
        self._custom_data_transfer_functions = {}
51
52
        self._opencl_queue = opencl_queue
        self._opencl_ctx = opencl_ctx
53

54
        if not array_handler:
55
56
57
            try:
                self.array_handler = PyCudaArrayHandler()
            except Exception:
58
                self.array_handler = PyCudaNotAvailableHandler()
59

Jan Hönig's avatar
Jan Hönig committed
60
            if default_target == Target.OPENCL or opencl_queue:
61
                self.array_handler = PyOpenClArrayHandler(opencl_queue)
62
63
        else:
            self.array_handler = array_handler
64

65
66
67
68
69
70
        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
71
        self._field_information = {}
Markus Holzer's avatar
Markus Holzer committed
72
        self._default_target = default_target
73
        self._start_time = time.perf_counter()
Martin Bauer's avatar
Martin Bauer committed
74

Markus Holzer's avatar
Markus Holzer committed
75
76
77
78
    @property
    def default_target(self):
        return self._default_target

Martin Bauer's avatar
Martin Bauer committed
79
80
81
82
    @property
    def dim(self):
        return len(self._domainSize)

Martin Bauer's avatar
Martin Bauer committed
83
84
85
86
87
88
89
90
    @property
    def shape(self):
        return self._domainSize

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

Martin Bauer's avatar
Martin Bauer committed
91
92
93
94
    @property
    def fields(self):
        return self._fields

Martin Bauer's avatar
Martin Bauer committed
95
96
    def ghost_layers_of_field(self, name):
        return self._field_information[name]['ghost_layers']
97

Martin Bauer's avatar
Martin Bauer committed
98
99
    def values_per_cell(self, name):
        return self._field_information[name]['values_per_cell']
100

Martin Bauer's avatar
Martin Bauer committed
101
    def add_array(self, name, values_per_cell=1, dtype=np.float64, latex_name=None, ghost_layers=None, layout=None,
102
                  cpu=True, gpu=None, alignment=False, field_type=FieldType.GENERIC):
Martin Bauer's avatar
Martin Bauer committed
103
        if ghost_layers is None:
Martin Bauer's avatar
Martin Bauer committed
104
            ghost_layers = self.default_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
105
        if layout is None:
Martin Bauer's avatar
Martin Bauer committed
106
            layout = self.default_layout
107
        if gpu is None:
108
            gpu = self.default_target in self._GPU_LIKE_TARGETS
Martin Bauer's avatar
Martin Bauer committed
109
110

        kwargs = {
Martin Bauer's avatar
Martin Bauer committed
111
            'shape': tuple(s + 2 * ghost_layers for s in self._domainSize),
Martin Bauer's avatar
Martin Bauer committed
112
113
            'dtype': dtype,
        }
114
115

        if not hasattr(values_per_cell, '__len__'):
Jan Hönig's avatar
Jan Hönig committed
116
            values_per_cell = (values_per_cell,)
117
118
119
        if len(values_per_cell) == 1 and values_per_cell[0] == 1:
            values_per_cell = ()

Martin Bauer's avatar
Martin Bauer committed
120
121
122
        self._field_information[name] = {
            'ghost_layers': ghost_layers,
            'values_per_cell': values_per_cell,
Martin Bauer's avatar
Martin Bauer committed
123
124
            'layout': layout,
            'dtype': dtype,
Martin Bauer's avatar
Martin Bauer committed
125
            'alignment': alignment,
126
            'field_type': field_type,
Martin Bauer's avatar
Martin Bauer committed
127
128
        }

129
130
131
132
133
        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
134
        else:
Martin Bauer's avatar
Martin Bauer committed
135
            layout_tuple = spatial_layout_string_to_tuple(layout, self.dim)
136

Martin Bauer's avatar
Martin Bauer committed
137
        # cpu_arr is always created - since there is no create_pycuda_array_with_layout()
Martin Bauer's avatar
Martin Bauer committed
138
139
140
        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)
141

Martin Bauer's avatar
Martin Bauer committed
142
143
144
        if alignment and gpu:
            raise NotImplementedError("Alignment for GPU fields not supported")

Martin Bauer's avatar
Martin Bauer committed
145
        if cpu:
Martin Bauer's avatar
Martin Bauer committed
146
            if name in self.cpu_arrays:
Martin Bauer's avatar
Martin Bauer committed
147
                raise ValueError("CPU Field with this name already exists")
Martin Bauer's avatar
Martin Bauer committed
148
            self.cpu_arrays[name] = cpu_arr
Martin Bauer's avatar
Martin Bauer committed
149
        if gpu:
Martin Bauer's avatar
Martin Bauer committed
150
            if name in self.gpu_arrays:
Martin Bauer's avatar
Martin Bauer committed
151
                raise ValueError("GPU Field with this name already exists")
152
            self.gpu_arrays[name] = self.array_handler.to_gpu(cpu_arr)
Martin Bauer's avatar
Martin Bauer committed
153

Martin Bauer's avatar
Martin Bauer committed
154
        assert all(f.name != name for f in self.fields.values()), "Symbolic field with this name already exists"
Michael Kuron's avatar
Michael Kuron committed
155
        self.fields[name] = Field.create_from_numpy_array(name, cpu_arr, index_dimensions=index_dimensions,
156
                                                          field_type=field_type)
Martin Bauer's avatar
Martin Bauer committed
157
        self.fields[name].latex_name = latex_name
158
        return self.fields[name]
Martin Bauer's avatar
Martin Bauer committed
159

Martin Bauer's avatar
Martin Bauer committed
160
161
    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):
162

Martin Bauer's avatar
Martin Bauer committed
163
164
        if cpu_creation_function and gpu_creation_function:
            if cpu_to_gpu_transfer_func is None or gpu_to_cpu_transfer_func is None:
165
                raise ValueError("For GPU data, both transfer functions have to be specified")
Martin Bauer's avatar
Martin Bauer committed
166
            self._custom_data_transfer_functions[name] = (cpu_to_gpu_transfer_func, gpu_to_cpu_transfer_func)
167

Martin Bauer's avatar
Martin Bauer committed
168
169
170
171
        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()
172

Martin Bauer's avatar
Martin Bauer committed
173
174
175
        if gpu_creation_function:
            assert name not in self.gpu_arrays
            self.custom_data_gpu[name] = gpu_creation_function()
176

Martin Bauer's avatar
Martin Bauer committed
177
    def has_data(self, name):
Martin Bauer's avatar
Martin Bauer committed
178
        return name in self.fields
179

Martin Bauer's avatar
Martin Bauer committed
180
181
182
183
184
185
    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
186
            ghost_layers = self.default_ghost_layers
Martin Bauer's avatar
Martin Bauer committed
187
188
189
190
191
192
193
194
195
196
197
198
199
        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()
200
        for name, arr in arrays.items():
Martin Bauer's avatar
Martin Bauer committed
201
202
            field_gls = self._field_information[name]['ghost_layers']
            if field_gls < ghost_layers:
203
                continue
Martin Bauer's avatar
Martin Bauer committed
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
            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]
228
229
230
231
        else:
            arr = arr.view()
        arr.flags.writeable = False
        return arr
Martin Bauer's avatar
Martin Bauer committed
232

233
234
    def swap(self, name1, name2, gpu=None):
        if gpu is None:
235
            gpu = self.default_target in self._GPU_LIKE_TARGETS
236
237
        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
238
239
240
241
242
243
244
245
246

    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
247
    def run_kernel(self, kernel_function, **kwargs):
248
        arrays = self.gpu_arrays if kernel_function.ast.backend in self._GPU_LIKE_BACKENDS else self.cpu_arrays
249
        kernel_function(**{**arrays, **kwargs})
Martin Bauer's avatar
Martin Bauer committed
250

251
252
    def get_kernel_kwargs(self, kernel_function, **kwargs):
        result = {}
253
        result.update(self.gpu_arrays if kernel_function.ast.backend in self._GPU_LIKE_BACKENDS else self.cpu_arrays)
254
255
256
        result.update(kwargs)
        return [result]

Martin Bauer's avatar
Martin Bauer committed
257
258
259
260
    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])
261
        else:
262
            self.array_handler.download(self.gpu_arrays[name], self.cpu_arrays[name])
Martin Bauer's avatar
Martin Bauer committed
263

Martin Bauer's avatar
Martin Bauer committed
264
265
266
267
    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])
268
        else:
269
            self.array_handler.upload(self.gpu_arrays[name], self.cpu_arrays[name])
270

Martin Bauer's avatar
Martin Bauer committed
271
272
    def is_on_gpu(self, name):
        return name in self.gpu_arrays
273

Martin Bauer's avatar
Martin Bauer committed
274
    def synchronization_function_cpu(self, names, stencil_name=None, **_):
Jan Hönig's avatar
Jan Hönig committed
275
        return self.synchronization_function(names, stencil_name, target=Target.CPU)
276

Martin Bauer's avatar
Martin Bauer committed
277
    def synchronization_function_gpu(self, names, stencil_name=None, **_):
Jan Hönig's avatar
Jan Hönig committed
278
        return self.synchronization_function(names, stencil_name, target=Target.GPU)
279

280
    def synchronization_function(self, names, stencil=None, target=None, functor=None, **_):
281
        if target is None:
Martin Bauer's avatar
Martin Bauer committed
282
            target = self.default_target
Jan Hönig's avatar
Jan Hönig committed
283
284
285
        if target == Target.OPENCL:  # TODO potential misuse between Target and Backend
            target = Target.GPU
        assert target in (Target.CPU, Target.GPU)
286
287
288
        if not hasattr(names, '__len__') or type(names) is str:
            names = [names]

Martin Bauer's avatar
Martin Bauer committed
289
        filtered_stencil = []
Martin Bauer's avatar
Martin Bauer committed
290
291
        neighbors = [-1, 0, 1]

292
        if (stencil is None and self.dim == 2) or (stencil is not None and stencil.startswith('D2')):
Martin Bauer's avatar
Martin Bauer committed
293
            directions = itertools.product(*[neighbors] * 2)
294
        elif (stencil is None and self.dim == 3) or (stencil is not None and stencil.startswith('D3')):
Martin Bauer's avatar
Martin Bauer committed
295
296
297
298
299
            directions = itertools.product(*[neighbors] * 3)
        else:
            raise ValueError("Invalid stencil")

        for direction in directions:
Martin Bauer's avatar
Martin Bauer committed
300
            use_direction = True
301
            if direction == (0, 0) or direction == (0, 0, 0):
Martin Bauer's avatar
Martin Bauer committed
302
                use_direction = False
303
304
            for component, periodicity in zip(direction, self._periodicity):
                if not periodicity and component != 0:
Martin Bauer's avatar
Martin Bauer committed
305
306
307
                    use_direction = False
            if use_direction:
                filtered_stencil.append(direction)
308

Martin Bauer's avatar
Martin Bauer committed
309
        result = []
310
        for name in names:
Martin Bauer's avatar
Martin Bauer committed
311
            gls = self._field_information[name]['ghost_layers']
312
313
            values_per_cell = self._field_information[name]['values_per_cell']
            if values_per_cell == ():
Jan Hönig's avatar
Jan Hönig committed
314
                values_per_cell = (1,)
315
316
317
            if len(values_per_cell) == 1:
                values_per_cell = values_per_cell[0]

Martin Bauer's avatar
Martin Bauer committed
318
            if len(filtered_stencil) > 0:
Jan Hönig's avatar
Jan Hönig committed
319
                if target == Target.CPU:
320
321
322
323
                    if functor is None:
                        from pystencils.slicing import get_periodic_boundary_functor
                        functor = get_periodic_boundary_functor
                    result.append(functor(filtered_stencil, ghost_layers=gls))
324
                else:
325
326
                    if functor is None:
                        from pystencils.gpucuda.periodicity import get_periodic_boundary_functor as functor
Jan Hönig's avatar
Jan Hönig committed
327
328
                        target = Target.GPU if not isinstance(self.array_handler,
                                                              PyOpenClArrayHandler) else Target.OPENCL
329
330
331
332
333
334
335
336
                    result.append(functor(filtered_stencil, self._domainSize,
                                          index_dimensions=self.fields[name].index_dimensions,
                                          index_dim_shape=values_per_cell,
                                          dtype=self.fields[name].dtype.numpy_dtype,
                                          ghost_layers=gls,
                                          target=target,
                                          opencl_queue=self._opencl_queue,
                                          opencl_ctx=self._opencl_ctx))
337

Jan Hönig's avatar
Jan Hönig committed
338
        if target == Target.CPU:
Martin Bauer's avatar
Martin Bauer committed
339
340
341
            def result_functor():
                for arr_name, func in zip(names, result):
                    func(pdfs=self.cpu_arrays[arr_name])
342
        else:
Martin Bauer's avatar
Martin Bauer committed
343
344
345
            def result_functor():
                for arr_name, func in zip(names, result):
                    func(pdfs=self.gpu_arrays[arr_name])
346

Martin Bauer's avatar
Martin Bauer committed
347
        return result_functor
Martin Bauer's avatar
Martin Bauer committed
348

349
    @property
Martin Bauer's avatar
Martin Bauer committed
350
    def array_names(self):
351
352
353
        return tuple(self.fields.keys())

    @property
Martin Bauer's avatar
Martin Bauer committed
354
355
    def custom_data_names(self):
        return tuple(self.custom_data_cpu.keys())
356

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

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

Martin Bauer's avatar
Martin Bauer committed
363
    def create_vtk_writer(self, file_name, data_names, ghost_layers=False):
364
        from pystencils.datahandling.vtk import image_to_vtk
Martin Bauer's avatar
Martin Bauer committed
365
366

        def writer(step):
Martin Bauer's avatar
Martin Bauer committed
367
368
369
370
            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
371
                if self.dim == 2:
372
                    field = field[:, :, np.newaxis]
Martin Bauer's avatar
Martin Bauer committed
373
                if len(field.shape) == 3:
Martin Bauer's avatar
Martin Bauer committed
374
                    cell_data[name] = np.ascontiguousarray(field)
Martin Bauer's avatar
Martin Bauer committed
375
                elif len(field.shape) == 4:
Martin Bauer's avatar
Martin Bauer committed
376
377
378
                    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
379
380
                        if len(field) == 2:
                            field.append(np.zeros_like(field[0]))
Martin Bauer's avatar
Martin Bauer committed
381
                        cell_data[name] = tuple(field)
Martin Bauer's avatar
Martin Bauer committed
382
                    else:
Martin Bauer's avatar
Martin Bauer committed
383
384
                        for i in range(values_per_cell):
                            cell_data["%s[%d]" % (name, i)] = np.ascontiguousarray(field[..., i])
Martin Bauer's avatar
Martin Bauer committed
385
                else:
386
387
                    raise NotImplementedError("VTK export for fields with more than one index "
                                              "coordinate not implemented")
Martin Bauer's avatar
Martin Bauer committed
388
            image_to_vtk(full_file_name, cell_data=cell_data)
Jan Hönig's avatar
Jan Hönig committed
389

Martin Bauer's avatar
Martin Bauer committed
390
391
        return writer

Martin Bauer's avatar
Martin Bauer committed
392
    def create_vtk_writer_for_flag_array(self, file_name, data_name, masks_to_name, ghost_layers=False):
393
        from pystencils.datahandling.vtk import image_to_vtk
Martin Bauer's avatar
Martin Bauer committed
394
395

        def writer(step):
Martin Bauer's avatar
Martin Bauer committed
396
397
            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
398
399
            if self.dim == 2:
                field = field[:, :, np.newaxis]
400
            cell_data = {name: np.ascontiguousarray(np.bitwise_and(field, field.dtype.type(mask)) > 0, dtype=np.uint8)
Martin Bauer's avatar
Martin Bauer committed
401
402
                         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
403
404
405

        return writer

Martin Bauer's avatar
Martin Bauer committed
406
407
408
409
    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
410

Martin Bauer's avatar
Martin Bauer committed
411
        gl_to_remove = actual_ghost_layers - ghost_layers
412
        ind_dims = len(self._field_information[name]['values_per_cell'])
Martin Bauer's avatar
Martin Bauer committed
413
        return remove_ghost_layers(self.cpu_arrays[name], ind_dims, gl_to_remove)
414
415
416
417
418
419
420

    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)))
421
        message = f"[{level: <8}]{spacing * '-'}({time_running:.3f} sec) {message} "
422
423
424
425
426
427
428
429
430
431
432
433
        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
434
435
436

    def save_all(self, file):
        np.savez_compressed(file, **self.cpu_arrays)
437
438

    def load_all(self, file):
439
440
        if '.npz' not in file:
            file += '.npz'
441
442
443
        file_contents = np.load(file)
        for arr_name, arr_contents in self.cpu_arrays.items():
            if arr_name not in file_contents:
444
                print(f"Skipping read data {arr_name} because there is no data with this name in data handling")
445
446
                continue
            if file_contents[arr_name].shape != arr_contents.shape:
447
448
                print(f"Skipping read data {arr_name} because shapes don't match. "
                      f"Read array shape {file_contents[arr_name].shape}, existing array shape {arr_contents.shape}")
449
450
                continue
            np.copyto(arr_contents, file_contents[arr_name])