cumulantbased.py 11 KB
Newer Older
1
from collections import OrderedDict
Martin Bauer's avatar
Martin Bauer committed
2
3
4
5

import sympy as sp

from lbmpy.cumulants import cumulant_as_function_of_raw_moments, raw_moment_as_function_of_cumulants
6
7
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
Martin Bauer's avatar
Martin Bauer committed
8
9
10
11
from lbmpy.moments import (
    MOMENT_SYMBOLS, extract_monomials, moment_matrix, monomial_to_polynomial_transformation_matrix)
from pystencils import Assignment
from pystencils.sympyextensions import fast_subs, subs_additive
12
13
14
15


class CumulantBasedLbMethod(AbstractLbMethod):

Martin Bauer's avatar
Martin Bauer committed
16
17
    def __init__(self, stencil, cumulant_to_relaxation_info_dict, conserved_quantity_computation, force_model=None):
        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
18
19
        super(CumulantBasedLbMethod, self).__init__(stencil)

Martin Bauer's avatar
Martin Bauer committed
20
21
22
        self._force_model = force_model
        self._cumulant_to_relaxation_info_dict = OrderedDict(cumulant_to_relaxation_info_dict.items())
        self._conserved_quantity_computation = conserved_quantity_computation
Martin Bauer's avatar
Martin Bauer committed
23
        self._weights = None
24
25

    @property
Martin Bauer's avatar
Martin Bauer committed
26
27
    def force_model(self):
        return self._force_model
28
29

    @property
Martin Bauer's avatar
Martin Bauer committed
30
31
    def relaxation_info_dict(self):
        return self._cumulant_to_relaxation_info_dict
32

Martin Bauer's avatar
Martin Bauer committed
33
    @property
Martin Bauer's avatar
Martin Bauer committed
34
35
    def zeroth_order_equilibrium_moment_symbol(self, ):
        return self._conserved_quantity_computation.zeroth_order_moment_symbol
Martin Bauer's avatar
Martin Bauer committed
36
37

    @property
Martin Bauer's avatar
Martin Bauer committed
38
39
    def first_order_equilibrium_moment_symbols(self, ):
        return self._conserved_quantity_computation.first_order_moment_symbols
Martin Bauer's avatar
Martin Bauer committed
40

Martin Bauer's avatar
Martin Bauer committed
41
    def set_first_moment_relaxation_rate(self, relaxation_rate):
42
        for e in MOMENT_SYMBOLS[:self.dim]:
Martin Bauer's avatar
Martin Bauer committed
43
44
            assert e in self._cumulant_to_relaxation_info_dict, \
                "First cumulants are not relaxed separately by this method"
45
        for e in MOMENT_SYMBOLS[:self.dim]:
Martin Bauer's avatar
Martin Bauer committed
46
47
48
            prev_entry = self._cumulant_to_relaxation_info_dict[e]
            new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
            self._cumulant_to_relaxation_info_dict[e] = new_entry
49
50

    @property
Martin Bauer's avatar
Martin Bauer committed
51
    def cumulants(self):
Martin Bauer's avatar
Martin Bauer committed
52
        return tuple(self._cumulant_to_relaxation_info_dict.keys())
53
54

    @property
Martin Bauer's avatar
Martin Bauer committed
55
    def cumulant_equilibrium_values(self):
Martin Bauer's avatar
Martin Bauer committed
56
        return tuple([e.equilibrium_value for e in self._cumulant_to_relaxation_info_dict.values()])
57
58

    @property
Martin Bauer's avatar
Martin Bauer committed
59
60
    def relaxation_rates(self):
        return tuple([e.relaxation_rate for e in self._cumulant_to_relaxation_info_dict.values()])
61

Martin Bauer's avatar
Martin Bauer committed
62
    @property
Martin Bauer's avatar
Martin Bauer committed
63
64
    def conserved_quantity_computation(self):
        return self._conserved_quantity_computation
Martin Bauer's avatar
Martin Bauer committed
65
66
67
68

    @property
    def weights(self):
        if self._weights is None:
Martin Bauer's avatar
Martin Bauer committed
69
            self._weights = self._compute_weights()
Martin Bauer's avatar
Martin Bauer committed
70
71
        return self._weights

Martin Bauer's avatar
Martin Bauer committed
72
    def override_weights(self, weights):
73
74
75
        assert len(weights) == len(self.stencil)
        self._weights = weights

76
77
78
79
80
81
    def _repr_html_(self):
        table = """
        <table style="border:none; width: 100%">
            <tr {nb}>
                <th {nb} >Cumulant</th>
                <th {nb} >Eq. Value </th>
82
                <th {nb} >Relaxation Rate</th>
83
84
85
86
87
            </tr>
            {content}
        </table>
        """
        content = ""
Martin Bauer's avatar
Martin Bauer committed
88
        for cumulant, (eq_value, rr) in self._cumulant_to_relaxation_info_dict.items():
89
90
91
            vals = {
                'rr': sp.latex(rr),
                'cumulant': sp.latex(cumulant),
Martin Bauer's avatar
Martin Bauer committed
92
                'eq_value': sp.latex(eq_value),
93
94
95
96
                'nb': 'style="border:none"',
            }
            content += """<tr {nb}>
                            <td {nb}>${cumulant}$</td>
Martin Bauer's avatar
Martin Bauer committed
97
                            <td {nb}>${eq_value}$</td>
98
99
100
101
                            <td {nb}>${rr}$</td>
                         </tr>\n""".format(**vals)
        return table.format(content=content, nb='style="border:none"')

Martin Bauer's avatar
Martin Bauer committed
102
103
104
105
    def get_equilibrium(self, conserved_quantity_equations=None):
        d = sp.eye(len(self.relaxation_rates))
        return self._get_collision_rule_with_relaxation_matrix(d, conserved_quantity_equations,
                                                               False, False, False, False)
Martin Bauer's avatar
Martin Bauer committed
106

Martin Bauer's avatar
Martin Bauer committed
107
108
    def get_equilibrium_terms(self):
        equilibrium = self.get_equilibrium()
Martin Bauer's avatar
Martin Bauer committed
109
        return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
110

Martin Bauer's avatar
Martin Bauer committed
111
    def get_collision_rule(self, conserved_quantity_equations=None, moment_subexpressions=False,
Martin Bauer's avatar
Martin Bauer committed
112
113
                           pre_collision_subexpressions=True, post_collision_subexpressions=False,
                           keep_rrs_symbolic=None):
Martin Bauer's avatar
Martin Bauer committed
114
115
116
117
118
119
120
121
122
123
124
125
126
127
        return self._get_collision_rule_with_relaxation_matrix(sp.diag(*self.relaxation_rates),
                                                               conserved_quantity_equations,
                                                               moment_subexpressions, pre_collision_subexpressions,
                                                               post_collision_subexpressions)

    def _compute_weights(self):
        replacements = self._conserved_quantity_computation.default_values
        eq = self.get_equilibrium()
        ac = eq.new_with_substitutions(replacements, substitute_on_lhs=False).new_without_subexpressions()
        new_main_eqs = [Assignment(e.lhs,
                                   subs_additive(e.rhs, sp.sympify(1), sum(self.pre_collision_pdf_symbols),
                                                 required_match_replacement=1.0))
                        for e in ac.main_assignments]
        ac = ac.copy(new_main_eqs)
Martin Bauer's avatar
Martin Bauer committed
128
129

        weights = []
Martin Bauer's avatar
Martin Bauer committed
130
        for eq in ac.main_assignments:
Martin Bauer's avatar
Martin Bauer committed
131
132
133
134
135
            value = eq.rhs.expand()
            assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights"
            weights.append(value)
        return weights

Martin Bauer's avatar
Martin Bauer committed
136
137
138
    def _get_collision_rule_with_relaxation_matrix(self, relaxation_matrix, conserved_quantity_equations=None,
                                                   moment_subexpressions=False, pre_collision_subexpressions=True,
                                                   post_collision_subexpressions=False, include_force_terms=True):
Martin Bauer's avatar
Martin Bauer committed
139
        def tuple_to_symbol(exp, prefix):
Martin Bauer's avatar
Martin Bauer committed
140
            dim = len(exp)
Martin Bauer's avatar
Martin Bauer committed
141
            format_string = prefix + "_" + "_".join(["%d"] * dim)
Martin Bauer's avatar
Martin Bauer committed
142
            return sp.Symbol(format_string % exp)
Martin Bauer's avatar
Martin Bauer committed
143

Martin Bauer's avatar
Martin Bauer committed
144
145
146
147
148
149
        def substitute_conserved_quantities(expressions, cqe):
            cqe = cqe.new_without_subexpressions()
            substitution_dict = {eq.rhs: eq.lhs for eq in cqe.main_assignments}
            density = cqe.main_assignments[0].lhs
            substitution_dict.update({density * eq.rhs: density * eq.lhs for eq in cqe.main_assignments[1:]})
            return [fast_subs(e, substitution_dict) for e in expressions]
Martin Bauer's avatar
Martin Bauer committed
150

Martin Bauer's avatar
Martin Bauer committed
151
        f = self.pre_collision_pdf_symbols
Martin Bauer's avatar
Martin Bauer committed
152
        if conserved_quantity_equations is None:
Martin Bauer's avatar
Martin Bauer committed
153
            conserved_quantity_equations = self._conserved_quantity_computation.equilibrium_input_equations_from_pdfs(f)
Martin Bauer's avatar
Martin Bauer committed
154

Martin Bauer's avatar
Martin Bauer committed
155
        subexpressions = conserved_quantity_equations.all_assignments
Martin Bauer's avatar
Martin Bauer committed
156

157
        # 1) Determine monomial indices, and arrange them such that the zeroth and first order indices come first
Martin Bauer's avatar
Martin Bauer committed
158
        indices = list(extract_monomials(self.cumulants, dim=len(self.stencil[0])))
Martin Bauer's avatar
Martin Bauer committed
159
160
161
162
163
        zeroth_moment_exponent = (0,) * self.dim
        first_moment_exponents = [tuple([1 if i == j else 0 for i in range(self.dim)]) for j in range(self.dim)]
        lower_order_indices = [zeroth_moment_exponent] + first_moment_exponents
        num_lower_order_indices = len(lower_order_indices)
        assert all(e in indices for e in lower_order_indices), \
Martin Bauer's avatar
Martin Bauer committed
164
            "Cumulant system does not contain relaxation rules for zeroth and first order cumulants"
Martin Bauer's avatar
Martin Bauer committed
165
166
        higher_order_indices = [e for e in indices if e not in lower_order_indices]
        indices = lower_order_indices + higher_order_indices  # reorder
Martin Bauer's avatar
Martin Bauer committed
167
168

        # 2) Transform pdfs to moments
Martin Bauer's avatar
Martin Bauer committed
169
        moment_transformation_matrix = moment_matrix(indices, self.stencil)
Martin Bauer's avatar
Martin Bauer committed
170
171
        moments = moment_transformation_matrix * sp.Matrix(f)
        moments = substitute_conserved_quantities(moments, conserved_quantity_equations)
Martin Bauer's avatar
Martin Bauer committed
172
        if moment_subexpressions:
Martin Bauer's avatar
Martin Bauer committed
173
            symbols = [tuple_to_symbol(t, "m") for t in higher_order_indices]
Martin Bauer's avatar
Martin Bauer committed
174
175
            subexpressions += [Assignment(sym, moment)
                               for sym, moment in zip(symbols, moments[num_lower_order_indices:])]
Martin Bauer's avatar
Martin Bauer committed
176
            moments = moments[:num_lower_order_indices] + symbols
Martin Bauer's avatar
Martin Bauer committed
177
178

        # 3) Transform moments to monomial cumulants
Martin Bauer's avatar
Martin Bauer committed
179
        moments_dict = {idx: m for idx, m in zip(indices, moments)}
Martin Bauer's avatar
Martin Bauer committed
180
        monomial_cumulants = [cumulant_as_function_of_raw_moments(idx, moments_dict) for idx in indices]
Martin Bauer's avatar
Martin Bauer committed
181

Martin Bauer's avatar
Martin Bauer committed
182
        if pre_collision_subexpressions:
Martin Bauer's avatar
Martin Bauer committed
183
            symbols = [tuple_to_symbol(t, "pre_c") for t in higher_order_indices]
184
            subexpressions += [Assignment(sym, c)
Martin Bauer's avatar
Martin Bauer committed
185
186
                               for sym, c in zip(symbols, monomial_cumulants[num_lower_order_indices:])]
            monomial_cumulants = monomial_cumulants[:num_lower_order_indices] + symbols
Martin Bauer's avatar
Martin Bauer committed
187
188

        # 4) Transform monomial to polynomial cumulants which are then relaxed and transformed back
Martin Bauer's avatar
Martin Bauer committed
189
        mon_to_poly = monomial_to_polynomial_transformation_matrix(indices, self.cumulants)
Martin Bauer's avatar
Martin Bauer committed
190
        poly_values = mon_to_poly * sp.Matrix(monomial_cumulants)
Martin Bauer's avatar
Martin Bauer committed
191
192
        eq_values = sp.Matrix(self.cumulant_equilibrium_values)
        collided_poly_values = poly_values + relaxation_matrix * (eq_values - poly_values)  # collision
Martin Bauer's avatar
Martin Bauer committed
193
        relaxed_monomial_cumulants = mon_to_poly.inv() * collided_poly_values
Martin Bauer's avatar
Martin Bauer committed
194

Martin Bauer's avatar
Martin Bauer committed
195
        if post_collision_subexpressions:
Martin Bauer's avatar
Martin Bauer committed
196
            symbols = [tuple_to_symbol(t, "post_c") for t in higher_order_indices]
197
            subexpressions += [Assignment(sym, c)
Martin Bauer's avatar
Martin Bauer committed
198
199
                               for sym, c in zip(symbols, relaxed_monomial_cumulants[num_lower_order_indices:])]
            relaxed_monomial_cumulants = relaxed_monomial_cumulants[:num_lower_order_indices] + symbols
Martin Bauer's avatar
Martin Bauer committed
200
201

        # 5) Transform post-collision cumulant back to moments and from there to pdfs
Martin Bauer's avatar
Martin Bauer committed
202
        cumulant_dict = {idx: value for idx, value in zip(indices, relaxed_monomial_cumulants)}
Martin Bauer's avatar
Martin Bauer committed
203
        collided_moments = [raw_moment_as_function_of_cumulants(idx, cumulant_dict) for idx in indices]
Martin Bauer's avatar
Martin Bauer committed
204
        result = moment_transformation_matrix.inv() * sp.Matrix(collided_moments)
Martin Bauer's avatar
Martin Bauer committed
205
        main_assignments = [Assignment(sym, val) for sym, val in zip(self.post_collision_pdf_symbols, result)]
Martin Bauer's avatar
Martin Bauer committed
206

Martin Bauer's avatar
Martin Bauer committed
207
        # 6) Add forcing terms
Martin Bauer's avatar
Martin Bauer committed
208
209
        if self._force_model is not None and include_force_terms:
            force_model_terms = self._force_model(self)
Martin Bauer's avatar
Martin Bauer committed
210
            force_term_symbols = sp.symbols("forceTerm_:%d" % (len(force_model_terms,)))
Martin Bauer's avatar
Martin Bauer committed
211
212
            force_subexpressions = [Assignment(sym, force_model_term)
                                    for sym, force_model_term in zip(force_term_symbols, force_model_terms)]
Martin Bauer's avatar
Martin Bauer committed
213
            subexpressions += force_subexpressions
Martin Bauer's avatar
Martin Bauer committed
214
215
            main_assignments = [Assignment(eq.lhs, eq.rhs + force_term_symbol)
                                for eq, force_term_symbol in zip(main_assignments, force_term_symbols)]
Martin Bauer's avatar
Martin Bauer committed
216

Martin Bauer's avatar
Martin Bauer committed
217
        sh = {'relaxation_rates': list(self.relaxation_rates)}
Martin Bauer's avatar
Martin Bauer committed
218
        return LbmCollisionRule(self, main_assignments, subexpressions, simplification_hints=sh)