moments.py 26 KB
Newer Older
Martin Bauer's avatar
Martin Bauer committed
1
# -*- coding: utf-8 -*-
2
r"""
3
4
Module Overview
~~~~~~~~~~~~~~~
5

6
This module provides functions to
7

8
9
10
11
- generate moments according to certain patterns
- compute moments of discrete probability distribution functions
- create transformation rules into moment space
- orthogonalize moment bases
12
13


14
15
Moment Representation
~~~~~~~~~~~~~~~~~~~~~
16

17
Moments can be represented in two ways:
18

19
20
21
- by an index :math:`i,j`: defined as :math:`m_{ij} := \sum_{\mathbf{d} \in stencil} <\mathbf{d}, \mathbf{x}> f_i`
- or by a polynomial in the variables x,y and z. For example the polynomial :math:`x^2 y^1 z^3 + x + 1` is
  describing the linear combination of moments: :math:`m_{213} + m_{100} + m_{000}`
22

23
24
The polynomial description is more powerful, since one polynomial can express a linear combination of single moments.
All moment polynomials have to use ``MOMENT_SYMBOLS`` (which is a module variable) as degrees of freedom.
25

26
Example ::
27

Martin Bauer's avatar
Martin Bauer committed
28
29
    >>> from lbmpy.moments import MOMENT_SYMBOLS
    >>> x, y, z = MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
30
    >>> second_order_moment = x*y + y*z
31

Martin Bauer's avatar
Martin Bauer committed
32
33
34
35

Functions
~~~~~~~~~

36
37
"""
import itertools
38
import math
Martin Bauer's avatar
Martin Bauer committed
39
from collections import Counter, defaultdict
Martin Bauer's avatar
Martin Bauer committed
40
41
42
from copy import copy
from typing import Iterable, List, Optional, Sequence, Tuple, TypeVar

43
import sympy as sp
44

45
from pystencils.cache import memorycache
Martin Bauer's avatar
Martin Bauer committed
46
from pystencils.sympyextensions import remove_higher_order_terms
47

Martin Bauer's avatar
Martin Bauer committed
48
49
MOMENT_SYMBOLS = sp.symbols('x y z')
T = TypeVar('T')
50
51


52
53
54
# ------------------------------ Discrete (Exponent Tuples) ------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
55
def moment_multiplicity(exponent_tuple):
56
57
58
59
    """
    Returns number of permutations of the given moment tuple

    Example:
Martin Bauer's avatar
Martin Bauer committed
60
    >>> moment_multiplicity((2,0,0))
61
    3
Martin Bauer's avatar
Martin Bauer committed
62
    >>> list(moment_permutations((2,0,0)))
63
64
    [(0, 0, 2), (0, 2, 0), (2, 0, 0)]
    """
Martin Bauer's avatar
Martin Bauer committed
65
66
    c = Counter(exponent_tuple)
    result = math.factorial(len(exponent_tuple))
67
68
69
70
71
    for d in c.values():
        result //= math.factorial(d)
    return result


Martin Bauer's avatar
Martin Bauer committed
72
def pick_representative_moments(moments):
73
    """Picks the representative i.e. of each permutation group only one is kept"""
Martin Bauer's avatar
Martin Bauer committed
74
    to_remove = []
75
    for m in moments:
Martin Bauer's avatar
Martin Bauer committed
76
77
78
        permutations = list(moment_permutations(m))
        to_remove += permutations[1:]
    return set(moments) - set(to_remove)
79
80


Martin Bauer's avatar
Martin Bauer committed
81
def moment_permutations(exponent_tuple):
82
    """Returns all (unique) permutations of the given tuple"""
Martin Bauer's avatar
Martin Bauer committed
83
    return __unique_permutations(exponent_tuple)
84
85


Martin Bauer's avatar
Martin Bauer committed
86
def moments_of_order(order, dim=3, include_permutations=True):
87
    """All tuples of length 'dim' which sum equals 'order'"""
Martin Bauer's avatar
Martin Bauer committed
88
    for item in __fixed_sum_tuples(dim, order, ordered=not include_permutations):
Michael Kuron's avatar
Michael Kuron committed
89
90
91
        assert(len(item) == dim)
        assert(sum(item) == order)
        yield item
92
93


Martin Bauer's avatar
Martin Bauer committed
94
def moments_up_to_order(order, dim=3, include_permutations=True):
95
    """All tuples of length 'dim' which sum is smaller than 'order' """
Martin Bauer's avatar
Martin Bauer committed
96
97
    single_moment_iterators = [moments_of_order(o, dim, include_permutations) for o in range(order + 1)]
    return tuple(itertools.chain(*single_moment_iterators))
98
99


Martin Bauer's avatar
Martin Bauer committed
100
def moments_up_to_component_order(order, dim=3):
101
    """All tuples of length 'dim' where each entry is smaller or equal to 'order' """
102
    return tuple(itertools.product(*[range(order + 1)] * dim))
103
104


Martin Bauer's avatar
Martin Bauer committed
105
def extend_moments_with_permutations(exponent_tuples):
106
    """Returns all permutations of the given exponent tuples"""
Martin Bauer's avatar
Martin Bauer committed
107
108
109
110
    all_moments = []
    for i in exponent_tuples:
        all_moments += list(moment_permutations(i))
    return __unique(all_moments)
111
112


113
114
115
# ------------------------------ Representation Conversions ------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
116
def exponent_to_polynomial_representation(exponent_tuple):
Martin Bauer's avatar
Martin Bauer committed
117
118
119
120
    """
    Converts an exponent tuple to corresponding polynomial representation

    Example:
Martin Bauer's avatar
Martin Bauer committed
121
        >>> exponent_to_polynomial_representation( (2,1,3) )
Martin Bauer's avatar
Martin Bauer committed
122
123
124
        x**2*y*z**3
    """
    poly = 1
Martin Bauer's avatar
Martin Bauer committed
125
126
    for sym, tuple_entry in zip(MOMENT_SYMBOLS[:len(exponent_tuple)], exponent_tuple):
        poly *= sym ** tuple_entry
Martin Bauer's avatar
Martin Bauer committed
127
128
129
    return poly


Martin Bauer's avatar
Martin Bauer committed
130
131
132
def exponents_to_polynomial_representations(sequence_of_exponent_tuples):
    """Applies :func:`exponent_to_polynomial_representation` to given sequence"""
    return tuple([exponent_to_polynomial_representation(t) for t in sequence_of_exponent_tuples])
Martin Bauer's avatar
Martin Bauer committed
133
134


Martin Bauer's avatar
Martin Bauer committed
135
def polynomial_to_exponent_representation(polynomial, dim=3):
Martin Bauer's avatar
Martin Bauer committed
136
137
138
139
    """
    Converts a linear combination of moments in polynomial representation into exponent representation

    :returns list of tuples where the first element is the coefficient and the second element is the exponent tuple
140

Martin Bauer's avatar
Martin Bauer committed
141
142
    Example:
        >>> x , y, z = MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
143
        >>> set(polynomial_to_exponent_representation(1 + (42 * x**2 * y**2 * z) )) == {(42, (2, 2, 1)), (1, (0, 0, 0))}
Martin Bauer's avatar
Martin Bauer committed
144
        True
Martin Bauer's avatar
Martin Bauer committed
145
146
147
148
    """
    assert dim <= 3
    x, y, z = MOMENT_SYMBOLS
    polynomial = polynomial.expand()
Martin Bauer's avatar
Martin Bauer committed
149
    coeff_exp_tuple_representation = []
150
151
152

    summands = [polynomial] if polynomial.func != sp.Add else polynomial.args
    for expr in summands:
Martin Bauer's avatar
Martin Bauer committed
153
154
        if len(expr.atoms(sp.Symbol) - set(MOMENT_SYMBOLS)) > 0:
            raise ValueError("Invalid moment polynomial: " + str(expr))
155
        c, x_exp, y_exp, z_exp = sp.Wild('c'), sp.Wild('xexp'), sp.Wild('yexp'), sp.Wild('zc')
Martin Bauer's avatar
Martin Bauer committed
156
157
158
        match_res = expr.match(c * x**x_exp * y**y_exp * z**z_exp)
        assert match_res[x_exp].is_integer and match_res[y_exp].is_integer and match_res[z_exp].is_integer
        exp_tuple = (int(match_res[x_exp]), int(match_res[y_exp]), int(match_res[z_exp]),)
Martin Bauer's avatar
Martin Bauer committed
159
160
        if dim < 3:
            for i in range(dim, 3):
Martin Bauer's avatar
Martin Bauer committed
161
162
163
164
                assert exp_tuple[i] == 0, "Used symbols in polynomial are not representable in that dimension"
            exp_tuple = exp_tuple[:dim]
        coeff_exp_tuple_representation.append((match_res[c], exp_tuple))
    return coeff_exp_tuple_representation
165

166

Martin Bauer's avatar
Martin Bauer committed
167
def moment_sort_key(moment):
168
169
    """Sort key function for moments to sort them by (in decreasing priority)
     order, number of occuring symbols, length of string representation, string representation"""
Martin Bauer's avatar
Martin Bauer committed
170
171
    mom_str = str(moment)
    return get_order(moment), len(moment.atoms(sp.Symbol)), len(mom_str), mom_str
172
173


Martin Bauer's avatar
Martin Bauer committed
174
def sort_moments_into_groups_of_same_order(moments):
175
176
177
    """Returns a dictionary mapping the order (int) to a list of moments with that order."""
    result = defaultdict(list)
    for i, moment in enumerate(moments):
Martin Bauer's avatar
Martin Bauer committed
178
        order = get_order(moment)
179
180
181
        result[order].append(moment)
    return result

182
183
184
# -------------------- Common Function working with exponent tuples and polynomial moments -----------------------------


Martin Bauer's avatar
Martin Bauer committed
185
def is_even(moment):
186
187
188
189
    """
    A moment is considered even when under sign reversal nothing changes i.e. :math:`m(-x,-y,-z) = m(x,y,z)`

    For the exponent tuple representation that means that the exponent sum is even  e.g.
Martin Bauer's avatar
Martin Bauer committed
190
        >>> x , y, z = MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
191
        >>> is_even(x**2 * y**2)
Martin Bauer's avatar
Martin Bauer committed
192
        True
Martin Bauer's avatar
Martin Bauer committed
193
        >>> is_even(x**2 * y)
Martin Bauer's avatar
Martin Bauer committed
194
        False
Martin Bauer's avatar
Martin Bauer committed
195
        >>> is_even((1,0,0))
Martin Bauer's avatar
Martin Bauer committed
196
        False
Martin Bauer's avatar
Martin Bauer committed
197
        >>> is_even(1)
Martin Bauer's avatar
Martin Bauer committed
198
        True
199
200
201
202
    """
    if type(moment) is tuple:
        return sum(moment) % 2 == 0
    else:
203
        moment = sp.sympify(moment)
204
205
206
        opposite = moment
        for s in MOMENT_SYMBOLS:
            opposite = opposite.subs(s, -s)
Martin Bauer's avatar
Martin Bauer committed
207
        return sp.expand(moment - opposite) == 0
208
209


Martin Bauer's avatar
Martin Bauer committed
210
def get_moment_indices(moment_exponent_tuple):
211
212
213
    """Returns indices for a given exponent tuple:
    
    Example:
Martin Bauer's avatar
Martin Bauer committed
214
        >>> get_moment_indices((2,1,0))
215
        [0, 0, 1]
Martin Bauer's avatar
Martin Bauer committed
216
        >>> get_moment_indices((0,0,3))
217
218
219
        [2, 2, 2]
    """
    result = []
Martin Bauer's avatar
Martin Bauer committed
220
    for i, element in enumerate(moment_exponent_tuple):
221
222
223
224
        result += [i] * element
    return result


Martin Bauer's avatar
Martin Bauer committed
225
def get_exponent_tuple_from_indices(moment_indices, dim):
226
    result = [0] * dim
Martin Bauer's avatar
Martin Bauer committed
227
    for i in moment_indices:
228
229
230
231
        result[i] += 1
    return tuple(result)


Martin Bauer's avatar
Martin Bauer committed
232
def get_order(moment):
Martin Bauer's avatar
Martin Bauer committed
233
    """Computes polynomial order of given moment.
234

235
    Examples:
Martin Bauer's avatar
Martin Bauer committed
236
        >>> x , y, z = MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
237
        >>> get_order(x**2 * y + x)
Martin Bauer's avatar
Martin Bauer committed
238
        3
Martin Bauer's avatar
Martin Bauer committed
239
        >>> get_order(z**4 * x**2)
Martin Bauer's avatar
Martin Bauer committed
240
        6
Martin Bauer's avatar
Martin Bauer committed
241
        >>> get_order((2,1,0))
Martin Bauer's avatar
Martin Bauer committed
242
        3
243
    """
Martin Bauer's avatar
Martin Bauer committed
244
245
    if isinstance(moment, tuple):
        return sum(moment)
246
247
    if len(moment.atoms(sp.Symbol)) == 0:
        return 0
Martin Bauer's avatar
Martin Bauer committed
248
249
250
251
252
    leading_coefficient = sp.polys.polytools.LM(moment)
    symbols_in_leading_coefficient = leading_coefficient.atoms(sp.Symbol)
    return sum([sp.degree(leading_coefficient, gen=m) for m in symbols_in_leading_coefficient])


Martin Bauer's avatar
Martin Bauer committed
253
def non_aliased_moment(moment_tuple: Sequence[int]) -> Tuple[int, ...]:
Martin Bauer's avatar
Martin Bauer committed
254
255
256
257
258
259
260
    """Takes a moment exponent tuple and returns the non-aliased version of it.

    For first neighborhood stencils, all moments with exponents 3, 5, 7... are equal to exponent 1,
    and exponents 4, 6, 8... are equal to exponent 2. This is because first neighborhood stencils only have values
    d ∈ {+1, 0, -1}. So for example d**5 is always the same as d**3 and d, and d**6 == d**4 == d**2

    Example:
Martin Bauer's avatar
Martin Bauer committed
261
         >>> non_aliased_moment((5, 4, 2))
Martin Bauer's avatar
Martin Bauer committed
262
         (1, 2, 2)
Martin Bauer's avatar
Martin Bauer committed
263
         >>> non_aliased_moment((9, 1, 2))
Martin Bauer's avatar
Martin Bauer committed
264
265
266
267
268
269
270
271
272
273
         (1, 1, 2)
    """
    moment = list(moment_tuple)
    result = []
    for element in moment:
        if element > 2:
            result.append(2 - (element % 2))
        else:
            result.append(element)
    return tuple(result)
274
275


Martin Bauer's avatar
Martin Bauer committed
276
def is_shear_moment(moment):
277
278
    """Shear moments in 3D are: x*y, x*z and y*z - in 2D its only x*y"""
    if type(moment) is tuple:
Martin Bauer's avatar
Martin Bauer committed
279
        moment = exponent_to_polynomial_representation(moment)
Martin Bauer's avatar
Martin Bauer committed
280
    return moment in is_shear_moment.shear_moments
Martin Bauer's avatar
Martin Bauer committed
281
282


Martin Bauer's avatar
Martin Bauer committed
283
is_shear_moment.shear_moments = set([c[0] * c[1] for c in itertools.combinations(MOMENT_SYMBOLS, 2)])
284
285


Michael Kuron's avatar
Michael Kuron committed
286
@memorycache(maxsize=512)
287
def discrete_moment(func, moment, stencil, shift_velocity=None):
288
    r"""
Martin Bauer's avatar
Martin Bauer committed
289
290
291
292
293
294
295
    Computes discrete moment of given distribution function

    .. math ::
        \sum_{d \in stencil} p(d) f_i

    where :math:`p(d)` is the moment polynomial where :math:`x, y, z` have been replaced with the components of the
    stencil direction, and :math:`f_i` is the i'th entry in the passed function sequence
296

297
298
299
300
    Args:
        func: list of distribution functions for each direction
        moment: can either be a exponent tuple, or a sympy polynomial expression
        stencil: sequence of directions
301
302
        shift_velocity: velocity vector u to compute central moments, the lattice velocity is replaced by
                        (lattice_velocity - shift_velocity)
303
    """
Martin Bauer's avatar
Martin Bauer committed
304
    assert len(stencil) == len(func)
305
306
    if shift_velocity is None:
        shift_velocity = (0,) * len(stencil[0])
307
    res = 0
Martin Bauer's avatar
Martin Bauer committed
308
    for factor, e in zip(func, stencil):
309
        if type(moment) is tuple:
310
            for vel, shift, exponent in zip(e, shift_velocity, moment):
311
                factor *= (vel - shift) ** exponent
312
313
            res += factor
        else:
314
            weight = moment
315
316
            for variable, e_i, shift in zip(MOMENT_SYMBOLS, e, shift_velocity):
                weight = weight.subs(variable, e_i - shift)
317
            res += weight * factor
318

Martin Bauer's avatar
Martin Bauer committed
319
    return res
320
321


322
def moment_matrix(moments, stencil, shift_velocity=None):
Martin Bauer's avatar
Martin Bauer committed
323
324
325
326
327
328
    """
    Returns transformation matrix to moment space

    each row corresponds to a moment, each column to a direction of the stencil
    The entry i,j is the i'th moment polynomial evaluated at direction j
    """
329
330
    if shift_velocity is None:
        shift_velocity = (0,) * len(stencil[0])
331
332
333
334

    if type(moments[0]) is tuple:
        def generator(row, column):
            result = sp.Rational(1, 1)
335
336
            for exponent, stencil_entry, shift in zip(moments[row], stencil[column], shift_velocity):
                result *= (sp.sympify(stencil_entry) - shift) ** exponent
337
338
339
340
            return result
    else:
        def generator(row, column):
            evaluated = moments[row]
341
342
            for var, stencil_entry, shift in zip(MOMENT_SYMBOLS, stencil[column], shift_velocity):
                evaluated = evaluated.subs(var, stencil_entry - shift)
343
344
345
346
347
            return evaluated

    return sp.Matrix(len(moments), len(stencil), generator)


Martin Bauer's avatar
Martin Bauer committed
348
def gram_schmidt(moments, stencil, weights=None):
349
    r"""
350
351
    Computes orthogonal set of moments using the method by Gram-Schmidt

352
353
354
355
356
357
358
359
    Args:
        moments: sequence of moments, either in tuple or polynomial form
        stencil: stencil as sequence of directions
        weights: optional weights, that define the scalar product which is used for normalization.
                 Scalar product :math:`< a,b > = \sum a_i b_i w_i` with weights :math:`w_i`.
                 Passing no weights sets all weights to 1.
    Returns:
        set of orthogonal moments in polynomial form
360
    """
361
362
363
364
365
366
367
    if weights is None:
        weights = sp.eye(len(stencil))
    if type(weights) is list:
        assert len(weights) == len(stencil)
        weights = sp.diag(*weights)

    if type(moments[0]) is tuple:
Martin Bauer's avatar
Martin Bauer committed
368
        moments = list(exponents_to_polynomial_representations(moments))
369
    else:
370
        moments = list(copy(moments))
371

Martin Bauer's avatar
Martin Bauer committed
372
373
374
375
376
    pdfs_to_moments = moment_matrix(moments, stencil).transpose()
    moment_matrix_columns = [pdfs_to_moments.col(i) for i in range(pdfs_to_moments.cols)]
    orthogonalized_vectors = []
    for i in range(len(moment_matrix_columns)):
        current_element = moment_matrix_columns[i]
377
        for j in range(i):
Martin Bauer's avatar
Martin Bauer committed
378
379
380
            prev_element = orthogonalized_vectors[j]
            denominator = prev_element.dot(weights * prev_element)
            if denominator == 0:
381
                raise ValueError("Not an independent set of vectors given: "
382
                                 "vector %d is dependent on previous vectors" % (i,))
Martin Bauer's avatar
Martin Bauer committed
383
384
            overlap = current_element.dot(weights * prev_element) / denominator
            current_element -= overlap * prev_element
385
            moments[i] -= overlap * moments[j]
Martin Bauer's avatar
Martin Bauer committed
386
        orthogonalized_vectors.append(current_element)
Martin Bauer's avatar
Martin Bauer committed
387
388

    return moments
389
390


Martin Bauer's avatar
Martin Bauer committed
391
def get_default_moment_set_for_stencil(stencil):
392
393
394
    """
    Returns a sequence of moments that are commonly used to construct a LBM equilibrium for the given stencil
    """
395
    from lbmpy.stencils import get_stencil
396
    from pystencils.stencil import have_same_entries
397

Martin Bauer's avatar
Martin Bauer committed
398
    to_poly = exponents_to_polynomial_representations
399

400
    if have_same_entries(stencil, get_stencil("D2Q9")):
Martin Bauer's avatar
Martin Bauer committed
401
        return sorted(to_poly(moments_up_to_component_order(2, dim=2)), key=moment_sort_key)
402

Martin Bauer's avatar
Martin Bauer committed
403
    all27_moments = moments_up_to_component_order(2, dim=3)
404
    if have_same_entries(stencil, get_stencil("D3Q27")):
Martin Bauer's avatar
Martin Bauer committed
405
        return to_poly(all27_moments)
406
    if have_same_entries(stencil, get_stencil("D3Q19")):
Martin Bauer's avatar
Martin Bauer committed
407
408
409
        non_matched_moments = [(1, 2, 2), (1, 1, 2), (2, 2, 2), (1, 1, 1)]
        moments19 = set(all27_moments) - set(extend_moments_with_permutations(non_matched_moments))
        return sorted(to_poly(moments19), key=moment_sort_key)
410
    if have_same_entries(stencil, get_stencil("D3Q15")):
411
        x, y, z = MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
412
413
414
415
416
417
418
419
        non_matched_moments = [(1, 2, 0), (2, 2, 0), (1, 1, 2), (1, 2, 2), (2, 2, 2)]
        additional_moments = (6 * (x ** 2 * y ** 2 + x ** 2 * z ** 2 + y ** 2 * z ** 2),
                              3 * (x * (y ** 2 + z ** 2)),
                              3 * (y * (x ** 2 + z ** 2)),
                              3 * (z * (x ** 2 + y ** 2)),
                              )
        to_remove = set(extend_moments_with_permutations(non_matched_moments))
        return sorted(to_poly(set(all27_moments) - to_remove) + additional_moments, key=moment_sort_key)
420
421
422
423

    raise NotImplementedError("No default moment system available for this stencil - define matched moments yourself")


Martin Bauer's avatar
Martin Bauer committed
424
425
426
def extract_monomials(sequence_of_polynomials, dim=3):
    """
    Returns a set of exponent tuples of all monomials contained in the given set of polynomials
427
428
429
430

    Args:
        sequence_of_polynomials: sequence of polynomials in the MOMENT_SYMBOLS
        dim: length of returned exponent tuples
Martin Bauer's avatar
Martin Bauer committed
431
432
433
434
435
436
437
438
439

    >>> x, y, z = MOMENT_SYMBOLS
    >>> extract_monomials([x**2 + y**2 + y, y + y**2])
    {(0, 2, 0), (0, 1, 0), (2, 0, 0)}
    >>> extract_monomials([x**2 + y**2 + y, y + y**2], dim=2)
    {(0, 1), (2, 0), (0, 2)}
    """
    monomials = set()
    for polynomial in sequence_of_polynomials:
Martin Bauer's avatar
Martin Bauer committed
440
441
        for factor, exponent_tuple in polynomial_to_exponent_representation(polynomial):
            monomials.add(exponent_tuple[:dim])
Martin Bauer's avatar
Martin Bauer committed
442
443
444
445
446
447
    return monomials


def monomial_to_polynomial_transformation_matrix(monomials, polynomials):
    """
    Returns a transformation matrix from a monomial to a polynomial representation
448
449
450
451

    Args:
        monomials: sequence of exponent tuples
        polynomials: sequence of polynomials in the MOMENT_SYMBOLS
Martin Bauer's avatar
Martin Bauer committed
452
453
454
455
456
457
458
459
460
461
462
463
464

    >>> x, y, z = MOMENT_SYMBOLS
    >>> polys = [7 * x**2 + 3 * x + 2 * y **2, \
                 9 * x**2 - 5 * x]
    >>> mons = list(extract_monomials(polys, dim=2))
    >>> monomial_to_polynomial_transformation_matrix(mons, polys)
    Matrix([
    [7,  3, 2],
    [9, -5, 0]])
    """
    dim = len(monomials[0])

    result = sp.zeros(len(polynomials), len(monomials))
Martin Bauer's avatar
Martin Bauer committed
465
    for polynomial_idx, polynomial in enumerate(polynomials):
Martin Bauer's avatar
Martin Bauer committed
466
467
        for factor, exponent_tuple in polynomial_to_exponent_representation(polynomial):
            exponent_tuple = exponent_tuple[:dim]
Martin Bauer's avatar
Martin Bauer committed
468
            result[polynomial_idx, monomials.index(exponent_tuple)] = factor
Martin Bauer's avatar
Martin Bauer committed
469
470
471
    return result


472
473
474
# ---------------------------------- Visualization ---------------------------------------------------------------------


Martin Bauer's avatar
Martin Bauer committed
475
def moment_equality_table(stencil, discrete_equilibrium=None, continuous_equilibrium=None,
Martin Bauer's avatar
Martin Bauer committed
476
                          max_order=4, truncate_order=3):
477
478
479
    """
    Creates a table showing which moments of a discrete stencil/equilibrium coincide with the
    corresponding continuous moments
480

481
482
483
484
485
486
487
488
489
490
    Args:
        stencil: list of stencil velocities
        discrete_equilibrium: list of sympy expr to compute discrete equilibrium for each direction, if left
                             to default the standard discrete Maxwellian equilibrium is used
        continuous_equilibrium: continuous equilibrium, if left to default, the continuous Maxwellian is used
        max_order: compare moments up to this order (number of rows in table)
        truncate_order: moments are considered equal if they match up to this order

    Returns:
        Object to display in an Jupyter notebook
491
492
    """
    import ipy_table
Martin Bauer's avatar
Martin Bauer committed
493
    from lbmpy.continuous_distribution_measures import continuous_moment
494

495
    dim = len(stencil[0])
496
    u = sp.symbols("u_:{dim}".format(dim=dim))
Martin Bauer's avatar
Martin Bauer committed
497
    if discrete_equilibrium is None:
Martin Bauer's avatar
Martin Bauer committed
498
499
500
        from lbmpy.maxwellian_equilibrium import discrete_maxwellian_equilibrium
        discrete_equilibrium = discrete_maxwellian_equilibrium(stencil, c_s_sq=sp.Rational(1, 3), compressible=True,
                                                               u=u, order=truncate_order)
Martin Bauer's avatar
Martin Bauer committed
501
    if continuous_equilibrium is None:
Martin Bauer's avatar
Martin Bauer committed
502
503
        from lbmpy.maxwellian_equilibrium import continuous_maxwellian_equilibrium
        continuous_equilibrium = continuous_maxwellian_equilibrium(dim=dim, u=u, c_s_sq=sp.Rational(1, 3))
504

505
    table = []
Martin Bauer's avatar
Martin Bauer committed
506
507
    matched_moments = 0
    non_matched_moments = 0
508

Martin Bauer's avatar
Martin Bauer committed
509
    moments_list = [list(moments_of_order(o, dim, include_permutations=False)) for o in range(max_order + 1)]
510
511

    colors = dict()
Martin Bauer's avatar
Martin Bauer committed
512
    nr_of_columns = max([len(v) for v in moments_list]) + 1
513

Martin Bauer's avatar
Martin Bauer committed
514
515
516
    header_row = [' '] * nr_of_columns
    header_row[0] = 'order'
    table.append(header_row)
517

Martin Bauer's avatar
Martin Bauer committed
518
519
    for order, moments in enumerate(moments_list):
        row = [' '] * nr_of_columns
520
        row[0] = '%d' % (order,)
Martin Bauer's avatar
Martin Bauer committed
521
        for moment, col_idx in zip(moments, range(1, len(row))):
Martin Bauer's avatar
Martin Bauer committed
522
523
524
            multiplicity = moment_multiplicity(moment)
            dm = discrete_moment(discrete_equilibrium, moment, stencil)
            cm = continuous_moment(continuous_equilibrium, moment, symbols=sp.symbols("v_0 v_1 v_2")[:dim])
525
            difference = sp.simplify(dm - cm)
Martin Bauer's avatar
Martin Bauer committed
526
527
            if truncate_order:
                difference = sp.simplify(remove_higher_order_terms(difference, symbols=u, order=truncate_order))
528
            if difference != 0:
Martin Bauer's avatar
Martin Bauer committed
529
                colors[(order + 1, col_idx)] = 'Orange'
Martin Bauer's avatar
Martin Bauer committed
530
                non_matched_moments += multiplicity
531
            else:
Martin Bauer's avatar
Martin Bauer committed
532
                colors[(order + 1, col_idx)] = 'lightGreen'
Martin Bauer's avatar
Martin Bauer committed
533
                matched_moments += multiplicity
534

Martin Bauer's avatar
Martin Bauer committed
535
            row[col_idx] = '%s  x %d' % (moment, moment_multiplicity(moment))
536
537
538

        table.append(row)

Martin Bauer's avatar
Martin Bauer committed
539
    table_display = ipy_table.make_table(table)
540
    ipy_table.set_row_style(0, color='#ddd')
Martin Bauer's avatar
Martin Bauer committed
541
542
    for cell_idx, color in colors.items():
        ipy_table.set_cell_style(cell_idx[0], cell_idx[1], color=color)
543
544

    print("Matched moments %d - non matched moments %d - total %d" %
Martin Bauer's avatar
Martin Bauer committed
545
          (matched_moments, non_matched_moments, matched_moments + non_matched_moments))
546

Martin Bauer's avatar
Martin Bauer committed
547
    return table_display
548
549


Martin Bauer's avatar
Martin Bauer committed
550
def moment_equality_table_by_stencil(name_to_stencil_dict, moments, truncate_order=3):
551
552
553
554
    """
    Creates a table for display in IPython notebooks that shows which moments agree between continuous and
    discrete equilibrium, group by stencils

555
556
557
558
559
    Args:
        name_to_stencil_dict: dict from stencil name to stencil
        moments: sequence of moments to compare - assumes that permutations have similar properties
                 so just one representative is shown labeled with its multiplicity
        truncate_order: compare up to this order
560
561
    """
    import ipy_table
Martin Bauer's avatar
Martin Bauer committed
562
563
564
    from lbmpy.maxwellian_equilibrium import discrete_maxwellian_equilibrium
    from lbmpy.maxwellian_equilibrium import continuous_maxwellian_equilibrium
    from lbmpy.continuous_distribution_measures import continuous_moment
565

Martin Bauer's avatar
Martin Bauer committed
566
    stencil_names = []
567
    stencils = []
Martin Bauer's avatar
Martin Bauer committed
568
    for key, value in name_to_stencil_dict.items():
Martin Bauer's avatar
Martin Bauer committed
569
        stencil_names.append(key)
570
571
        stencils.append(value)

Martin Bauer's avatar
Martin Bauer committed
572
    moments = list(pick_representative_moments(moments))
573
574

    colors = {}
Martin Bauer's avatar
Martin Bauer committed
575
    for stencil_idx, stencil in enumerate(stencils):
576
        dim = len(stencil[0])
577
        u = sp.symbols("u_:{dim}".format(dim=dim))
Martin Bauer's avatar
Martin Bauer committed
578
579
580
        discrete_equilibrium = discrete_maxwellian_equilibrium(stencil, c_s_sq=sp.Rational(1, 3), compressible=True,
                                                               u=u, order=truncate_order)
        continuous_equilibrium = continuous_maxwellian_equilibrium(dim=dim, u=u, c_s_sq=sp.Rational(1, 3))
581

Martin Bauer's avatar
Martin Bauer committed
582
        for moment_idx, moment in enumerate(moments):
583
            moment = moment[:dim]
Martin Bauer's avatar
Martin Bauer committed
584
585
            dm = discrete_moment(discrete_equilibrium, moment, stencil)
            cm = continuous_moment(continuous_equilibrium, moment, symbols=sp.symbols("v_0 v_1 v_2")[:dim])
586
            difference = sp.simplify(dm - cm)
Martin Bauer's avatar
Martin Bauer committed
587
588
            if truncate_order:
                difference = sp.simplify(remove_higher_order_terms(difference, symbols=u, order=truncate_order))
Martin Bauer's avatar
Martin Bauer committed
589
            colors[(moment_idx + 1, stencil_idx + 2)] = 'Orange' if difference != 0 else 'lightGreen'
590
591

    table = []
Martin Bauer's avatar
Martin Bauer committed
592
593
    header_row = [' ', '#'] + stencil_names
    table.append(header_row)
594
    for moment in moments:
Martin Bauer's avatar
Martin Bauer committed
595
        row = [str(moment), str(moment_multiplicity(moment))] + [' '] * len(stencils)
596
597
        table.append(row)

Martin Bauer's avatar
Martin Bauer committed
598
    table_display = ipy_table.make_table(table)
599
    ipy_table.set_row_style(0, color='#ddd')
Martin Bauer's avatar
Martin Bauer committed
600
601
    for cell_idx, color in colors.items():
        ipy_table.set_cell_style(cell_idx[0], cell_idx[1], color=color)
602

Martin Bauer's avatar
Martin Bauer committed
603
    return table_display
Martin Bauer's avatar
Martin Bauer committed
604
605


Martin Bauer's avatar
Martin Bauer committed
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
# --------------------------------------- Internal Functions -----------------------------------------------------------

def __unique(seq: Sequence[T]) -> List[T]:
    """Removes duplicates from a sequence in an order preserving way.

    Example:
        >>> __unique((1, 2, 3, 1, 4, 6, 3))
        [1, 2, 3, 4, 6]
    """
    seen = {}
    result = []
    for item in seq:
        if item in seen:
            continue
        seen[item] = 1
        result.append(item)
    return result


def __unique_permutations(elements: Sequence[T]) -> Iterable[T]:
    """Generates all unique permutations of the passed sequence.

    Example:
        >>> list(__unique_permutations([1, 1, 2]))
        [(1, 1, 2), (1, 2, 1), (2, 1, 1)]

    """
    if len(elements) == 1:
        yield (elements[0],)
    else:
        unique_elements = set(elements)
        for first_element in unique_elements:
            remaining_elements = list(elements)
            remaining_elements.remove(first_element)
            for sub_permutation in __unique_permutations(remaining_elements):
                yield (first_element,) + sub_permutation


def __fixed_sum_tuples(tuple_length: int, tuple_sum: int,
                       allowed_values: Optional[Sequence[int]] = None,
                       ordered: bool = False) -> Iterable[Tuple[int, ...]]:
    """Generates all possible tuples of positive integers with a fixed sum of all entries.

    Args:
        tuple_length: length of the returned tuples
        tuple_sum: summing over the entries of a yielded tuple should give this number
        allowed_values: optional sequence of positive integers that are considered as tuple entries
                        zero has to be in the set of allowed values
                        if None, all possible positive integers are allowed
        ordered: if True, only tuples are returned where the entries are descending, i.e. where the entries are ordered

    Examples:
        Generate all 2-tuples where the sum of entries is 3
        >>> list(__fixed_sum_tuples(tuple_length=2, tuple_sum=3))
        [(0, 3), (1, 2), (2, 1), (3, 0)]

        Same with ordered tuples only
        >>> list(__fixed_sum_tuples(tuple_length=2, tuple_sum=3, ordered=True))
        [(2, 1), (3, 0)]

        Restricting the allowed values, note that zero has to be in the allowed values!
        >>> list(__fixed_sum_tuples(tuple_length=3, tuple_sum=4, allowed_values=(0, 1, 3)))
        [(0, 1, 3), (0, 3, 1), (1, 0, 3), (1, 3, 0), (3, 0, 1), (3, 1, 0)]
    """
    if not allowed_values:
        allowed_values = set(range(0, tuple_sum + 1))

    assert 0 in allowed_values and all(i >= 0 for i in allowed_values)

    def recursive_helper(current_list, position, total_sum):
        new_position = position + 1
        if new_position < len(current_list):
            for i in allowed_values:
                current_list[position] = i
                new_sum = total_sum - i
                if new_sum < 0:
                    continue
                for item in recursive_helper(current_list, new_position, new_sum):
                    yield item
        else:
            if total_sum in allowed_values:
                current_list[-1] = total_sum
                if not ordered:
                    yield tuple(current_list)
                if ordered and current_list == sorted(current_list, reverse=True):
                    yield tuple(current_list)

Martin Bauer's avatar
Martin Bauer committed
693
    return recursive_helper([0] * tuple_length, 0, tuple_sum)