lattice.py 19.3 KB
Newer Older
MischaD's avatar
MischaD committed
1
import numpy as np
MischaD's avatar
MischaD committed
2
import sympy as sp
MischaD's avatar
MischaD committed
3
import random
MischaD's avatar
MischaD committed
4
from sympy import oo
Mischa Dombrowski's avatar
Mischa Dombrowski committed
5
from .utils.mylog import logger
MischaD's avatar
MischaD committed
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from .exceptions import (
    OrderNotImplementedException,
    ImpossibleVelocityShellException,
    ReducibleShellException,
    NoSolutionException,
    InfiniteSolutionsException,
    UninitializedAttributeException,
)
from .utils.utils import (
    analyse_tensor_dimension,
    get_group,
    get_subshells,
    lattice_sum,
    approximate_ratio,
    timer,
MischaD's avatar
MischaD committed
21
    ps_installed,
MischaD's avatar
MischaD committed
22
23
24
    get_random_vector,
    velocities_for_shell,
)
25
from .shell import Shell
MischaD's avatar
MischaD committed
26
27


MischaD's avatar
MischaD committed
28
class Lattice:
MischaD's avatar
MischaD committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
    BY_NAME = {
        "D2Q9": {"dimension": 2, "order": 4, "shell_list": [1, 2, 4], "seed": 44},
        "D2V17": {
            "dimension": 2,
            "order": 6,
            "shell_list": [1, 2, 4, 8, 9],
            "seed": 44,
        },
        "D2V37": {
            "dimension": 2,
            "order": 8,
            "shell_list": [1, 2, 4, 5, 8, 9, 10, 16],
            "seed": 44,
        },
        "D3Q19": {"dimension": 3, "order": 4, "shell_list": [1, 2, 4], "seed": 44},
        "D3Q15": {"dimension": 3, "order": 4, "shell_list": [1, 3, 4], "seed": 44},
        "D3Q41-ZOT": {
            "dimension": 3,
            "order": 6,
            "shell_list": [1, 2, 3, 9, 16, 27],
            "boundary": "sup",
            "unwanted_subshells": ["221", "511"],
            "seed": 44,
        },
    }

    def __init__(
        self,
        dimension=2,
        order=4,
        shell_list=[1, 2, 4],
        seed=None,
        boundary=None,
        unwanted_subshells=[],
        svd_tolerance=1e-8,
    ):
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
        """Create an instance of the Lattice class without calculating the weights.
        Default values are chosen to lead to the standard D2Q9 lattice.
        Not all input values are valid. :func:`Lattice.init_by_order`

        Args:
            name: Initialize Lattice with the correct values for some known velocities. Possible names are
                D2Q9, D2V17, D2V37, D3Q17,
            dimension (int): Dimension of the resulting lattice. Default 2.
            order (int): Order of tensor isotropy. Only odd orders are
            shell_list (list): List of the Squared lengths of the velocities.
            unwanted_subshells (list) : List of strings that correspond to the type of a subshell that isn't supposed to
                to be part of lattice. Some shells are ambiguous and in the default case all possiblities are taken,
                but some lattice, like the D3Q41-ZOT need certain subshells to be remove. E.g. 2 Dimensions and
                shell 25 is added. Bother type = "50" and type = "43" could be used. Add one of them to the list of
                unwanted subshells to remove them.
            seed (float): Seed for the random number generator.
            svd_tolerance (float): tolerance to set the singular values to zero after svd.
MischaD's avatar
MischaD committed
82

MischaD's avatar
MischaD committed
83
        """
MischaD's avatar
MischaD committed
84
        self._dimension = dimension
MischaD's avatar
MischaD committed
85
86
87
        self._order = order
        self._shell_list = shell_list
        self._seed = seed
MischaD's avatar
MischaD committed
88
        self._svd_tolerance = svd_tolerance
MischaD's avatar
MischaD committed
89
        random.seed(self._seed)
MischaD's avatar
MischaD committed
90

MischaD's avatar
MischaD committed
91
        # Preperation
MischaD's avatar
MischaD committed
92
        self._tensor_space_dimension = 0
MischaD's avatar
MischaD committed
93
        self._all_velocity_vectors = []
MischaD's avatar
MischaD committed
94
        self._possible_tensors = []
MischaD's avatar
MischaD committed
95
        self._tensor_dimensions = []
MischaD's avatar
MischaD committed
96
        self._shells = 0
97
        self._unwanted_subshells = unwanted_subshells
MischaD's avatar
MischaD committed
98

MischaD's avatar
MischaD committed
99
100
101
102
        # LES
        self._lhs = []
        self._rhs = []

MischaD's avatar
MischaD committed
103
104
        # SVD
        self._singular_values = []
MischaD's avatar
MischaD committed
105
        self._solution = []
MischaD's avatar
MischaD committed
106

107
        # Polynomials
MischaD's avatar
MischaD committed
108
        self._coefficients = []
MischaD's avatar
MischaD committed
109
110
        self._weight_polynomials = []
        self._interval = None
111
        self._boundary = boundary if boundary == "sup" else "inf"
MischaD's avatar
MischaD committed
112

113
114
115
116
117
        # Final values
        self._discrete_velocities = []
        self._weights = []
        self._c_s_sq = None

MischaD's avatar
MischaD committed
118
119
    def __str__(self):
        string = "D{} - Order: {} - Shells: {}".format(
MischaD's avatar
MischaD committed
120
            self._dimension, self._order, str(self._shell_list)
MischaD's avatar
MischaD committed
121
        )
MischaD's avatar
MischaD committed
122
123
124
        if self._boundary:
            string += " - Boundary: {}".format(self._boundary)
        if self._unwanted_subshells:
MischaD's avatar
MischaD committed
125
126
127
            string += " - Unwanted Subshells: {}".format(
                "".join([s + ", " for s in self._unwanted_subshells])
            )
MischaD's avatar
MischaD committed
128
129
130
        return string

    @classmethod
MischaD's avatar
MischaD committed
131
    def from_name(cls, name):
132
133
134
135
136
        """Initiate the Lattice with the correct values corresponding to a know Lattice in literature.
        A complete list can be seen in Lattice.BY_NAME.keys()"""
        kwargs = cls.BY_NAME.get(name)
        if not kwargs:
            raise ValueError("No Lattice with this name is known or implemented")
MischaD's avatar
MischaD committed
137
        return cls(**kwargs)
138
139

    @classmethod
MischaD's avatar
MischaD committed
140
    def from_order(cls, dimension, order):
MischaD's avatar
MischaD committed
141
142
143
144
145
        """
        :param dimension:
        :param order:
        :return:
        """
MischaD's avatar
MischaD committed
146
147
148
149
150
151
152
153
        if order % 2 != 0:
            order -= 1

        for key, value in cls.BY_NAME.items():
            if value.get("dimension") == dimension and value.get("order") == order:
                return cls(**cls.BY_NAME.get(key))

        raise ValueError(
MischaD's avatar
MischaD committed
154
            "No Lattice known with order {} and dimension {}.".format(order, dimension)
MischaD's avatar
MischaD committed
155
        )
MischaD's avatar
MischaD committed
156
        return cls(dimension=dimension, order=order, shell_list=shell_list, seed=seed)
MischaD's avatar
MischaD committed
157

MischaD's avatar
MischaD committed
158
159
160
161
    @property
    def tensor_space_dimension(self):
        return np.array([len(x) for x in self._possible_tensors]).sum()

MischaD's avatar
MischaD committed
162
163
    @property
    def shells(self):
164
        return self._shells
MischaD's avatar
MischaD committed
165
166
167

    @property
    def velocities(self):
MischaD's avatar
MischaD committed
168
169
        """Get a list of all types of velocities. Velocities with zero weight are included
        For D2O9 this equals [00, 10, 11, 20]"""
170
171
        if not self._discrete_velocities and self._shells:
            self._discrete_velocities = [shell.type for shell in self._shells]
172
        return self._discrete_velocities
MischaD's avatar
MischaD committed
173

174
175
176
177
178
179
    @property
    def weights(self):
        if self._weights:
            return self._weights
        return self.calculate_weights()

MischaD's avatar
MischaD committed
180
181
182
183
184
185
    @property
    def c_s_sq(self):
        if self._c_s_sq:
            return self._c_s_sq
        return self.velocity_set()[0]

MischaD's avatar
MischaD committed
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
    @property
    @ps_installed
    def assignment_collection(self):
        print("okay")
        return

        if self._dimension == 3:
            return self._assignment_collection3d()
        Q = len(self.weights)
        weights = self.weights
        c_s_sq = float(self.c_s_sq)
        c_ix = self.velocities[0]
        c_iy = self.velocities[1]
        src, dst = ps.fields("src, dst({Q}): double[2D]".format(Q=Q))
        density, u_x, u_y, omega = sp.symbols("density, u_x, u_y, omega")

        density_formula = sum([src[0, 0](i) for i in range(Q)])
        vel_x_formula = (sum([src[0, 0](i) * c_ix[i] for i in range(Q)])) / density
        vel_y_formula = (sum([src[0, 0](i) * c_iy[i] for i in range(Q)])) / density

        symbolic_description = [ps.Assignment(density, density_formula),
                                ps.Assignment(u_x, vel_x_formula),
                                ps.Assignment(u_y, vel_y_formula), ]

        for i in range(Q):
            feq_formula = weights[i] * density * (
                    1
                    + (u_x * c_ix[i] + u_y * c_iy[i]) / c_s_sq
                    + (u_x * c_ix[i] + u_y * c_iy[i]) ** 2 / (2 * c_s_sq ** 2)
                    - (u_x ** 2 + u_y ** 2) / (2 * c_s_sq)
            )
            if len(c_ix) >= 17:
                feq_formula += weights[i] * density * (
                        -2 * (1 / (2 * c_s_sq)) ** 2 * (
                            c_ix[i] * u_x ** 3 + c_ix[i] * u_x * u_y ** 2 + c_iy[i] * u_y * u_x ** 2 + c_iy[
                        i] * u_y ** 3)
                        + 4 / 3 * (1 / (2 * c_s_sq)) ** 3 * (u_x * c_ix[i] + u_y * c_iy[i]) ** 3
                )
            if len(c_ix) >= 37:
                feq_formula += weights[i] * density * \
                               (1 / (6 * c_s_sq ** 2) * (
                                       (u_x * c_ix[i] + u_y * c_iy[i]) ** 4 / (4 * c_s_sq ** 2) - (
                                           3 * (u_x ** 2 + u_y ** 2) *
                                           (u_x * c_ix[i] + u_y * c_iy[i]) ** 2) / (2 * c_s_sq) + 3 * (
                                                   u_x ** 4 + u_y ** 4)))

        for i in range(Q):
            symbolic_description.append(
                ps.Assignment(dst[-1 * c_iy[i], c_ix[i]](i), omega * feq_formula + (1 - omega) * src[0, 0](i))
            )
        return ps.AssignmentCollection(symbolic_description)

    def _assignment_collection3d(self):
        pass

MischaD's avatar
MischaD committed
241
242
243
244
245
    def shell_from_type(self, type):
        for shell in self._shells:
            if shell.type == type:
                return shell

MischaD's avatar
MischaD committed
246
    def _calculate_velocity_vectors(self):
MischaD's avatar
MischaD committed
247
        velocity_vectors = []
MischaD's avatar
MischaD committed
248

MischaD's avatar
MischaD committed
249
        group = get_group(self._dimension)
MischaD's avatar
MischaD committed
250
        for shell in range(len(self._shell_list)):
MischaD's avatar
MischaD committed
251
252
253
            velocities = velocities_for_shell(
                dimension=self._dimension, shell=self._shell_list[shell]
            )
MischaD's avatar
MischaD committed
254
            if len(velocities) == 0:
MischaD's avatar
MischaD committed
255
256
257
                raise ImpossibleVelocityShellException(
                    "No velocity sets found for velocity shell {}.".format(shell)
                )
MischaD's avatar
MischaD committed
258

259
            # If e.g. dimension = 2, then shell 25 in ambiguous. Both shell (3,4) and (5,0) are added
MischaD's avatar
MischaD committed
260
261
            subshells = get_subshells(velocities, group)
            if len(subshells) > 1:
MischaD's avatar
MischaD committed
262
                # 3 4, 0 5 can only be given together
263
                velocities = []
MischaD's avatar
MischaD committed
264
                for i_subs, subshell in enumerate(subshells):
265
                    # get type for subshell. Example shell: all vectors of type [0, -1, 0] --> 100
MischaD's avatar
MischaD committed
266
267
268
                    type = "".join(
                        [str(int(x)) for x in sorted(abs(subshell[0]), reverse=True)]
                    )
269
270
                    if type not in self._unwanted_subshells:
                        velocities.append(subshell)
MischaD's avatar
MischaD committed
271
272
            else:
                velocities = [velocities]
MischaD's avatar
MischaD committed
273
            velocity_vectors.extend(velocities)
MischaD's avatar
MischaD committed
274

MischaD's avatar
MischaD committed
275
        self._shells = [Shell([np.array([0] * self._dimension)])]
276
277
        for final_velocities in velocity_vectors:
            self._shells.append(Shell(final_velocities))
MischaD's avatar
MischaD committed
278
279
        self._all_velocity_vectors = velocity_vectors
        return velocity_vectors
MischaD's avatar
MischaD committed
280

MischaD's avatar
MischaD committed
281
    def _fill_lhs(self):
MischaD's avatar
MischaD committed
282
283
284
285
        lhs = []
        for i, rank in enumerate(np.arange(2, self._order + 2, 2)):
            tensor_space_dimension = self._tensor_dimensions[i]
            for j in range(tensor_space_dimension):
MischaD's avatar
MischaD committed
286
                vector = get_random_vector(self._dimension)
MischaD's avatar
MischaD committed
287
288
                # vector = 2 * np.random.rand(self._dimension) - 1
                # vector = vector / np.linalg.norm(vector)
MischaD's avatar
MischaD committed
289
290

                rows = []
MischaD's avatar
MischaD committed
291
                for shell_velocities in self._all_velocity_vectors:
MischaD's avatar
MischaD committed
292
293
294
                    shell_sum = lattice_sum(vector, shell_velocities, rank)
                    rows.append(shell_sum)
                lhs.append(rows)
MischaD's avatar
MischaD committed
295
        self._lhs = np.array(lhs)
MischaD's avatar
MischaD committed
296
        return np.array(lhs)
MischaD's avatar
MischaD committed
297

MischaD's avatar
MischaD committed
298
    def _fill_rhs(self):
MischaD's avatar
MischaD committed
299
300
301
302
303
304
305
306
307
308
309
310
        rhs = []
        cols = self._order // 2
        for k, rank in enumerate(np.arange(2, self._order + 2, 2)):
            for j in range(self._tensor_dimensions[k]):
                rows = []
                for i in range(cols):
                    c_s_power = 2 * i + 2
                    rows.append(1 if c_s_power == rank else 0)
                rhs.append(rows)
        self._rhs = np.array(rhs)
        return np.array(rhs)

MischaD's avatar
MischaD committed
311
    def _svd(self):
MischaD's avatar
MischaD committed
312
313
314
        U, s, V = np.linalg.svd(self._lhs)
        rows = self._lhs.shape[0]
        cols = self._lhs.shape[1]
MischaD's avatar
MischaD committed
315
316
317
318

        self._singular_values = len(s)
        s = np.array([sv if sv > self._svd_tolerance else 0 for sv in s])
        rank = np.count_nonzero(s)
MischaD's avatar
MischaD committed
319
        rhs = np.dot(np.transpose(U), self._rhs)
MischaD's avatar
MischaD committed
320
321
322
323
        if rank < rows:
            overlap = rows - rank
            if np.linalg.norm(rhs[-overlap:]) > self._svd_tolerance:
                # NO Solution
MischaD's avatar
MischaD committed
324
325
326
                raise NoSolutionException(
                    "Add shells or lower the order"
                )  # change shells ?
MischaD's avatar
MischaD committed
327
328
            rows = rank
            s.resize(rows)
MischaD's avatar
MischaD committed
329
            rhs.resize((rows, rhs.shape[1]), refcheck=False)
MischaD's avatar
MischaD committed
330
331
332
333
334
335
336
337
338
339
340
341
            self._singular_values = len(s)

        reduced_rhs = np.zeros((rows, rhs.shape[1]))
        for i, SingularValue in enumerate(s):
            for j in range(rhs.shape[1]):
                reduced_rhs[i, j] = rhs[i, j] / SingularValue

        reduced_rhs = np.zeros((rows, rhs.shape[1]))
        for i, singular_value in enumerate(s):
            reduced_rhs[i, :] = rhs[i, :] / singular_value

        if cols - rows > 0:
MischaD's avatar
MischaD committed
342
343
344
345
            raise InfiniteSolutionsException(
                "Reduce order, remove shells, or use continue.py in "
                "https://github.com/BDuenweg/Lattice-Boltzmann-weights"
            )
MischaD's avatar
MischaD committed
346
347
        self._solution = np.dot(np.transpose(V), reduced_rhs)
        return self._solution
MischaD's avatar
MischaD committed
348

MischaD's avatar
MischaD committed
349
    def _calculate_coefficients(self):
MischaD's avatar
MischaD committed
350
351
352
353
        """Use _all_velocity_vectors and _solution to calculate the coefficients of the
        weigth polynomials. Writes them into _coefficients
        """

MischaD's avatar
MischaD committed
354
355
356
357
358
        solution_columns = self._order // 2
        coefficients = np.zeros((1 + solution_columns))
        coefficients[0] = 1
        for j in range(solution_columns):
            tmp = 0
MischaD's avatar
MischaD committed
359
360
361
            for i in range(len(self._all_velocity_vectors)):
                tmp -= self._solution[i, j] * len(self._all_velocity_vectors[i])
            coefficients[j + 1] = tmp
MischaD's avatar
MischaD committed
362

MischaD's avatar
MischaD committed
363
        coeffs = np.zeros((self._solution.shape[0] + 1, self._solution.shape[1] + 1))
MischaD's avatar
Bugfix    
MischaD committed
364
        coeffs[0, :] = coefficients[-1::-1]
MischaD's avatar
MischaD committed
365
366
367
368
        coeffs[1:, :-1] = self._solution[:, -1::-1]
        self._coefficients = coeffs
        return self._coefficients

MischaD's avatar
MischaD committed
369
    def _calculate_valid_interval(self):
MischaD's avatar
MischaD committed
370
        if not self._weight_polynomials:
MischaD's avatar
MischaD committed
371
372
373
            raise UninitializedAttributeException(
                "Weight Polynomials not give. Calculate them first!"
            )
MischaD's avatar
MischaD committed
374
375
376
377
378
379
380
381
382

        interval = sp.Interval(-oo, oo)
        for equation in self._weight_polynomials:
            roots = equation.real_roots()
            if not roots:
                continue
            roots = [roots[0] - 1] + roots + [roots[-1] + 1]
            cur_interval = sp.EmptySet()
            for i in range(len(roots) - 1):
MischaD's avatar
MischaD committed
383
                mesh_point = (roots[i + 1] + roots[i]) / 2
MischaD's avatar
MischaD committed
384
385
386
387
                if equation.eval(mesh_point) < 0:
                    continue
                # positive weights
                if i == 0:
MischaD's avatar
MischaD committed
388
389
390
                    cur_interval = sp.Union(
                        cur_interval, sp.Interval(-oo, roots[i + 1])
                    )
MischaD's avatar
MischaD committed
391
392
393
                elif i == (len(roots) - 2):
                    cur_interval = sp.Union(cur_interval, sp.Interval(roots[i], oo))
                else:
MischaD's avatar
MischaD committed
394
395
396
                    cur_interval = sp.Union(
                        cur_interval, sp.Interval(roots[i], roots[i + 1])
                    )
MischaD's avatar
MischaD committed
397
398
399
400
            interval = interval.intersect(cur_interval)
        self._interval = sp.Complement(interval, sp.FiniteSet(0))
        return self._interval

MischaD's avatar
MischaD committed
401
402
403
404
405
406
407
    def solution_expected(self):
        """Evaluate if the order, dimension and amount of shells entered are expected to retrieve a solution.
        This is in no way a guarantee on the existence of a solution or whether or not the program
        will finish successful, but it can be used as a quick hint on the existence of a solution.

        Changes the state of self._solution and self._velocity vectors for future calculations."""

MischaD's avatar
MischaD committed
408
409
410
411
412
413
414
415
416
        if self._order % 2 == 1:
            self._order -= 1
            logger.warning("Only odd order valid. Decrease by one")

        for k in np.arange(2, self._order + 2, 2):
            possible_tensors = analyse_tensor_dimension(k)
            self._possible_tensors.append(possible_tensors)
        self._tensor_dimensions = [len(x) for x in self._possible_tensors]

MischaD's avatar
MischaD committed
417
418
419
        self._calculate_velocity_vectors()
        self._fill_lhs()
        self._fill_rhs()
MischaD's avatar
MischaD committed
420

MischaD's avatar
MischaD committed
421
        self._svd()
MischaD's avatar
MischaD committed
422

MischaD's avatar
MischaD committed
423
        return self._solution is not None
MischaD's avatar
Bugfix    
MischaD committed
424

MischaD's avatar
MischaD committed
425
426
    def calculate_polynomials(self, approximate=True):
        """Main algorithm as proposed by D. Spiller and B Duenweg.
MischaD's avatar
MischaD committed
427

MischaD's avatar
MischaD committed
428
429
430
        Divided into the part where the existence of a solution and calculated together with all the velocity
        vectors. If one is expected the coefficients are calculated, the interval evaluated and the weight
        polynomials returned.
MischaD's avatar
MischaD committed
431

MischaD's avatar
MischaD committed
432
433
        :param approximate: Use python Fraction to approximate the calculated polynomial coefficients as rational
         number. This should give better results, since they should have a rational representation.
434

MischaD's avatar
MischaD committed
435
        :return: List of Sympy Polynomials."""
436

MischaD's avatar
MischaD committed
437
438
439
440
441
        self.solution_expected()
        self._calculate_coefficients()

        c_s_sq = sp.Symbol("c_s_sq")
        if approximate:
MischaD's avatar
MischaD committed
442
443
444
445
            self._weight_polynomials = [
                sp.Poly([approximate_ratio(x) for x in pol_coffs], c_s_sq)
                for pol_coffs in self._coefficients
            ]
MischaD's avatar
MischaD committed
446
        else:
MischaD's avatar
MischaD committed
447
448
449
            self._weight_polynomials = [
                sp.Poly(pol_coffs, c_s_sq) for pol_coffs in self._coefficients
            ]
MischaD's avatar
MischaD committed
450
451
452

        self._interval = self._calculate_valid_interval()
        return [] if self._interval == sp.EmptySet else self._weight_polynomials
453

454
    def calculate_weights(self, c_s_sq=None, boundary=None):
455
        """
456
457
458
459
460
461
        Calculate the weights for this lattice for a given speed of sound value c_s_sq, or for
        a given boundary.
        It is first checked, whether a c_s_sq value is given and then boundary.
        If both are not giving, self._boundary is taken which is defined in the
        init function.

462
463
464
465
466
        If no value is given, the infimum is used, i.e. the smallest possible value.

        :param c_s_sq: Float value for the speed of sound, has to be inside of the valid interval
        :param boundary: Get reduced model by using the infimum or supremum of the interval. Fallback value if
            no c_s_sq is given
MischaD's avatar
MischaD committed
467
        :return: weights With zero weight(s) included
468
        """
469

470
471
472
473
474
        if not self._weight_polynomials:
            self.calculate_polynomials()
        if not self._interval:  # Empty Interval
            return []
        if not c_s_sq:
475
476
            if not boundary:
                boundary = self._boundary
477
478
479
480
            if boundary == "inf":
                c_s_sq = self._interval.inf
            else:
                c_s_sq = self._interval.sup
481

482
483
484
        if self._interval.contains(c_s_sq) is sp.EmptySet:  # Invalid c_s_sq input
            return []

MischaD's avatar
MischaD committed
485
        weights = [x.eval(c_s_sq) for x in self._weight_polynomials]
MischaD's avatar
MischaD committed
486
        weights = [x if abs(x) >= 1e-10 else 0 for x in weights]
487
488
489
490
        for i, shell in enumerate(self.shells):
            shell.set_weight(c_s_sq, weights[i])
        self._c_s_sq = c_s_sq
        self._weights = weights
MischaD's avatar
MischaD committed
491
        return weights
MischaD's avatar
MischaD committed
492

493
494
    def velocity_set(self):
        """
MischaD's avatar
MischaD committed
495
496
        Removes unwanted shells, i.e. weight == = and
        returns the values that are important for the Lattice Boltzmann Model, i.e.
497
        a list of the velocity values their corresponding weights and the value for the speed of sound.
MischaD's avatar
MischaD committed
498
499

        :return: (c_s_sq, weights, velocities)
500
        """
MischaD's avatar
MischaD committed
501
502
        if not self.weights:
            self.calculate_weights()
503
504
505
506
507
508

        weights = []
        for shell in self.shells:
            if shell.weight == 0:
                continue
            weights += [shell.weight] * shell.size
MischaD's avatar
MischaD committed
509
        velocities = np.zeros((self._dimension, len(weights)), dtype=np.int32)
510
511
512
513
514
        i = 0
        for shell in self.shells:
            if shell.weight == 0:
                continue
            for velocity in shell.velocities:
MischaD's avatar
MischaD committed
515
                velocities[:, i] = velocity
516
517
518
                i += 1

        return self._c_s_sq, weights, velocities