__init__.py 2.55 KB
Newer Older
1
2
3
from typing import Tuple, Union
from .serial_datahandling import SerialDataHandling
from .datahandling_interface import DataHandling
4
5

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


17
18
19
def create_data_handling(domain_size: Tuple[int, ...],
                         periodicity: Union[bool, Tuple[bool, ...]] = False,
                         default_layout: str = 'SoA',
20
                         default_target: str = 'cpu',
21
22
23
24
25
                         parallel: bool = False,
                         default_ghost_layers: int = 1) -> DataHandling:
    """Creates a data handling instance.

    Args:
26
27
28
29
30
31
32
        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'
        default_target: either 'cpu' or 'gpu'
        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'
33
    """
34
    if parallel:
Martin Bauer's avatar
Martin Bauer committed
35
36
        if wlb is None:
            raise ValueError("Cannot create parallel data handling because walberla module is not available")
37
38
39
40
41
42

        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
43
            periodicity = tuple(int(bool(x)) for x in periodicity)
44
45
46
            if len(periodicity) == 2:
                periodicity += (1,)

Martin Bauer's avatar
Martin Bauer committed
47
        if len(domain_size) == 2:
48
            dim = 2
Martin Bauer's avatar
Martin Bauer committed
49
            domain_size = (domain_size[0], domain_size[1], 1)
50
51
52
        else:
            dim = 3

Martin Bauer's avatar
Martin Bauer committed
53
54
        # noinspection PyArgumentList
        block_storage = wlb.createUniformBlockGrid(cells=domain_size, periodic=periodicity)
55
        return ParallelDataHandling(blocks=block_storage, dim=dim, default_target=default_target,
Martin Bauer's avatar
Martin Bauer committed
56
                                    default_layout=default_layout, default_ghost_layers=default_ghost_layers)
57
    else:
58
        return SerialDataHandling(domain_size, periodicity=periodicity, default_target=default_target,
Martin Bauer's avatar
Martin Bauer committed
59
                                  default_layout=default_layout, default_ghost_layers=default_ghost_layers)
60
61
62


__all__ = ['create_data_handling']