Skip to content

XLB lattice

Bases: object

This class represents a lattice in the Lattice Boltzmann Method.

It stores the properties of the lattice, including the dimensions, the number of velocities, the velocity vectors, the weights, the moments, and the indices of the opposite, main, right, and left velocities.

The class also provides methods to construct these properties based on the name of the lattice.

Parameters

name: str The name of the lattice, which specifies the dimensions and the number of velocities. For example, “D2Q9” represents a 2D lattice with 9 velocities. precision: str, optional The precision of the computations. It can be “f32/f32”, “f32/f16”, “f64/f64”, “f64/f32”, or “f64/f16”. The first part before the slash is the precision of the computations, and the second part after the slash is the precision of the outputs.

Source code in src/lattice.py 
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 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
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
class Lattice(object):
    """
    This class represents a lattice in the Lattice Boltzmann Method.

    It stores the properties of the lattice, including the dimensions, the number of 
    velocities, the velocity vectors, the weights, the moments, and the indices of the 
    opposite, main, right, and left velocities.

    The class also provides methods to construct these properties based on the name of the 
    lattice.

    Parameters
    ----------
    name: str
        The name of the lattice, which specifies the dimensions and the number of velocities.
        For example, "D2Q9" represents a 2D lattice with 9 velocities.
    precision: str, optional
        The precision of the computations. It can be "f32/f32", "f32/f16", "f64/f64", 
        "f64/f32", or "f64/f16". The first part before the slash is the precision of the 
        computations, and the second part after the slash is the precision of the outputs.
    """
    def __init__(self, name, precision="f32/f32") -> None:
        self.name = name
        dq = re.findall(r"\d+", name)
        self.precision = precision
        self.d = int(dq[0])
        self.q = int(dq[1])
        if precision == "f32/f32" or precision == "f32/f16":
            self.precisionPolicy = jnp.float32
        elif precision == "f64/f64" or precision == "f64/f32" or precision == "f64/f16":
            self.precisionPolicy = jnp.float64
        elif precision == "f16/f16":
            self.precisionPolicy = jnp.float16
        else:
            raise ValueError("precision not supported")

        # Construct the properties of the lattice
        self.c = jnp.array(self.construct_lattice_velocity(), dtype=jnp.int8)
        self.w = jnp.array(self.construct_lattice_weight(), dtype=self.precisionPolicy)
        self.cc = jnp.array(self.construct_lattice_moment(), dtype=self.precisionPolicy)
        self.opp_indices = jnp.array(self.construct_opposite_indices(), dtype=jnp.int8)
        self.main_indices = jnp.array(self.construct_main_indices(), dtype=jnp.int8)
        self.right_indices = np.array(self.construct_right_indices(), dtype=jnp.int8)
        self.left_indices = np.array(self.construct_left_indices(), dtype=jnp.int8)

    def construct_opposite_indices(self):
        """
        This function constructs the indices of the opposite velocities for each velocity.

        The opposite velocity of a velocity is the velocity that has the same magnitude but the 
        opposite direction.

        Returns
        -------
        opposite: numpy.ndarray
            The indices of the opposite velocities.
        """
        c = self.c.T
        opposite = np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)])
        return opposite

    def construct_right_indices(self):
        """
        This function constructs the indices of the velocities that point in the positive 
        x-direction.

        Returns
        -------
        numpy.ndarray
            The indices of the right velocities.
        """
        c = self.c.T
        return np.nonzero(c[:, 0] == 1)[0]

    def construct_left_indices(self):
        """
        This function constructs the indices of the velocities that point in the negative 
        x-direction.

        Returns
        -------
        numpy.ndarray
            The indices of the left velocities.
        """
        c = self.c.T
        return np.nonzero(c[:, 0] == -1)[0]

    def construct_main_indices(self):
        """
        This function constructs the indices of the main velocities.

        The main velocities are the velocities that have a magnitude of 1 in lattice units.

        Returns
        -------
        numpy.ndarray
            The indices of the main velocities.
        """
        c = self.c.T
        if self.d == 2:
            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0]

        elif self.d == 3:
            return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0]

    def construct_lattice_velocity(self):
        """
        This function constructs the velocity vectors of the lattice.

        The velocity vectors are defined based on the name of the lattice. For example, for a D2Q9 
        lattice, there are 9 velocities: (0,0), (1,0), (-1,0), (0,1), (0,-1), (1,1), (-1,-1), 
        (1,-1), and (-1,1).

        Returns
        -------
        c.T: numpy.ndarray
            The velocity vectors of the lattice.
        """
        if self.name == "D2Q9":  # D2Q9
            cx = [0, 0, 0, 1, -1, 1, -1, 1, -1]
            cy = [0, 1, -1, 0, 1, -1, 0, 1, -1]
            c = np.array(tuple(zip(cx, cy)))
        elif self.name == "D3Q19":  # D3Q19
            c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
            c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
        elif self.name == "D3Q27":  # D3Q27
            c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
            # c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
            c = np.array(c)
        else:
            raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")

        return c.T

    def construct_lattice_weight(self):
        """
        This function constructs the weights of the lattice.

        The weights are defined based on the name of the lattice. For example, for a D2Q9 lattice, 
        the weights are 4/9 for the rest velocity, 1/9 for the main velocities, and 1/36 for the 
        diagonal velocities.

        Returns
        -------
        w: numpy.ndarray
            The weights of the lattice.
        """
        # Get the transpose of the lattice vector
        c = self.c.T

        # Initialize the weights to be 1/36
        w = 1.0 / 36.0 * np.ones(self.q)

        # Update the weights for 2D and 3D lattices
        if self.name == "D2Q9":
            w[np.linalg.norm(c, axis=1) < 1.1] = 1.0 / 9.0
            w[0] = 4.0 / 9.0
        elif self.name == "D3Q19":
            w[np.linalg.norm(c, axis=1) < 1.1] = 2.0 / 36.0
            w[0] = 1.0 / 3.0
        elif self.name == "D3Q27":
            cl = np.linalg.norm(c, axis=1)
            w[np.isclose(cl, 1.0, atol=1e-8)] = 2.0 / 27.0
            w[(cl > 1) & (cl <= np.sqrt(2))] = 1.0 / 54.0
            w[(cl > np.sqrt(2)) & (cl <= np.sqrt(3))] = 1.0 / 216.0
            w[0] = 8.0 / 27.0
        else:
            raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")

        # Return the weights
        return w

    def construct_lattice_moment(self):
        """
        This function constructs the moments of the lattice.

        The moments are the products of the velocity vectors, which are used in the computation of 
        the equilibrium distribution functions and the collision operator in the Lattice Boltzmann 
        Method (LBM).

        Returns
        -------
        cc: numpy.ndarray
            The moments of the lattice.
        """
        c = self.c.T
        # Counter for the loop
        cntr = 0

        # nt: number of independent elements of a symmetric tensor
        nt = self.d * (self.d + 1) // 2

        cc = np.zeros((self.q, nt))
        for a in range(0, self.d):
            for b in range(a, self.d):
                cc[:, cntr] = c[:, a] * c[:, b]
                cntr += 1

        return cc

    def __str__(self):
        return self.name

construct_lattice_moment 

1
construct_lattice_moment()

This function constructs the moments of the lattice.

The moments are the products of the velocity vectors, which are used in the computation of the equilibrium distribution functions and the collision operator in the Lattice Boltzmann Method (LBM).

Returns

cc: numpy.ndarray The moments of the lattice.

Source code in src/lattice.py 
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def construct_lattice_moment(self):
    """
    This function constructs the moments of the lattice.

    The moments are the products of the velocity vectors, which are used in the computation of 
    the equilibrium distribution functions and the collision operator in the Lattice Boltzmann 
    Method (LBM).

    Returns
    -------
    cc: numpy.ndarray
        The moments of the lattice.
    """
    c = self.c.T
    # Counter for the loop
    cntr = 0

    # nt: number of independent elements of a symmetric tensor
    nt = self.d * (self.d + 1) // 2

    cc = np.zeros((self.q, nt))
    for a in range(0, self.d):
        for b in range(a, self.d):
            cc[:, cntr] = c[:, a] * c[:, b]
            cntr += 1

    return cc

construct_lattice_velocity 

1
construct_lattice_velocity()

This function constructs the velocity vectors of the lattice.

The velocity vectors are defined based on the name of the lattice. For example, for a D2Q9 lattice, there are 9 velocities: (0,0), (1,0), (-1,0), (0,1), (0,-1), (1,1), (-1,-1), (1,-1), and (-1,1).

Returns

c.T: numpy.ndarray The velocity vectors of the lattice.

Source code in src/lattice.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def construct_lattice_velocity(self):
    """
    This function constructs the velocity vectors of the lattice.

    The velocity vectors are defined based on the name of the lattice. For example, for a D2Q9 
    lattice, there are 9 velocities: (0,0), (1,0), (-1,0), (0,1), (0,-1), (1,1), (-1,-1), 
    (1,-1), and (-1,1).

    Returns
    -------
    c.T: numpy.ndarray
        The velocity vectors of the lattice.
    """
    if self.name == "D2Q9":  # D2Q9
        cx = [0, 0, 0, 1, -1, 1, -1, 1, -1]
        cy = [0, 1, -1, 0, 1, -1, 0, 1, -1]
        c = np.array(tuple(zip(cx, cy)))
    elif self.name == "D3Q19":  # D3Q19
        c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
        c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
    elif self.name == "D3Q27":  # D3Q27
        c = [(x, y, z) for x in [0, -1, 1] for y in [0, -1, 1] for z in [0, -1, 1]]
        # c = np.array([ci for ci in c if np.linalg.norm(ci) < 1.5])
        c = np.array(c)
    else:
        raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")

    return c.T

construct_lattice_weight 

1
construct_lattice_weight()

This function constructs the weights of the lattice.

The weights are defined based on the name of the lattice. For example, for a D2Q9 lattice, the weights are 4/9 for the rest velocity, 1/9 for the main velocities, and 1/36 for the diagonal velocities.

Returns

w: numpy.ndarray The weights of the lattice.

Source code in src/lattice.py 
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def construct_lattice_weight(self):
    """
    This function constructs the weights of the lattice.

    The weights are defined based on the name of the lattice. For example, for a D2Q9 lattice, 
    the weights are 4/9 for the rest velocity, 1/9 for the main velocities, and 1/36 for the 
    diagonal velocities.

    Returns
    -------
    w: numpy.ndarray
        The weights of the lattice.
    """
    # Get the transpose of the lattice vector
    c = self.c.T

    # Initialize the weights to be 1/36
    w = 1.0 / 36.0 * np.ones(self.q)

    # Update the weights for 2D and 3D lattices
    if self.name == "D2Q9":
        w[np.linalg.norm(c, axis=1) < 1.1] = 1.0 / 9.0
        w[0] = 4.0 / 9.0
    elif self.name == "D3Q19":
        w[np.linalg.norm(c, axis=1) < 1.1] = 2.0 / 36.0
        w[0] = 1.0 / 3.0
    elif self.name == "D3Q27":
        cl = np.linalg.norm(c, axis=1)
        w[np.isclose(cl, 1.0, atol=1e-8)] = 2.0 / 27.0
        w[(cl > 1) & (cl <= np.sqrt(2))] = 1.0 / 54.0
        w[(cl > np.sqrt(2)) & (cl <= np.sqrt(3))] = 1.0 / 216.0
        w[0] = 8.0 / 27.0
    else:
        raise ValueError("Supported Lattice types are D2Q9, D3Q19 and D3Q27")

    # Return the weights
    return w

construct_left_indices 

1
construct_left_indices()

This function constructs the indices of the velocities that point in the negative x-direction.

Returns

numpy.ndarray The indices of the left velocities.

Source code in src/lattice.py 
80
81
82
83
84
85
86
87
88
89
90
91
def construct_left_indices(self):
    """
    This function constructs the indices of the velocities that point in the negative 
    x-direction.

    Returns
    -------
    numpy.ndarray
        The indices of the left velocities.
    """
    c = self.c.T
    return np.nonzero(c[:, 0] == -1)[0]

construct_main_indices 

1
construct_main_indices()

This function constructs the indices of the main velocities.

The main velocities are the velocities that have a magnitude of 1 in lattice units.

Returns

numpy.ndarray The indices of the main velocities.

Source code in src/lattice.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def construct_main_indices(self):
    """
    This function constructs the indices of the main velocities.

    The main velocities are the velocities that have a magnitude of 1 in lattice units.

    Returns
    -------
    numpy.ndarray
        The indices of the main velocities.
    """
    c = self.c.T
    if self.d == 2:
        return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) == 1))[0]

    elif self.d == 3:
        return np.nonzero((np.abs(c[:, 0]) + np.abs(c[:, 1]) + np.abs(c[:, 2]) == 1))[0]

construct_opposite_indices 

1
construct_opposite_indices()

This function constructs the indices of the opposite velocities for each velocity.

The opposite velocity of a velocity is the velocity that has the same magnitude but the opposite direction.

Returns

opposite: numpy.ndarray The indices of the opposite velocities.

Source code in src/lattice.py 
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
def construct_opposite_indices(self):
    """
    This function constructs the indices of the opposite velocities for each velocity.

    The opposite velocity of a velocity is the velocity that has the same magnitude but the 
    opposite direction.

    Returns
    -------
    opposite: numpy.ndarray
        The indices of the opposite velocities.
    """
    c = self.c.T
    opposite = np.array([c.tolist().index((-c[i]).tolist()) for i in range(self.q)])
    return opposite

construct_right_indices 

1
construct_right_indices()

This function constructs the indices of the velocities that point in the positive x-direction.

Returns

numpy.ndarray The indices of the right velocities.

Source code in src/lattice.py 
67
68
69
70
71
72
73
74
75
76
77
78
def construct_right_indices(self):
    """
    This function constructs the indices of the velocities that point in the positive 
    x-direction.

    Returns
    -------
    numpy.ndarray
        The indices of the right velocities.
    """
    c = self.c.T
    return np.nonzero(c[:, 0] == 1)[0]

Bases: Lattice

Lattice class for 2D D2Q9 lattice.

D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the Lat tice Boltzmann Method for simulating fluid flows in two dimensions.

Parameters

precision: str, optional The precision of the lattice. The default is “f32/f32”

Source code in src/lattice.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class LatticeD2Q9(Lattice):
    """
    Lattice class for 2D D2Q9 lattice.

    D2Q9 stands for two-dimensional nine-velocity model. It is a common model used in the 
    Lat tice Boltzmann Method for simulating fluid flows in two dimensions.

    Parameters
    ----------
    precision: str, optional
        The precision of the lattice. The default is "f32/f32"
    """
    def __init__(self, precision="f32/f32"):
        super().__init__("D2Q9", precision)
        self._set_constants()

    def _set_constants(self):
        self.cs = jnp.sqrt(3) / 3.0
        self.cs2 = 1.0 / 3.0
        self.inv_cs2 = 3.0
        self.i_s = jnp.asarray(list(range(9)))
        self.im = 3  # Number of imiddles (includes center)
        self.ik = 3  # Number of iknowns or iunknowns

Bases: Lattice

Lattice class for 3D D3Q19 lattice.

D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the Lattice Boltzmann Method for simulating fluid flows in three dimensions.

Parameters

precision: str, optional The precision of the lattice. The default is “f32/f32”

Source code in src/lattice.py 
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
class LatticeD3Q19(Lattice):
    """
    Lattice class for 3D D3Q19 lattice.

    D3Q19 stands for three-dimensional nineteen-velocity model. It is a common model used in the 
    Lattice Boltzmann Method for simulating fluid flows in three dimensions.

    Parameters
    ----------
    precision: str, optional
        The precision of the lattice. The default is "f32/f32"
    """
    def __init__(self, precision="f32/f32"):
        super().__init__("D3Q19", precision)
        self._set_constants()

    def _set_constants(self):
        self.cs = jnp.sqrt(3) / 3.0
        self.cs2 = 1.0 / 3.0
        self.inv_cs2 = 3.0
        self.i_s = jnp.asarray(list(range(19)), dtype=jnp.int8)

        self.im = 9  # Number of imiddles (includes center)
        self.ik = 5  # Number of iknowns or iunknowns

Bases: Lattice

Lattice class for 3D D3Q27 lattice.

D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the Lattice Boltzmann Method for simulating fluid flows in three dimensions.

Parameters

precision: str, optional The precision of the lattice. The default is “f32/f32”

Source code in src/lattice.py 
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
class LatticeD3Q27(Lattice):
    """
    Lattice class for 3D D3Q27 lattice.

    D3Q27 stands for three-dimensional twenty-seven-velocity model. It is a common model used in the 
    Lattice Boltzmann Method for simulating fluid flows in three dimensions.

    Parameters
    ----------
    precision: str, optional
        The precision of the lattice. The default is "f32/f32"
    """

    def __init__(self, precision="f32/f32"):
        super().__init__("D3Q27", precision)
        self._set_constants()

    def _set_constants(self):
        self.cs = jnp.sqrt(3) / 3.0
        self.cs2 = 1.0 / 3.0
        self.inv_cs2 = 3.0
        self.i_s = jnp.asarray(list(range(27)), dtype=jnp.int8)