creationfunctions.py 29.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
r"""
Creating LBM kernels
====================

The following list describes common parameters for the creation functions. They have to be passed as keyword parameters.

Method parameters
^^^^^^^^^^^^^^^^^

General:

Martin Bauer's avatar
Martin Bauer committed
12
- ``stencil='D2Q9'``: stencil name e.g. 'D2Q9', 'D3Q19'. See :func:`pystencils.stencils.get_stencil` for details
13
- ``method='srt'``: name of lattice Boltzmann method. This determines the selection and relaxation pattern of
Martin Bauer's avatar
Martin Bauer committed
14
  moments/cumulants, i.e. which moment/cumulant basis is chosen, and which of the basis vectors are relaxed together
Martin Bauer's avatar
Martin Bauer committed
15

Martin Bauer's avatar
Martin Bauer committed
16
    - ``srt``: single relaxation time (:func:`lbmpy.methods.create_srt`)
17
18
    - ``trt``: two relaxation time, first relaxation rate is for even moments and determines the viscosity (as in SRT),
      the second relaxation rate is used for relaxing odd moments, and controls the bulk viscosity.
Martin Bauer's avatar
Martin Bauer committed
19
      (:func:`lbmpy.methods.create_trt`)
Martin Bauer's avatar
Martin Bauer committed
20
21
22
23
24
25
26
27
28
29
    - ``mrt``: orthogonal multi relaxation time model, relaxation rates are used in this order for :
      shear modes, bulk modes, 3rd order modes, 4th order modes, etc.
      Requires also a parameter 'weighted' that should be True if the moments should be orthogonal w.r.t. weighted
      scalar product using the lattice weights. If `False` the normal scalar product is used.
      For custom definition of the method, a 'nested_moments' can be passed.
      For example: [ [1, x, y], [x*y, x**2, y**2], ... ] that groups all moments together that should be relaxed with
      the same rate. Literature values of this list can be obtained through
      :func:`lbmpy.methods.creationfunctions.mrt_orthogonal_modes_literature`.
      See also :func:`lbmpy.methods.create_mrt_orthogonal`
    - ``mrt3``: deprecated
30
31
    - ``mrt_raw``: non-orthogonal MRT where all relaxation rates can be specified independently i.e. there are as many
      relaxation rates as stencil entries. Look at the generated method in Jupyter to see which moment<->relaxation rate
Martin Bauer's avatar
Martin Bauer committed
32
      mapping (:func:`lbmpy.methods.create_mrt_raw`)
33
    - ``trt-kbc-n<N>`` where <N> is 1,2,3 or 4. Special two-relaxation rate method. This is not the entropic method
34
      yet, only the relaxation pattern. To get the entropic method, see parameters below!
Martin Bauer's avatar
Martin Bauer committed
35
      (:func:`lbmpy.methods.create_trt_kbc`)
Martin Bauer's avatar
Martin Bauer committed
36

Martin Bauer's avatar
Martin Bauer committed
37
- ``relaxation_rates``: sequence of relaxation rates, number depends on selected method. If you specify more rates than
38
  method needs, the additional rates are ignored. For SRT and TRT models it is possible ot define a single
Martin Bauer's avatar
Martin Bauer committed
39
  ``relaxation_rate`` instead of a list, the second rate for TRT is then determined via magic number.
40
41
42
- ``compressible=False``: affects the selection of equilibrium moments. Both options approximate the *incompressible*
  Navier Stokes Equations. However when chosen as False, the approximation is better, the standard LBM derivation is
  compressible.
Martin Bauer's avatar
Martin Bauer committed
43
- ``equilibrium_order=2``: order in velocity, at which the equilibrium moment/cumulant approximation is
44
  truncated. Order 2 is sufficient to approximate Navier-Stokes
Martin Bauer's avatar
Martin Bauer committed
45
- ``force_model=None``: possible values: ``None``, ``'simple'``, ``'luo'``, ``'guo'`` ``'buick'``, or an instance of a
46
  class implementing the same methods as the classes in :mod:`lbmpy.forcemodels`. For details, see
47
  :mod:`lbmpy.forcemodels`
48
- ``force=(0,0,0)``: either constant force or a symbolic expression depending on field value
Martin Bauer's avatar
Martin Bauer committed
49
- ``maxwellian_moments=True``: way to compute equilibrium moments/cumulants, if False the standard
50
51
52
  discretized LBM equilibrium is used, otherwise the equilibrium moments are computed from the continuous Maxwellian.
  This makes only a difference if sparse stencils are used e.g. D2Q9 and D3Q27 are not affected, D319 and DQ15 are
- ``cumulant=False``: use cumulants instead of moments
Martin Bauer's avatar
Martin Bauer committed
53
- ``initial_velocity=None``: initial velocity in domain, can either be a tuple (x,y,z) velocity to set a constant
54
55
  velocity everywhere, or a numpy array with the same size of the domain, with a last coordinate of shape dim to set
  velocities on cell level
56
- ``output={}``: a dictionary mapping macroscopic quantites e.g. the strings 'density' and 'velocity' to pystencils
57
  fields. In each timestep the corresponding quantities are written to the given fields.
Martin Bauer's avatar
Martin Bauer committed
58
- ``velocity_input``: symbolic field where the velocities are read from (for advection diffusion LBM)
Martin Bauer's avatar
Martin Bauer committed
59
- ``density_input``: symbolic field or field access where to read density from. When passing this parameter,
Martin Bauer's avatar
Martin Bauer committed
60
  ``velocity_input`` has to be passed as well
Markus Holzer's avatar
Markus Holzer committed
61
62
- ``kernel_type``: supported values: 'stream_pull_collide' (default), 'collide_only', stream_pull_only,
   collide_stream_push, esotwist_even, esotwist_odd, aa_even, aa_odd
63
64
65
66
67
68

Entropic methods:

- ``entropic=False``: In case there are two distinct relaxation rate in a method, one of them (usually the one, not
  determining the viscosity) can be automatically chosen w.r.t an entropy condition. For details see
  :mod:`lbmpy.methods.entropic`
Martin Bauer's avatar
Martin Bauer committed
69
- ``entropic_newton_iterations=None``: For moment methods the entropy optimum can be calculated in closed form.
70
71
  For cumulant methods this is not possible, in that case it is computed using Newton iterations. This parameter can be
  used to force Newton iterations and specify how many should be done
Martin Bauer's avatar
Martin Bauer committed
72
73
- ``omega_output_field=None``: you can pass a pystencils Field here, where the calculated free relaxation rate of
  an entropic or Smagorinsky method is written to
74

75
76
LES methods:

Martin Bauer's avatar
Martin Bauer committed
77
78
- ``smagorinsky=False``: set to Smagorinsky constant to activate turbulence model, ``omega_output_field`` can be set to
  write out adapted relaxation rates
79

80
81
Fluctuating LB:

Martin Bauer's avatar
Martin Bauer committed
82
83
- ``fluctuating``: enables fluctuating lattice Boltzmann by randomizing collision process.
  Pass dictionary with parameters to  ``lbmpy.fluctuatinglb.add_fluctuations_to_collision_rule``
84

85
86
87
88
89
90

Optimization Parameters
^^^^^^^^^^^^^^^^^^^^^^^

Simplifications / Transformations:

Martin Bauer's avatar
Martin Bauer committed
91
92
- ``cse_pdfs=False``: run common subexpression elimination for opposing stencil directions
- ``cse_global=False``: run common subexpression elimination after all other simplifications have been executed
93
94
- ``split=False``: split innermost loop, to handle only 2 directions per loop. This reduces the number of parallel
  load/store streams and thus speeds up the kernel on most architectures
Martin Bauer's avatar
Martin Bauer committed
95
- ``builtin_periodicity=(False,False,False)``: instead of handling periodicity by copying ghost layers, the periodicity
96
97
  is built into the kernel. This parameters specifies if the domain is periodic in (x,y,z) direction. Even if the
  periodicity is built into the kernel, the fields have one ghost layer to be consistent with other functions. 
98
    
99
100
101

Field size information:

102
103
- ``pdf_arr=None``: pass a numpy array here to create kernels with fixed size and create the loop nest according
  to layout of this array
Martin Bauer's avatar
Martin Bauer committed
104
- ``field_size=None``: create kernel for fixed field size
Martin Bauer's avatar
Martin Bauer committed
105
106
- ``field_layout='c'``:   ``'c'`` or ``'numpy'`` for standard numpy layout, ``'reverse_numpy'`` or ``'f'`` for fortran
  layout, this does not apply when pdf_arr was given, then the same layout as pdf_arr is used
107
108
- ``symbolic_field``: pystencils field for source (pdf field that is read)
- ``symbolic_temporary_field``: pystencils field for temporary (pdf field that is written in stream, or stream-collide)
109

110
111
112
113
114
115
116
117
118
119
120

CPU:

- ``openmp=True``: Can be a boolean to turn multi threading on/off, or an integer
  specifying the number of threads. If True is specified OpenMP chooses the number of threads
- ``vectorization=False``: controls manual vectorization using SIMD instrinsics. If True default vectorization settings
  are use. Alternatively a dictionary with parameters for vectorize function can be passed. For example
  ``{'instruction_set': 'avx', 'assume_aligned': True, 'nontemporal': True}``. Nontemporal stores are only used if
  assume_aligned is also activated.


121
122
123
124
GPU:

- ``target='cpu'``: ``'cpu'`` or ``'gpu'``, last option requires a CUDA enabled graphics card
  and installed *pycuda* package
Martin Bauer's avatar
Martin Bauer committed
125
126
- ``gpu_indexing='block'``: determines mapping of CUDA threads to cells. Can be either ``'block'`` or ``'line'``
- ``gpu_indexing_params='block'``: parameters passed to init function of gpu indexing.
Martin Bauer's avatar
Martin Bauer committed
127
  For ``'block'`` indexing one can e.g. specify the block size ``{'block_size' : (128, 4, 1)}``
128

129

130
131
Other:

Martin Bauer's avatar
Martin Bauer committed
132
- ``openmp=True``: only applicable for cpu simulations. Can be a boolean to turn multi threading on/off, or an integer
133
  specifying the number of threads. If True is specified OpenMP chooses the number of threads
Martin Bauer's avatar
Martin Bauer committed
134
- ``double_precision=True``:  by default simulations run with double precision floating point numbers, by setting this
135
  parameter to False, single precision is used, which is much faster, especially on GPUs
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164




Terminology and creation pipeline
---------------------------------

Kernel functions are created in three steps:

1. *Method*:
         the method defines the collision process. Currently there are two big categories:
         moment and cumulant based methods. A method defines how each moment or cumulant is relaxed by
         storing the equilibrium value and the relaxation rate for each moment/cumulant.
2. *Collision/Update Rule*:
         Methods can generate a "collision rule" which is an equation collection that define the
         post collision values as a function of the pre-collision values. On these equation collection
         simplifications are applied to reduce the number of floating point operations.
         At this stage an entropic optimization step can also be added to determine one relaxation rate by an
         entropy condition.
         Then a streaming rule is added which transforms the collision rule into an update rule.
         The streaming step depends on the pdf storage (source/destination, AABB pattern, EsoTwist).
         Currently only the simple source/destination  pattern is supported.
3. *AST*:
        The abstract syntax tree describes the structure of the kernel, including loops and conditionals.
        The ast can be modified e.g. to add OpenMP pragmas, reorder loops or apply other optimizations.
4. *Function*:
        This step compiles the AST into an executable function, either for CPU or GPUs. This function
        behaves like a normal Python function and runs one LBM time step.

Martin Bauer's avatar
Martin Bauer committed
165
The function :func:`create_lb_function` runs the whole pipeline, the other functions in this module
166
167
168
169
execute this pipeline only up to a certain step. Each function optionally also takes the result of the previous step.

For example, to modify the AST one can run::

Martin Bauer's avatar
Martin Bauer committed
170
    ast = create_lb_ast(...)
171
    # modify ast here
Martin Bauer's avatar
Martin Bauer committed
172
    func = create_lb_function(ast=ast, ...)
173

174
175
"""
from copy import copy
176

Martin Bauer's avatar
Martin Bauer committed
177
178
179
180
import sympy as sp

import lbmpy.forcemodels as forcemodels
from lbmpy.fieldaccess import (
Martin Bauer's avatar
Martin Bauer committed
181
182
183
    AAEvenTimeStepAccessor, AAOddTimeStepAccessor, CollideOnlyInplaceAccessor,
    EsoTwistEvenTimeStepAccessor, EsoTwistOddTimeStepAccessor, PdfFieldAccessor,
    PeriodicTwoFieldsAccessor, StreamPullTwoFieldsAccessor, StreamPushTwoFieldsAccessor)
184
from lbmpy.fluctuatinglb import add_fluctuations_to_collision_rule
Martin Bauer's avatar
Martin Bauer committed
185
186
from lbmpy.methods import (
    create_mrt3, create_mrt_orthogonal, create_mrt_raw, create_srt, create_trt, create_trt_kbc)
Martin Bauer's avatar
Martin Bauer committed
187
188
from lbmpy.methods.creationfunctions import create_generic_mrt
from lbmpy.methods.cumulantbased import CumulantBasedLbMethod
Martin Bauer's avatar
Martin Bauer committed
189
from lbmpy.methods.entropic import add_entropy_condition, add_iterative_entropy_condition
Martin Bauer's avatar
Martin Bauer committed
190
from lbmpy.methods.entropic_eq_srt import create_srt_entropic
Martin Bauer's avatar
Martin Bauer committed
191
from lbmpy.moments import get_order
Martin Bauer's avatar
Martin Bauer committed
192
from lbmpy.relaxationrates import relaxation_rate_from_magic_number
Martin Bauer's avatar
Martin Bauer committed
193
from lbmpy.simplificationfactory import create_simplification_strategy
194
from lbmpy.stencils import get_stencil
Martin Bauer's avatar
Martin Bauer committed
195
196
197
198
199
200
from lbmpy.turbulence_models import add_smagorinsky_model
from lbmpy.updatekernels import create_lbm_kernel, create_stream_pull_with_output_kernel
from pystencils import Assignment, AssignmentCollection, create_kernel
from pystencils.cache import disk_cache_no_fallback
from pystencils.data_types import collate_types
from pystencils.field import Field, get_layout_of_array
Martin Bauer's avatar
Martin Bauer committed
201
from pystencils.simp import sympy_cse
202
from pystencils.stencil import have_same_entries
203
204


Martin Bauer's avatar
Martin Bauer committed
205
def create_lb_function(ast=None, optimization={}, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
206
    """Creates a Python function for the LB method"""
Martin Bauer's avatar
Martin Bauer committed
207
    params, opt_params = update_with_default_parameters(kwargs, optimization)
208

209
    if ast is None:
Martin Bauer's avatar
Martin Bauer committed
210
        params['optimization'] = opt_params
Martin Bauer's avatar
Martin Bauer committed
211
        ast = create_lb_ast(**params)
212

Martin Bauer's avatar
Martin Bauer committed
213
    res = ast.compile()
214

215
    res.method = ast.method
Martin Bauer's avatar
Martin Bauer committed
216
    res.update_rule = ast.update_rule
Martin Bauer's avatar
Martin Bauer committed
217
    res.ast = ast
218
219
    return res

220

Martin Bauer's avatar
Martin Bauer committed
221
def create_lb_ast(update_rule=None, optimization={}, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
222
    """Creates a pystencils AST for the LB method"""
Martin Bauer's avatar
Martin Bauer committed
223
    params, opt_params = update_with_default_parameters(kwargs, optimization)
224

Martin Bauer's avatar
Martin Bauer committed
225
226
227
    if update_rule is None:
        params['optimization'] = optimization
        update_rule = create_lb_update_rule(**params)
228

Martin Bauer's avatar
Martin Bauer committed
229
230
231
232
233
    field_types = set(fa.field.dtype for fa in update_rule.defined_symbols if isinstance(fa, Field.Access))
    res = create_kernel(update_rule, target=opt_params['target'], data_type=collate_types(field_types),
                        cpu_openmp=opt_params['openmp'], cpu_vectorize_info=opt_params['vectorization'],
                        gpu_indexing=opt_params['gpu_indexing'], gpu_indexing_params=opt_params['gpu_indexing_params'],
                        ghost_layers=1)
234

Martin Bauer's avatar
Martin Bauer committed
235
236
    res.method = update_rule.method
    res.update_rule = update_rule
237
    return res
238

239

Martin Bauer's avatar
Martin Bauer committed
240
241
@disk_cache_no_fallback
def create_lb_update_rule(collision_rule=None, optimization={}, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
242
    """Creates an update rule (list of Assignments) for a LB method that describe a full sweep"""
Martin Bauer's avatar
Martin Bauer committed
243
    params, opt_params = update_with_default_parameters(kwargs, optimization)
244

Martin Bauer's avatar
Martin Bauer committed
245
    if collision_rule is None:
Martin Bauer's avatar
Martin Bauer committed
246
        collision_rule = create_lb_collision_rule(**params, optimization=opt_params)
247

Martin Bauer's avatar
Martin Bauer committed
248
    lb_method = collision_rule.method
249

Martin Bauer's avatar
Martin Bauer committed
250
    field_data_type = 'float64' if opt_params['double_precision'] else 'float32'
251
    q = len(collision_rule.method.stencil)
Martin Bauer's avatar
Martin Bauer committed
252

Martin Bauer's avatar
Martin Bauer committed
253
254
255
    if opt_params['symbolic_field'] is not None:
        src_field = opt_params['symbolic_field']
    elif opt_params['field_size']:
256
        field_size = [s + 2 for s in opt_params['field_size']] + [q]
Martin Bauer's avatar
Martin Bauer committed
257
258
        src_field = Field.create_fixed_size(params['field_name'], field_size, index_dimensions=1,
                                            layout=opt_params['field_layout'], dtype=field_data_type)
259
    else:
Martin Bauer's avatar
Martin Bauer committed
260
        src_field = Field.create_generic(params['field_name'], spatial_dimensions=collision_rule.method.dim,
261
                                         index_shape=(q,), layout=opt_params['field_layout'], dtype=field_data_type)
262

263
264
265
266
    if opt_params['symbolic_temporary_field'] is not None:
        dst_field = opt_params['symbolic_temporary_field']
    else:
        dst_field = src_field.new_field_with_different_name(params['temporary_field_name'])
267

268
269
270
271
272
    kernel_type = params['kernel_type']
    if isinstance(kernel_type, PdfFieldAccessor):
        accessor = kernel_type
        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
    elif params['kernel_type'] == 'stream_pull_collide':
273
        accessor = StreamPullTwoFieldsAccessor
Martin Bauer's avatar
Martin Bauer committed
274
275
276
277
278
        if any(opt_params['builtin_periodicity']):
            accessor = PeriodicTwoFieldsAccessor(opt_params['builtin_periodicity'], ghost_layers=1)
        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
    elif params['kernel_type'] == 'stream_pull_only':
        return create_stream_pull_with_output_kernel(lb_method, src_field, dst_field, params['output'])
279
    else:
280
        kernel_type_to_accessor = {
281
            'collide_only': CollideOnlyInplaceAccessor,
282
283
284
285
286
287
288
289
290
291
292
            'collide_stream_push': StreamPushTwoFieldsAccessor,
            'esotwist_even': EsoTwistEvenTimeStepAccessor,
            'esotwist_odd': EsoTwistOddTimeStepAccessor,
            'aa_even': AAEvenTimeStepAccessor,
            'aa_odd': AAOddTimeStepAccessor,
        }
        try:
            accessor = kernel_type_to_accessor[kernel_type]()
        except KeyError:
            raise ValueError("Invalid value of parameter 'kernel_type'", params['kernel_type'])
        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
293
294


Martin Bauer's avatar
Martin Bauer committed
295
@disk_cache_no_fallback
Martin Bauer's avatar
Martin Bauer committed
296
def create_lb_collision_rule(lb_method=None, optimization={}, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
297
    """Creates a collision rule (list of Assignments) for a LB method describing the collision operator (no stream)"""
Martin Bauer's avatar
Martin Bauer committed
298
    params, opt_params = update_with_default_parameters(kwargs, optimization)
299

Martin Bauer's avatar
Martin Bauer committed
300
301
    if lb_method is None:
        lb_method = create_lb_method(**params)
302

Martin Bauer's avatar
Martin Bauer committed
303
304
    split_inner_loop = 'split' in opt_params and opt_params['split']
    cqc = lb_method.conserved_quantity_computation
305

Martin Bauer's avatar
Martin Bauer committed
306
307
308
309
310
311
312
313
    rho_in = params['density_input']
    u_in = params['velocity_input']

    if u_in is not None and isinstance(u_in, Field):
        u_in = u_in.center_vector
    if rho_in is not None and isinstance(rho_in, Field):
        rho_in = rho_in.center

Martin Bauer's avatar
Martin Bauer committed
314
    keep_rrs_symbolic = opt_params['keep_rrs_symbolic']
Martin Bauer's avatar
Martin Bauer committed
315
316
317
318
    if u_in is not None:
        density_rhs = sum(lb_method.pre_collision_pdf_symbols) if rho_in is None else rho_in
        eqs = [Assignment(cqc.zeroth_order_moment_symbol, density_rhs)]
        eqs += [Assignment(u_sym, u_in[i]) for i, u_sym in enumerate(cqc.first_order_moment_symbols)]
319
        eqs = AssignmentCollection(eqs, [])
Martin Bauer's avatar
Martin Bauer committed
320
321
        collision_rule = lb_method.get_collision_rule(conserved_quantity_equations=eqs,
                                                      keep_rrs_symbolic=keep_rrs_symbolic)
Martin Bauer's avatar
Martin Bauer committed
322
323
    elif u_in is None and rho_in is not None:
        raise ValueError("When setting 'density_input' parameter, 'velocity_input' has to be specified as well.")
324
    else:
Martin Bauer's avatar
Martin Bauer committed
325
        collision_rule = lb_method.get_collision_rule(keep_rrs_symbolic=keep_rrs_symbolic)
326

Martin Bauer's avatar
Martin Bauer committed
327
328
329
330
    if opt_params['simplification'] == 'auto':
        simplification = create_simplification_strategy(lb_method, split_inner_loop=split_inner_loop)
    else:
        simplification = opt_params['simplification']
331
332
    collision_rule = simplification(collision_rule)

333
    if params['fluctuating']:
Martin Bauer's avatar
Martin Bauer committed
334
        add_fluctuations_to_collision_rule(collision_rule, **params['fluctuating'])
335

336
    if params['entropic']:
337
338
        if params['smagorinsky']:
            raise ValueError("Choose either entropic or smagorinsky")
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
        if params['entropic_newton_iterations']:
            if isinstance(params['entropic_newton_iterations'], bool):
                iterations = 3
            else:
                iterations = params['entropic_newton_iterations']
            collision_rule = add_iterative_entropy_condition(collision_rule, newton_iterations=iterations,
                                                             omega_output_field=params['omega_output_field'])
        else:
            collision_rule = add_entropy_condition(collision_rule, omega_output_field=params['omega_output_field'])
    elif params['smagorinsky']:
        smagorinsky_constant = 0.12 if params['smagorinsky'] is True else params['smagorinsky']
        collision_rule = add_smagorinsky_model(collision_rule, smagorinsky_constant,
                                               omega_output_field=params['omega_output_field'])
        if 'split_groups' in collision_rule.simplification_hints:
            collision_rule.simplification_hints['split_groups'][0].append(sp.Symbol("smagorinsky_omega"))

355
356
357
358
359
360
361
362
    cse_pdfs = False if 'cse_pdfs' not in opt_params else opt_params['cse_pdfs']
    cse_global = False if 'cse_global' not in opt_params else opt_params['cse_global']
    if cse_pdfs:
        from lbmpy.methods.momentbasedsimplifications import cse_in_opposing_directions
        collision_rule = cse_in_opposing_directions(collision_rule)
    if cse_global:
        collision_rule = sympy_cse(collision_rule)

363
364
365
366
367
    if params['output'] and params['kernel_type'] == 'stream_pull_collide':
        cqc = lb_method.conserved_quantity_computation
        output_eqs = cqc.output_equations_from_pdfs(lb_method.pre_collision_pdf_symbols, params['output'])
        collision_rule = collision_rule.new_merged(output_eqs)

368
    return collision_rule
369
370


Martin Bauer's avatar
Martin Bauer committed
371
def create_lb_method(**params):
Martin Bauer's avatar
Martin Bauer committed
372
    """Creates a LB method, defined by moments/cumulants for collision space, equilibrium and relaxation rates."""
Martin Bauer's avatar
Martin Bauer committed
373
    params, _ = update_with_default_parameters(params, {}, fail_on_unknown_parameter=False)
374

375
    if isinstance(params['stencil'], tuple) or isinstance(params['stencil'], list):
Martin Bauer's avatar
Martin Bauer committed
376
        stencil_entries = params['stencil']
377
    else:
Martin Bauer's avatar
Martin Bauer committed
378
        stencil_entries = get_stencil(params['stencil'])
379

Martin Bauer's avatar
Martin Bauer committed
380
    dim = len(stencil_entries[0])
381

382
383
384
    if isinstance(params['force'], Field):
        params['force'] = tuple(params['force'](i) for i in range(dim))

Martin Bauer's avatar
Martin Bauer committed
385
    force_is_zero = True
386
387
    for f_i in params['force']:
        if f_i != 0:
Martin Bauer's avatar
Martin Bauer committed
388
            force_is_zero = False
389

Martin Bauer's avatar
Martin Bauer committed
390
391
392
    no_force_model = 'force_model' not in params or params['force_model'] == 'none' or params['force_model'] is None
    if not force_is_zero and no_force_model:
        params['force_model'] = 'guo'
393

Martin Bauer's avatar
Martin Bauer committed
394
395
    if 'force_model' in params:
        force_model = force_model_from_string(params['force_model'], params['force'][:dim])
396
    else:
Martin Bauer's avatar
Martin Bauer committed
397
        force_model = None
398

Martin Bauer's avatar
Martin Bauer committed
399
    common_params = {
400
        'compressible': params['compressible'],
Martin Bauer's avatar
Martin Bauer committed
401
402
403
        'equilibrium_order': params['equilibrium_order'],
        'force_model': force_model,
        'maxwellian_moments': params['maxwellian_moments'],
Martin Bauer's avatar
Martin Bauer committed
404
        'cumulant': params['cumulant'],
405
        'c_s_sq': params['c_s_sq'],
406
    }
Martin Bauer's avatar
Martin Bauer committed
407
408
409
410
411
412
413
414
415
416
417
418
    method_name = params['method']
    relaxation_rates = params['relaxation_rates']

    if method_name.lower() == 'srt':
        assert len(relaxation_rates) >= 1, "Not enough relaxation rates"
        method = create_srt(stencil_entries, relaxation_rates[0], **common_params)
    elif method_name.lower() == 'trt':
        assert len(relaxation_rates) >= 2, "Not enough relaxation rates"
        method = create_trt(stencil_entries, relaxation_rates[0], relaxation_rates[1], **common_params)
    elif method_name.lower() == 'mrt':
        next_relaxation_rate = [0]

Martin Bauer's avatar
Martin Bauer committed
419
        def relaxation_rate_getter(moments):
420
            try:
Martin Bauer's avatar
Martin Bauer committed
421
                if all(get_order(m) < 2 for m in moments):
Markus Holzer's avatar
Markus Holzer committed
422
423
424
425
                    if params['entropic']:
                        return relaxation_rates[0]
                    else:
                        return 0
426
427
428
429
                res = relaxation_rates[next_relaxation_rate[0]]
                next_relaxation_rate[0] += 1
            except IndexError:
                raise ValueError("Too few relaxation rates specified")
430
            return res
Martin Bauer's avatar
Martin Bauer committed
431
        weighted = params['weighted'] if 'weighted' in params else True
432
433
434
        nested_moments = params['nested_moments'] if 'nested_moments' in params else None
        method = create_mrt_orthogonal(stencil_entries, relaxation_rate_getter, weighted=weighted,
                                       nested_moments=nested_moments, **common_params)
Martin Bauer's avatar
Martin Bauer committed
435
436
437
438
439
    elif method_name.lower() == 'mrt_raw':
        method = create_mrt_raw(stencil_entries, relaxation_rates, **common_params)
    elif method_name.lower() == 'mrt3':
        method = create_mrt3(stencil_entries, relaxation_rates, **common_params)
    elif method_name.lower().startswith('trt-kbc-n'):
440
        if have_same_entries(stencil_entries, get_stencil("D2Q9")):
441
            dim = 2
442
        elif have_same_entries(stencil_entries, get_stencil("D3Q27")):
443
444
445
            dim = 3
        else:
            raise NotImplementedError("KBC type TRT methods can only be constructed for D2Q9 and D3Q27 stencils")
Martin Bauer's avatar
Martin Bauer committed
446
447
448
449
        method_nr = method_name[-1]
        method = create_trt_kbc(dim, relaxation_rates[0], relaxation_rates[1], 'KBC-N' + method_nr, **common_params)
    elif method_name.lower() == 'entropic-srt':
        method = create_srt_entropic(stencil_entries, relaxation_rates[0], force_model, params['compressible'])
450
    else:
Martin Bauer's avatar
Martin Bauer committed
451
        raise ValueError("Unknown method %s" % (method_name,))
452
453
454

    return method

455

Martin Bauer's avatar
Martin Bauer committed
456
457
458
459
460
def create_lb_method_from_existing(method, modification_function):
    """Creates a new method based on an existing method by modifying its collision table.

    Args:
        method: old method
Martin Bauer's avatar
Martin Bauer committed
461
        modification_function: function receiving (moment, equilibrium_value, relaxation_rate) as arguments,
Martin Bauer's avatar
Martin Bauer committed
462
463
464
465
466
467
468
469
                               i.e. one row of the relaxation table, returning a modified version
    """
    relaxation_table = (modification_function(m, eq, rr)
                        for m, eq, rr in zip(method.moments, method.moment_equilibrium_values, method.relaxation_rates))
    compressible = method.conserved_quantity_computation.compressible
    cumulant = isinstance(method, CumulantBasedLbMethod)
    return create_generic_mrt(method.stencil, relaxation_table, compressible, method.force_model, cumulant)

470
471
472
# ----------------------------------------------------------------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
473
474
475
476
def force_model_from_string(force_model_name, force_values):
    if type(force_model_name) is not str:
        return force_model_name
    if force_model_name == 'none':
Martin Bauer's avatar
Martin Bauer committed
477
478
        return None

Martin Bauer's avatar
Martin Bauer committed
479
    force_model_dict = {
Martin Bauer's avatar
Martin Bauer committed
480
481
482
483
484
485
486
        'simple': forcemodels.Simple,
        'luo': forcemodels.Luo,
        'guo': forcemodels.Guo,
        'buick': forcemodels.Buick,
        'silva': forcemodels.Buick,
        'edm': forcemodels.EDM,
    }
Martin Bauer's avatar
Martin Bauer committed
487
488
    if force_model_name.lower() not in force_model_dict:
        raise ValueError("Unknown force model %s" % (force_model_name,))
Martin Bauer's avatar
Martin Bauer committed
489

Martin Bauer's avatar
Martin Bauer committed
490
491
    force_model_class = force_model_dict[force_model_name.lower()]
    return force_model_class(force_values)
492
493


Martin Bauer's avatar
Martin Bauer committed
494
def switch_to_symbolic_relaxation_rates_for_omega_adapting_methods(method_parameters, kernel_params, force=False):
495
496
497
498
    """
    For entropic kernels the relaxation rate has to be a variable. If a constant was passed a
    new dummy variable is inserted and the value of this variable is later on passed to the kernel
    """
Martin Bauer's avatar
Martin Bauer committed
499
    if method_parameters['entropic'] or method_parameters['smagorinsky'] or force:
Martin Bauer's avatar
Martin Bauer committed
500
501
502
        value_to_symbol_map = {}
        new_relaxation_rates = []
        for rr in method_parameters['relaxation_rates']:
503
            if not isinstance(rr, sp.Symbol):
Martin Bauer's avatar
Martin Bauer committed
504
505
506
507
508
                if rr not in value_to_symbol_map:
                    value_to_symbol_map[rr] = sp.Dummy()
                dummy_var = value_to_symbol_map[rr]
                new_relaxation_rates.append(dummy_var)
                kernel_params[dummy_var.name] = rr
509
            else:
Martin Bauer's avatar
Martin Bauer committed
510
511
512
513
                new_relaxation_rates.append(rr)
        if len(new_relaxation_rates) < 2:
            new_relaxation_rates.append(sp.Dummy())
        method_parameters['relaxation_rates'] = new_relaxation_rates
514
515


Martin Bauer's avatar
Martin Bauer committed
516
517
518
519
def update_with_default_parameters(params, opt_params=None, fail_on_unknown_parameter=True):
    opt_params = opt_params if opt_params is not None else {}

    default_method_description = {
520
521
        'stencil': 'D2Q9',
        'method': 'srt',  # can be srt, trt or mrt
Martin Bauer's avatar
Martin Bauer committed
522
        'relaxation_rates': None,
523
        'compressible': False,
Martin Bauer's avatar
Martin Bauer committed
524
        'equilibrium_order': 2,
525
        'c_s_sq': sp.Rational(1, 3),
Martin Bauer's avatar
Martin Bauer committed
526
        'weighted': True,
Michael Kuron's avatar
Michael Kuron committed
527
        'nested_moments': None,
528

Martin Bauer's avatar
Martin Bauer committed
529
        'force_model': 'none',
530
        'force': (0, 0, 0),
Martin Bauer's avatar
Martin Bauer committed
531
        'maxwellian_moments': True,
532
        'cumulant': False,
Martin Bauer's avatar
Martin Bauer committed
533
        'initial_velocity': None,
534
535

        'entropic': False,
Martin Bauer's avatar
Martin Bauer committed
536
537
        'entropic_newton_iterations': None,
        'omega_output_field': None,
538
        'smagorinsky': False,
539
        'fluctuating': False,
540
        'temperature': None,
541
542

        'output': {},
Martin Bauer's avatar
Martin Bauer committed
543
        'velocity_input': None,
Martin Bauer's avatar
Martin Bauer committed
544
        'density_input': None,
545

Martin Bauer's avatar
Martin Bauer committed
546
        'kernel_type': 'stream_pull_collide',
547

Martin Bauer's avatar
Martin Bauer committed
548
549
        'field_name': 'src',
        'temporary_field_name': 'dst',
550

Martin Bauer's avatar
Martin Bauer committed
551
552
553
        'lb_method': None,
        'collision_rule': None,
        'update_rule': None,
554
555
556
        'ast': None,
    }

Martin Bauer's avatar
Martin Bauer committed
557
558
559
    default_optimization_description = {
        'cse_pdfs': False,
        'cse_global': False,
Martin Bauer's avatar
Martin Bauer committed
560
561
        'simplification': 'auto',
        'keep_rrs_symbolic': True,
562
563
        'split': False,

Martin Bauer's avatar
Martin Bauer committed
564
        'field_size': None,
Martin Bauer's avatar
Martin Bauer committed
565
        'field_layout': 'fzyx',  # can be 'numpy' (='c'), 'reverse_numpy' (='f'), 'fzyx', 'zyxf'
Martin Bauer's avatar
Martin Bauer committed
566
        'symbolic_field': None,
567
        'symbolic_temporary_field': None,
568
569

        'target': 'cpu',
Martin Bauer's avatar
Martin Bauer committed
570
571
        'openmp': False,
        'double_precision': True,
572

Martin Bauer's avatar
Martin Bauer committed
573
574
        'gpu_indexing': 'block',
        'gpu_indexing_params': {},
575
576
577

        'vectorization': None,

Martin Bauer's avatar
Martin Bauer committed
578
        'builtin_periodicity': (False, False, False),
579
    }
Martin Bauer's avatar
Martin Bauer committed
580
581
    if 'relaxation_rate' in params:
        if 'relaxation_rates' not in params:
582
            if 'entropic' in params and params['entropic']:
Martin Bauer's avatar
Martin Bauer committed
583
                params['relaxation_rates'] = [params['relaxation_rate']]
584
            else:
Martin Bauer's avatar
Martin Bauer committed
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
                params['relaxation_rates'] = [params['relaxation_rate'],
                                              relaxation_rate_from_magic_number(params['relaxation_rate'])]

            del params['relaxation_rate']

    if 'pdf_arr' in opt_params:
        arr = opt_params['pdf_arr']
        opt_params['field_size'] = tuple(e - 2 for e in arr.shape[:-1])
        opt_params['field_layout'] = get_layout_of_array(arr)
        del opt_params['pdf_arr']

    if fail_on_unknown_parameter:
        unknown_params = [k for k in params.keys() if k not in default_method_description]
        unknown_opt_params = [k for k in opt_params.keys() if k not in default_optimization_description]
        if unknown_params:
            raise ValueError("Unknown parameter(s): " + ", ".join(unknown_params))
        if unknown_opt_params:
            raise ValueError("Unknown optimization parameter(s): " + ",".join(unknown_opt_params))

    params_result = copy(default_method_description)
    params_result.update(params)
    opt_params_result = copy(default_optimization_description)
    opt_params_result.update(opt_params)

    if params_result['relaxation_rates'] is None:
        stencil_param = params_result['stencil']
        if isinstance(stencil_param, tuple) or isinstance(stencil_param, list):
            stencil_entries = stencil_param
613
        else:
Martin Bauer's avatar
Martin Bauer committed
614
615
            stencil_entries = get_stencil(params_result['stencil'])
        params_result['relaxation_rates'] = sp.symbols("omega_:%d" % len(stencil_entries))
616

Martin Bauer's avatar
Martin Bauer committed
617
    return params_result, opt_params_result