momentbasedmethod.py 14.4 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
            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

Markus Holzer's avatar
Markus Holzer committed
61
62
63
64
65
66
67
68
    @property
    def compressible(self):
        return self._conservedQuantityComputation.compressible

    @property
    def zero_centered_pdfs(self):
        return self._conservedQuantityComputation.zero_centered_pdfs

69
70
    @property
    def moments(self):
71
72
73
        return tuple(self._momentToRelaxationInfoDict.keys())

    @property
Martin Bauer's avatar
Martin Bauer committed
74
    def moment_equilibrium_values(self):
Martin Bauer's avatar
Martin Bauer committed
75
        return tuple([e.equilibrium_value for e in self._momentToRelaxationInfoDict.values()])
76
77

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

81
    @property
Markus Holzer's avatar
Markus Holzer committed
82
    def zeroth_order_equilibrium_moment_symbol(self):
Martin Bauer's avatar
Martin Bauer committed
83
        return self._conservedQuantityComputation.zeroth_order_moment_symbol
84
85

    @property
Markus Holzer's avatar
Markus Holzer committed
86
    def first_order_equilibrium_moment_symbols(self):
Martin Bauer's avatar
Martin Bauer committed
87
        return self._conservedQuantityComputation.first_order_moment_symbols
88
89
90

    @property
    def weights(self):
91
        if self._weights is None:
Martin Bauer's avatar
Martin Bauer committed
92
            self._weights = self._compute_weights()
93
94
        return self._weights

Martin Bauer's avatar
Martin Bauer committed
95
    def override_weights(self, weights):
96
97
98
        assert len(weights) == len(self.stencil)
        self._weights = weights

Frederik Hennig's avatar
Frederik Hennig committed
99
100
    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
101
        relaxation_matrix = sp.eye(len(self.relaxation_rates))
Frederik Hennig's avatar
Frederik Hennig committed
102
103
104
105
106
107
108
109
110
111
112
113
        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
114

Martin Bauer's avatar
Martin Bauer committed
115
116
    def get_equilibrium_terms(self):
        equilibrium = self.get_equilibrium()
Martin Bauer's avatar
Martin Bauer committed
117
        return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments])
118

Frederik Hennig's avatar
Frederik Hennig committed
119
    def get_collision_rule(self, conserved_quantity_equations=None, pre_simplification=True):
120
        relaxation_rate_sub_expressions, d = self._generate_symbolic_relaxation_matrix()
Martin Bauer's avatar
Martin Bauer committed
121
        ac = self._collision_rule_with_relaxation_matrix(d, relaxation_rate_sub_expressions,
Frederik Hennig's avatar
Frederik Hennig committed
122
123
                                                         True, conserved_quantity_equations,
                                                         pre_simplification=pre_simplification)
Martin Bauer's avatar
Martin Bauer committed
124
        return ac
125

Martin Bauer's avatar
Martin Bauer committed
126
    def set_zeroth_moment_relaxation_rate(self, relaxation_rate):
127
128
        one = sp.Rational(1, 1)
        prev_entry = self._momentToRelaxationInfoDict[one]
Martin Bauer's avatar
Martin Bauer committed
129
        new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
130
        self._momentToRelaxationInfoDict[one] = new_entry
Martin Bauer's avatar
Martin Bauer committed
131

Martin Bauer's avatar
Martin Bauer committed
132
    def set_first_moment_relaxation_rate(self, relaxation_rate):
133
134
135
        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
136
137
138
            prev_entry = self._momentToRelaxationInfoDict[e]
            new_entry = RelaxationInfo(prev_entry[0], relaxation_rate)
            self._momentToRelaxationInfoDict[e] = new_entry
139

Markus Holzer's avatar
Markus Holzer committed
140
141
142
143
144
145
146
    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

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

    @property
Martin Bauer's avatar
Martin Bauer committed
154
155
    def inverse_collision_matrix(self):
        pdfs_to_moments = self.moment_matrix
Markus Holzer's avatar
Markus Holzer committed
156
        inverse_relaxation_matrix = self.relaxation_matrix.inv()
Martin Bauer's avatar
Martin Bauer committed
157
        return pdfs_to_moments.inv() * inverse_relaxation_matrix * pdfs_to_moments
158

159
    @property
Martin Bauer's avatar
Martin Bauer committed
160
161
    def moment_matrix(self):
        return moment_matrix(self.moments, self.stencil)
162

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
    def _bound_symbols_cqc(self, conserved_quantity_equations=None):
        f = self.pre_collision_pdf_symbols
        cqe = conserved_quantity_equations

        if cqe is None:
227
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f, False)
Frederik Hennig's avatar
Frederik Hennig committed
228
229
230

        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)
234
        moment_polynomials = list(self.moments)
Frederik Hennig's avatar
Frederik Hennig committed
235
236
237

        cqe = conserved_quantity_equations
        if cqe is None:
238
            cqe = self._conservedQuantityComputation.equilibrium_input_equations_from_pdfs(f, False)
Frederik Hennig's avatar
Frederik Hennig committed
239
240
241
242
243
244
245

        if self._forceModel is None:
            include_force_terms = False

        moment_space_forcing = False

        if include_force_terms and self._moment_transform_class:
246
247
248
249
250
251
            if self._forceModel is not None:
                moment_space_forcing = self._forceModel.has_moment_space_forcing

        forcing_subexpressions = []
        if self._forceModel is not None:
            forcing_subexpressions = AssignmentCollection(self._forceModel.subs_dict_force).all_assignments
Frederik Hennig's avatar
Frederik Hennig committed
252
253
254

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

Frederik Hennig's avatar
Frederik Hennig committed
257
258
259
260
261
        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)

262
263
            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
264
265
266
267
268
269
270
271

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

273
            if include_force_terms and moment_space_forcing:
Frederik Hennig's avatar
Frederik Hennig committed
274
                collision_rule += self._forceModel.moment_space_forcing(self)
275

Frederik Hennig's avatar
Frederik Hennig committed
276
277
278
279
280
            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]
281
282
            subexpressions = list(additional_subexpressions) + forcing_subexpressions + [ac.all_assignments for ac in
                                                                                         all_acs]
Frederik Hennig's avatar
Frederik Hennig committed
283
284
285
286
287
288
289
290
            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)]
291
            subexpressions = list(additional_subexpressions) + forcing_subexpressions + cqe.all_assignments
Frederik Hennig's avatar
Frederik Hennig committed
292
293
294
            main_assignments = collision_eqs

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

Frederik Hennig's avatar
Frederik Hennig committed
298
        if include_force_terms and not moment_space_forcing:
Martin Bauer's avatar
Martin Bauer committed
299
            force_model_terms = self._forceModel(self)
Markus Holzer's avatar
Markus Holzer committed
300
            force_term_symbols = sp.symbols(f"forceTerm_:{len(force_model_terms)}")
Martin Bauer's avatar
Martin Bauer committed
301
302
            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
303
304
305
            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
306
            simplification_hints['force_terms'] = force_term_symbols
Martin Bauer's avatar
Martin Bauer committed
307

Frederik Hennig's avatar
Frederik Hennig committed
308
309
310
        ac = LbmCollisionRule(self, main_assignments, subexpressions, simplification_hints)
        ac.topological_sort()
        return ac