momentbasedmethod.py 15.2 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
from pystencils.sympyextensions import subs_additive
Frederik Hennig's avatar
Frederik Hennig committed
10
11
12
from pystencils import Assignment, AssignmentCollection

from lbmpy.moment_transforms import PdfsToMomentsByChimeraTransform
13

Martin Bauer's avatar
Martin Bauer committed
14

Martin Bauer's avatar
Martin Bauer committed
15
class MomentBasedLbMethod(AbstractLbMethod):
Frederik Hennig's avatar
Frederik Hennig committed
16
17
    def __init__(self, stencil, moment_to_relaxation_info_dict, conserved_quantity_computation=None, force_model=None,
                 moment_transform_class=PdfsToMomentsByChimeraTransform):
18
        """
19
20
21
22
23
        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.

24
25
26
        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
27
28
                                            to a RelaxationInfo, which consists of the corresponding equilibrium moment
                                            and a relaxation rate
29
            conserved_quantity_computation: instance of :class:`lbmpy.methods.AbstractConservedQuantityComputation`.
Frederik Hennig's avatar
Frederik Hennig committed
30
31
                                            This determines how conserved quantities are computed, and defines
                                            the symbols used in the equilibrium moments like e.g. density and velocity
32
            force_model: force model instance, or None if no forcing terms are required
Markus Holzer's avatar
Markus Holzer committed
33
            moment_transform_class: transformation class to transform PDFs to the moment space.
34
        """
Martin Bauer's avatar
Martin Bauer committed
35
        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
36
        super(MomentBasedLbMethod, self).__init__(stencil)
37

Martin Bauer's avatar
Martin Bauer committed
38
39
40
        self._forceModel = force_model
        self._momentToRelaxationInfoDict = OrderedDict(moment_to_relaxation_info_dict.items())
        self._conservedQuantityComputation = conserved_quantity_computation
41
        self._weights = None
Frederik Hennig's avatar
Frederik Hennig committed
42
        self._moment_transform_class = moment_transform_class
43

Martin Bauer's avatar
Martin Bauer committed
44
    @property
Martin Bauer's avatar
Martin Bauer committed
45
    def force_model(self):
46
47
        return self._forceModel

Frederik Hennig's avatar
Frederik Hennig committed
48
49
50
51
52
    @property
    def moment_space_collision(self):
        """Returns whether collision is derived in terms of moments or in terms of populations only."""
        return (self._moment_transform_class is not None)

53
    @property
Martin Bauer's avatar
Martin Bauer committed
54
    def relaxation_info_dict(self):
Martin Bauer's avatar
Martin Bauer committed
55
56
        return self._momentToRelaxationInfoDict

57
    @property
Martin Bauer's avatar
Martin Bauer committed
58
    def conserved_quantity_computation(self):
59
        return self._conservedQuantityComputation
60

61
62
    @property
    def moments(self):
63
64
65
        return tuple(self._momentToRelaxationInfoDict.keys())

    @property
Martin Bauer's avatar
Martin Bauer committed
66
    def moment_equilibrium_values(self):
Martin Bauer's avatar
Martin Bauer committed
67
        return tuple([e.equilibrium_value for e in self._momentToRelaxationInfoDict.values()])
68
69

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

73
    @property
Martin Bauer's avatar
Martin Bauer committed
74
75
    def zeroth_order_equilibrium_moment_symbol(self, ):
        return self._conservedQuantityComputation.zeroth_order_moment_symbol
76
77

    @property
Martin Bauer's avatar
Martin Bauer committed
78
79
    def first_order_equilibrium_moment_symbols(self, ):
        return self._conservedQuantityComputation.first_order_moment_symbols
80
81
82

    @property
    def weights(self):
83
        if self._weights is None:
Martin Bauer's avatar
Martin Bauer committed
84
            self._weights = self._compute_weights()
85
86
        return self._weights

Martin Bauer's avatar
Martin Bauer committed
87
    def override_weights(self, weights):
88
89
90
        assert len(weights) == len(self.stencil)
        self._weights = weights

Frederik Hennig's avatar
Frederik Hennig committed
91
92
    def get_equilibrium(self, conserved_quantity_equations=None, include_force_terms=False,
                        pre_simplification=False, subexpressions=False, keep_cqc_subexpressions=True):
Martin Bauer's avatar
Martin Bauer committed
93
        relaxation_matrix = sp.eye(len(self.relaxation_rates))
Frederik Hennig's avatar
Frederik Hennig committed
94
95
96
97
98
99
100
101
102
103
104
105
        ac = self._collision_rule_with_relaxation_matrix(relaxation_matrix,
                                                         conserved_quantity_equations=conserved_quantity_equations,
                                                         include_force_terms=include_force_terms,
                                                         pre_simplification=pre_simplification)
        if not subexpressions:
            if keep_cqc_subexpressions:
                bs = self._bound_symbols_cqc(conserved_quantity_equations)
                return ac.new_without_subexpressions(subexpressions_to_keep=bs)
            else:
                return ac.new_without_subexpressions()
        else:
            return ac
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

Frederik Hennig's avatar
Frederik Hennig committed
111
    def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=True):
Markus Holzer's avatar
Markus Holzer committed
112
        d = self.relaxation_matrix
Frederik Hennig's avatar
Frederik Hennig committed
113
        relaxation_rate_sub_expressions, d = self._generate_relaxation_matrix(d, True)
Martin Bauer's avatar
Martin Bauer committed
114
        ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
Frederik Hennig's avatar
Frederik Hennig committed
115
116
                                                         True, conserved_quantity_equations,
                                                         pre_simplification=pre_simplification)
Martin Bauer's avatar
Martin Bauer committed
117
        return ac
118

Martin Bauer's avatar
Martin Bauer committed
119
    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
Markus Holzer's avatar
Markus Holzer committed
120
121
        one = sp.Rational(1, 1)
        prev_entry = self._momentToRelaxationInfoDict[one]
Martin Bauer's avatar
Martin Bauer committed
122
        new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
Markus Holzer's avatar
Markus Holzer committed
123
        self._momentToRelaxationInfoDict[one] = new_entry
Martin Bauer's avatar
Martin Bauer committed
124

Martin Bauer's avatar
Martin Bauer committed
125
    def set_first_moment_relaxation_rate(self, relaxation_rate):
126
127
128
        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
129
130
131
            prev_entry = self._momentToRelaxationInfoDict[e]
            new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
            self._momentToRelaxationInfoDict[e] = new_entry
132

Markus Holzer's avatar
Markus Holzer committed
133
134
135
136
137
138
139
    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

140
    @property
Martin Bauer's avatar
Martin Bauer committed
141
142
    def collision_matrix(self):
        pdfs_to_moments = self.moment_matrix
Markus Holzer's avatar
Markus Holzer committed
143
144
        d = self.relaxation_matrix
        return pdfs_to_moments.inv() * d * pdfs_to_moments
145
146

    @property
Martin Bauer's avatar
Martin Bauer committed
147
148
    def inverse_collision_matrix(self):
        pdfs_to_moments = self.moment_matrix
Markus Holzer's avatar
Markus Holzer committed
149
        inverse_relaxation_matrix = self.relaxation_matrix.inv()
Martin Bauer's avatar
Martin Bauer committed
150
        return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
151

152
    @property
Martin Bauer's avatar
Martin Bauer committed
153
154
    def moment_matrix(self):
        return moment_matrix(self.moments, self.stencil)
155

Markus Holzer's avatar
Markus Holzer committed
156
157
158
159
160
161
162
    @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

163
164
165
166
167
168
169
170
171
172
    @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()

173
174
175
176
177
    def __getstate__(self):
        # Workaround for a bug in joblib
        self._momentToRelaxationInfoDictToPickle = [i for i in self._momentToRelaxationInfoDict.items()]
        return self.__dict__

178
179
180
181
182
183
    def _repr_html_(self):
        table = """
        <table style="border:none; width: 100%">
            <tr {nb}>
                <th {nb} >Moment</th>
                <th {nb} >Eq. Value </th>
184
                <th {nb} >Relaxation Rate</th>
185
186
187
188
189
            </tr>
            {content}
        </table>
        """
        content = ""
Martin Bauer's avatar
Martin Bauer committed
190
        for moment, (eq_value, rr) in self._momentToRelaxationInfoDict.items():
191
192
193
            vals = {
                'rr': sp.latex(rr),
                'moment': sp.latex(moment),
Martin Bauer's avatar
Martin Bauer committed
194
                'eq_value': sp.latex(eq_value),
195
196
197
198
                'nb': 'style="border:none"',
            }
            content += """<tr {nb}>
                            <td {nb}>${moment}$</td>
Martin Bauer's avatar
Martin Bauer committed
199
                            <td {nb}>${eq_value}$</td>
200
201
202
203
                            <td {nb}>${rr}$</td>
                         </tr>\n""".format(**vals)
        return table.format(content=content, nb='style="border:none"')

Martin Bauer's avatar
Martin Bauer committed
204
205
206
207
    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()
208

Martin Bauer's avatar
Martin Bauer committed
209
210
211
212
213
        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)
214
215

        weights = []
Martin Bauer's avatar
Martin Bauer committed
216
        for eq in ac.main_assignments:
217
            value = eq.rhs.expand()
218
            assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights " + str(value)
219
220
            weights.append(value)
        return weights
221

Frederik Hennig's avatar
Frederik Hennig committed
222
223
224
225
226
227
228
229
230
    def _bound_symbols_cqc(self, conserved_quantity_equations=None):
        f = self.pre_collision_pdf_symbols
        cqe = conserved_quantity_equations

        if cqe is None:
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f)

        return cqe.bound_symbols

Martin Bauer's avatar
Martin Bauer committed
231
    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
Frederik Hennig's avatar
Frederik Hennig committed
232
                                               conserved_quantity_equations=None, pre_simplification=False):
Martin Bauer's avatar
Martin Bauer committed
233
        f = sp.Matrix(self.pre_collision_pdf_symbols)
Markus Holzer's avatar
Markus Holzer committed
234
        moment_polynomials = list(self.moments)
Frederik Hennig's avatar
Frederik Hennig committed
235
236
237
238
239
240
241

        cqe = conserved_quantity_equations
        if cqe is None:
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f)

        if self._forceModel is None:
            include_force_terms = False
Markus Holzer's avatar
Markus Holzer committed
242
243
244
245
246
247
248

        moment_space_forcing = False

        if include_force_terms and self._moment_transform_class:
            if self._forceModel is not None:
                moment_space_forcing = self._forceModel.has_moment_space_forcing

Frederik Hennig's avatar
Frederik Hennig committed
249
250
        rho = self.zeroth_order_equilibrium_moment_symbol
        u = self.first_order_equilibrium_moment_symbols
Martin Bauer's avatar
Martin Bauer committed
251
        m_eq = sp.Matrix(self.moment_equilibrium_values)
252

Frederik Hennig's avatar
Frederik Hennig committed
253
254
255
256
257
        if self._moment_transform_class:
            #   Derive equations in moment space if a transform is given
            pdf_to_m_transform = self._moment_transform_class(self.stencil, moment_polynomials, rho, u,
                                                              conserved_quantity_equations=cqe)

258
259
            m_pre = pdf_to_m_transform.pre_collision_symbols
            m_post = pdf_to_m_transform.post_collision_symbols
Frederik Hennig's avatar
Frederik Hennig committed
260
261
262
263
264
265
266
267

            pdf_to_m_eqs = pdf_to_m_transform.forward_transform(self.pre_collision_pdf_symbols,
                                                                simplification=pre_simplification)
            m_post_to_f_post_eqs = pdf_to_m_transform.backward_transform(self.post_collision_pdf_symbols,
                                                                         simplification=pre_simplification)

            m_pre_vec = sp.Matrix(m_pre)
            collision_rule = m_pre_vec + d * (m_eq - m_pre_vec)
268

Markus Holzer's avatar
Markus Holzer committed
269
            if include_force_terms and moment_space_forcing:
Frederik Hennig's avatar
Frederik Hennig committed
270
                collision_rule += self._forceModel.moment_space_forcing(self)
271

Frederik Hennig's avatar
Frederik Hennig committed
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
            collision_eqs = [Assignment(lhs, rhs) for lhs, rhs in zip(m_post, collision_rule)]
            collision_eqs = AssignmentCollection(collision_eqs)

            all_acs = [] if pdf_to_m_transform.absorbs_conserved_quantity_equations else [cqe]
            all_acs += [pdf_to_m_eqs, collision_eqs]
            subexpressions = list(additional_subexpressions) + [ac.all_assignments for ac in all_acs]
            subexpressions += m_post_to_f_post_eqs.subexpressions
            main_assignments = m_post_to_f_post_eqs.main_assignments
        else:
            #   For SRT, TRT by default, and whenever customly required, derive equations entirely in
            #   population space
            pdf_to_moment = self.moment_matrix
            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)]
            subexpressions = list(additional_subexpressions) + cqe.all_assignments
            main_assignments = collision_eqs

        simplification_hints = cqe.simplification_hints.copy()
Martin Bauer's avatar
Martin Bauer committed
290
        simplification_hints.update(self._conservedQuantityComputation.defined_symbols())
Martin Bauer's avatar
Martin Bauer committed
291
        simplification_hints['relaxation_rates'] = [d[i, i] for i in range(d.rows)]
292

Frederik Hennig's avatar
Frederik Hennig committed
293
        if include_force_terms and not moment_space_forcing:
Martin Bauer's avatar
Martin Bauer committed
294
            force_model_terms = self._forceModel(self)
Markus Holzer's avatar
Markus Holzer committed
295
            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
Martin Bauer's avatar
Martin Bauer committed
296
297
            force_subexpressions = [Assignment(sym, force_model_term)
                                    for sym, force_model_term in zip(force_term_symbols, force_model_terms)]
Frederik Hennig's avatar
Frederik Hennig committed
298
299
300
            subexpressions += force_subexpressions
            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
301
            simplification_hints['force_terms'] = force_term_symbols
Martin Bauer's avatar
Martin Bauer committed
302

Frederik Hennig's avatar
Frederik Hennig committed
303
304
305
        ac = LbmCollisionRule(self, main_assignments, subexpressions, simplification_hints)
        ac.topological_sort()
        return ac
306

307
    @staticmethod
Martin Bauer's avatar
Martin Bauer committed
308
    def _generate_relaxation_matrix(relaxation_matrix, keep_rr_symbolic):
309
310
311
        """
        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
312
        the subexpressions, that assign the number to the newly introduced symbol
313
        """
Martin Bauer's avatar
Martin Bauer committed
314
        rr = [relaxation_matrix[i, i] for i in range(relaxation_matrix.rows)]
Frederik Hennig's avatar
Frederik Hennig committed
315
        if keep_rr_symbolic:
Martin Bauer's avatar
Martin Bauer committed
316
            unique_relaxation_rates = set(rr)
317
            subexpressions = {}
Martin Bauer's avatar
Martin Bauer committed
318
            for rt in unique_relaxation_rates:
319
320
                rt = sp.sympify(rt)
                if not isinstance(rt, sp.Symbol):
Markus Holzer's avatar
Markus Holzer committed
321
                    rt_symbol = sp.Symbol(f"rr_{len(subexpressions)}")
Martin Bauer's avatar
Martin Bauer committed
322
                    subexpressions[rt] = rt_symbol
323

Martin Bauer's avatar
Martin Bauer committed
324
325
            new_rr = [subexpressions[sp.sympify(e)] if sp.sympify(e) in subexpressions else e
                      for e in rr]
326
            substitutions = [Assignment(e[1], e[0]) for e in subexpressions.items()]
Martin Bauer's avatar
Martin Bauer committed
327
            return substitutions, sp.diag(*new_rr)
328
        else:
Martin Bauer's avatar
Martin Bauer committed
329
            return [], relaxation_matrix