utility.py 3.67 KB
Newer Older
Frederik Hennig's avatar
Frederik Hennig committed
1
2
from lbmpy.fieldaccess import PdfFieldAccessor, \
    StreamPullTwoFieldsAccessor, \
3
4
5
6
7
8
    StreamPushTwoFieldsAccessor, \
    AAEvenTimeStepAccessor, \
    AAOddTimeStepAccessor, \
    EsoTwistEvenTimeStepAccessor, \
    EsoTwistOddTimeStepAccessor

9
import numpy as np
Frederik Hennig's avatar
Frederik Hennig committed
10
import pystencils as ps
Frederik Hennig's avatar
Frederik Hennig committed
11
12
from enum import IntEnum

13

Frederik Hennig's avatar
Frederik Hennig committed
14
15
16
17
18
19
20
21
22
23
24
25
class Timestep(IntEnum):
    EVEN = 0
    ODD = 1
    BOTH = 2

    def next(self):
        return self if self == Timestep.BOTH else Timestep((self + 1) % 2)

    @property
    def idx(self):
        """To use this timestep as an array index"""
        return self % 2
26

27

28
streaming_patterns = ['push', 'pull', 'aa', 'esotwist']
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

even_accessors = {
    'pull': StreamPullTwoFieldsAccessor,
    'push': StreamPushTwoFieldsAccessor,
    'aa': AAEvenTimeStepAccessor,
    'esotwist': EsoTwistEvenTimeStepAccessor
}

odd_accessors = {
    'pull': StreamPullTwoFieldsAccessor,
    'push': StreamPushTwoFieldsAccessor,
    'aa': AAOddTimeStepAccessor,
    'esotwist': EsoTwistOddTimeStepAccessor
}

44

45
def get_accessor(streaming_pattern: str, timestep: Timestep) -> PdfFieldAccessor:
46
    if streaming_pattern not in streaming_patterns:
47
        raise ValueError(
48
            "Invalid value of parameter 'streaming_pattern'.", streaming_pattern)
49

Frederik Hennig's avatar
Frederik Hennig committed
50
    if timestep == Timestep.EVEN:
51
        return even_accessors[streaming_pattern]
52
    else:
53
54
        return odd_accessors[streaming_pattern]

Frederik Hennig's avatar
Frederik Hennig committed
55

56
57
58
59
def is_inplace(streaming_pattern):
    if streaming_pattern not in streaming_patterns:
        raise ValueError('Invalid streaming pattern', streaming_pattern)

Frederik Hennig's avatar
Frederik Hennig committed
60
61
    return streaming_pattern in ['aa', 'esotwist']

62

Frederik Hennig's avatar
Frederik Hennig committed
63
64
def get_timesteps(streaming_pattern):
    return (Timestep.EVEN, Timestep.ODD) if is_inplace(streaming_pattern) else (Timestep.BOTH, )
65

66

Frederik Hennig's avatar
Frederik Hennig committed
67
68
def numeric_offsets(field_access: ps.Field.Access):
    return tuple(int(o) for o in field_access.offsets)
69
70


Frederik Hennig's avatar
Frederik Hennig committed
71
72
73
74
def numeric_index(field_access: ps.Field.Access):
    return tuple(int(i) for i in field_access.index)


Frederik Hennig's avatar
Frederik Hennig committed
75
76
def inverse_dir_index(stencil, direction):
    return stencil.index(tuple(-d for d in stencil[direction]))
77
78


79
class AccessPdfValues:
Frederik Hennig's avatar
Frederik Hennig committed
80
81
82
    """Allows to access values from a PDF array correctly depending on 
    the streaming pattern."""

83
84
85
86
87
88
    def __init__(self, pdf_field, stencil,
                 streaming_pattern='pull', timestep=Timestep.BOTH, streaming_dir='out',
                 accessor=None):
        if streaming_dir not in ['in', 'out']:
            raise ValueError('Invalid streaming direction.', streaming_dir)

89
        if accessor is None:
90
            accessor = get_accessor(streaming_pattern, timestep)
91
92
93
        self.accs = accessor.read(pdf_field, stencil) \
            if streaming_dir == 'in' \
            else accessor.write(pdf_field, stencil)
94
95

    def write_pdf(self, pdf_arr, pos, d, value):
96
        offsets = numeric_offsets(self.accs[d])
97
        pos = tuple(p + o for p, o in zip(pos, offsets))
98
        i = numeric_index(self.accs[d])[0]
99
100
101
        pdf_arr[pos + (i,)] = value

    def read_pdf(self, pdf_arr, pos, d):
102
        offsets = numeric_offsets(self.accs[d])
103
        pos = tuple(p + o for p, o in zip(pos, offsets))
104
        i = numeric_index(self.accs[d])[0]
105
        return pdf_arr[pos + (i,)]
106

107
108
109
110
    def read_multiple(self, pdf_arr, indices):
        """Returns PDF values for a list of index tuples (x, y, [z,] dir)"""
        return np.array([self.read_pdf(pdf_arr, idx[:-1], idx[-1]) for idx in indices])

111
    def collect_from_index_list(self, pdf_arr, index_list):
112
113
114
115
116
        """To collect PDF values according to an pystencils boundary handling index list"""
        def to_index_tuple(idx):
            return tuple(idx[v] for v in ('x', 'y', 'z')[:len(idx) - 1] + ('dir',))

        return self.read_multiple(pdf_arr, (to_index_tuple(idx) for idx in index_list))