momentbasedmethod.py 14.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
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):
112
        relaxation_rate_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
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):
Markus Holzer's avatar
Markus Holzer committed
119
120
        one = sp.Rational(1, 1)
        prev_entry = self._momentToRelaxationInfoDict[one]
Martin Bauer's avatar
Martin Bauer committed
121
        new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
Markus Holzer's avatar
Markus Holzer committed
122
        self._momentToRelaxationInfoDict[one] = 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

155
156
157
158
159
160
161
162
163
164
    @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()

165
166
167
168
169
    def __getstate__(self):
        # Workaround for a bug in joblib
        self._momentToRelaxationInfoDictToPickle = [i for i in self._momentToRelaxationInfoDict.items()]
        return self.__dict__

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

Martin Bauer's avatar
Martin Bauer committed
196
197
198
199
    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()
200

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

        weights = []
Martin Bauer's avatar
Martin Bauer committed
208
        for eq in ac.main_assignments:
209
            value = eq.rhs.expand()
210
            assert len(value.atoms(sp.Symbol)) == 0, "Failed to compute weights " + str(value)
211
212
            weights.append(value)
        return weights
213

Frederik Hennig's avatar
Frederik Hennig committed
214
215
216
217
218
    def _bound_symbols_cqc(self, conserved_quantity_equations=None):
        f = self.pre_collision_pdf_symbols
        cqe = conserved_quantity_equations

        if cqe is None:
219
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f, False)
Frederik Hennig's avatar
Frederik Hennig committed
220
221
222

        return cqe.bound_symbols

Martin Bauer's avatar
Martin Bauer committed
223
    def _collision_rule_with_relaxation_matrix(self, d, additional_subexpressions=(), include_force_terms=True,
Frederik Hennig's avatar
Frederik Hennig committed
224
                                               conserved_quantity_equations=None, pre_simplification=False):
Martin Bauer's avatar
Martin Bauer committed
225
        f = sp.Matrix(self.pre_collision_pdf_symbols)
Markus Holzer's avatar
Markus Holzer committed
226
        moment_polynomials = list(self.moments)
Frederik Hennig's avatar
Frederik Hennig committed
227
228
229

        cqe = conserved_quantity_equations
        if cqe is None:
230
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f, False)
Frederik Hennig's avatar
Frederik Hennig committed
231
232
233

        if self._forceModel is None:
            include_force_terms = False
Markus Holzer's avatar
Markus Holzer committed
234
235
236
237
238
239
240

        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

241
        forcing_subexpressions = []
Frederik Hennig's avatar
Frederik Hennig committed
242
        if self._forceModel is not None:
243
244
            forcing_subexpressions = AssignmentCollection(self._forceModel.subs_dict_force).all_assignments

Frederik Hennig's avatar
Frederik Hennig committed
245
246
        rho = self.zeroth_order_equilibrium_moment_symbol
        u = self.first_order_equilibrium_moment_symbols
Martin Bauer's avatar
Martin Bauer committed
247
        m_eq = sp.Matrix(self.moment_equilibrium_values)
248

Frederik Hennig's avatar
Frederik Hennig committed
249
250
251
252
253
        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)

254
255
            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
256
257
258
259
260
261
262
263

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

Markus Holzer's avatar
Markus Holzer committed
265
            if include_force_terms and moment_space_forcing:
Frederik Hennig's avatar
Frederik Hennig committed
266
                collision_rule += self._forceModel.moment_space_forcing(self)
267

Frederik Hennig's avatar
Frederik Hennig committed
268
269
270
271
272
            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]
273
274
            subexpressions = list(additional_subexpressions) + forcing_subexpressions + [ac.all_assignments for ac in
                                                                                         all_acs]
Frederik Hennig's avatar
Frederik Hennig committed
275
276
277
278
279
280
281
282
            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)]
283
            subexpressions = list(additional_subexpressions) + forcing_subexpressions + cqe.all_assignments
Frederik Hennig's avatar
Frederik Hennig committed
284
285
286
            main_assignments = collision_eqs

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

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

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