macroscopic_value_kernels.py 12.2 KB
Newer Older
1
import functools
2
from copy import deepcopy
Martin Bauer's avatar
Martin Bauer committed
3

Martin Bauer's avatar
Martin Bauer committed
4
from lbmpy.simplificationfactory import create_simplification_strategy
Martin Bauer's avatar
Martin Bauer committed
5
from pystencils.field import Field, get_layout_of_array
Martin Bauer's avatar
Martin Bauer committed
6

7
8
from lbmpy.fieldaccess import StreamPullTwoFieldsAccessor

Martin Bauer's avatar
Martin Bauer committed
9

10
11
12
13
14
15
16
17
18
19
def pdf_initialization_assignments(lb_method, density, velocity, pdfs):
    """Assignments to initialize the pdf field with equilibrium"""
    cqc = lb_method.conserved_quantity_computation
    inp_eqs = cqc.equilibrium_input_equations_from_init_values(density, velocity)
    setter_eqs = lb_method.get_equilibrium(conserved_quantity_equations=inp_eqs)
    setter_eqs = setter_eqs.new_with_substitutions({sym: pdfs[i]
                                                    for i, sym in enumerate(lb_method.post_collision_pdf_symbols)})
    return setter_eqs


20
21
22
23
24
25
26
27
28
29
30
31
32
def macroscopic_values_getter(lb_method, density, velocity, pdfs):
    cqc = lb_method.conserved_quantity_computation
    assert not (velocity is None and density is None)
    output_spec = {}
    if velocity is not None:
        output_spec['velocity'] = velocity
    if density is not None:
        output_spec['density'] = density
    return cqc.output_equations_from_pdfs(pdfs, output_spec)


macroscopic_values_setter = pdf_initialization_assignments

33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
#   Extensions for any streaming patterns


def flexible_macroscopic_values_setter(lb_method, density, velocity,
                                       pdf_field,
                                       previous_step_accessor=StreamPullTwoFieldsAccessor):
    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
    return pdf_initialization_assignments(lb_method, density, velocity, write_accesses)


def flexible_macroscopic_values_getter(lb_method, density, velocity,
                                       pdf_field,
                                       previous_step_accessor=StreamPullTwoFieldsAccessor):
    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)
    return macroscopic_values_getter(lb_method, density, velocity, write_accesses)
48

49
50
51
52
53

def compile_macroscopic_values_getter(lb_method, output_quantities, pdf_arr=None,
                                      ghost_layers=1, iteration_slice=None,
                                      field_layout='numpy', target='cpu',
                                      previous_step_accessor=StreamPullTwoFieldsAccessor):
Martin Bauer's avatar
Martin Bauer committed
54
    """
55
    Create kernel to compute macroscopic value(s) from a pdf field (e.g. density or velocity)
Martin Bauer's avatar
Martin Bauer committed
56

57
58
59
60
    Args:
        lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
        output_quantities: sequence of quantities to compute e.g. ['density', 'velocity']
        pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
61
62
63
64
        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
                      that should be excluded from the iteration. If None, the number of ghost layers 
                      is determined automatically and assumed to be equal for all dimensions.        
        iteration_slice: if not None, iteration is done only over this slice of the field
65
66
        field_layout: layout for output field, also used for pdf field if pdf_arr is not given
        target: 'cpu' or 'gpu'
67
        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
68
69
70

    Returns:
        a function to compute macroscopic values:
Martin Bauer's avatar
Martin Bauer committed
71
        - pdf_array
Martin Bauer's avatar
Martin Bauer committed
72
        - keyword arguments from name of conserved quantity (as in output_quantities) to numpy field
Martin Bauer's avatar
Martin Bauer committed
73
    """
Martin Bauer's avatar
Martin Bauer committed
74
75
    if not (isinstance(output_quantities, list) or isinstance(output_quantities, tuple)):
        output_quantities = [output_quantities]
76

Martin Bauer's avatar
Martin Bauer committed
77
    cqc = lb_method.conserved_quantity_computation
Martin Bauer's avatar
Martin Bauer committed
78
79
    unknown_quantities = [oq for oq in output_quantities if oq not in cqc.conserved_quantities]
    if unknown_quantities:
Martin Bauer's avatar
Martin Bauer committed
80
        raise ValueError("No such conserved quantity: %s, conserved quantities are %s" %
Martin Bauer's avatar
Martin Bauer committed
81
                         (str(unknown_quantities), str(cqc.conserved_quantities.keys())))
Martin Bauer's avatar
Martin Bauer committed
82

Martin Bauer's avatar
Martin Bauer committed
83
84
    if pdf_arr is None:
        pdf_field = Field.create_generic('pdfs', lb_method.dim, index_dimensions=1, layout=field_layout)
Martin Bauer's avatar
Martin Bauer committed
85
    else:
Martin Bauer's avatar
Martin Bauer committed
86
        pdf_field = Field.create_from_numpy_array('pdfs', pdf_arr, index_dimensions=1)
Martin Bauer's avatar
Martin Bauer committed
87

Martin Bauer's avatar
Martin Bauer committed
88
89
90
91
    output_mapping = {}
    for output_quantity in output_quantities:
        number_of_elements = cqc.conserved_quantities[output_quantity]
        assert number_of_elements >= 1
92

Martin Bauer's avatar
Martin Bauer committed
93
94
95
96
        ind_dims = 0 if number_of_elements <= 1 else 1
        if pdf_arr is None:
            output_field = Field.create_generic(output_quantity, lb_method.dim, layout=field_layout,
                                                index_dimensions=ind_dims)
97
        else:
Martin Bauer's avatar
Martin Bauer committed
98
99
100
101
            output_field_shape = pdf_arr.shape[:-1]
            if ind_dims > 0:
                output_field_shape += (number_of_elements,)
                field_layout = get_layout_of_array(pdf_arr)
102
            else:
Martin Bauer's avatar
Martin Bauer committed
103
104
105
                field_layout = get_layout_of_array(pdf_arr, index_dimension_ids=[len(pdf_field.shape) - 1])
            output_field = Field.create_fixed_size(output_quantity, output_field_shape, ind_dims, pdf_arr.dtype,
                                                   field_layout)
Martin Bauer's avatar
Martin Bauer committed
106

Martin Bauer's avatar
Martin Bauer committed
107
108
109
        output_mapping[output_quantity] = [output_field(i) for i in range(number_of_elements)]
        if len(output_mapping[output_quantity]) == 1:
            output_mapping[output_quantity] = output_mapping[output_quantity][0]
Martin Bauer's avatar
Martin Bauer committed
110

Martin Bauer's avatar
Martin Bauer committed
111
    stencil = lb_method.stencil
112
113
114
115
116
    # pdf_symbols = [pdf_field(i) for i in range(len(stencil))]

    #   Extension for any kind of streaming rule
    pdf_symbols = previous_step_accessor.write(pdf_field, stencil)

Martin Bauer's avatar
Martin Bauer committed
117
    eqs = cqc.output_equations_from_pdfs(pdf_symbols, output_mapping).all_assignments
Martin Bauer's avatar
Martin Bauer committed
118
119

    if target == 'cpu':
120
        import pystencils.cpu as cpu
121
122
        kernel = cpu.make_python_function(cpu.create_kernel(
            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
Martin Bauer's avatar
Martin Bauer committed
123
    elif target == 'gpu':
124
        import pystencils.gpucuda as gpu
125
126
        kernel = gpu.make_python_function(gpu.create_cuda_kernel(
            eqs, ghost_layers=ghost_layers, iteration_slice=iteration_slice))
Martin Bauer's avatar
Martin Bauer committed
127
128
129
    else:
        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))

130
    def getter(pdfs, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
131
132
133
134
        if pdf_arr is not None:
            assert pdfs.shape == pdf_arr.shape and pdfs.strides == pdf_arr.strides, \
                "Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdf_arr.shape)
        if not set(output_quantities).issubset(kwargs.keys()):
135
            raise ValueError("You have to specify the output field for each of the following quantities: %s"
Martin Bauer's avatar
Martin Bauer committed
136
                             % (str(output_quantities),))
137
        kernel(pdfs=pdfs, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
138
139
140
141

    return getter


142
143
144
145
def compile_macroscopic_values_setter(lb_method, quantities_to_set, pdf_arr=None, 
                                      ghost_layers=1, iteration_slice=None,
                                      field_layout='numpy', target='cpu', 
                                      previous_step_accessor=StreamPullTwoFieldsAccessor):
Martin Bauer's avatar
Martin Bauer committed
146
    """
147
148
149
    Creates a function that sets a pdf field to specified macroscopic quantities
    The returned function can be called with the pdf field to set as single argument

150
151
152
153
    Args:
        lb_method: instance of :class:`lbmpy.methods.AbstractLbMethod`
        quantities_to_set: map from conserved quantity name to fixed value or numpy array
        pdf_arr: optional numpy array for pdf field - used to get optimal loop structure for kernel
154
155
156
157
        ghost_layers: a sequence of pairs for each coordinate with lower and upper nr of ghost layers
                      that should be excluded from the iteration. If None, the number of ghost layers 
                      is determined automatically and assumed to be equal for all dimensions.        
        iteration_slice: if not None, iteration is done only over this slice of the field
158
159
        field_layout: layout of the pdf field if pdf_arr was not given
        target: 'cpu' or 'gpu'
160
        previous_step_accessor: The accessor used by the streaming pattern of the previous timestep
161
162
163

    Returns:
        function taking pdf array as single argument and which sets the field to the given values
Martin Bauer's avatar
Martin Bauer committed
164
    """
Martin Bauer's avatar
Martin Bauer committed
165
166
    if pdf_arr is not None:
        pdf_field = Field.create_from_numpy_array('pdfs', pdf_arr, index_dimensions=1)
167
    else:
Martin Bauer's avatar
Martin Bauer committed
168
        pdf_field = Field.create_generic('pdfs', lb_method.dim, index_dimensions=1, layout=field_layout)
Martin Bauer's avatar
Martin Bauer committed
169

Martin Bauer's avatar
Martin Bauer committed
170
    fixed_kernel_parameters = {}
Martin Bauer's avatar
Martin Bauer committed
171
    cqc = lb_method.conserved_quantity_computation
Martin Bauer's avatar
Martin Bauer committed
172

Martin Bauer's avatar
Martin Bauer committed
173
174
175
    value_map = {}
    at_least_one_field_input = False
    for quantity_name, value in quantities_to_set.items():
176
        if hasattr(value, 'shape'):
Martin Bauer's avatar
Martin Bauer committed
177
178
179
            fixed_kernel_parameters[quantity_name] = value
            at_least_one_field_input = True
            num_components = cqc.conserved_quantities[quantity_name]
180
181
            field = Field.create_from_numpy_array(quantity_name, value,
                                                  index_dimensions=0 if num_components <= 1 else 1)
Martin Bauer's avatar
Martin Bauer committed
182
            if num_components == 1:
183
184
                value = field(0)
            else:
Martin Bauer's avatar
Martin Bauer committed
185
                value = [field(i) for i in range(num_components)]
186

Martin Bauer's avatar
Martin Bauer committed
187
        value_map[quantity_name] = value
Martin Bauer's avatar
Martin Bauer committed
188

Martin Bauer's avatar
Martin Bauer committed
189
    cq_equations = cqc.equilibrium_input_equations_from_init_values(**value_map)
Martin Bauer's avatar
Martin Bauer committed
190

Martin Bauer's avatar
Martin Bauer committed
191
192
    eq = lb_method.get_equilibrium(conserved_quantity_equations=cq_equations)
    if at_least_one_field_input:
Martin Bauer's avatar
Martin Bauer committed
193
        simplification = create_simplification_strategy(lb_method)
194
195
        eq = simplification(eq)
    else:
Martin Bauer's avatar
Martin Bauer committed
196
        eq = eq.new_without_subexpressions()
Martin Bauer's avatar
Martin Bauer committed
197

198
199
200
201
    #   Extension for any streaming rule
    write_accesses = previous_step_accessor.write(pdf_field, lb_method.stencil)

    substitutions = {sym: write_accesses[i] for i, sym in enumerate(lb_method.post_collision_pdf_symbols)}
Martin Bauer's avatar
Martin Bauer committed
202
    eq = eq.new_with_substitutions(substitutions).all_assignments
Martin Bauer's avatar
Martin Bauer committed
203

204
205
    if target == 'cpu':
        import pystencils.cpu as cpu
206
207
        kernel = cpu.make_python_function(cpu.create_kernel(eq))
        kernel = functools.partial(kernel, **fixed_kernel_parameters)
208
209
    elif target == 'gpu':
        import pystencils.gpucuda as gpu
210
211
        kernel = gpu.make_python_function(gpu.create_cuda_kernel(eq))
        kernel = functools.partial(kernel, **fixed_kernel_parameters)
212
213
    else:
        raise ValueError("Unknown target '%s'. Possible targets are 'cpu' and 'gpu'" % (target,))
Martin Bauer's avatar
Martin Bauer committed
214

215
    def setter(pdfs, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
216
217
218
        if pdf_arr is not None:
            assert pdfs.shape == pdf_arr.shape and pdfs.strides == pdf_arr.strides, \
                "Pdf array not matching blueprint which was used to compile" + str(pdfs.shape) + str(pdf_arr.shape)
219
        kernel(pdfs=pdfs, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
220

221
    return setter
Martin Bauer's avatar
Martin Bauer committed
222
223


224
225
def create_advanced_velocity_setter_collision_rule(method, velocity_field: Field, velocity_relaxation_rate=0.8):
    """Advanced initialization of velocity field through iteration procedure.
Martin Bauer's avatar
Martin Bauer committed
226

227
    by Mei, Luo, Lallemand and Humieres: Consistent initial conditions for LBM simulations, 2005
228

229
230
231
232
    Args:
        method: lattice boltzmann method object
        velocity_field: pystencils field
        velocity_relaxation_rate: relaxation rate for the velocity moments - determines convergence behaviour
233
                                  of the initialization scheme
234
235
236
237
238

    Returns:
        LB collision rule
    """
    cqc = method.conserved_quantity_computation
Martin Bauer's avatar
Martin Bauer committed
239
240
    density_symbol = cqc.defined_symbols(order=0)[1]
    velocity_symbols = cqc.defined_symbols(order=1)[1]
241
242

    # density is computed from pdfs
243
    eq_input_from_pdfs = cqc.equilibrium_input_equations_from_pdfs(method.pre_collision_pdf_symbols)
Martin Bauer's avatar
Martin Bauer committed
244
    eq_input_from_pdfs = eq_input_from_pdfs.new_filtered([density_symbol])
245
    # velocity is read from input field
246
    vel_symbols = [velocity_field(i) for i in range(method.dim)]
Martin Bauer's avatar
Martin Bauer committed
247
248
    eq_input_from_field = cqc.equilibrium_input_equations_from_init_values(velocity=vel_symbols)
    eq_input_from_field = eq_input_from_field.new_filtered(velocity_symbols)
249
    # then both are merged together
Martin Bauer's avatar
Martin Bauer committed
250
    eq_input = eq_input_from_pdfs.new_merged(eq_input_from_field)
Martin Bauer's avatar
Martin Bauer committed
251

252
    # set first order relaxation rate
253
254
    method = deepcopy(method)
    method.set_first_moment_relaxation_rate(velocity_relaxation_rate)
Martin Bauer's avatar
Martin Bauer committed
255

256
257
    simplification_strategy = create_simplification_strategy(method)
    new_collision_rule = simplification_strategy(method.get_collision_rule(eq_input))
Martin Bauer's avatar
Martin Bauer committed
258

Martin Bauer's avatar
Martin Bauer committed
259
    return new_collision_rule