momentbasedmethod.py 11.5 KB
Newer Older
1
from collections import OrderedDict
Martin Bauer's avatar
Martin Bauer committed
2
3
4

import sympy as sp

5
from lbmpy.maxwellian_equilibrium import get_weights
6
from lbmpy.methods.abstractlbmethod import AbstractLbMethod, LbmCollisionRule, RelaxationInfo
Martin Bauer's avatar
Martin Bauer committed
7
from lbmpy.methods.conservedquantitycomputation import AbstractConservedQuantityComputation
Martin Bauer's avatar
Martin Bauer committed
8
from lbmpy.moments import MOMENT_SYMBOLS, moment_matrix
Martin Bauer's avatar
Martin Bauer committed
9
10
from pystencils import Assignment
from pystencils.sympyextensions import subs_additive
11

Martin Bauer's avatar
Martin Bauer committed
12

Martin Bauer's avatar
Martin Bauer committed
13
class MomentBasedLbMethod(AbstractLbMethod):
Martin Bauer's avatar
Martin Bauer committed
14
    def __init__(self, stencil, moment_to_relaxation_info_dict, conserved_quantity_computation=None, force_model=None):
15
        """
16
17
18
19
20
        Moment based LBM is a class to represent the single (SRT), two (TRT) and multi relaxation time (MRT) methods.
        These methods work by transforming the pdfs into moment space using a linear transformation. In the moment
        space each component (moment) is relaxed to an equilibrium moment by a certain relaxation rate. These
        equilibrium moments can e.g. be determined by taking the equilibrium moments of the continuous Maxwellian.

21
22
23
        Args:
            stencil: see :func:`lbmpy.stencils.get_stencil`
            moment_to_relaxation_info_dict: a dictionary mapping moments in either tuple or polynomial formulation
Frederik Hennig's avatar
Frederik Hennig committed
24
25
                                            to a RelaxationInfo, which consists of the corresponding equilibrium moment
                                            and a relaxation rate
26
            conserved_quantity_computation: instance of :class:`lbmpy.methods.AbstractConservedQuantityComputation`.
Frederik Hennig's avatar
Frederik Hennig committed
27
28
                                            This determines how conserved quantities are computed, and defines
                                            the symbols used in the equilibrium moments like e.g. density and velocity
29
            force_model: force model instance, or None if no forcing terms are required
30
        """
Martin Bauer's avatar
Martin Bauer committed
31
        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
32
        super(MomentBasedLbMethod, self).__init__(stencil)
33

Martin Bauer's avatar
Martin Bauer committed
34
35
36
        self._forceModel = force_model
        self._momentToRelaxationInfoDict = OrderedDict(moment_to_relaxation_info_dict.items())
        self._conservedQuantityComputation = conserved_quantity_computation
37
        self._weights = None
38

Martin Bauer's avatar
Martin Bauer committed
39
    @property
Martin Bauer's avatar
Martin Bauer committed
40
    def force_model(self):
41
42
43
        return self._forceModel

    @property
Martin Bauer's avatar
Martin Bauer committed
44
    def relaxation_info_dict(self):
Martin Bauer's avatar
Martin Bauer committed
45
46
        return self._momentToRelaxationInfoDict

47
    @property
Martin Bauer's avatar
Martin Bauer committed
48
    def conserved_quantity_computation(self):
49
        return self._conservedQuantityComputation
50

51
52
    @property
    def moments(self):
53
54
55
        return tuple(self._momentToRelaxationInfoDict.keys())

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

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

63
    @property
Martin Bauer's avatar
Martin Bauer committed
64
65
    def zeroth_order_equilibrium_moment_symbol(self, ):
        return self._conservedQuantityComputation.zeroth_order_moment_symbol
66
67

    @property
Martin Bauer's avatar
Martin Bauer committed
68
69
    def first_order_equilibrium_moment_symbols(self, ):
        return self._conservedQuantityComputation.first_order_moment_symbols
70
71
72

    @property
    def weights(self):
73
        if self._weights is None:
Martin Bauer's avatar
Martin Bauer committed
74
            self._weights = self._compute_weights()
75
76
        return self._weights

Martin Bauer's avatar
Martin Bauer committed
77
    def override_weights(self, weights):
78
79
80
        assert len(weights) == len(self.stencil)
        self._weights = weights

Martin Bauer's avatar
Martin Bauer committed
81
82
    def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False):
        relaxation_matrix = sp.eye(len(self.relaxation_rates))
Martin Bauer's avatar
Martin Bauer committed
83
84
85
        return self._collision_rule_with_relaxation_matrix(relaxation_matrix,
                                                           conserved_quantity_equations=conserved_quantity_equations,
                                                           include_force_terms=include_force_terms)
86

Martin Bauer's avatar
Martin Bauer committed
87
88
    def get_equilibrium_terms(self):
        equilibrium = self.get_equilibrium()
Martin Bauer's avatar
Martin Bauer committed
89
        return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
90

Frederik Hennig's avatar
Frederik Hennig committed
91
    def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=True):
Markus Holzer's avatar
Markus Holzer committed
92
        d = self.relaxation_matrix
Frederik Hennig's avatar
Frederik Hennig committed
93
        relaxation_rate_sub_expressions, d = self._generate_relaxation_matrix(d, pre_simplification)
Martin Bauer's avatar
Martin Bauer committed
94
95
        ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
                                                         True, conserved_quantity_equations)
Martin Bauer's avatar
Martin Bauer committed
96
        return ac
97

Martin Bauer's avatar
Martin Bauer committed
98
    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
Martin Bauer's avatar
Martin Bauer committed
99
        e = sp.Rational(1, 1)
Martin Bauer's avatar
Martin Bauer committed
100
101
102
        prev_entry = self._momentToRelaxationInfoDict[e]
        new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
        self._momentToRelaxationInfoDict[e] = new_entry
Martin Bauer's avatar
Martin Bauer committed
103

Martin Bauer's avatar
Martin Bauer committed
104
    def set_first_moment_relaxation_rate(self, relaxation_rate):
105
106
107
        for e in MOMENT_SYMBOLS[:self.dim]:
            assert e in self._momentToRelaxationInfoDict, "First moments are not relaxed separately by this method"
        for e in MOMENT_SYMBOLS[:self.dim]:
Martin Bauer's avatar
Martin Bauer committed
108
109
110
            prev_entry = self._momentToRelaxationInfoDict[e]
            new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
            self._momentToRelaxationInfoDict[e] = new_entry
111

Markus Holzer's avatar
Markus Holzer committed
112
113
114
115
116
117
118
    def set_conserved_moments_relaxation_rate(self, relaxation_rate):
        self.set_zeroth_moment_relaxation_rate(relaxation_rate)
        self.set_first_moment_relaxation_rate(relaxation_rate)

    def set_force_model(self, force_model):
        self._forceModel = force_model

119
    @property
Martin Bauer's avatar
Martin Bauer committed
120
121
    def collision_matrix(self):
        pdfs_to_moments = self.moment_matrix
Markus Holzer's avatar
Markus Holzer committed
122
123
        d = self.relaxation_matrix
        return pdfs_to_moments.inv() * d * pdfs_to_moments
124
125

    @property
Martin Bauer's avatar
Martin Bauer committed
126
127
    def inverse_collision_matrix(self):
        pdfs_to_moments = self.moment_matrix
Markus Holzer's avatar
Markus Holzer committed
128
        inverse_relaxation_matrix = self.relaxation_matrix.inv()
Martin Bauer's avatar
Martin Bauer committed
129
        return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
130

131
    @property
Martin Bauer's avatar
Martin Bauer committed
132
133
    def moment_matrix(self):
        return moment_matrix(self.moments, self.stencil)
134

Markus Holzer's avatar
Markus Holzer committed
135
136
137
138
139
140
141
    @property
    def relaxation_matrix(self):
        d = sp.zeros(len(self.relaxation_rates))
        for i in range(0, len(self.relaxation_rates)):
            d[i, i] = self.relaxation_rates[i]
        return d

142
143
144
145
146
147
148
149
150
151
    @property
    def is_orthogonal(self):
        return (self.moment_matrix * self.moment_matrix.T).is_diagonal()

    @property
    def is_weighted_orthogonal(self):
        w = get_weights(self.stencil, sp.Rational(1, 3))
        return (sp.matrix_multiply_elementwise(self.moment_matrix, sp.Matrix([w] * len(w))) * self.moment_matrix.T
                ).is_diagonal()

152
153
154
155
156
    def __getstate__(self):
        # Workaround for a bug in joblib
        self._momentToRelaxationInfoDictToPickle = [i for i in self._momentToRelaxationInfoDict.items()]
        return self.__dict__

157
158
159
160
161
162
    def _repr_html_(self):
        table = """
        <table style="border:none; width: 100%">
            <tr {nb}>
                <th {nb} >Moment</th>
                <th {nb} >Eq. Value </th>
163
                <th {nb} >Relaxation Rate</th>
164
165
166
167
168
            </tr>
            {content}
        </table>
        """
        content = ""
Martin Bauer's avatar
Martin Bauer committed
169
        for moment, (eq_value, rr) in self._momentToRelaxationInfoDict.items():
170
171
172
            vals = {
                'rr': sp.latex(rr),
                'moment': sp.latex(moment),
Martin Bauer's avatar
Martin Bauer committed
173
                'eq_value': sp.latex(eq_value),
174
175
176
177
                'nb': 'style="border:none"',
            }
            content += """<tr {nb}>
                            <td {nb}>${moment}$</td>
Martin Bauer's avatar
Martin Bauer committed
178
                            <td {nb}>${eq_value}$</td>
179
180
181
182
                            <td {nb}>${rr}$</td>
                         </tr>\n""".format(**vals)
        return table.format(content=content, nb='style="border:none"')

Martin Bauer's avatar
Martin Bauer committed
183
184
185
186
    def _compute_weights(self):
        replacements = self._conservedQuantityComputation.default_values
        ac = self.get_equilibrium(include_force_terms=False)
        ac = ac.new_with_substitutions(replacements, substitute_on_lhs=False).new_without_subexpressions()
187

Martin Bauer's avatar
Martin Bauer committed
188
189
190
191
192
        new_assignments = [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_assignments)
193
194

        weights = []
Martin Bauer's avatar
Martin Bauer committed
195
        for eq in ac.main_assignments:
196
            value = eq.rhs.expand()
197
            assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights " + str(value)
198
199
            weights.append(value)
        return weights
200

Martin Bauer's avatar
Martin Bauer committed
201
202
    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
                                               conserved_quantity_equations=None):
Martin Bauer's avatar
Martin Bauer committed
203
204
205
        f = sp.Matrix(self.pre_collision_pdf_symbols)
        pdf_to_moment = self.moment_matrix
        m_eq = sp.Matrix(self.moment_equilibrium_values)
206

Martin Bauer's avatar
Martin Bauer committed
207
208
        collision_rule = f + pdf_to_moment.inv() * d * (m_eq - pdf_to_moment * f)
        collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(self.post_collision_pdf_symbols, collision_rule)]
209

Martin Bauer's avatar
Martin Bauer committed
210
        if conserved_quantity_equations is None:
Martin Bauer's avatar
Martin Bauer committed
211
            conserved_quantity_equations = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f)
212

Martin Bauer's avatar
Martin Bauer committed
213
214
        simplification_hints = conserved_quantity_equations.simplification_hints.copy()
        simplification_hints.update(self._conservedQuantityComputation.defined_symbols())
Martin Bauer's avatar
Martin Bauer committed
215
        simplification_hints['relaxation_rates'] = [d[i, i] for i in range(d.rows)]
216

Martin Bauer's avatar
Martin Bauer committed
217
        all_subexpressions = list(additional_subexpressions) + conserved_quantity_equations.all_assignments
218

Martin Bauer's avatar
Martin Bauer committed
219
        if self._forceModel is not None and include_force_terms:
Martin Bauer's avatar
Martin Bauer committed
220
            force_model_terms = self._forceModel(self)
Markus Holzer's avatar
Markus Holzer committed
221
            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
Martin Bauer's avatar
Martin Bauer committed
222
223
            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
224
            all_subexpressions += force_subexpressions
Martin Bauer's avatar
Martin Bauer committed
225
226
227
            collision_eqs = [Assignment(eq.lhs, eq.rhs + force_term_symbol)
                             for eq, force_term_symbol in zip(collision_eqs, force_term_symbols)]
            simplification_hints['force_terms'] = force_term_symbols
Martin Bauer's avatar
Martin Bauer committed
228
229
230

        return LbmCollisionRule(self, collision_eqs, all_subexpressions,
                                simplification_hints)
231

232
    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
233
    def _generate_relaxation_matrix(relaxation_matrix, keep_rr_symbolic):
234
235
236
        """
        For SRT and TRT the equations can be easier simplified if the relaxation times are symbols, not numbers.
        This function replaces the numbers in the relaxation matrix with symbols in this case, and returns also
Markus Holzer's avatar
Markus Holzer committed
237
        the subexpressions, that assign the number to the newly introduced symbol
238
        """
Martin Bauer's avatar
Martin Bauer committed
239
        rr = [relaxation_matrix[i, i] for i in range(relaxation_matrix.rows)]
Martin Bauer's avatar
Martin Bauer committed
240
241
        if keep_rr_symbolic <= 2:
            unique_relaxation_rates = set(rr)
242
            subexpressions = {}
Martin Bauer's avatar
Martin Bauer committed
243
            for rt in unique_relaxation_rates:
244
245
                rt = sp.sympify(rt)
                if not isinstance(rt, sp.Symbol):
Markus Holzer's avatar
Markus Holzer committed
246
                    rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
Martin Bauer's avatar
Martin Bauer committed
247
                    subexpressions[rt] = rt_symbol
248

Martin Bauer's avatar
Martin Bauer committed
249
250
            new_rr = [subexpressions[sp.sympify(e)] if sp.sympify(e) in subexpressions else e
                      for e in rr]
251
            substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
Martin Bauer's avatar
Martin Bauer committed
252
            return substitutions, sp.diag(*new_rr)
253
        else:
Martin Bauer's avatar
Martin Bauer committed
254
            return [], relaxation_matrix