creationfunctions.py 35.5 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
    - ``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`
29
30
    - ``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
31
      mapping (:func:`lbmpy.methods.create_mrt_raw`)
32
    - ``trt-kbc-n<N>`` where <N> is 1,2,3 or 4. Special two-relaxation rate method. This is not the entropic method
33
      yet, only the relaxation pattern. To get the entropic method, see parameters below!
Martin Bauer's avatar
Martin Bauer committed
34
      (:func:`lbmpy.methods.create_trt_kbc`)
Frederik Hennig's avatar
Frederik Hennig committed
35
36
37
38
    - ``cumulant``: cumulant-based lb method (:func:`lbmpy.methods.create_with_default_polynomial_cumulants`) which 
      relaxes groups of polynomial cumulants chosen to optimize rotational invariance.
    - ``monomial_cumulant``: cumulant-based lb method (:func:`lbmpy.methods.create_with_monomial_cumulants`) which
      relaxes monomial cumulants.
Martin Bauer's avatar
Martin Bauer committed
39

Martin Bauer's avatar
Martin Bauer committed
40
- ``relaxation_rates``: sequence of relaxation rates, number depends on selected method. If you specify more rates than
Frederik Hennig's avatar
Frederik Hennig committed
41
42
43
  method needs, the additional rates are ignored. For SRT, TRT and polynomial cumulant models it is possible ot define 
  a single ``relaxation_rate`` instead of a list. The second rate for TRT is then determined via magic number. For the
  cumulant model, it sets only the relaxation rate corresponding to shear viscosity, setting all others to unity.
44
45
46
- ``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
47
- ``equilibrium_order=2``: order in velocity, at which the equilibrium moment/cumulant approximation is
48
  truncated. Order 2 is sufficient to approximate Navier-Stokes
49
50
51
- ``force_model=None``: possible values: ``None``, ``'simple'``, ``'luo'``, ``'guo'``/``'schiller'``,
  ``'buick'``/``'silva'``, ``'edm'``/``'kupershtokh'``, ``'he'``, ``'shanchen'``, ``'cumulant'``,
  or an instance of a class implementing the same methods as the classes in :mod:`lbmpy.forcemodels`. 
Frederik Hennig's avatar
Frederik Hennig committed
52
  For details, see :mod:`lbmpy.forcemodels`
53
- ``force=(0,0,0)``: either constant force or a symbolic expression depending on field value
Martin Bauer's avatar
Martin Bauer committed
54
- ``maxwellian_moments=True``: way to compute equilibrium moments/cumulants, if False the standard
55
56
  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
Martin Bauer's avatar
Martin Bauer committed
57
- ``initial_velocity=None``: initial velocity in domain, can either be a tuple (x,y,z) velocity to set a constant
58
59
  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
60
- ``output={}``: a dictionary mapping macroscopic quantites e.g. the strings 'density' and 'velocity' to pystencils
61
  fields. In each timestep the corresponding quantities are written to the given fields.
Martin Bauer's avatar
Martin Bauer committed
62
- ``velocity_input``: symbolic field where the velocities are read from (for advection diffusion LBM)
Martin Bauer's avatar
Martin Bauer committed
63
- ``density_input``: symbolic field or field access where to read density from. When passing this parameter,
Martin Bauer's avatar
Martin Bauer committed
64
  ``velocity_input`` has to be passed as well
Frederik Hennig's avatar
Frederik Hennig committed
65
66
67
68
69
70
71
72
73
- ``kernel_type``: supported values: 'default_stream_collide' (default), 'collide_only', 'stream_pull_only'. 
  With 'default_stream_collide', streaming pattern and even/odd time-step (for in-place patterns) can be specified
  by the ``streaming_pattern`` and ``timestep`` arguments. For backwards compatibility, ``kernel_type`` also accepts
  'stream_pull_collide', 'collide_stream_push', 'esotwist_even', 'esotwist_odd', 'aa_even' and 'aa_odd' for selection 
  of the streaming pattern. 
- ``streaming_pattern``: The streaming pattern to be used with a 'default_stream_collide' kernel. Accepted values are
  'pull', 'push', 'aa' and 'esotwist'.
- ``timestep``: Timestep modulus for the streaming pattern. For two-fields patterns, this argument is irrelevant and
  by default set to ``Timestep.BOTH``. For in-place patterns, ``Timestep.EVEN`` or ``Timestep.ODD`` must be speficied.
74
75
76
77
78

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
Frederik Hennig's avatar
Frederik Hennig committed
79
  :mod:`lbmpy.methods.momentbased.entropic`
Martin Bauer's avatar
Martin Bauer committed
80
- ``entropic_newton_iterations=None``: For moment methods the entropy optimum can be calculated in closed form.
81
82
  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
83
84
- ``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
85

Frederik Hennig's avatar
Frederik Hennig committed
86
87
88
89
90
Cumulant methods:

- ``galilean_correction=False``: Special correction for D3Q27 cumulant LBMs. For Details see
  :mod:`lbmpy.methods.centeredcumulant.galilean_correction`

91
92
LES methods:

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

96
97
Fluctuating LB:

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

101
102
103
104
105
106

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

Simplifications / Transformations:

Martin Bauer's avatar
Martin Bauer committed
107
108
- ``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
109
110
- ``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
111
- ``builtin_periodicity=(False,False,False)``: instead of handling periodicity by copying ghost layers, the periodicity
112
  is built into the kernel. This parameters specifies if the domain is periodic in (x,y,z) direction. Even if the
113
114
  periodicity is built into the kernel, the fields have one ghost layer to be consistent with other functions.
- ``pre_simplification=True``: Simplifications applied during the derivaton of the collision rule for cumulant LBMs
Frederik Hennig's avatar
Frederik Hennig committed
115
  For details see :mod:`lbmpy.moment_transforms`.
116
    
117
118
119

Field size information:

120
121
- ``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
122
- ``field_size=None``: create kernel for fixed field size
Martin Bauer's avatar
Martin Bauer committed
123
124
- ``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
125
126
- ``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)
127

128
129
130
131
132
133
134
135
136
137
138

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.


139
140
141
142
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
143
144
- ``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
145
  For ``'block'`` indexing one can e.g. specify the block size ``{'block_size' : (128, 4, 1)}``
146

147

148
149
Other:

Martin Bauer's avatar
Martin Bauer committed
150
- ``openmp=True``: only applicable for cpu simulations. Can be a boolean to turn multi threading on/off, or an integer
151
  specifying the number of threads. If True is specified OpenMP chooses the number of threads
Martin Bauer's avatar
Martin Bauer committed
152
- ``double_precision=True``:  by default simulations run with double precision floating point numbers, by setting this
153
  parameter to False, single precision is used, which is much faster, especially on GPUs
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
181
182




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
183
The function :func:`create_lb_function` runs the whole pipeline, the other functions in this module
184
185
186
187
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
188
    ast = create_lb_ast(...)
189
    # modify ast here
Martin Bauer's avatar
Martin Bauer committed
190
    func = create_lb_function(ast=ast, ...)
191

192
"""
193
from collections import OrderedDict
194
from copy import copy
195

Martin Bauer's avatar
Martin Bauer committed
196
197
198
import sympy as sp

import lbmpy.forcemodels as forcemodels
Frederik Hennig's avatar
Frederik Hennig committed
199
import lbmpy.methods.centeredcumulant.force_model as cumulant_force_model
Frederik Hennig's avatar
Frederik Hennig committed
200
from lbmpy.fieldaccess import CollideOnlyInplaceAccessor, PdfFieldAccessor, PeriodicTwoFieldsAccessor
201
from lbmpy.fluctuatinglb import add_fluctuations_to_collision_rule
Markus Holzer's avatar
Markus Holzer committed
202
203
from lbmpy.methods import (create_mrt_orthogonal, create_mrt_raw, create_central_moment,
                           create_srt, create_trt, create_trt_kbc)
204
from lbmpy.methods.abstractlbmethod import RelaxationInfo
Frederik Hennig's avatar
Frederik Hennig committed
205
from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
Frederik Hennig's avatar
Frederik Hennig committed
206
from lbmpy.moment_transforms import PdfsToMomentsByChimeraTransform, PdfsToCentralMomentsByShiftMatrix
Frederik Hennig's avatar
Frederik Hennig committed
207
208
209
from lbmpy.methods.centeredcumulant.cumulant_transform import CentralMomentsToCumulantsByGeneratingFunc
from lbmpy.methods.creationfunctions import (
    create_with_monomial_cumulants, create_with_polynomial_cumulants, create_with_default_polynomial_cumulants)
Martin Bauer's avatar
Martin Bauer committed
210
from lbmpy.methods.creationfunctions import create_generic_mrt
Frederik Hennig's avatar
Frederik Hennig committed
211
212
from lbmpy.methods.momentbased.entropic import add_entropy_condition, add_iterative_entropy_condition
from lbmpy.methods.momentbased.entropic_eq_srt import create_srt_entropic
Martin Bauer's avatar
Martin Bauer committed
213
from lbmpy.moments import get_order
Martin Bauer's avatar
Martin Bauer committed
214
from lbmpy.relaxationrates import relaxation_rate_from_magic_number
Martin Bauer's avatar
Martin Bauer committed
215
from lbmpy.simplificationfactory import create_simplification_strategy
216
from lbmpy.stencils import get_stencil
Martin Bauer's avatar
Martin Bauer committed
217
218
from lbmpy.turbulence_models import add_smagorinsky_model
from lbmpy.updatekernels import create_lbm_kernel, create_stream_pull_with_output_kernel
Frederik Hennig's avatar
Frederik Hennig committed
219
from lbmpy.advanced_streaming.utility import Timestep, get_accessor
Martin Bauer's avatar
Martin Bauer committed
220
221
222
223
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
Frederik Hennig's avatar
Frederik Hennig committed
224
from pystencils.simp import sympy_cse, SimplificationStrategy
225
from pystencils.stencil import have_same_entries
226
227


Frederik Hennig's avatar
Frederik Hennig committed
228
def create_lb_function(ast=None, optimization=None, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
229
    """Creates a Python function for the LB method"""
Martin Bauer's avatar
Martin Bauer committed
230
    params, opt_params = update_with_default_parameters(kwargs, optimization)
231

232
    if ast is None:
Martin Bauer's avatar
Martin Bauer committed
233
        params['optimization'] = opt_params
Martin Bauer's avatar
Martin Bauer committed
234
        ast = create_lb_ast(**params)
235

Martin Bauer's avatar
Martin Bauer committed
236
    res = ast.compile()
237

238
    res.method = ast.method
Martin Bauer's avatar
Martin Bauer committed
239
    res.update_rule = ast.update_rule
Martin Bauer's avatar
Martin Bauer committed
240
    res.ast = ast
241
242
    return res

243

Frederik Hennig's avatar
Frederik Hennig committed
244
def create_lb_ast(update_rule=None, optimization=None, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
245
    """Creates a pystencils AST for the LB method"""
Martin Bauer's avatar
Martin Bauer committed
246
    params, opt_params = update_with_default_parameters(kwargs, optimization)
247

Martin Bauer's avatar
Martin Bauer committed
248
249
250
    if update_rule is None:
        params['optimization'] = optimization
        update_rule = create_lb_update_rule(**params)
251

Martin Bauer's avatar
Martin Bauer committed
252
253
254
255
256
    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)
257

Martin Bauer's avatar
Martin Bauer committed
258
259
    res.method = update_rule.method
    res.update_rule = update_rule
260
    return res
261

262

Martin Bauer's avatar
Martin Bauer committed
263
@disk_cache_no_fallback
Frederik Hennig's avatar
Frederik Hennig committed
264
def create_lb_update_rule(collision_rule=None, optimization=None, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
265
    """Creates an update rule (list of Assignments) for a LB method that describe a full sweep"""
Martin Bauer's avatar
Martin Bauer committed
266
    params, opt_params = update_with_default_parameters(kwargs, optimization)
267

Martin Bauer's avatar
Martin Bauer committed
268
    if collision_rule is None:
Martin Bauer's avatar
Martin Bauer committed
269
        collision_rule = create_lb_collision_rule(**params, optimization=opt_params)
270

Martin Bauer's avatar
Martin Bauer committed
271
    lb_method = collision_rule.method
272

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

Martin Bauer's avatar
Martin Bauer committed
276
277
278
    if opt_params['symbolic_field'] is not None:
        src_field = opt_params['symbolic_field']
    elif opt_params['field_size']:
279
        field_size = [s + 2 for s in opt_params['field_size']] + [q]
Martin Bauer's avatar
Martin Bauer committed
280
281
        src_field = Field.create_fixed_size(params['field_name'], field_size, index_dimensions=1,
                                            layout=opt_params['field_layout'], dtype=field_data_type)
282
    else:
Martin Bauer's avatar
Martin Bauer committed
283
        src_field = Field.create_generic(params['field_name'], spatial_dimensions=collision_rule.method.dim,
284
                                         index_shape=(q,), layout=opt_params['field_layout'], dtype=field_data_type)
285

286
287
288
289
    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'])
290

291
    kernel_type = params['kernel_type']
Frederik Hennig's avatar
Frederik Hennig committed
292
    if kernel_type == 'stream_pull_only':
Martin Bauer's avatar
Martin Bauer committed
293
        return create_stream_pull_with_output_kernel(lb_method, src_field, dst_field, params['output'])
294
    else:
Frederik Hennig's avatar
Frederik Hennig committed
295
296
297
298
299
300
301
302
303
304
        if kernel_type == 'default_stream_collide':
            if params['streaming_pattern'] == 'pull' and any(opt_params['builtin_periodicity']):
                accessor = PeriodicTwoFieldsAccessor(opt_params['builtin_periodicity'], ghost_layers=1)
            else:
                accessor = get_accessor(params['streaming_pattern'], params['timestep'])
        elif kernel_type == 'collide_only':
            accessor = CollideOnlyInplaceAccessor
        elif isinstance(kernel_type, PdfFieldAccessor):
            accessor = kernel_type
        else:
305
306
            raise ValueError("Invalid value of parameter 'kernel_type'", params['kernel_type'])
        return create_lbm_kernel(collision_rule, src_field, dst_field, accessor)
307
308


Martin Bauer's avatar
Martin Bauer committed
309
@disk_cache_no_fallback
Frederik Hennig's avatar
Frederik Hennig committed
310
def create_lb_collision_rule(lb_method=None, optimization=None, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
311
    """Creates a collision rule (list of Assignments) for a LB method describing the collision operator (no stream)"""
Martin Bauer's avatar
Martin Bauer committed
312
    params, opt_params = update_with_default_parameters(kwargs, optimization)
313

Martin Bauer's avatar
Martin Bauer committed
314
315
    if lb_method is None:
        lb_method = create_lb_method(**params)
316

Martin Bauer's avatar
Martin Bauer committed
317
318
    split_inner_loop = 'split' in opt_params and opt_params['split']
    cqc = lb_method.conserved_quantity_computation
319

Martin Bauer's avatar
Martin Bauer committed
320
321
322
323
324
325
326
327
    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

Frederik Hennig's avatar
Frederik Hennig committed
328
    pre_simplification = opt_params['pre_simplification']
Martin Bauer's avatar
Martin Bauer committed
329
330
331
332
    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)]
333
        eqs = AssignmentCollection(eqs, [])
Martin Bauer's avatar
Martin Bauer committed
334
        collision_rule = lb_method.get_collision_rule(conserved_quantity_equations=eqs,
Frederik Hennig's avatar
Frederik Hennig committed
335
336
                                                      pre_simplification=pre_simplification)

Martin Bauer's avatar
Martin Bauer committed
337
338
    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.")
339
    else:
Frederik Hennig's avatar
Frederik Hennig committed
340
        collision_rule = lb_method.get_collision_rule(pre_simplification=pre_simplification)
341

342
    if params['entropic']:
343
344
        if params['smagorinsky']:
            raise ValueError("Choose either entropic or smagorinsky")
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
        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"))

Frederik Hennig's avatar
Frederik Hennig committed
361
    if params['output']:
362
363
364
365
366
367
        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)

    if opt_params['simplification'] == 'auto':
        simplification = create_simplification_strategy(lb_method, split_inner_loop=split_inner_loop)
Frederik Hennig's avatar
Frederik Hennig committed
368
    elif callable(opt_params['simplification']):
369
        simplification = opt_params['simplification']
Frederik Hennig's avatar
Frederik Hennig committed
370
371
    else:
        simplification = SimplificationStrategy()
372
373
374
375
376
    collision_rule = simplification(collision_rule)

    if params['fluctuating']:
        add_fluctuations_to_collision_rule(collision_rule, **params['fluctuating'])

377
378
379
    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:
Frederik Hennig's avatar
Frederik Hennig committed
380
        from lbmpy.methods.momentbased.momentbasedsimplifications import cse_in_opposing_directions
381
382
383
384
        collision_rule = cse_in_opposing_directions(collision_rule)
    if cse_global:
        collision_rule = sympy_cse(collision_rule)

385
    return collision_rule
386
387


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

Frederik Hennig's avatar
Frederik Hennig committed
392
393
394
    method_name = params['method']
    relaxation_rates = params['relaxation_rates']

395
    if isinstance(params['stencil'], tuple) or isinstance(params['stencil'], list):
Martin Bauer's avatar
Martin Bauer committed
396
        stencil_entries = params['stencil']
397
    else:
Martin Bauer's avatar
Martin Bauer committed
398
        stencil_entries = get_stencil(params['stencil'])
399

Martin Bauer's avatar
Martin Bauer committed
400
    dim = len(stencil_entries[0])
401

402
403
404
    if isinstance(params['force'], Field):
        params['force'] = tuple(params['force'](i) for i in range(dim))

Martin Bauer's avatar
Martin Bauer committed
405
    force_is_zero = True
406
407
    for f_i in params['force']:
        if f_i != 0:
Martin Bauer's avatar
Martin Bauer committed
408
            force_is_zero = False
409

Martin Bauer's avatar
Martin Bauer committed
410
411
    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:
412
        params['force_model'] = 'cumulant' if method_name.lower().endswith('cumulant') else 'guo'
413

Martin Bauer's avatar
Martin Bauer committed
414
415
    if 'force_model' in params:
        force_model = force_model_from_string(params['force_model'], params['force'][:dim])
416
    else:
Martin Bauer's avatar
Martin Bauer committed
417
        force_model = None
418

Martin Bauer's avatar
Martin Bauer committed
419
    common_params = {
420
        'compressible': params['compressible'],
Martin Bauer's avatar
Martin Bauer committed
421
422
423
        'equilibrium_order': params['equilibrium_order'],
        'force_model': force_model,
        'maxwellian_moments': params['maxwellian_moments'],
Frederik Hennig's avatar
Frederik Hennig committed
424
425
426
        'c_s_sq': params['c_s_sq'],
        'moment_transform_class': params['moment_transform_class'],
        'central_moment_transform_class': params['central_moment_transform_class'],
427
    }
Frederik Hennig's avatar
Frederik Hennig committed
428
429
430
431
432
433
434
435
436

    cumulant_params = {
        'equilibrium_order': params['equilibrium_order'],
        'force_model': force_model,
        'c_s_sq': params['c_s_sq'],
        'galilean_correction': params['galilean_correction'],
        'central_moment_transform_class': params['central_moment_transform_class'],
        'cumulant_transform_class': params['cumulant_transform_class'],
    }
Martin Bauer's avatar
Martin Bauer committed
437
438
439
440
441
442
443
444
445
446

    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
447
        def relaxation_rate_getter(moments):
448
            try:
Martin Bauer's avatar
Martin Bauer committed
449
                if all(get_order(m) < 2 for m in moments):
Markus Holzer's avatar
Markus Holzer committed
450
451
452
453
                    if params['entropic']:
                        return relaxation_rates[0]
                    else:
                        return 0
454
455
456
457
                res = relaxation_rates[next_relaxation_rate[0]]
                next_relaxation_rate[0] += 1
            except IndexError:
                raise ValueError("Too few relaxation rates specified")
458
            return res
459

Martin Bauer's avatar
Martin Bauer committed
460
        weighted = params['weighted'] if 'weighted' in params else True
461
462
463
        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)
Markus Holzer's avatar
Markus Holzer committed
464
    elif method_name.lower() == 'central_moment':
465
466
467
        nested_moments = params['nested_moments'] if 'nested_moments' in params else None
        method = create_central_moment(stencil_entries, relaxation_rates,
                                       nested_moments=nested_moments, **common_params)
Martin Bauer's avatar
Martin Bauer committed
468
469
470
    elif method_name.lower() == 'mrt_raw':
        method = create_mrt_raw(stencil_entries, relaxation_rates, **common_params)
    elif method_name.lower().startswith('trt-kbc-n'):
471
        if have_same_entries(stencil_entries, get_stencil("D2Q9")):
472
            dim = 2
473
        elif have_same_entries(stencil_entries, get_stencil("D3Q27")):
474
475
476
            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
477
478
479
480
        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'])
Frederik Hennig's avatar
Frederik Hennig committed
481
482
483
484
485
486
487
488
489
    elif method_name.lower() == 'cumulant':
        nested_moments = params['nested_moments'] if 'nested_moments' in params else None
        if nested_moments is not None:
            method = create_with_polynomial_cumulants(
                stencil_entries, relaxation_rates, nested_moments, **cumulant_params)
        else:
            method = create_with_default_polynomial_cumulants(stencil_entries, relaxation_rates, **cumulant_params)
    elif method_name.lower() == 'monomial_cumulant':
        method = create_with_monomial_cumulants(stencil_entries, relaxation_rates, **cumulant_params)
490
    else:
Martin Bauer's avatar
Martin Bauer committed
491
        raise ValueError("Unknown method %s" % (method_name,))
492
493
494

    return method

495

Martin Bauer's avatar
Martin Bauer committed
496
497
498
499
500
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
501
        modification_function: function receiving (moment, equilibrium_value, relaxation_rate) as arguments,
Martin Bauer's avatar
Martin Bauer committed
502
503
504
                               i.e. one row of the relaxation table, returning a modified version
    """
    compressible = method.conserved_quantity_computation.compressible
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
    if isinstance(method, CenteredCumulantBasedLbMethod):
        rr_dict = OrderedDict()
        relaxation_table = (modification_function(m, eq, rr)
                            for m, eq, rr in
                            zip(method.cumulants, method.cumulant_equilibrium_values, method.relaxation_rates))

        for cumulant, eq_value, rr in relaxation_table:
            cumulant = sp.sympify(cumulant)
            rr_dict[cumulant] = RelaxationInfo(eq_value, rr)

        return CenteredCumulantBasedLbMethod(method.stencil, rr_dict, method.conserved_quantity_computation,
                                             method.force_model,
                                             galilean_correction=method.galilean_correction,
                                             central_moment_transform_class=method.central_moment_transform_class,
                                             cumulant_transform_class=method.cumulant_transform_class)
    else:
        relaxation_table = (modification_function(m, eq, rr)
                            for m, eq, rr in
                            zip(method.moments, method.moment_equilibrium_values, method.relaxation_rates))
        return create_generic_mrt(method.stencil, relaxation_table, compressible, method.force_model)
Martin Bauer's avatar
Martin Bauer committed
525

526

527
528
529
# ----------------------------------------------------------------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
530
531
532
533
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
534
535
        return None

Martin Bauer's avatar
Martin Bauer committed
536
    force_model_dict = {
Martin Bauer's avatar
Martin Bauer committed
537
538
539
        'simple': forcemodels.Simple,
        'luo': forcemodels.Luo,
        'guo': forcemodels.Guo,
540
        'schiller': forcemodels.Guo,
Martin Bauer's avatar
Martin Bauer committed
541
542
543
        'buick': forcemodels.Buick,
        'silva': forcemodels.Buick,
        'edm': forcemodels.EDM,
544
        'kupershtokh': forcemodels.EDM,
Markus Holzer's avatar
Markus Holzer committed
545
        'cumulant': cumulant_force_model.CenteredCumulantForceModel,
546
547
        'he': forcemodels.He,
        'shanchen': forcemodels.ShanChen
Martin Bauer's avatar
Martin Bauer committed
548
    }
Martin Bauer's avatar
Martin Bauer committed
549
550
    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
551

Martin Bauer's avatar
Martin Bauer committed
552
553
    force_model_class = force_model_dict[force_model_name.lower()]
    return force_model_class(force_values)
554
555


Martin Bauer's avatar
Martin Bauer committed
556
def switch_to_symbolic_relaxation_rates_for_omega_adapting_methods(method_parameters, kernel_params, force=False):
557
558
559
560
    """
    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
561
    if method_parameters['entropic'] or method_parameters['smagorinsky'] or force:
Martin Bauer's avatar
Martin Bauer committed
562
563
564
        value_to_symbol_map = {}
        new_relaxation_rates = []
        for rr in method_parameters['relaxation_rates']:
565
            if not isinstance(rr, sp.Symbol):
Martin Bauer's avatar
Martin Bauer committed
566
567
568
569
570
                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
571
            else:
Martin Bauer's avatar
Martin Bauer committed
572
573
574
575
                new_relaxation_rates.append(rr)
        if len(new_relaxation_rates) < 2:
            new_relaxation_rates.append(sp.Dummy())
        method_parameters['relaxation_rates'] = new_relaxation_rates
576
577


Martin Bauer's avatar
Martin Bauer committed
578
579
580
581
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 = {
582
        'stencil': 'D2Q9',
Frederik Hennig's avatar
Frederik Hennig committed
583
        'method': 'srt',  # can be srt, trt, mrt or cumulant
Martin Bauer's avatar
Martin Bauer committed
584
        'relaxation_rates': None,
585
        'compressible': False,
Martin Bauer's avatar
Martin Bauer committed
586
        'equilibrium_order': 2,
587
        'c_s_sq': sp.Rational(1, 3),
Martin Bauer's avatar
Martin Bauer committed
588
        'weighted': True,
Michael Kuron's avatar
Michael Kuron committed
589
        'nested_moments': None,
590

Martin Bauer's avatar
Martin Bauer committed
591
        'force_model': 'none',
592
        'force': (0, 0, 0),
Martin Bauer's avatar
Martin Bauer committed
593
594
        'maxwellian_moments': True,
        'initial_velocity': None,
595

Frederik Hennig's avatar
Frederik Hennig committed
596
        'galilean_correction': False,  # only available for D3Q27 cumulant methods
Frederik Hennig's avatar
Frederik Hennig committed
597
598

        'moment_transform_class': PdfsToMomentsByChimeraTransform,
Frederik Hennig's avatar
Frederik Hennig committed
599
        'central_moment_transform_class': PdfsToCentralMomentsByShiftMatrix,
Frederik Hennig's avatar
Frederik Hennig committed
600
601
        'cumulant_transform_class': CentralMomentsToCumulantsByGeneratingFunc,

602
        'entropic': False,
Martin Bauer's avatar
Martin Bauer committed
603
604
        'entropic_newton_iterations': None,
        'omega_output_field': None,
605
        'smagorinsky': False,
606
        'fluctuating': False,
607
        'temperature': None,
608
609

        'output': {},
Martin Bauer's avatar
Martin Bauer committed
610
        'velocity_input': None,
Martin Bauer's avatar
Martin Bauer committed
611
        'density_input': None,
612

Frederik Hennig's avatar
Frederik Hennig committed
613
614
615
        'kernel_type': 'default_stream_collide',
        'streaming_pattern': 'pull',
        'timestep': Timestep.BOTH,
616

Martin Bauer's avatar
Martin Bauer committed
617
618
        'field_name': 'src',
        'temporary_field_name': 'dst',
619

Martin Bauer's avatar
Martin Bauer committed
620
621
622
        'lb_method': None,
        'collision_rule': None,
        'update_rule': None,
623
624
625
        'ast': None,
    }

Martin Bauer's avatar
Martin Bauer committed
626
627
628
    default_optimization_description = {
        'cse_pdfs': False,
        'cse_global': False,
Martin Bauer's avatar
Martin Bauer committed
629
        'simplification': 'auto',
Frederik Hennig's avatar
Frederik Hennig committed
630
        'pre_simplification': True,
631
632
        'split': False,

Martin Bauer's avatar
Martin Bauer committed
633
        'field_size': None,
Martin Bauer's avatar
Martin Bauer committed
634
        'field_layout': 'fzyx',  # can be 'numpy' (='c'), 'reverse_numpy' (='f'), 'fzyx', 'zyxf'
Martin Bauer's avatar
Martin Bauer committed
635
        'symbolic_field': None,
636
        'symbolic_temporary_field': None,
637
638

        'target': 'cpu',
Martin Bauer's avatar
Martin Bauer committed
639
640
        'openmp': False,
        'double_precision': True,
641

Martin Bauer's avatar
Martin Bauer committed
642
643
        'gpu_indexing': 'block',
        'gpu_indexing_params': {},
644
645
646

        'vectorization': None,

Martin Bauer's avatar
Martin Bauer committed
647
        'builtin_periodicity': (False, False, False),
648
    }
Martin Bauer's avatar
Martin Bauer committed
649
650
    if 'relaxation_rate' in params:
        if 'relaxation_rates' not in params:
651
            if 'entropic' in params and params['entropic']:
Martin Bauer's avatar
Martin Bauer committed
652
                params['relaxation_rates'] = [params['relaxation_rate']]
653
654
            elif 'method' in params and (params['method'].endswith('cumulant')
                                         or params['method'].endswith('central_moment')):
Frederik Hennig's avatar
Frederik Hennig committed
655
                params['relaxation_rates'] = [params['relaxation_rate']]
656
            else:
Martin Bauer's avatar
Martin Bauer committed
657
658
659
660
661
                params['relaxation_rates'] = [params['relaxation_rate'],
                                              relaxation_rate_from_magic_number(params['relaxation_rate'])]

            del params['relaxation_rate']

Frederik Hennig's avatar
Frederik Hennig committed
662
663
664
665
666
667
668
669
670
671
672
673
    #   By default, do not derive moment equations for SRT and TRT methods
    if 'method' not in params or params['method'][:3].lower() in ('srt', 'trt'):
        if 'moment_transform_class' not in params:
            params['moment_transform_class'] = None

    #   Workaround until entropic method supports relaxation in subexpressions
    #   and the problem with RNGs in the assignment collection has been solved
    if (('entropic' in params and params['entropic'])
            or ('fluctuating' in params and params['fluctuating'])):
        if 'moment_transform_class' not in params:
            params['moment_transform_class'] = None

Frederik Hennig's avatar
Frederik Hennig committed
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
    #   for backwards compatibility
    if 'kernel_type' in params:
        kernel_type_to_streaming_pattern = {
            'stream_pull_collide': ('pull', Timestep.BOTH),
            'collide_stream_push': ('push', Timestep.BOTH),
            'aa_even': ('aa', Timestep.EVEN),
            'aa_odd': ('aa', Timestep.ODD),
            'esotwist_even': ('esotwist', Timestep.EVEN),
            'esotwist_odd': ('esotwist', Timestep.ODD)
        }

        kernel_type = params['kernel_type']
        if kernel_type in kernel_type_to_streaming_pattern.keys():
            params['kernel_type'] = 'default_stream_collide'
            params['streaming_pattern'], params['timestep'] = kernel_type_to_streaming_pattern[kernel_type]

Martin Bauer's avatar
Martin Bauer committed
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
    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
713
        else:
Martin Bauer's avatar
Martin Bauer committed
714
            stencil_entries = get_stencil(params_result['stencil'])
Markus Holzer's avatar
Markus Holzer committed
715
        params_result['relaxation_rates'] = sp.symbols(f"omega_:{len(stencil_entries)}")
716

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