__init__.py 3.15 KB
Newer Older
Jan Hönig's avatar
Jan Hönig committed
1
2
import warnings

3
from typing import Tuple, Union
Martin Bauer's avatar
Martin Bauer committed
4

5
from .datahandling_interface import DataHandling
Jan Hönig's avatar
Jan Hönig committed
6
from ..enums import Target
Martin Bauer's avatar
Martin Bauer committed
7
from .serial_datahandling import SerialDataHandling
8
9

try:
Martin Bauer's avatar
Martin Bauer committed
10
11
12
    # noinspection PyPep8Naming
    import waLBerla as wlb
    if wlb.cpp_available:
13
        from pystencils.datahandling.parallel_datahandling import ParallelDataHandling
14
    else:
Martin Bauer's avatar
Martin Bauer committed
15
        wlb = None
16
except ImportError:
Martin Bauer's avatar
Martin Bauer committed
17
    wlb = None
18
19
20
    ParallelDataHandling = None


21
22
23
def create_data_handling(domain_size: Tuple[int, ...],
                         periodicity: Union[bool, Tuple[bool, ...]] = False,
                         default_layout: str = 'SoA',
Jan Hönig's avatar
Jan Hönig committed
24
                         default_target: Target = Target.CPU,
25
                         parallel: bool = False,
26
27
                         default_ghost_layers: int = 1,
                         opencl_queue=None) -> DataHandling:
28
29
30
    """Creates a data handling instance.

    Args:
31
32
33
34
        domain_size: size of the rectangular domain
        periodicity: either True, False for full or no periodicity or a tuple of booleans indicating periodicity
                     for each coordinate
        default_layout: default array layout, that is used if not explicitly specified in 'add_array'
Jan Hönig's avatar
Jan Hönig committed
35
        default_target: `Target`
36
37
        parallel: if True a parallel domain is created using walberla - each MPI process gets a part of the domain
        default_ghost_layers: default number of ghost layers if not overwritten in 'add_array'
38
    """
Jan Hönig's avatar
Jan Hönig committed
39
40
41
42
43
44
    if isinstance(default_target, str):
        new_target = Target[default_target.upper()]
        warnings.warn(f'Target "{default_target}" as str is deprecated. Use {new_target} instead',
                      category=DeprecationWarning)
        default_target = new_target

45
    if parallel:
46
        assert not opencl_queue, "OpenCL is only supported for SerialDataHandling"
Martin Bauer's avatar
Martin Bauer committed
47
48
        if wlb is None:
            raise ValueError("Cannot create parallel data handling because walberla module is not available")
49
50
51
52
53
54

        if periodicity is False or periodicity is None:
            periodicity = (0, 0, 0)
        elif periodicity is True:
            periodicity = (1, 1, 1)
        else:
Martin Bauer's avatar
Martin Bauer committed
55
            periodicity = tuple(int(bool(x)) for x in periodicity)
56
57
58
            if len(periodicity) == 2:
                periodicity += (1,)

Martin Bauer's avatar
Martin Bauer committed
59
        if len(domain_size) == 2:
60
            dim = 2
Martin Bauer's avatar
Martin Bauer committed
61
            domain_size = (domain_size[0], domain_size[1], 1)
62
63
64
        else:
            dim = 3

Martin Bauer's avatar
Martin Bauer committed
65
66
        # noinspection PyArgumentList
        block_storage = wlb.createUniformBlockGrid(cells=domain_size, periodic=periodicity)
67
        return ParallelDataHandling(blocks=block_storage, dim=dim, default_target=default_target,
Martin Bauer's avatar
Martin Bauer committed
68
                                    default_layout=default_layout, default_ghost_layers=default_ghost_layers)
69
    else:
70
71
72
73
74
75
        return SerialDataHandling(domain_size,
                                  periodicity=periodicity,
                                  default_target=default_target,
                                  default_layout=default_layout,
                                  default_ghost_layers=default_ghost_layers,
                                  opencl_queue=opencl_queue)
76
77
78


__all__ = ['create_data_handling']