flexible_boundaries_prototype.py 10.2 KB
Newer Older
Frederik Hennig's avatar
Frederik Hennig committed
1
2
3
4
5
6
7
8
9
10
11
12
import sympy as sp
import numpy as np

from lbmpy.boundaries.boundaryconditions import Boundary
from lbmpy.boundaries.boundaryhandling import LbmWeightInfo
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor, \
    StreamPushTwoFieldsAccessor, \
    AAEvenTimeStepAccessor, \
    AAOddTimeStepAccessor, \
    EsoTwistEvenTimeStepAccessor, \
    EsoTwistOddTimeStepAccessor

Frederik Hennig's avatar
Frederik Hennig committed
13
from pystencils import Assignment, AssignmentCollection
Frederik Hennig's avatar
Frederik Hennig committed
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
from pystencils.data_types import TypedSymbol, create_type
from pystencils.backends.cbackend import CustomCodeNode
from pystencils import create_indexed_kernel


class AdvancedStreamingBoundaryOffsetInfo(CustomCodeNode):
    # --------------------------- Functions to be used by boundaries --------------------------

    @staticmethod
    def outward_offset_from_dir(dir_idx, dim):
        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
                      for symbol in AdvancedStreamingBoundaryOffsetInfo._outward_offset_symbols(dim)])

    @staticmethod
    def inverse_inward_offset_from_dir(dir_idx, dim):
        return tuple([sp.IndexedBase(symbol, shape=(1,))[dir_idx]
                      for symbol in AdvancedStreamingBoundaryOffsetInfo._inverse_inward_offset_symbols(dim)])

    @staticmethod
    def outward_index_from_dir(dir_idx):
        return sp.IndexedBase(AdvancedStreamingBoundaryOffsetInfo.OUTWARD_INDEX_SYMBOL, shape=(1,))[dir_idx]

    @staticmethod
    def inverse_inward_index_from_dir(dir_idx):
        return sp.IndexedBase(AdvancedStreamingBoundaryOffsetInfo.INV_INWARD_INDEX_SYMBOL, shape=(1,))[dir_idx]

Frederik Hennig's avatar
Frederik Hennig committed
40
41
42
43
44
45
46
    @staticmethod
    def stream_out_pdfs(stencil):
        return sp.symbols('f_out_:{}'.format(len(stencil)))

    @staticmethod
    def stream_in_pdfs(stencil):
        return sp.symbols('f_in_:{}'.format(len(stencil)))
Frederik Hennig's avatar
Frederik Hennig committed
47

Frederik Hennig's avatar
Frederik Hennig committed
48
49
50
51
52
53
54
55
56
57
58
59
    @staticmethod
    def stream_out_pdf_symbolic_dir_indexed(dir_index):
        return sp.IndexedBase('f_out', shape=(1,))[dir_index]

    @staticmethod
    def stream_in_pdf_symbolic_inverse_dir_indexed(dir_index):
        return sp.IndexedBase('f_in', shape=(1,))[dir_index]

    @staticmethod
    def inverse_dir_symbol(dir_symbol):
        symbol_type = dir_symbol.dtype if isinstance(dir_symbol, TypedSymbol) else create_type(np.int64)
        return TypedSymbol("{}_inv".format(dir_symbol.name), symbol_type)
Frederik Hennig's avatar
Frederik Hennig committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

    def __init__(self, pdf_field, stencil, between_timesteps='both', kernel_type='pull'):
        if between_timesteps not in ['both', 'odd_to_even', 'even_to_odd']:
            raise ValueError(
                "Invalid value of parameter 'between_timesteps'. Must be one of 'both', 'odd_to_even' or 'even_to_odd'.",
                between_timesteps)

        if between_timesteps == 'both' and kernel_type in ['aa', 'esotwist']:
            raise ValueError('Cannot create an offset info for both kinds of timesteps for kernel type ' + kernel_type)

        odd_to_even = (between_timesteps == 'odd_to_even')

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

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

        try:
            even_accessor = even_accessors[kernel_type]
            odd_accessor = odd_accessors[kernel_type]
        except KeyError:
            raise ValueError(
                "Invalid value of parameter 'kernel_type'.", kernel_type)

        if between_timesteps == 'both':
            assert even_accessor == odd_accessor

Frederik Hennig's avatar
Frederik Hennig committed
96
97
98
        self.pdf_field = pdf_field
        self.stencil = stencil

Frederik Hennig's avatar
Frederik Hennig committed
99
100
101
102
103
        write_field_accesses = (
            odd_accessor if odd_to_even else even_accessor).write(pdf_field, stencil)
        read_field_accesses = (
            even_accessor if odd_to_even else odd_accessor).read(pdf_field, stencil)

Frederik Hennig's avatar
Frederik Hennig committed
104
105
106
        self.write_field_accesses = write_field_accesses
        self.read_field_accesses = read_field_accesses

Frederik Hennig's avatar
Frederik Hennig committed
107
108
109
110
111
112
113
114
115
116
117
118
        dim = len(stencil[0])
        code = "\n"

        #   Offsets and Indices for the stream-out-populations
        outward_offset_sym = self._outward_offset_symbols(dim)
        for i in range(dim):
            offset_str = ", ".join([str(a.offsets[i])
                                    for a in write_field_accesses])
            code += "const int64_t %s [] = { %s };\n" % (
                outward_offset_sym[i].name, offset_str)

        out_acc_indices = [str(a.index[0]) for a in write_field_accesses]
Frederik Hennig's avatar
Frederik Hennig committed
119
        code += "const int64_t %s [] = { %s };\n" % (
Frederik Hennig's avatar
Frederik Hennig committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
            self.OUTWARD_INDEX_SYMBOL.name, ", ".join(out_acc_indices))

        #   Offsets and Indices for the stream-in-populations
        #   Permute the read_field_accesses for the inverse directions
        inverse_indices = [self._inverse_dir_index(stencil, i)
                           for i in range(len(stencil))]
        inv_read_field_accesses = [read_field_accesses[idx]
                                   for idx in inverse_indices]

        inv_inward_offset_sym = self._inverse_inward_offset_symbols(dim)
        for i in range(dim):
            offset_str = ", ".join([str(a.offsets[i])
                                    for a in inv_read_field_accesses])
            code += "const int64_t %s [] = { %s };\n" % (
                inv_inward_offset_sym[i].name, offset_str)

        inv_in_acc_indices = [str(a.index[0]) for a in inv_read_field_accesses]
Frederik Hennig's avatar
Frederik Hennig committed
137
        code += "const int64_t %s [] = { %s };\n" % (
Frederik Hennig's avatar
Frederik Hennig committed
138
139
140
141
142
143
144
145
            self.INV_INWARD_INDEX_SYMBOL.name, ", ".join(inv_in_acc_indices))

        defined_symbols = set(outward_offset_sym + inv_inward_offset_sym +
                              [self.OUTWARD_INDEX_SYMBOL, self.INV_INWARD_INDEX_SYMBOL])

        super(AdvancedStreamingBoundaryOffsetInfo, self).__init__(
            code, symbols_read=set(), symbols_defined=defined_symbols)

Frederik Hennig's avatar
Frederik Hennig committed
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
    def substitute_pdf_proxies(self, boundary_assignments, dir_symbol):
        if not isinstance(boundary_assignments, AssignmentCollection):
            boundary_assignments = AssignmentCollection(boundary_assignments)

        out_pdfs = self.stream_out_pdfs(self.stencil)
        in_pdfs = self.stream_in_pdfs(self.stencil)
        dir_indexed_out_pdf = self.stream_out_pdf_symbolic_dir_indexed(dir_symbol)
        inv_dir_symbol = self.inverse_dir_symbol(dir_symbol)
        inv_dir_indexed_in_pdf = self.stream_in_pdf_symbolic_inverse_dir_indexed(inv_dir_symbol)

        q = len(self.stencil)
        dim = len(self.stencil[0])

        pdf_subs = dict()

        #   Direct accesses
        for i in range(q):
            pdf_subs[out_pdfs[i]] = self.write_field_accesses[i]
            pdf_subs[in_pdfs[i]] = self.read_field_accesses[i]

        #   Symbolically indexed accesses
        outward_offset = self.outward_offset_from_dir(dir_symbol, dim)
        outward_idx = self.outward_index_from_dir(dir_symbol)

        inv_inward_offset = self.inverse_inward_offset_from_dir(dir_symbol, dim)
        inv_inward_idx = self.inverse_inward_index_from_dir(dir_symbol)

        pdf_out_access_by_dir = self.pdf_field[outward_offset](outward_idx)
        pdf_subs[dir_indexed_out_pdf] = pdf_out_access_by_dir

        pdf_inv_in_access_by_dir = self.pdf_field[inv_inward_offset](inv_inward_idx)
        pdf_subs[inv_dir_indexed_in_pdf] = pdf_inv_in_access_by_dir

        return boundary_assignments.new_with_substitutions(pdf_subs)

Frederik Hennig's avatar
Frederik Hennig committed
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    @staticmethod
    def _inverse_dir_index(stencil, dir):
        return stencil.index(tuple(-d for d in stencil[dir]))

    @staticmethod
    def _outward_offset_symbols(dim):
        return [TypedSymbol(f"OutOffset_{d}", create_type(np.int64)) for d in ['x', 'y', 'z'][:dim]]

    @staticmethod
    def _inverse_inward_offset_symbols(dim):
        return [TypedSymbol(f"InvInOffset_{d}", create_type(np.int64)) for d in ['x', 'y', 'z'][:dim]]

    OUTWARD_INDEX_SYMBOL = TypedSymbol("outIdx", int)
    INV_INWARD_INDEX_SYMBOL = TypedSymbol("invInIdx", int)

# end class AdvancedStreamingBoundaryOffsetInfo


class FlexibleNoSlip(Boundary):

    def __init__(self, name=None):
        """Set an optional name here, to mark boundaries, for example for force evaluations"""
        super(FlexibleNoSlip, self).__init__(name)

    """
        No-Slip, (half-way) simple bounce back boundary condition, enforcing zero velocity at obstacle.
        Extended for use with any streaming pattern.
    """

    def __call__(self, pdf_field, direction_symbol, lb_method, **kwargs):
        outward_offset = AdvancedStreamingBoundaryOffsetInfo.outward_offset_from_dir(direction_symbol, lb_method.dim)
        outward_idx = AdvancedStreamingBoundaryOffsetInfo.outward_index_from_dir(direction_symbol)

Frederik Hennig's avatar
Frederik Hennig committed
214
215
        inv_inward_offset = AdvancedStreamingBoundaryOffsetInfo.inverse_inward_offset_from_dir(
            direction_symbol, lb_method.dim)
Frederik Hennig's avatar
Frederik Hennig committed
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
        inv_inward_idx = AdvancedStreamingBoundaryOffsetInfo.inverse_inward_index_from_dir(direction_symbol)

        return [Assignment(pdf_field[inv_inward_offset](inv_inward_idx), pdf_field[outward_offset](outward_idx))]

    def __hash__(self):
        # All boundaries of these class behave equal -> should also be equal (as long as name is equal)
        return hash(self.name)

    def __eq__(self, other):
        if not isinstance(other, FlexibleNoSlip):
            return False
        return self.name == other.name

# end class FlexibleNoSlip

Frederik Hennig's avatar
Frederik Hennig committed
231

Frederik Hennig's avatar
Frederik Hennig committed
232
233
234
def create_flexible_lbm_boundary_kernel(pdf_field, index_field, lb_method, boundary_functor,
                                        between_timesteps='both', kernel_type='pull',
                                        target='cpu', openmp=True):
Frederik Hennig's avatar
Frederik Hennig committed
235
236

    elements = [AdvancedStreamingBoundaryOffsetInfo(pdf_field, lb_method.stencil, between_timesteps, kernel_type),
Frederik Hennig's avatar
Frederik Hennig committed
237
238
239
240
241
242
243
                LbmWeightInfo(lb_method)]
    index_arr_dtype = index_field.dtype.numpy_dtype
    dir_symbol = TypedSymbol("dir", index_arr_dtype.fields['dir'][0])
    elements += [Assignment(dir_symbol, index_field[0]('dir'))]
    elements += boundary_functor(pdf_field=pdf_field, direction_symbol=dir_symbol,
                                 lb_method=lb_method, index_field=index_field)
    return create_indexed_kernel(elements, [index_field], target=target, cpu_openmp=openmp)