conservedquantitycomputation.py 15.6 KB
Newer Older
1
import abc
2
from collections import OrderedDict
Martin Bauer's avatar
Martin Bauer committed
3
4
5
6

import sympy as sp

from pystencils import Assignment, AssignmentCollection, Field
7
8


Martin Bauer's avatar
Martin Bauer committed
9
class AbstractConservedQuantityComputation(abc.ABC):
10
    r"""
11

12
    This class defines how conserved quantities are computed as functions of the pdfs.
13
    Conserved quantities are used for output and as input to the equilibrium in the collision step
14
15
16
17
18
19
20

    Depending on the method they might also be computed slightly different, e.g. due to a force model.

    An additional method describes how to get the conserved quantities for the equilibrium for initialization.
    In most cases the inputs can be used directly, but for some methods they have to be altered slightly.
    For example in zero centered hydrodynamic schemes with force model, the density has
    to be decreased by one, and the given velocity has to be shifted dependent on the force.
21

Martin Bauer's avatar
Martin Bauer committed
22
    .. image:: /img/moment_shift.svg
23

24
25
    """

26
27
    @property
    @abc.abstractmethod
Martin Bauer's avatar
Martin Bauer committed
28
    def conserved_quantities(self):
29
30
31
        """
        Dict, mapping names (symbol) to dimensionality (int)
        For example: {'density' : 1, 'velocity' : 3}
Martin Bauer's avatar
Martin Bauer committed
32
33
        The naming strings can be used in :func:`output_equations_from_pdfs`
        and :func:`equilibrium_input_equations_from_init_values`
34
35
        """

Martin Bauer's avatar
Martin Bauer committed
36
    def defined_symbols(self, order='all'):
37
38
39
40
        """
        Returns a dict, mapping names of conserved quantities to their symbols
        """

41
42
    @property
    @abc.abstractmethod
Martin Bauer's avatar
Martin Bauer committed
43
    def default_values(self):
44
45
46
47
48
49
50
        """
        Returns a dict of symbol to default value, where "default" means that
        the equilibrium simplifies to the weights if these values are inserted.
        Hydrodynamic example: rho=1, u_i = 0
        """

    @abc.abstractmethod
Martin Bauer's avatar
Martin Bauer committed
51
    def equilibrium_input_equations_from_pdfs(self, pdfs):
52
53
54
55
56
        """
        Returns an equation collection that defines all necessary quantities to compute the equilibrium as functions
        of the pdfs.
        For hydrodynamic LBM schemes this is usually the density and velocity.

57
58
        Args:
            pdfs: values or symbols for the pdf values
59
60
61
        """

    @abc.abstractmethod
Martin Bauer's avatar
Martin Bauer committed
62
    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
63
64
65
        """
        Returns an equation collection that defines conserved quantities for output. These conserved quantities might
        be slightly different that the ones used as input for the equilibrium e.g. due to a force model.
66

67
68
69
70
        Args:
            pdfs: values for the pdf entries
            output_quantity_names_to_symbols: dict mapping of conserved quantity names
             (See :func:`conserved_quantities`) to symbols or field accesses where they should be written to
71
72
73
        """

    @abc.abstractmethod
Martin Bauer's avatar
Martin Bauer committed
74
    def equilibrium_input_equations_from_init_values(self, **kwargs):
75
76
77
        """
        Returns an equation collection that defines all necessary quantities to compute the equilibrium as function of
        given conserved quantities. Parameters can be names that are given by
Martin Bauer's avatar
Martin Bauer committed
78
        symbol names of :func:`conserved_quantities`.
79
80
81
82
83
84
        For all parameters not specified each implementation should use sensible defaults. For example hydrodynamic
        schemes use density=1 and velocity=0.
        """


class DensityVelocityComputation(AbstractConservedQuantityComputation):
Martin Bauer's avatar
Martin Bauer committed
85
86
    def __init__(self, stencil, compressible, force_model=None,
                 zeroth_order_moment_symbol=sp.Symbol("rho"),
87
88
                 first_order_moment_symbols=sp.symbols("u_:3"),
                 second_order_moment_symbols=sp.symbols("p_:9")):
89
90
91
        dim = len(stencil[0])
        self._stencil = stencil
        self._compressible = compressible
Martin Bauer's avatar
Martin Bauer committed
92
93
94
        self._forceModel = force_model
        self._symbolOrder0 = zeroth_order_moment_symbol
        self._symbolsOrder1 = first_order_moment_symbols[:dim]
itischler's avatar
itischler committed
95
        self._symbolsOrder2 = second_order_moment_symbols[:(dim * dim)]
96
97

    @property
Martin Bauer's avatar
Martin Bauer committed
98
    def conserved_quantities(self):
99
        return {'density': 1,
100
                'velocity': len(self._stencil[0])}
101

102
103
104
105
    @property
    def compressible(self):
        return self._compressible

Martin Bauer's avatar
Martin Bauer committed
106
    def defined_symbols(self, order='all'):
107
108
109
110
111
112
113
114
115
116
        if order == 'all':
            return {'density': self._symbolOrder0,
                    'velocity': self._symbolsOrder1}
        elif order == 0:
            return 'density', self._symbolOrder0
        elif order == 1:
            return 'velocity', self._symbolsOrder1
        else:
            return None

117
    @property
Martin Bauer's avatar
Martin Bauer committed
118
    def zero_centered_pdfs(self):
119
120
        return not self._compressible

Martin Bauer's avatar
Martin Bauer committed
121
    @property
Martin Bauer's avatar
Martin Bauer committed
122
    def zeroth_order_moment_symbol(self):
Martin Bauer's avatar
Martin Bauer committed
123
124
125
        return self._symbolOrder0

    @property
Martin Bauer's avatar
Martin Bauer committed
126
    def first_order_moment_symbols(self):
Martin Bauer's avatar
Martin Bauer committed
127
128
        return self._symbolsOrder1

129
    @property
Martin Bauer's avatar
Martin Bauer committed
130
    def default_values(self):
131
132
133
134
135
        result = {self._symbolOrder0: 1}
        for s in self._symbolsOrder1:
            result[s] = 0
        return result

Martin Bauer's avatar
Martin Bauer committed
136
    def equilibrium_input_equations_from_pdfs(self, pdfs):
137
        dim = len(self._stencil[0])
Martin Bauer's avatar
Martin Bauer committed
138
        eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs, self._symbolOrder0,
139
                                                                  self._symbolsOrder1[:dim])
140
        if self._compressible:
Martin Bauer's avatar
Martin Bauer committed
141
            eq_coll = divide_first_order_moments_by_rho(eq_coll, dim)
142

Martin Bauer's avatar
Martin Bauer committed
143
144
        eq_coll = apply_force_model_shift('equilibrium_velocity_shift', dim, eq_coll,
                                          self._forceModel, self._compressible)
145
        return eq_coll
146

Martin Bauer's avatar
Martin Bauer committed
147
    def equilibrium_input_equations_from_init_values(self, density=1, velocity=(0, 0, 0)):
148
        dim = len(self._stencil[0])
Martin Bauer's avatar
Martin Bauer committed
149
        zeroth_order_moment = density
150
151
        first_order_moments = velocity[:dim]
        vel_offset = [0] * dim
152

153
        if self._compressible:
Martin Bauer's avatar
Martin Bauer committed
154
155
            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
                vel_offset = self._forceModel.macroscopic_velocity_shift(zeroth_order_moment)
156
        else:
Martin Bauer's avatar
Martin Bauer committed
157
158
159
160
            if self._forceModel and hasattr(self._forceModel, 'macroscopic_velocity_shift'):
                vel_offset = self._forceModel.macroscopic_velocity_shift(sp.Rational(1, 1))
            zeroth_order_moment -= sp.Rational(1, 1)
        eqs = [Assignment(self._symbolOrder0, zeroth_order_moment)]
161

162
163
        first_order_moments = [a - b for a, b in zip(first_order_moments, vel_offset)]
        eqs += [Assignment(l, r) for l, r in zip(self._symbolsOrder1, first_order_moments)]
164

165
        return AssignmentCollection(eqs, [])
166

Martin Bauer's avatar
Martin Bauer committed
167
    def output_equations_from_pdfs(self, pdfs, output_quantity_names_to_symbols):
168
        dim = len(self._stencil[0])
Martin Bauer's avatar
Martin Bauer committed
169

Martin Bauer's avatar
Martin Bauer committed
170
        ac = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
itischler's avatar
itischler committed
171
172
                                                             self._symbolOrder0, self._symbolsOrder1,
                                                             self._symbolsOrder2)
Martin Bauer's avatar
Martin Bauer committed
173

174
        if self._compressible:
Martin Bauer's avatar
Martin Bauer committed
175
            ac = divide_first_order_moments_by_rho(ac, dim)
176
        else:
Martin Bauer's avatar
Martin Bauer committed
177
            ac = add_density_offset(ac)
178

Martin Bauer's avatar
Martin Bauer committed
179
        ac = apply_force_model_shift('macroscopic_velocity_shift', dim, ac, self._forceModel, self._compressible)
180

181
        main_assignments = []
Martin Bauer's avatar
Martin Bauer committed
182
        eqs = OrderedDict([(eq.lhs, eq.rhs) for eq in ac.all_assignments])
Martin Bauer's avatar
Martin Bauer committed
183

Martin Bauer's avatar
Martin Bauer committed
184
185
        if 'density' in output_quantity_names_to_symbols:
            density_output_symbol = output_quantity_names_to_symbols['density']
186
            if isinstance(density_output_symbol, Field):
187
                density_output_symbol = density_output_symbol.center
188
189
            if density_output_symbol != self._symbolOrder0:
                main_assignments.append(Assignment(density_output_symbol, self._symbolOrder0))
Martin Bauer's avatar
Martin Bauer committed
190
            else:
191
                main_assignments.append(Assignment(self._symbolOrder0, eqs[self._symbolOrder0]))
Martin Bauer's avatar
Martin Bauer committed
192
                del eqs[self._symbolOrder0]
Martin Bauer's avatar
Martin Bauer committed
193
194
        if 'velocity' in output_quantity_names_to_symbols:
            vel_output_symbols = output_quantity_names_to_symbols['velocity']
195
            if isinstance(vel_output_symbols, Field):
196
                vel_output_symbols = vel_output_symbols.center_vector
197
198
            if tuple(vel_output_symbols) != tuple(self._symbolsOrder1):
                main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._symbolsOrder1)]
Martin Bauer's avatar
Martin Bauer committed
199
            else:
200
                for u_i in self._symbolsOrder1:
201
                    main_assignments.append(Assignment(u_i, eqs[u_i]))
202
                    del eqs[u_i]
Martin Bauer's avatar
Martin Bauer committed
203
        if 'momentum_density' in output_quantity_names_to_symbols:
Martin Bauer's avatar
Martin Bauer committed
204
205
206
            # get zeroth and first moments again - force-shift them if necessary
            # and add their values directly to the main equations assuming that subexpressions are already in
            # main equation collection
Martin Bauer's avatar
Martin Bauer committed
207
208
209
210
            # Is not optimal when velocity and momentum_density are calculated together,
            # but this is usually not the case
            momentum_density_output_symbols = output_quantity_names_to_symbols['momentum_density']
            mom_density_eq_coll = get_equations_for_zeroth_and_first_order_moment(self._stencil, pdfs,
Martin Bauer's avatar
Martin Bauer committed
211
                                                                                  self._symbolOrder0,
212
                                                                                  self._symbolsOrder1)
213
214
            mom_density_eq_coll = apply_force_model_shift('macroscopic_momentum_density_shift', dim, 
                                                          mom_density_eq_coll, self._forceModel, self._compressible)
215

Martin Bauer's avatar
Martin Bauer committed
216
            for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]):
217
                main_assignments.append(Assignment(sym, val.rhs))
218
219
220
221
222
223
224
225
226
227
228
        if 'moment0' in output_quantity_names_to_symbols:
            moment0_output_symbol = output_quantity_names_to_symbols['moment0']
            if isinstance(moment0_output_symbol, Field):
                moment0_output_symbol = moment0_output_symbol.center
            main_assignments.append(Assignment(moment0_output_symbol, sum(pdfs)))
        if 'moment1' in output_quantity_names_to_symbols:
            moment1_output_symbol = output_quantity_names_to_symbols['moment1']
            if isinstance(moment1_output_symbol, Field):
                moment1_output_symbol = moment1_output_symbol.center_vector
            main_assignments.extend([Assignment(lhs, sum(d[i] * pdf for d, pdf in zip(self._stencil, pdfs)))
                                     for i, lhs in enumerate(moment1_output_symbol)])
229
        if 'moment2' in output_quantity_names_to_symbols:
230
231
232
233
234
235
            moment2_output_symbol = output_quantity_names_to_symbols['moment2']
            if isinstance(moment2_output_symbol, Field):
                moment2_output_symbol = moment2_output_symbol.center_vector
            for i, p in enumerate(moment2_output_symbol):
                main_assignments.append(Assignment(p, eqs[self._symbolsOrder2[i]]))
                del eqs[self._symbolsOrder2[i]]
Martin Bauer's avatar
Martin Bauer committed
236

237
        ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()])
Martin Bauer's avatar
Martin Bauer committed
238
        return ac.new_without_unused_subexpressions()
239
240

    def __repr__(self):
Martin Bauer's avatar
Martin Bauer committed
241
        return "ConservedValueComputation for %s" % (", " .join(self.conserved_quantities.keys()),)
242

243

244
245
246
# -----------------------------------------  Helper functions ----------------------------------------------------------


itischler's avatar
itischler committed
247
def get_equations_for_zeroth_and_first_order_moment(stencil, symbolic_pdfs, symbolic_zeroth_moment,
248
                                                    symbolic_first_moments, symbolic_second_moments=None):
249
    r"""
250
251
252
253
254
255
256
257
    Returns an equation system that computes the zeroth and first order moments with the least amount of operations

    The first equation of the system is equivalent to

    .. math :

        \rho = \sum_{d \in S} f_d
        u_j = \sum_{d \in S} f_d u_jd
258
        p_j = \sum_{d \in S} {d \in S} f_d u_jd
259

Martin Bauer's avatar
Martin Bauer committed
260
261
262
263
264
    Args:
        stencil: called :math:`S` above
        symbolic_pdfs: called :math:`f` above
        symbolic_zeroth_moment:  called :math:`\rho` above
        symbolic_first_moments: called :math:`u` above
265
        symbolic_second_moments: called :math:`p` above
266
    """
267
    def filter_out_plus_terms(expr):
268
269
270
271
272
273
274
275
276
        result = 0
        for term in expr.args:
            if not type(term) is sp.Mul:
                result += term
        return result

    dim = len(stencil[0])

    subexpressions = []
Martin Bauer's avatar
Martin Bauer committed
277
    pdf_sum = sum(symbolic_pdfs)
278
    u = [0] * dim
Martin Bauer's avatar
Martin Bauer committed
279
    for f, offset in zip(symbolic_pdfs, stencil):
280
281
282
        for i in range(dim):
            u[i] += f * int(offset[i])

283
284
    p = [0] * dim * dim
    for f, offset in zip(symbolic_pdfs, stencil):
itischler's avatar
itischler committed
285
286
287
        for i in range(dim):
            for j in range(dim):
                p[dim * i + j] += f * int(offset[i]) * int(offset[j])
288

289
    plus_terms = [set(filter_out_plus_terms(u_i).args) for u_i in u]
290
    for i in range(dim):
291
        rhs = plus_terms[i]
292
        for j in range(i):
293
294
            rhs -= plus_terms[j]
        eq = Assignment(sp.Symbol("vel%dTerm" % (i,)), sum(rhs))
295
296
297
        subexpressions.append(eq)

    for subexpression in subexpressions:
298
        pdf_sum = pdf_sum.subs(subexpression.rhs, subexpression.lhs)
299
300
301
302
303

    for i in range(dim):
        u[i] = u[i].subs(subexpressions[i].rhs, subexpressions[i].lhs)

    equations = []
Martin Bauer's avatar
Martin Bauer committed
304
305
    equations += [Assignment(symbolic_zeroth_moment, pdf_sum)]
    equations += [Assignment(u_i_sym, u_i) for u_i_sym, u_i in zip(symbolic_first_moments, u)]
306
307
    if symbolic_second_moments:
        equations += [Assignment(symbolic_second_moments[i], p[i]) for i in range(dim**2)]
308

309
    return AssignmentCollection(equations, subexpressions)
310
311


Martin Bauer's avatar
Martin Bauer committed
312
def divide_first_order_moments_by_rho(assignment_collection, dim):
313
    r"""
314
315
316
317
318
319
320
    Assumes that the equations of the passed equation collection are the following
        - rho = f_0  + f_1 + ...
        - u_0 = ...
        - u_1 = ...
    Returns a new equation collection where the u terms (first order moments) are divided by rho.
    The dim parameter specifies the number of first order moments. All subsequent equations are just copied over.
    """
Martin Bauer's avatar
Martin Bauer committed
321
322
    old_eqs = assignment_collection.main_assignments
    rho = old_eqs[0].lhs
Martin Bauer's avatar
Martin Bauer committed
323
324
    new_first_order_moment_eq = [Assignment(eq.lhs, eq.rhs / rho) for eq in old_eqs[1:dim + 1]]
    new_eqs = [old_eqs[0]] + new_first_order_moment_eq + old_eqs[dim + 1:]
325
    return assignment_collection.copy(new_eqs)
326
327


Martin Bauer's avatar
Martin Bauer committed
328
def add_density_offset(assignment_collection, offset=sp.Rational(1, 1)):
329
    r"""
330
331
    Assumes that first equation is the density (zeroth moment). Changes the density equations by adding offset to it.
    """
Martin Bauer's avatar
Martin Bauer committed
332
333
334
    old_eqs = assignment_collection.main_assignments
    new_density = Assignment(old_eqs[0].lhs, old_eqs[0].rhs + offset)
    return assignment_collection.copy([new_density] + old_eqs[1:])
335
336


Martin Bauer's avatar
Martin Bauer committed
337
def apply_force_model_shift(shift_member_name, dim, assignment_collection, force_model, compressible, reverse=False):
338
    """
339
    Modifies the first order moment equations in assignment collection according to the force model shift.
Martin Bauer's avatar
Martin Bauer committed
340
    It is applied if force model has a method named shift_member_name. The equations 1: dim+1 of the passed
341
342
    equation collection are assumed to be the velocity equations.
    """
Martin Bauer's avatar
Martin Bauer committed
343
    if force_model is not None and hasattr(force_model, shift_member_name):
Martin Bauer's avatar
Martin Bauer committed
344
        old_eqs = assignment_collection.main_assignments
345
346
        density = old_eqs[0].lhs if compressible else sp.Rational(1, 1)
        old_vel_eqs = old_eqs[1:dim + 1]
Martin Bauer's avatar
Martin Bauer committed
347
        shift_func = getattr(force_model, shift_member_name)
348
        vel_offsets = shift_func(density)
349
        if reverse:
350
            vel_offsets = [-v for v in vel_offsets]
Martin Bauer's avatar
Martin Bauer committed
351
352
        shifted_velocity_eqs = [Assignment(old_eq.lhs, old_eq.rhs + offset)
                                for old_eq, offset in zip(old_vel_eqs, vel_offsets)]
353
354
        new_eqs = [old_eqs[0]] + shifted_velocity_eqs + old_eqs[dim + 1:]
        return assignment_collection.copy(new_eqs)
355
    else:
356
        return assignment_collection