momentbasedmethod.py 15.1 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
33
        """
Martin Bauer's avatar
Martin Bauer committed
34
        assert isinstance(conserved_quantity_computation, AbstractConservedQuantityComputation)
35
        super(MomentBasedLbMethod, self).__init__(stencil)
36

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

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

Frederik Hennig's avatar
Frederik Hennig committed
47
48
49
50
51
    @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)

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

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

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

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

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

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

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

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

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

Frederik Hennig's avatar
Frederik Hennig committed
90
91
    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
92
        relaxation_matrix = sp.eye(len(self.relaxation_rates))
Frederik Hennig's avatar
Frederik Hennig committed
93
94
95
96
97
98
99
100
101
102
103
104
        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
105

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

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

Martin Bauer's avatar
Martin Bauer committed
118
    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
Martin Bauer's avatar
Martin Bauer committed
119
        e = sp.Rational(1, 1)
Martin Bauer's avatar
Martin Bauer committed
120
121
122
        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
123

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

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

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

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

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

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

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

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

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

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

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

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

Frederik Hennig's avatar
Frederik Hennig committed
221
222
223
224
225
226
227
228
229
    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
230
    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
Frederik Hennig's avatar
Frederik Hennig committed
231
                                               conserved_quantity_equations=None, pre_simplification=False):
Martin Bauer's avatar
Martin Bauer committed
232
        f = sp.Matrix(self.pre_collision_pdf_symbols)
Frederik Hennig's avatar
Frederik Hennig committed
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
        moment_polynomials = list(self._momentToRelaxationInfoDict.keys())

        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

        moment_space_forcing = False

        if include_force_terms and self._moment_transform_class:
            moment_space_forcing = hasattr(self.force_model, 'moment_space_forcing')

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

Frederik Hennig's avatar
Frederik Hennig committed
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
        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)

            m_pre = pdf_to_m_transform.pre_collision_moment_symbols
            m_post = pdf_to_m_transform.post_collision_moment_symbols

            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)
266

Frederik Hennig's avatar
Frederik Hennig committed
267
268
            if moment_space_forcing:
                collision_rule += self._forceModel.moment_space_forcing(self)
269

Frederik Hennig's avatar
Frederik Hennig committed
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
            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
288
        simplification_hints.update(self._conservedQuantityComputation.defined_symbols())
Martin Bauer's avatar
Martin Bauer committed
289
        simplification_hints['relaxation_rates'] = [d[i, i] for i in range(d.rows)]
290

Frederik Hennig's avatar
Frederik Hennig committed
291
        if include_force_terms and not moment_space_forcing:
Martin Bauer's avatar
Martin Bauer committed
292
            force_model_terms = self._forceModel(self)
Markus Holzer's avatar
Markus Holzer committed
293
            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
Martin Bauer's avatar
Martin Bauer committed
294
295
            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
296
297
298
            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
299
            simplification_hints['force_terms'] = force_term_symbols
Martin Bauer's avatar
Martin Bauer committed
300

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

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

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