creationfunctions.py 41.1 KB
Newer Older
1
2
3
4
from warnings import warn
from dataclasses import dataclass
from typing import Type

Martin Bauer's avatar
Martin Bauer committed
5
6
7
8
import itertools
import operator
from collections import OrderedDict
from functools import reduce
Martin Bauer's avatar
Martin Bauer committed
9

Martin Bauer's avatar
Martin Bauer committed
10
import sympy as sp
Martin Bauer's avatar
Martin Bauer committed
11

12
13
from lbmpy.maxwellian_equilibrium import get_weights
from lbmpy.equilibrium import ContinuousHydrodynamicMaxwellian, DiscreteHydrodynamicMaxwellian
Frederik Hennig's avatar
Frederik Hennig committed
14

Markus Holzer's avatar
Markus Holzer committed
15
from lbmpy.methods.default_moment_sets import cascaded_moment_sets_literature
Markus Holzer's avatar
Markus Holzer committed
16
17

from lbmpy.methods.centeredcumulant import CenteredCumulantBasedLbMethod
Frederik Hennig's avatar
Frederik Hennig committed
18
19
from lbmpy.methods.centeredcumulant.cumulant_transform import CentralMomentsToCumulantsByGeneratingFunc

Martin Bauer's avatar
Martin Bauer committed
20
from lbmpy.methods.conservedquantitycomputation import DensityVelocityComputation
Markus Holzer's avatar
Markus Holzer committed
21

Frederik Hennig's avatar
Frederik Hennig committed
22
from lbmpy.methods.momentbased.momentbasedmethod import MomentBasedLbMethod
Markus Holzer's avatar
Markus Holzer committed
23
from lbmpy.methods.momentbased.centralmomentbasedmethod import CentralMomentBasedLbMethod
24
25
26
27
from lbmpy.moment_transforms import (
    AbstractMomentTransform, PdfsToCentralMomentsByShiftMatrix, PdfsToMomentsByChimeraTransform)
from lbmpy.moment_transforms.rawmomenttransforms import AbstractRawMomentTransform
from lbmpy.moment_transforms.centralmomenttransforms import AbstractCentralMomentTransform
Frederik Hennig's avatar
Frederik Hennig committed
28

Martin Bauer's avatar
Martin Bauer committed
29
from lbmpy.moments import (
30
    MOMENT_SYMBOLS, exponents_to_polynomial_representations,
Markus Holzer's avatar
Markus Holzer committed
31
    get_default_moment_set_for_stencil, gram_schmidt, is_even, moments_of_order,
Frederik Hennig's avatar
Frederik Hennig committed
32
    moments_up_to_component_order, sort_moments_into_groups_of_same_order,
33
    is_bulk_moment, is_shear_moment, get_order)
Frederik Hennig's avatar
Frederik Hennig committed
34

Martin Bauer's avatar
Martin Bauer committed
35
from lbmpy.relaxationrates import relaxation_rate_from_magic_number
36
from lbmpy.enums import Stencil, CollisionSpace
Markus Holzer's avatar
Markus Holzer committed
37
from lbmpy.stencils import LBStencil
Martin Bauer's avatar
Martin Bauer committed
38
from pystencils.sympyextensions import common_denominator
Martin Bauer's avatar
Martin Bauer committed
39
40


41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@dataclass
class CollisionSpaceInfo:
    """
    This class encapsulates necessary details about a method's collision space.
    """
    collision_space: CollisionSpace
    """
    The method's collision space.
    """
    raw_moment_transform_class: Type[AbstractRawMomentTransform] = None
    """
    Python class that determines how PDFs are transformed to raw moment space. If left as 'None', this parameter
    will be inferred from `collision_space`, defaulting to 
    :class:`lbmpy.moment_transforms.PdfsToMomentsByChimeraTransform`
    if `CollisionSpace.RAW_MOMENTS` is given, or `None` otherwise.
    """
    central_moment_transform_class: Type[AbstractCentralMomentTransform] = None
    """
    Python class that determines how PDFs are transformed to central moment space. If left as 'None', this parameter
    will be inferred from `collision_space`, defaulting to 
    :class:`lbmpy.moment_transforms.PdfsToCentralMomentsByShiftMatrix`
    if `CollisionSpace.CENTRAL_MOMENTS` or `CollisionSpace.CUMULANTS` is given, or `None` otherwise.
    """
    cumulant_transform_class: Type[AbstractMomentTransform] = None
    """
    Python class that determines how central moments are transformed to cumulant space. If left as 'None', this 
    parameter will be inferred from `collision_space`, defaulting to 
    :class:`lbmpy.methods.centeredcumulant.cumulant_transform.CentralMomentsToCumulantsByGeneratingFunc`
    if `CollisionSpace.CUMULANTS` is given, or `None` otherwise.
    """

    def __post_init__(self):
        if self.collision_space == CollisionSpace.RAW_MOMENTS and self.raw_moment_transform_class is None:
            self.raw_moment_transform_class = PdfsToMomentsByChimeraTransform
        if self.collision_space in (CollisionSpace.CENTRAL_MOMENTS, CollisionSpace.CUMULANTS) \
                and self.central_moment_transform_class is None:
            self.central_moment_transform_class = PdfsToCentralMomentsByShiftMatrix
        if self.collision_space == CollisionSpace.CUMULANTS and self.cumulant_transform_class is None:
            self.cumulant_transform_class = CentralMomentsToCumulantsByGeneratingFunc

Martin Bauer's avatar
Martin Bauer committed
81

82
83
84
85
86
87
88
89
90
91
def create_with_discrete_maxwellian_equilibrium(stencil, moment_to_relaxation_rate_dict,
                                                compressible=False, zero_centered=False, delta_equilibrium=False,
                                                force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
                                                **kwargs):
    r"""Creates a moment-based LBM by taking a dictionary of moments with corresponding relaxation rates.

    These moments are relaxed against the moments of the discrete Maxwellian distribution
    (see :class:`lbmpy.equilibrium.DiscreteHydrodynamicMaxwellian`).

    Internally, this method calls :func:`lbmpy.methods.create_from_equilibrium`.
Martin Bauer's avatar
Martin Bauer committed
92

Martin Bauer's avatar
Martin Bauer committed
93
    Args:
Markus Holzer's avatar
Markus Holzer committed
94
        stencil: instance of :class:`lbmpy.stencils.LBStenil`
Martin Bauer's avatar
Martin Bauer committed
95
        moment_to_relaxation_rate_dict: dict that has as many entries as the stencil. Each moment, which can be
Martin Bauer's avatar
Martin Bauer committed
96
97
                                        represented by an exponent tuple or in polynomial form
                                        (see `lbmpy.moments`), is mapped to a relaxation rate.
98
99
100
101
102
103
104
105
        compressible: If `False`, the incompressible equilibrium formulation is used to better approximate the
                      incompressible Navier-Stokes equations. Otherwise, the default textbook equilibrium is used.
        zero_centered: If `True`, the zero-centered storage format for PDFs is used, storing only their deviation from
                       the background distribution (given by the lattice weights).
        delta_equilibrium: Takes effect only if zero-centered storage is used. If `True`, the equilibrium distribution
                           is also given only as its deviation from the background distribution.
        force_model: instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces are 
                     present.
Martin Bauer's avatar
Martin Bauer committed
106
107
        equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
        c_s_sq: Speed of sound squared
108
        kwargs: See :func:`lbmpy.methods.create_from_equilibrium`
Martin Bauer's avatar
Martin Bauer committed
109
110

    Returns:
111
        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
Martin Bauer's avatar
Martin Bauer committed
112
    """
113
114
115
116
117
118
119
120
121
122
123
124
125
    cqc = DensityVelocityComputation(stencil, compressible, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
    equilibrium = DiscreteHydrodynamicMaxwellian(stencil, compressible=compressible,
                                                 deviation_only=delta_equilibrium,
                                                 order=equilibrium_order,
                                                 c_s_sq=c_s_sq)
    return create_from_equilibrium(stencil, equilibrium, cqc, moment_to_relaxation_rate_dict,
                                   zero_centered=zero_centered, force_model=force_model, **kwargs)


def create_with_continuous_maxwellian_equilibrium(stencil, moment_to_relaxation_rate_dict,
                                                  compressible=False, zero_centered=False, delta_equilibrium=False,
                                                  force_model=None, equilibrium_order=2, c_s_sq=sp.Rational(1, 3),
                                                  **kwargs):
Martin Bauer's avatar
Martin Bauer committed
126
    r"""
127
128
129
130
131
    Creates a moment-based LBM by taking a dictionary of moments with corresponding relaxation rates. 
    These moments are relaxed against the moments of the continuous Maxwellian distribution
    (see :class:`lbmpy.equilibrium.ContinuousHydrodynamicMaxwellian`).

    Internally, this method calls :func:`lbmpy.methods.create_from_equilibrium`.
132
133

    Args:
134
        stencil: instance of :class:`lbmpy.stencils.LBStenil`
135
136
137
        moment_to_relaxation_rate_dict: dict that has as many entries as the stencil. Each moment, which can be
                                        represented by an exponent tuple or in polynomial form
                                        (see `lbmpy.moments`), is mapped to a relaxation rate.
138
139
140
141
142
143
144
145
        compressible: If `False`, the incompressible equilibrium formulation is used to better approximate the
                      incompressible Navier-Stokes equations. Otherwise, the default textbook equilibrium is used.
        zero_centered: If `True`, the zero-centered storage format for PDFs is used, storing only their deviation from 
                       the background distribution (given by the lattice weights).
        delta_equilibrium: Takes effect only if zero-centered storage is used. If `True`, the equilibrium 
                           distribution is also given only as its deviation from the background distribution.
        force_model: Instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces 
                     are present.
146
147
        equilibrium_order: approximation order of macroscopic velocity :math:`\mathbf{u}` in the equilibrium
        c_s_sq: Speed of sound squared
148
        kwargs: See :func:`lbmpy.methods.create_from_equilibrium`
149
150

    Returns:
151
        Instance of a subclass of :class:`lbmpy.methods.AbstractLbMethod`.
Martin Bauer's avatar
Martin Bauer committed
152
    """
153
154
155
156
157
158
159
160
161
162
163
164
165
    cqc = DensityVelocityComputation(stencil, compressible, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
    equilibrium = ContinuousHydrodynamicMaxwellian(dim=stencil.D, compressible=compressible,
                                                   deviation_only=delta_equilibrium,
                                                   order=equilibrium_order,
                                                   c_s_sq=c_s_sq)
    return create_from_equilibrium(stencil, equilibrium, cqc, moment_to_relaxation_rate_dict,
                                   zero_centered=zero_centered, force_model=force_model, **kwargs)


def create_from_equilibrium(stencil, equilibrium, conserved_quantity_computation,
                            moment_to_relaxation_rate_dict,
                            collision_space_info=CollisionSpaceInfo(CollisionSpace.POPULATIONS),
                            zero_centered=False, force_model=None):
166
    r"""
167
168
    Creates a lattice Boltzmann method in either population, moment, or central moment space, from a given
    discrete velocity set and equilibrium distribution. 
169

170
171
172
173
    This function takes a stencil, an equilibrium distribution, an appropriate conserved quantity computation
    instance, a dictionary mapping moment polynomials to their relaxation rates, and a collision space info
    determining the desired collision space. It returns a method instance relaxing the given moments against
    their equilibrium values computed from the given distribution, in the given collision space.
Martin Bauer's avatar
Martin Bauer committed
174
175

    Args:
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        stencil: Instance of :class:`lbmpy.stencils.LBStencil`
        equilibrium: Instance of a subclass of :class:`lbmpy.equilibrium.AbstractEquilibrium`.
        conserved_quantity_computation: Instance of a subclass of 
                                        :class:`lbmpy.methods.AbstractConservedQuantityComputation`,
                                        which must provide equations for the conserved quantities used in
                                        the equilibrium.
        moment_to_relaxation_rate_dict: Dictionary mapping moment polynomials to relaxation rates.
        collision_space_info: Instance of :class:`CollisionSpaceInfo`, defining the method's desired collision space
                              and the manner of transformation to this space. Cumulant-based methods are not supported
                              yet.
        zero_centered: Whether or not the zero-centered storage format should be used. If `True`, the given equilibrium
                       must either be a deviation-only formulation, or must provide a background distribution for PDFs
                       to be centered around.
        force_model: Instance of :class:`lbmpy.forcemodels.AbstractForceModel`, or None if no external forces are
                     present.
191
    """
Martin Bauer's avatar
Martin Bauer committed
192
193
194
    if any(isinstance(moment_to_relaxation_rate_dict, t) for t in (sp.Symbol, float, int)):
        moment_to_relaxation_rate_dict = {m: moment_to_relaxation_rate_dict
                                          for m in get_default_moment_set_for_stencil(stencil)}
195

Martin Bauer's avatar
Martin Bauer committed
196
    mom_to_rr_dict = OrderedDict(moment_to_relaxation_rate_dict)
Markus Holzer's avatar
Markus Holzer committed
197
    assert len(mom_to_rr_dict) == stencil.Q, "The number of moments has to be equal to the number of stencil entries"
198

199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
    cqc = conserved_quantity_computation
    cspace = collision_space_info

    if cspace.collision_space == CollisionSpace.POPULATIONS:
        return MomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                   force_model=force_model, zero_centered=zero_centered,
                                   moment_transform_class=None)
    elif cspace.collision_space == CollisionSpace.RAW_MOMENTS:
        return MomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                   force_model=force_model, zero_centered=zero_centered,
                                   moment_transform_class=cspace.raw_moment_transform_class)
    elif cspace.collision_space == CollisionSpace.CENTRAL_MOMENTS:
        return CentralMomentBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
                                          force_model=force_model, zero_centered=zero_centered,
                                          central_moment_transform_class=cspace.central_moment_transform_class)
    elif cspace.collision_space == CollisionSpace.CUMULANTS:
        raise NotImplementedError("Creating a cumulant method from general equilibria is not supported yet.")
        # return CenteredCumulantBasedLbMethod(stencil, equilibrium, mom_to_rr_dict, conserved_quantity_computation=cqc,
        #                                      force_model=force_model, zero_centered=zero_centered)
218
219


Martin Bauer's avatar
Martin Bauer committed
220
221
222
# ------------------------------------ SRT / TRT/ MRT Creators ---------------------------------------------------------


223
def create_srt(stencil, relaxation_rate, continuous_equilibrium=True, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
224
    r"""Creates a single relaxation time (SRT) lattice Boltzmann model also known as BGK model.
Martin Bauer's avatar
Martin Bauer committed
225

226
227
228
229
230
    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in population space.

Martin Bauer's avatar
Martin Bauer committed
231
    Args:
Markus Holzer's avatar
Markus Holzer committed
232
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
Martin Bauer's avatar
Martin Bauer committed
233
234
        relaxation_rate: relaxation rate (inverse of the relaxation time)
                        usually called :math:`\omega` in LBM literature
235
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
Martin Bauer's avatar
Martin Bauer committed
236
237
238
                        used to compute the equilibrium moments

    Returns:
Frederik Hennig's avatar
Frederik Hennig committed
239
        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
Martin Bauer's avatar
Martin Bauer committed
240
    """
241
242
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
Martin Bauer's avatar
Martin Bauer committed
243
244
    moments = get_default_moment_set_for_stencil(stencil)
    rr_dict = OrderedDict([(m, relaxation_rate) for m in moments])
245
246
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
247
    else:
248
        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
249
250


Martin Bauer's avatar
Martin Bauer committed
251
def create_trt(stencil, relaxation_rate_even_moments, relaxation_rate_odd_moments,
252
               continuous_equilibrium=True, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
253
254
255
256
257
    """
    Creates a two relaxation time (TRT) lattice Boltzmann model, where even and odd moments are relaxed differently.
    In the SRT model the exact wall position of no-slip boundaries depends on the viscosity, the TRT method does not
    have this problem.

Martin Bauer's avatar
Martin Bauer committed
258
    Parameters are similar to :func:`lbmpy.methods.create_srt`, but instead of one relaxation rate there are
Martin Bauer's avatar
Martin Bauer committed
259
    two relaxation rates: one for even moments (determines viscosity) and one for odd moments.
Martin Bauer's avatar
Martin Bauer committed
260
    If unsure how to choose the odd relaxation rate, use the function :func:`lbmpy.methods.create_trt_with_magic_number`
261
262
263
264
265
266
267
268

    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in population space.

    Returns:
        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
Martin Bauer's avatar
Martin Bauer committed
269
    """
270
271
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
Martin Bauer's avatar
Martin Bauer committed
272
273
274
    moments = get_default_moment_set_for_stencil(stencil)
    rr_dict = OrderedDict([(m, relaxation_rate_even_moments if is_even(m) else relaxation_rate_odd_moments)
                           for m in moments])
275
276
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
277
    else:
278
        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
279
280


Martin Bauer's avatar
Martin Bauer committed
281
def create_trt_with_magic_number(stencil, relaxation_rate, magic_number=sp.Rational(3, 16), **kwargs):
282
    r"""
Martin Bauer's avatar
Martin Bauer committed
283
284
    Creates a two relaxation time (TRT) lattice Boltzmann method, where the relaxation time for odd moments is
    determines from the even moment relaxation time and a "magic number".
Martin Bauer's avatar
Martin Bauer committed
285
    For possible parameters see :func:`lbmpy.methods.create_trt`
Markus Holzer's avatar
Markus Holzer committed
286

287
288
289
290
291
    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in population space.

292
    Args:
Markus Holzer's avatar
Markus Holzer committed
293
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
294
295
296
297
298
299
        relaxation_rate: relaxation rate (inverse of the relaxation time)
                        usually called :math:`\omega` in LBM literature
        magic_number: magic number which is used to calculate the relxation rate for the odd moments.

    Returns:
        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
Martin Bauer's avatar
Martin Bauer committed
300
    """
Martin Bauer's avatar
Martin Bauer committed
301
302
303
    rr_odd = relaxation_rate_from_magic_number(relaxation_rate, magic_number)
    return create_trt(stencil, relaxation_rate_even_moments=relaxation_rate,
                      relaxation_rate_odd_moments=rr_odd, **kwargs)
Martin Bauer's avatar
Martin Bauer committed
304
305


306
def create_mrt_raw(stencil, relaxation_rates, continuous_equilibrium=True, **kwargs):
307
308
    r"""
    Creates a MRT method using non-orthogonalized moments.
Markus Holzer's avatar
Markus Holzer committed
309

310
311
312
313
314
    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in raw moment space.

315
    Args:
Markus Holzer's avatar
Markus Holzer committed
316
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
317
        relaxation_rates: relaxation rates (inverse of the relaxation times) for each moment
318
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
319
320
321
322
                        used to compute the equilibrium moments.
    Returns:
        :class:`lbmpy.methods.momentbased.MomentBasedLbMethod` instance
    """
323
324
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.RAW_MOMENTS))
Martin Bauer's avatar
Martin Bauer committed
325
    moments = get_default_moment_set_for_stencil(stencil)
Markus Holzer's avatar
Markus Holzer committed
326
327
    nested_moments = [(c,) for c in moments]
    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D)
328
329
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
330
    else:
331
        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
332
333


334
def create_central_moment(stencil, relaxation_rates, nested_moments=None,
335
                          continuous_equilibrium=True, **kwargs):
Markus Holzer's avatar
Markus Holzer committed
336
337
338
    r"""
    Creates moment based LB method where the collision takes place in the central moment space.

339
340
341
    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

Markus Holzer's avatar
Markus Holzer committed
342
    Args:
Markus Holzer's avatar
Markus Holzer committed
343
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
Markus Holzer's avatar
Markus Holzer committed
344
        relaxation_rates: relaxation rates (inverse of the relaxation times) for each moment
Markus Holzer's avatar
Markus Holzer committed
345
        nested_moments: a list of lists of modes, grouped by common relaxation times.
346
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
Markus Holzer's avatar
Markus Holzer committed
347
348
349
350
                        used to compute the equilibrium moments.
    Returns:
        :class:`lbmpy.methods.momentbased.CentralMomentBasedLbMethod` instance
    """
351
352
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.CENTRAL_MOMENTS))
353
354
355
    if nested_moments and not isinstance(nested_moments[0], list):
        nested_moments = list(sort_moments_into_groups_of_same_order(nested_moments).values())
        second_order_moments = nested_moments[2]
Markus Holzer's avatar
Markus Holzer committed
356
357
        bulk_moment = [m for m in second_order_moments if is_bulk_moment(m, stencil.D)]
        shear_moments = [m for m in second_order_moments if is_shear_moment(m, stencil.D)]
358
359
360
361
362
363
364
        assert len(shear_moments) + len(bulk_moment) == len(second_order_moments)
        nested_moments[2] = shear_moments
        nested_moments.insert(3, bulk_moment)

    if not nested_moments:
        nested_moments = cascaded_moment_sets_literature(stencil)

Markus Holzer's avatar
Markus Holzer committed
365
    rr_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D)
Markus Holzer's avatar
Markus Holzer committed
366

367
368
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Markus Holzer's avatar
Markus Holzer committed
369
    else:
370
        return create_with_discrete_maxwellian_equilibrium(stencil, rr_dict, **kwargs)
Markus Holzer's avatar
Markus Holzer committed
371
372


Martin Bauer's avatar
Martin Bauer committed
373
def create_trt_kbc(dim, shear_relaxation_rate, higher_order_relaxation_rate, method_name='KBC-N4',
374
                   continuous_equilibrium=True, **kwargs):
375
376
377
    """
    Creates a method with two relaxation rates, one for lower order moments which determines the viscosity and
    one for higher order moments. In entropic models this second relaxation rate is chosen subject to an entropy
Martin Bauer's avatar
Martin Bauer committed
378
    condition. Which moments are relaxed by which rate is determined by the method_name
379

380
381
382
383
384
    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in population space.

385
386
387
388
389
390
    Args:
        dim: 2 or 3, leads to stencil D2Q9 or D3Q27
        shear_relaxation_rate: relaxation rate that determines viscosity
        higher_order_relaxation_rate: relaxation rate for higher order moments
        method_name: string 'KBC-Nx' where x can be an number from 1 to 4, for details see
                    "Karlin 2015: Entropic multi relaxation lattice Boltzmann models for turbulent flows"
391
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
392
                        used to compute the equilibrium moments.
393
    """
394
395
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.POPULATIONS))
Markus Holzer's avatar
Markus Holzer committed
396

397
398
399
    def product(iterable):
        return reduce(operator.mul, iterable, 1)

Martin Bauer's avatar
Martin Bauer committed
400
    the_moment = MOMENT_SYMBOLS[:dim]
401
402

    rho = [sp.Rational(1, 1)]
Martin Bauer's avatar
Martin Bauer committed
403
    velocity = list(the_moment)
404

Martin Bauer's avatar
Martin Bauer committed
405
406
407
408
    shear_tensor_off_diagonal = [product(t) for t in itertools.combinations(the_moment, 2)]
    shear_tensor_diagonal = [m_i * m_i for m_i in the_moment]
    shear_tensor_trace = sum(shear_tensor_diagonal)
    shear_tensor_trace_free_diagonal = [dim * d - shear_tensor_trace for d in shear_tensor_diagonal]
409

Martin Bauer's avatar
Martin Bauer committed
410
411
412
    poly_repr = exponents_to_polynomial_representations
    energy_transport_tensor = list(poly_repr([a for a in moments_of_order(3, dim, True)
                                              if 3 not in a]))
413

Martin Bauer's avatar
Martin Bauer committed
414
415
    explicitly_defined = set(rho + velocity + shear_tensor_off_diagonal
                             + shear_tensor_diagonal + energy_transport_tensor)
Martin Bauer's avatar
Martin Bauer committed
416
    rest = list(set(poly_repr(moments_up_to_component_order(2, dim))) - explicitly_defined)
Markus Holzer's avatar
Markus Holzer committed
417
    assert len(rest) + len(explicitly_defined) == 3 ** dim
418

Martin Bauer's avatar
Martin Bauer committed
419
    # naming according to paper Karlin 2015: Entropic multi relaxation lattice Boltzmann models for turbulent flows
Martin Bauer's avatar
Martin Bauer committed
420
421
422
423
    d = shear_tensor_off_diagonal + shear_tensor_trace_free_diagonal[:-1]
    t = [shear_tensor_trace]
    q = energy_transport_tensor
    if method_name == 'KBC-N1':
Martin Bauer's avatar
Martin Bauer committed
424
        decomposition = [d, t + q + rest]
Martin Bauer's avatar
Martin Bauer committed
425
    elif method_name == 'KBC-N2':
Martin Bauer's avatar
Martin Bauer committed
426
        decomposition = [d + t, q + rest]
Martin Bauer's avatar
Martin Bauer committed
427
    elif method_name == 'KBC-N3':
Martin Bauer's avatar
Martin Bauer committed
428
        decomposition = [d + q, t + rest]
Martin Bauer's avatar
Martin Bauer committed
429
    elif method_name == 'KBC-N4':
Martin Bauer's avatar
Martin Bauer committed
430
        decomposition = [d + t + q, rest]
431
432
433
    else:
        raise ValueError("Unknown model. Supported models KBC-Nx where x in (1,2,3,4)")

Martin Bauer's avatar
Martin Bauer committed
434
435
    omega_s, omega_h = shear_relaxation_rate, higher_order_relaxation_rate
    shear_part, rest_part = decomposition
436

Martin Bauer's avatar
Martin Bauer committed
437
438
439
440
    relaxation_rates = [omega_s] + \
                       [omega_s] * len(velocity) + \
                       [omega_s] * len(shear_part) + \
                       [omega_h] * len(rest_part)
441

Markus Holzer's avatar
Markus Holzer committed
442
    stencil = LBStencil(Stencil.D2Q9) if dim == 2 else LBStencil(Stencil.D3Q27)
Martin Bauer's avatar
Martin Bauer committed
443
444
    all_moments = rho + velocity + shear_part + rest_part
    moment_to_rr = OrderedDict((m, rr) for m, rr in zip(all_moments, relaxation_rates))
445

446
447
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil, moment_to_rr, **kwargs)
448
    else:
449
        return create_with_discrete_maxwellian_equilibrium(stencil, moment_to_rr, **kwargs)
450
451


452
def create_mrt_orthogonal(stencil, relaxation_rates, continuous_equilibrium=True, weighted=None,
453
                          nested_moments=None, **kwargs):
Martin Bauer's avatar
Martin Bauer committed
454
    r"""
455
    Returns an orthogonal multi-relaxation time model for the stencils D2Q9, D3Q15, D3Q19 and D3Q27.
Martin Bauer's avatar
Martin Bauer committed
456
457
    These MRT methods are just one specific version - there are many MRT methods possible for all these stencils
    which differ by the linear combination of moments and the grouping into equal relaxation times.
458
459
460
461
462
463
    To create a generic MRT method use `create_with_discrete_maxwellian_equilibrium`

    Internally calls either :func:`create_with_discrete_maxwellian_equilibrium`
    or :func:`create_with_continuous_maxwellian_equilibrium`, depending on the value of `continuous_equilibrium`.

    If not specified otherwise, collision equations will be derived in raw moment space.
Martin Bauer's avatar
Martin Bauer committed
464

Martin Bauer's avatar
Martin Bauer committed
465
    Args:
Markus Holzer's avatar
Markus Holzer committed
466
467
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
        relaxation_rates: relaxation rates for the moments
468
        continuous_equilibrium: determines if the discrete or continuous maxwellian equilibrium is
469
                                               used to compute the equilibrium moments
470
        weighted: whether to use weighted or unweighted orthogonality
Markus Holzer's avatar
Markus Holzer committed
471
472
473
        nested_moments: a list of lists of modes, grouped by common relaxation times. If this argument is not provided,
                        Gram-Schmidt orthogonalization of the default modes is performed. The default modes equal the
                        raw moments except for the separation of the shear and bulk viscosity.
Martin Bauer's avatar
Martin Bauer committed
474
    """
475
476
477
    continuous_equilibrium = _deprecate_maxwellian_moments(continuous_equilibrium, kwargs)
    kwargs.setdefault('collision_space_info', CollisionSpaceInfo(CollisionSpace.RAW_MOMENTS))

Markus Holzer's avatar
Markus Holzer committed
478
    if weighted:
479
480
481
        weights = get_weights(stencil, sp.Rational(1, 3))
    else:
        weights = None
Martin Bauer's avatar
Martin Bauer committed
482

483
    if not nested_moments:
Martin Bauer's avatar
Martin Bauer committed
484
        moments = get_default_moment_set_for_stencil(stencil)
485
        x, y, z = MOMENT_SYMBOLS
Markus Holzer's avatar
Markus Holzer committed
486
        if stencil.D == 2:
487
488
489
            diagonal_viscous_moments = [x ** 2 + y ** 2, x ** 2]
        else:
            diagonal_viscous_moments = [x ** 2 + y ** 2 + z ** 2, x ** 2, y ** 2 - z ** 2]
Markus Holzer's avatar
Markus Holzer committed
490

Markus Holzer's avatar
Markus Holzer committed
491
        for i, d in enumerate(MOMENT_SYMBOLS[:stencil.D]):
Markus Holzer's avatar
Markus Holzer committed
492
493
            if d ** 2 in moments:
                moments[moments.index(d ** 2)] = diagonal_viscous_moments[i]
494
        orthogonal_moments = gram_schmidt(moments, stencil, weights)
Martin Bauer's avatar
Martin Bauer committed
495
496
        orthogonal_moments_scaled = [e * common_denominator(e) for e in orthogonal_moments]
        nested_moments = list(sort_moments_into_groups_of_same_order(orthogonal_moments_scaled).values())
Martin Bauer's avatar
Martin Bauer committed
497
498
        # second order moments: separate bulk from shear
        second_order_moments = nested_moments[2]
Markus Holzer's avatar
Markus Holzer committed
499
500
        bulk_moment = [m for m in second_order_moments if is_bulk_moment(m, stencil.D)]
        shear_moments = [m for m in second_order_moments if is_shear_moment(m, stencil.D)]
Martin Bauer's avatar
Martin Bauer committed
501
502
503
        assert len(shear_moments) + len(bulk_moment) == len(second_order_moments)
        nested_moments[2] = shear_moments
        nested_moments.insert(3, bulk_moment)
Markus Holzer's avatar
Markus Holzer committed
504
505

    moment_to_relaxation_rate_dict = _get_relaxation_info_dict(relaxation_rates, nested_moments, stencil.D)
506

507
508
509
    if continuous_equilibrium:
        return create_with_continuous_maxwellian_equilibrium(stencil,
                                                             moment_to_relaxation_rate_dict, **kwargs)
510
    else:
511
512
        return create_with_discrete_maxwellian_equilibrium(stencil,
                                                           moment_to_relaxation_rate_dict, **kwargs)
513
514


Frederik Hennig's avatar
Frederik Hennig committed
515
516
517
518
# ----------------------------------------- Cumulant method creators ---------------------------------------------------


def create_centered_cumulant_model(stencil, cumulant_to_rr_dict, force_model=None,
519
520
                                   zero_centered=True,
                                   c_s_sq=sp.Rational(1, 3),
Frederik Hennig's avatar
Frederik Hennig committed
521
                                   galilean_correction=False,
522
                                   collision_space_info=CollisionSpaceInfo(CollisionSpace.CUMULANTS)):
Frederik Hennig's avatar
Frederik Hennig committed
523
524
525
    r"""Creates a cumulant lattice Boltzmann model.

    Args:
Markus Holzer's avatar
Markus Holzer committed
526
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
Frederik Hennig's avatar
Frederik Hennig committed
527
        cumulant_to_rr_dict: dict that has as many entries as the stencil. Each cumulant, which can be
Markus Holzer's avatar
Markus Holzer committed
528
                             represented by an exponent tuple or in polynomial form is mapped to a relaxation rate.
Markus Holzer's avatar
Markus Holzer committed
529
                             See :func:`lbmpy.methods.default_moment_sets.cascaded_moment_sets_literature`
Frederik Hennig's avatar
Frederik Hennig committed
530
        force_model: force model used for the collision. For cumulant LB method a good choice is
Markus Holzer's avatar
Markus Holzer committed
531
                     `lbmpy.methods.centeredcumulant.CenteredCumulantForceModel`
532
533
        zero_centered: If `True`, the zero-centered storage format for PDFs is used, storing only their deviation from 
                       the background distribution (given by the lattice weights).
Frederik Hennig's avatar
Frederik Hennig committed
534
535
536
537
        c_s_sq: Speed of sound squared
        galilean_correction: special correction for D3Q27 cumulant collisions. See Appendix H in
                             :cite:`geier2015`. Implemented in :mod:`lbmpy.methods.centeredcumulant.galilean_correction`
        central_moment_transform_class: Class which defines the transformation to the central moment space
Frederik Hennig's avatar
Frederik Hennig committed
538
                                        (see :mod:`lbmpy.moment_transforms`)
Frederik Hennig's avatar
Frederik Hennig committed
539
540
541
542
543
544
545
        cumulant_transform_class: Class which defines the transformation from the central moment space to the
                                  cumulant space (see :mod:`lbmpy.methods.centeredcumulant.cumulant_transform`)

    Returns:
        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
    """

Markus Holzer's avatar
Markus Holzer committed
546
547
    assert len(cumulant_to_rr_dict) == stencil.Q, \
        "The number of moments has to be equal to the number of stencil entries"
548
    assert collision_space_info.collision_space == CollisionSpace.CUMULANTS
Frederik Hennig's avatar
Frederik Hennig committed
549
550

    # CQC
551
    cqc = DensityVelocityComputation(stencil, True, zero_centered, force_model=force_model, c_s_sq=c_s_sq)
Frederik Hennig's avatar
Frederik Hennig committed
552

553
554
555
556
    equilibrium = ContinuousHydrodynamicMaxwellian(dim=stencil.D, compressible=True,
                                                   deviation_only=False,
                                                   order=None,
                                                   c_s_sq=c_s_sq)
Frederik Hennig's avatar
Frederik Hennig committed
557

558
559
560
561
    cspace = collision_space_info
    return CenteredCumulantBasedLbMethod(stencil, equilibrium, cumulant_to_rr_dict,
                                         conserved_quantity_computation=cqc, force_model=force_model,
                                         zero_centered=zero_centered,
Frederik Hennig's avatar
Frederik Hennig committed
562
                                         galilean_correction=galilean_correction,
563
564
                                         central_moment_transform_class=cspace.central_moment_transform_class,
                                         cumulant_transform_class=cspace.cumulant_transform_class)
Frederik Hennig's avatar
Frederik Hennig committed
565
566
567
568
569
570


def create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs):
    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.

    Args:
Markus Holzer's avatar
Markus Holzer committed
571
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
Frederik Hennig's avatar
Frederik Hennig committed
572
573
574
575
576
577
578
579
580
581
582
583
        relaxation_rates: relaxation rates for each cumulant group. If None are provided a list of symbolic relaxation
                          rates is created and used. If only a list with one entry is provided this relaxation rate is
                          used for determine the viscosity of the simulation. All other cumulants are relaxed with one.
                          If a cumulant force model is provided the first order cumulants are relaxed with two to ensure
                          that the force is applied correctly to the momentum groups
        cumulant_groups: Nested sequence of polynomial expressions defining the cumulants to be relaxed. All cumulants 
                         within one group are relaxed with the same relaxation rate.
        kwargs: See :func:`create_centered_cumulant_model`

    Returns:
        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
    """
Markus Holzer's avatar
Markus Holzer committed
584
    cumulant_to_rr_dict = _get_relaxation_info_dict(relaxation_rates, cumulant_groups, stencil.D)
Frederik Hennig's avatar
Frederik Hennig committed
585
586
587
588
589
590
591
    return create_centered_cumulant_model(stencil, cumulant_to_rr_dict, **kwargs)


def create_with_monomial_cumulants(stencil, relaxation_rates, **kwargs):
    r"""Creates a cumulant lattice Boltzmann model based on a default polinomial set.

    Args:
Markus Holzer's avatar
Markus Holzer committed
592
        stencil: instance of :class:`lbmpy.stencils.LBStencil`
Frederik Hennig's avatar
Frederik Hennig committed
593
594
595
596
597
598
599
600
601
602
603
604
605
606
        relaxation_rates: relaxation rates for each cumulant group. If None are provided a list of symbolic relaxation
                          rates is created and used. If only a list with one entry is provided this relaxation rate is
                          used for determine the viscosity of the simulation. All other cumulants are relaxed with one.
                          If a cumulant force model is provided the first order cumulants are relaxed with two to ensure
                          that the force is applied correctly to the momentum groups
        kwargs: See :func:`create_centered_cumulant_model`

    Returns:
        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
    """
    # Get monomial moments
    cumulants = get_default_moment_set_for_stencil(stencil)
    cumulant_groups = [(c,) for c in cumulants]

Markus Holzer's avatar
Markus Holzer committed
607
    return create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs)
Frederik Hennig's avatar
Frederik Hennig committed
608
609
610
611
612
613
614
615
616
617
618


def create_with_default_polynomial_cumulants(stencil, relaxation_rates, **kwargs):
    r"""Creates a cumulant lattice Boltzmann model based on a default polynomial set.

    Args: See :func:`create_with_polynomial_cumulants`.

    Returns:
        :class:`lbmpy.methods.centeredcumulant.CenteredCumulantBasedLbMethod` instance
    """
    # Get polynomial groups
619
    cumulant_groups = cascaded_moment_sets_literature(stencil)
Markus Holzer's avatar
Markus Holzer committed
620
    return create_with_polynomial_cumulants(stencil, relaxation_rates, cumulant_groups, **kwargs)
Frederik Hennig's avatar
Frederik Hennig committed
621
622


Markus Holzer's avatar
Markus Holzer committed
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
def _get_relaxation_info_dict(relaxation_rates, nested_moments, dim):
    r"""Creates a dictionary where each moment is mapped to a relaxation rate.

    Args:
        relaxation_rates: list of relaxation rates which should be used. This can also be a function which
                          takes a moment group in the list of nested moments and returns a list of relaxation rates.
                          This list has to have the length of the moment group and is then used for the moments
                          in the moment group.
        nested_moments: list of lists containing the moments.
        dim: dimension
    """
    result = OrderedDict()

    if callable(relaxation_rates):
        for group in nested_moments:
            rr = iter(relaxation_rates(group))
            for moment in group:
                result[moment] = next(rr)
Frederik Hennig's avatar
Frederik Hennig committed
641

Markus Holzer's avatar
Markus Holzer committed
642
        return result
Martin Bauer's avatar
Martin Bauer committed
643

Markus Holzer's avatar
Markus Holzer committed
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
    number_of_moments = 0
    shear_moments = 0
    bulk_moments = 0

    for group in nested_moments:
        for moment in group:
            number_of_moments += 1
            if is_shear_moment(moment, dim):
                shear_moments += 1
            if is_bulk_moment(moment, dim):
                bulk_moments += 1

    # if only one relaxation rate is specified it is used as the shear relaxation rate
    if len(relaxation_rates) == 1:
        for group in nested_moments:
            for moment in group:
                if get_order(moment) <= 1:
Markus Holzer's avatar
Markus Holzer committed
661
                    result[moment] = 0.0
Markus Holzer's avatar
Markus Holzer committed
662
663
664
                elif is_shear_moment(moment, dim):
                    result[moment] = relaxation_rates[0]
                else:
Markus Holzer's avatar
Markus Holzer committed
665
                    result[moment] = 1.0
Markus Holzer's avatar
Markus Holzer committed
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689

    # if relaxation rate for each moment is specified they are all used
    if len(relaxation_rates) == number_of_moments:
        rr_iter = iter(relaxation_rates)
        for group in nested_moments:
            for moment in group:
                rr = next(rr_iter)
                result[moment] = rr

    # Fallback case, relaxes each group with the same relaxation rate and separates shear and bulk moments
    next_rr = True
    if len(relaxation_rates) != 1 and len(relaxation_rates) != number_of_moments:
        try:
            rr_iter = iter(relaxation_rates)
            if shear_moments > 0:
                shear_rr = next(rr_iter)
            if bulk_moments > 0:
                bulk_rr = next(rr_iter)
            for group in nested_moments:
                if next_rr:
                    rr = next(rr_iter)
                next_rr = False
                for moment in group:
                    if get_order(moment) <= 1:
Markus Holzer's avatar
Markus Holzer committed
690
                        result[moment] = 0.0
Markus Holzer's avatar
Markus Holzer committed
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
                    elif is_shear_moment(moment, dim):
                        result[moment] = shear_rr
                    elif is_bulk_moment(moment, dim):
                        result[moment] = bulk_rr
                    else:
                        next_rr = True
                        result[moment] = rr
        except StopIteration:
            raise ValueError("Not enough relaxation rates are specified. You can either specify one relaxation rate, "
                             "which is used as the shear relaxation rate. In this case, conserved moments are "
                             "relaxed with 0, and higher-order moments are relaxed with 1. Another "
                             "possibility is to specify a relaxation rate for shear, bulk and one for each order "
                             "moment. In this case, conserved moments are also "
                             "relaxed with 0. The last possibility is to specify a relaxation rate for each moment, "
                             "including conserved moments")
    return result
Martin Bauer's avatar
Martin Bauer committed
707
708
709
# ----------------------------------------- Comparison view for notebooks ----------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
710
def compare_moment_based_lb_methods(reference, other, show_deviations_only=False):
Martin Bauer's avatar
Martin Bauer committed
711
712
    import ipy_table
    table = []
Martin Bauer's avatar
Martin Bauer committed
713
    caption_rows = [len(table)]
Martin Bauer's avatar
Martin Bauer committed
714
715
    table.append(['Shared Moment', 'ref', 'other', 'difference'])

Martin Bauer's avatar
Martin Bauer committed
716
717
718
    reference_moments = set(reference.moments)
    other_moments = set(other.moments)
    for moment in reference_moments.intersection(other_moments):
Martin Bauer's avatar
Martin Bauer committed
719
720
        reference_value = reference.relaxation_info_dict[moment].equilibrium_value
        other_value = other.relaxation_info_dict[moment].equilibrium_value
Martin Bauer's avatar
Martin Bauer committed
721
722
        diff = sp.simplify(reference_value - other_value)
        if show_deviations_only and diff == 0:
Martin Bauer's avatar
Martin Bauer committed
723
724
            pass
        else:
Markus Holzer's avatar
Markus Holzer committed
725
726
727
728
            table.append([f"${sp.latex(moment)}$",
                          f"${sp.latex(reference_value)}$",
                          f"${sp.latex(other_value)}$",
                          f"${sp.latex(diff)}$"])
Martin Bauer's avatar
Martin Bauer committed
729

Martin Bauer's avatar
Martin Bauer committed
730
731
732
    only_in_ref = reference_moments - other_moments
    if only_in_ref:
        caption_rows.append(len(table))
Martin Bauer's avatar
Martin Bauer committed
733
        table.append(['Only in Ref', 'value', '', ''])
Martin Bauer's avatar
Martin Bauer committed
734
        for moment in only_in_ref:
Martin Bauer's avatar
Martin Bauer committed
735
            val = reference.relaxation_info_dict[moment].equilibrium_value
Markus Holzer's avatar
Markus Holzer committed
736
737
            table.append([f"${sp.latex(moment)}$",
                          f"${sp.latex(val)}$",
Martin Bauer's avatar
Martin Bauer committed
738
739
                          " ", " "])

Martin Bauer's avatar
Martin Bauer committed
740
741
742
    only_in_other = other_moments - reference_moments
    if only_in_other:
        caption_rows.append(len(table))
Martin Bauer's avatar
Martin Bauer committed
743
        table.append(['Only in Other', '', 'value', ''])
Martin Bauer's avatar
Martin Bauer committed
744
        for moment in only_in_other:
Martin Bauer's avatar
Martin Bauer committed
745
            val = other.relaxation_info_dict[moment].equilibrium_value
Markus Holzer's avatar
Markus Holzer committed
746
            table.append([f"${sp.latex(moment)}$",
Martin Bauer's avatar
Martin Bauer committed
747
                          " ",
Markus Holzer's avatar
Markus Holzer committed
748
                          f"${sp.latex(val)}$",
Martin Bauer's avatar
Martin Bauer committed
749
750
                          " "])

Martin Bauer's avatar
Martin Bauer committed
751
752
    table_display = ipy_table.make_table(table)
    for row_idx in caption_rows:
Martin Bauer's avatar
Martin Bauer committed
753
        for col in range(4):
Martin Bauer's avatar
Martin Bauer committed
754
755
            ipy_table.set_cell_style(row_idx, col, color='#bbbbbb')
    return table_display
756
757
758
759
760
761
762
763
764
765
766
767

# ----------------------------------------- Deprecation of Maxwellian Moments -----------------------------------------


def _deprecate_maxwellian_moments(continuous_equilibrium, kwargs):
    if 'maxwellian_moments' in kwargs:
        warn("Argument 'maxwellian_moments' is deprecated and will be removed by version 0.5."
             "Use `continuous_equilibrium` instead.",
             stacklevel=2)
        return kwargs['maxwellian_moments']
    else:
        return continuous_equilibrium