Skip to content

XLB boundary conditions

Bases: object

Base class for boundary conditions in a LBM simulation.

This class provides a general structure for implementing boundary conditions. It includes methods for preparing the boundary attributes and for applying the boundary condition. Specific boundary conditions should be implemented as subclasses of this class, with the apply method overridden as necessary.

Attributes

lattice : Lattice The lattice used in the simulation. nx: The number of nodes in the x direction. ny: The number of nodes in the y direction. nz: The number of nodes in the z direction. dim : int The number of dimensions in the simulation (2 or 3). precision_policy : PrecisionPolicy The precision policy used in the simulation. indices : array-like The indices of the boundary nodes. name : str or None The name of the boundary condition. This should be set in subclasses. isSolid : bool Whether the boundary condition is for a solid boundary. This should be set in subclasses. isDynamic : bool Whether the boundary condition is dynamic (changes over time). This should be set in subclasses. needsExtraConfiguration : bool Whether the boundary condition requires extra configuration. This should be set in subclasses. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. This should be set in subclasses.

Source code in src/boundary_conditions.py 
  5
  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
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
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
class BoundaryCondition(object):
    """
    Base class for boundary conditions in a LBM simulation.

    This class provides a general structure for implementing boundary conditions. It includes methods for preparing the
    boundary attributes and for applying the boundary condition. Specific boundary conditions should be implemented as
    subclasses of this class, with the `apply` method overridden as necessary.

    Attributes
    ----------
    lattice : Lattice
        The lattice used in the simulation.
    nx:
        The number of nodes in the x direction.
    ny:
        The number of nodes in the y direction.
    nz:
        The number of nodes in the z direction.
    dim : int
        The number of dimensions in the simulation (2 or 3).
    precision_policy : PrecisionPolicy
        The precision policy used in the simulation.
    indices : array-like
        The indices of the boundary nodes.
    name : str or None
        The name of the boundary condition. This should be set in subclasses.
    isSolid : bool
        Whether the boundary condition is for a solid boundary. This should be set in subclasses.
    isDynamic : bool
        Whether the boundary condition is dynamic (changes over time). This should be set in subclasses.
    needsExtraConfiguration : bool
        Whether the boundary condition requires extra configuration. This should be set in subclasses.
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. This should be set in subclasses.
    """

    def __init__(self, indices, gridInfo, precision_policy):
        self.lattice = gridInfo["lattice"]
        self.nx = gridInfo["nx"]
        self.ny = gridInfo["ny"]
        self.nz = gridInfo["nz"]
        self.dim = gridInfo["dim"]
        self.precisionPolicy = precision_policy
        self.indices = indices
        self.name = None
        self.isSolid = False
        self.isDynamic = False
        self.needsExtraConfiguration = False
        self.implementationStep = "PostStreaming"

    def create_local_mask_and_normal_arrays(self, grid_mask):

        """
        Creates local mask and normal arrays for the boundary condition.

        Parameters
        ----------
        grid_mask : array-like
            The grid mask for the lattice.

        Returns
        -------
        None

        Notes
        -----
        This method creates local mask and normal arrays for the boundary condition based on the grid mask.
        If the boundary condition requires extra configuration, the `configure` method is called.
        """

        if self.needsExtraConfiguration:
            boundaryMask = self.get_boundary_mask(grid_mask)
            self.configure(boundaryMask)
            self.needsExtraConfiguration = False

        boundaryMask = self.get_boundary_mask(grid_mask)
        self.normals = self.get_normals(boundaryMask)
        self.imissing, self.iknown = self.get_missing_indices(boundaryMask)
        self.imissingMask, self.iknownMask, self.imiddleMask = self.get_missing_mask(boundaryMask)

        return

    def get_boundary_mask(self, grid_mask):  
        """
        Add jax.device_count() to the self.indices in x-direction, and 1 to the self.indices other directions
        This is to make sure the boundary condition is applied to the correct nodes as grid_mask is
        expanded by (jax.device_count(), 1, 1)

        Parameters
        ----------
        grid_mask : array-like
            The grid mask for the lattice.

        Returns
        -------
        boundaryMask : array-like
        """   
        shifted_indices = np.array(self.indices)
        shifted_indices[0] += device_count()
        shifted_indices[1:] += 1
        # Convert back to tuple
        shifted_indices = tuple(shifted_indices)
        boundaryMask = np.array(grid_mask[shifted_indices])

        return boundaryMask

    def configure(self, boundaryMask):
        """
        Configures the boundary condition.

        Parameters
        ----------
        boundaryMask : array-like
            The grid mask for the boundary voxels.

        Returns
        -------
        None

        Notes
        -----
        This method should be overridden in subclasses if the boundary condition requires extra configuration.
        """
        return

    @partial(jit, static_argnums=(0, 3), inline=True)
    def prepare_populations(self, fout, fin, implementation_step):
        """
        Prepares the distribution functions for the boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The incoming distribution functions.
        fin : jax.numpy.ndarray
            The outgoing distribution functions.
        implementation_step : str
            The step in the lattice Boltzmann method algorithm at which the preparation is applied.

        Returns
        -------
        jax.numpy.ndarray
            The prepared distribution functions.

        Notes
        -----
        This method should be overridden in subclasses if the boundary condition requires preparation of the distribution functions during post-collision or post-streaming. See ExtrapolationBoundaryCondition for an example.
        """   
        return fout

    def get_normals(self, boundaryMask):
        """
        Calculates the normal vectors at the boundary nodes.

        Parameters
        ----------
        boundaryMask : array-like
            The boundary mask for the lattice.

        Returns
        -------
        array-like
            The normal vectors at the boundary nodes.

        Notes
        -----
        This method calculates the normal vectors by dotting the boundary mask with the main lattice directions.
        """
        main_c = self.lattice.c.T[self.lattice.main_indices]
        m = boundaryMask[..., self.lattice.main_indices]
        normals = -np.dot(m, main_c)
        return normals

    def get_missing_indices(self, boundaryMask):
        """
        Returns two int8 arrays the same shape as boundaryMask. The non-zero entries of these arrays indicate missing
        directions that require BCs (imissing) as well as their corresponding opposite directions (iknown).

        Parameters
        ----------
        boundaryMask : array-like
            The boundary mask for the lattice.

        Returns
        -------
        tuple of array-like
            The missing and known indices for the boundary condition.

        Notes
        -----
        This method calculates the missing and known indices based on the boundary mask. The missing indices are the
        non-zero entries of the boundary mask, and the known indices are their corresponding opposite directions.
        """

        # Find imissing, iknown 1-to-1 corresponding indices
        # Note: the "zero" index is used as default value here and won't affect BC computations
        nbd = len(self.indices[0])
        imissing = np.vstack([np.arange(self.lattice.q, dtype='uint8')] * nbd)
        iknown = np.vstack([self.lattice.opp_indices] * nbd)
        imissing[~boundaryMask] = 0
        iknown[~boundaryMask] = 0
        return imissing, iknown

    def get_missing_mask(self, boundaryMask):
        """
        Returns three boolean arrays the same shape as boundaryMask.
        Note: these boundary masks are useful for reduction (eg. summation) operators of selected q-directions.

        Parameters
        ----------
        boundaryMask : array-like
            The boundary mask for the lattice.

        Returns
        -------
        tuple of array-like
            The missing, known, and middle masks for the boundary condition.

        Notes
        -----
        This method calculates the missing, known, and middle masks based on the boundary mask. The missing mask
        is the boundary mask, the known mask is the opposite directions of the missing mask, and the middle mask
        is the directions that are neither missing nor known.
        """
        # Find masks for imissing, iknown and imiddle
        imissingMask = boundaryMask
        iknownMask = imissingMask[:, self.lattice.opp_indices]
        imiddleMask = ~(imissingMask | iknownMask)
        return imissingMask, iknownMask, imiddleMask

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        None

        Notes
        -----
        This method should be overridden in subclasses to implement the specific boundary condition. The method should
        modify the output distribution functions in place to apply the boundary condition.
        """
        pass

    @partial(jit, static_argnums=(0,))
    def equilibrium(self, rho, u):
        """
        Compute equilibrium distribution function.

        Parameters
        ----------
        rho : jax.numpy.ndarray
            The density at each node in the lattice.
        u : jax.numpy.ndarray
            The velocity at each node in the lattice.

        Returns
        -------
        jax.numpy.ndarray
            The equilibrium distribution function at each node in the lattice.

        Notes
        -----
        This method computes the equilibrium distribution function based on the density and velocity. The computation is
        performed in the compute precision specified by the precision policy. The result is not cast to the output precision as
        this is function is used inside other functions that require the compute precision.
        """
        rho, u = self.precisionPolicy.cast_to_compute((rho, u))
        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
        cu = 3.0 * jnp.dot(u, c)
        usqr = 1.5 * jnp.sum(u**2, axis=-1, keepdims=True)
        feq = rho * self.lattice.w * (1.0 + 1.0 * cu + 0.5 * cu**2 - usqr)

        return feq

    @partial(jit, static_argnums=(0,))
    def momentum_flux(self, fneq):
        """
        Compute the momentum flux.

        Parameters
        ----------
        fneq : jax.numpy.ndarray
            The non-equilibrium distribution function at each node in the lattice.

        Returns
        -------
        jax.numpy.ndarray
            The momentum flux at each node in the lattice.

        Notes
        -----
        This method computes the momentum flux by dotting the non-equilibrium distribution function with the lattice
        direction vectors.
        """
        return jnp.dot(fneq, self.lattice.cc)

    @partial(jit, static_argnums=(0,))
    def momentum_exchange_force(self, f_poststreaming, f_postcollision):
        """
        Using the momentum exchange method to compute the boundary force vector exerted on the solid geometry
        based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a
        bounce-back method (e.g. Bouzidi) that accounts for curved boundaries.
        NOTE: this function should be called after BC's are imposed.
        [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation.
            Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339.
        [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving
            curved geometry, Phys. Rev. E 65 (2002) 041203.
        [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange
            algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423.

        Parameters
        ----------
        f_poststreaming : jax.numpy.ndarray
            The post-streaming distribution function at each node in the lattice.
        f_postcollision : jax.numpy.ndarray
            The post-collision distribution function at each node in the lattice.

        Returns
        -------
        jax.numpy.ndarray
            The force exerted on the solid geometry at each boundary node.

        Notes
        -----
        This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. 
        The force is computed based on the post-streaming and post-collision distribution functions. This method
        should be called after the boundary conditions are imposed.
        """
        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
        nbd = len(self.indices[0])
        bindex = np.arange(nbd)[:, None]
        phi = f_postcollision[self.indices][bindex, self.iknown] + \
              f_poststreaming[self.indices][bindex, self.imissing]
        force = jnp.sum(c[:, self.iknown] * phi, axis=-1).T
        return force

apply 

1
apply(fout, fin)

Applies the boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

None

Notes

This method should be overridden in subclasses to implement the specific boundary condition. The method should modify the output distribution functions in place to apply the boundary condition.

Source code in src/boundary_conditions.py 
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    None

    Notes
    -----
    This method should be overridden in subclasses to implement the specific boundary condition. The method should
    modify the output distribution functions in place to apply the boundary condition.
    """
    pass

configure 

1
configure(boundaryMask)

Configures the boundary condition.

Parameters

boundaryMask : array-like The grid mask for the boundary voxels.

Returns

None

Notes

This method should be overridden in subclasses if the boundary condition requires extra configuration.

Source code in src/boundary_conditions.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def configure(self, boundaryMask):
    """
    Configures the boundary condition.

    Parameters
    ----------
    boundaryMask : array-like
        The grid mask for the boundary voxels.

    Returns
    -------
    None

    Notes
    -----
    This method should be overridden in subclasses if the boundary condition requires extra configuration.
    """
    return

create_local_mask_and_normal_arrays 

1
create_local_mask_and_normal_arrays(grid_mask)

Creates local mask and normal arrays for the boundary condition.

Parameters

grid_mask : array-like The grid mask for the lattice.

Returns

None

Notes

This method creates local mask and normal arrays for the boundary condition based on the grid mask. If the boundary condition requires extra configuration, the configure method is called.

Source code in src/boundary_conditions.py 
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
def create_local_mask_and_normal_arrays(self, grid_mask):

    """
    Creates local mask and normal arrays for the boundary condition.

    Parameters
    ----------
    grid_mask : array-like
        The grid mask for the lattice.

    Returns
    -------
    None

    Notes
    -----
    This method creates local mask and normal arrays for the boundary condition based on the grid mask.
    If the boundary condition requires extra configuration, the `configure` method is called.
    """

    if self.needsExtraConfiguration:
        boundaryMask = self.get_boundary_mask(grid_mask)
        self.configure(boundaryMask)
        self.needsExtraConfiguration = False

    boundaryMask = self.get_boundary_mask(grid_mask)
    self.normals = self.get_normals(boundaryMask)
    self.imissing, self.iknown = self.get_missing_indices(boundaryMask)
    self.imissingMask, self.iknownMask, self.imiddleMask = self.get_missing_mask(boundaryMask)

    return

equilibrium 

1
equilibrium(rho, u)

Compute equilibrium distribution function.

Parameters

rho : jax.numpy.ndarray The density at each node in the lattice. u : jax.numpy.ndarray The velocity at each node in the lattice.

Returns

jax.numpy.ndarray The equilibrium distribution function at each node in the lattice.

Notes

This method computes the equilibrium distribution function based on the density and velocity. The computation is performed in the compute precision specified by the precision policy. The result is not cast to the output precision as this is function is used inside other functions that require the compute precision.

Source code in src/boundary_conditions.py 
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
@partial(jit, static_argnums=(0,))
def equilibrium(self, rho, u):
    """
    Compute equilibrium distribution function.

    Parameters
    ----------
    rho : jax.numpy.ndarray
        The density at each node in the lattice.
    u : jax.numpy.ndarray
        The velocity at each node in the lattice.

    Returns
    -------
    jax.numpy.ndarray
        The equilibrium distribution function at each node in the lattice.

    Notes
    -----
    This method computes the equilibrium distribution function based on the density and velocity. The computation is
    performed in the compute precision specified by the precision policy. The result is not cast to the output precision as
    this is function is used inside other functions that require the compute precision.
    """
    rho, u = self.precisionPolicy.cast_to_compute((rho, u))
    c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
    cu = 3.0 * jnp.dot(u, c)
    usqr = 1.5 * jnp.sum(u**2, axis=-1, keepdims=True)
    feq = rho * self.lattice.w * (1.0 + 1.0 * cu + 0.5 * cu**2 - usqr)

    return feq

get_boundary_mask 

1
get_boundary_mask(grid_mask)

Add jax.device_count() to the self.indices in x-direction, and 1 to the self.indices other directions This is to make sure the boundary condition is applied to the correct nodes as grid_mask is expanded by (jax.device_count(), 1, 1)

Parameters

grid_mask : array-like The grid mask for the lattice.

Returns

boundaryMask : array-like

Source code in src/boundary_conditions.py 
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def get_boundary_mask(self, grid_mask):  
    """
    Add jax.device_count() to the self.indices in x-direction, and 1 to the self.indices other directions
    This is to make sure the boundary condition is applied to the correct nodes as grid_mask is
    expanded by (jax.device_count(), 1, 1)

    Parameters
    ----------
    grid_mask : array-like
        The grid mask for the lattice.

    Returns
    -------
    boundaryMask : array-like
    """   
    shifted_indices = np.array(self.indices)
    shifted_indices[0] += device_count()
    shifted_indices[1:] += 1
    # Convert back to tuple
    shifted_indices = tuple(shifted_indices)
    boundaryMask = np.array(grid_mask[shifted_indices])

    return boundaryMask

get_missing_indices 

1
get_missing_indices(boundaryMask)

Returns two int8 arrays the same shape as boundaryMask. The non-zero entries of these arrays indicate missing directions that require BCs (imissing) as well as their corresponding opposite directions (iknown).

Parameters

boundaryMask : array-like The boundary mask for the lattice.

Returns

tuple of array-like The missing and known indices for the boundary condition.

Notes

This method calculates the missing and known indices based on the boundary mask. The missing indices are the non-zero entries of the boundary mask, and the known indices are their corresponding opposite directions.

Source code in src/boundary_conditions.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
205
206
def get_missing_indices(self, boundaryMask):
    """
    Returns two int8 arrays the same shape as boundaryMask. The non-zero entries of these arrays indicate missing
    directions that require BCs (imissing) as well as their corresponding opposite directions (iknown).

    Parameters
    ----------
    boundaryMask : array-like
        The boundary mask for the lattice.

    Returns
    -------
    tuple of array-like
        The missing and known indices for the boundary condition.

    Notes
    -----
    This method calculates the missing and known indices based on the boundary mask. The missing indices are the
    non-zero entries of the boundary mask, and the known indices are their corresponding opposite directions.
    """

    # Find imissing, iknown 1-to-1 corresponding indices
    # Note: the "zero" index is used as default value here and won't affect BC computations
    nbd = len(self.indices[0])
    imissing = np.vstack([np.arange(self.lattice.q, dtype='uint8')] * nbd)
    iknown = np.vstack([self.lattice.opp_indices] * nbd)
    imissing[~boundaryMask] = 0
    iknown[~boundaryMask] = 0
    return imissing, iknown

get_missing_mask 

1
get_missing_mask(boundaryMask)

Returns three boolean arrays the same shape as boundaryMask. Note: these boundary masks are useful for reduction (eg. summation) operators of selected q-directions.

Parameters

boundaryMask : array-like The boundary mask for the lattice.

Returns

tuple of array-like The missing, known, and middle masks for the boundary condition.

Notes

This method calculates the missing, known, and middle masks based on the boundary mask. The missing mask is the boundary mask, the known mask is the opposite directions of the missing mask, and the middle mask is the directions that are neither missing nor known.

Source code in src/boundary_conditions.py 
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
def get_missing_mask(self, boundaryMask):
    """
    Returns three boolean arrays the same shape as boundaryMask.
    Note: these boundary masks are useful for reduction (eg. summation) operators of selected q-directions.

    Parameters
    ----------
    boundaryMask : array-like
        The boundary mask for the lattice.

    Returns
    -------
    tuple of array-like
        The missing, known, and middle masks for the boundary condition.

    Notes
    -----
    This method calculates the missing, known, and middle masks based on the boundary mask. The missing mask
    is the boundary mask, the known mask is the opposite directions of the missing mask, and the middle mask
    is the directions that are neither missing nor known.
    """
    # Find masks for imissing, iknown and imiddle
    imissingMask = boundaryMask
    iknownMask = imissingMask[:, self.lattice.opp_indices]
    imiddleMask = ~(imissingMask | iknownMask)
    return imissingMask, iknownMask, imiddleMask

get_normals 

1
get_normals(boundaryMask)

Calculates the normal vectors at the boundary nodes.

Parameters

boundaryMask : array-like The boundary mask for the lattice.

Returns

array-like The normal vectors at the boundary nodes.

Notes

This method calculates the normal vectors by dotting the boundary mask with the main lattice directions.

Source code in src/boundary_conditions.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
def get_normals(self, boundaryMask):
    """
    Calculates the normal vectors at the boundary nodes.

    Parameters
    ----------
    boundaryMask : array-like
        The boundary mask for the lattice.

    Returns
    -------
    array-like
        The normal vectors at the boundary nodes.

    Notes
    -----
    This method calculates the normal vectors by dotting the boundary mask with the main lattice directions.
    """
    main_c = self.lattice.c.T[self.lattice.main_indices]
    m = boundaryMask[..., self.lattice.main_indices]
    normals = -np.dot(m, main_c)
    return normals

momentum_exchange_force 

1
momentum_exchange_force(f_poststreaming, f_postcollision)

Using the momentum exchange method to compute the boundary force vector exerted on the solid geometry based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a bounce-back method (e.g. Bouzidi) that accounts for curved boundaries. NOTE: this function should be called after BC’s are imposed. [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation. Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339. [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving curved geometry, Phys. Rev. E 65 (2002) 041203. [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423.

Parameters

f_poststreaming : jax.numpy.ndarray The post-streaming distribution function at each node in the lattice. f_postcollision : jax.numpy.ndarray The post-collision distribution function at each node in the lattice.

Returns

jax.numpy.ndarray The force exerted on the solid geometry at each boundary node.

Notes

This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. The force is computed based on the post-streaming and post-collision distribution functions. This method should be called after the boundary conditions are imposed.

Source code in src/boundary_conditions.py 
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
@partial(jit, static_argnums=(0,))
def momentum_exchange_force(self, f_poststreaming, f_postcollision):
    """
    Using the momentum exchange method to compute the boundary force vector exerted on the solid geometry
    based on [1] as described in [3]. Ref [2] shows how [1] is applicable to curved geometries only by using a
    bounce-back method (e.g. Bouzidi) that accounts for curved boundaries.
    NOTE: this function should be called after BC's are imposed.
    [1] A.J.C. Ladd, Numerical simulations of particular suspensions via a discretized Boltzmann equation.
        Part 2 (numerical results), J. Fluid Mech. 271 (1994) 311-339.
    [2] R. Mei, D. Yu, W. Shyy, L.-S. Luo, Force evaluation in the lattice Boltzmann method involving
        curved geometry, Phys. Rev. E 65 (2002) 041203.
    [3] Caiazzo, A., & Junk, M. (2008). Boundary forces in lattice Boltzmann: Analysis of momentum exchange
        algorithm. Computers & Mathematics with Applications, 55(7), 1415-1423.

    Parameters
    ----------
    f_poststreaming : jax.numpy.ndarray
        The post-streaming distribution function at each node in the lattice.
    f_postcollision : jax.numpy.ndarray
        The post-collision distribution function at each node in the lattice.

    Returns
    -------
    jax.numpy.ndarray
        The force exerted on the solid geometry at each boundary node.

    Notes
    -----
    This method computes the force exerted on the solid geometry at each boundary node using the momentum exchange method. 
    The force is computed based on the post-streaming and post-collision distribution functions. This method
    should be called after the boundary conditions are imposed.
    """
    c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
    nbd = len(self.indices[0])
    bindex = np.arange(nbd)[:, None]
    phi = f_postcollision[self.indices][bindex, self.iknown] + \
          f_poststreaming[self.indices][bindex, self.imissing]
    force = jnp.sum(c[:, self.iknown] * phi, axis=-1).T
    return force

momentum_flux 

1
momentum_flux(fneq)

Compute the momentum flux.

Parameters

fneq : jax.numpy.ndarray The non-equilibrium distribution function at each node in the lattice.

Returns

jax.numpy.ndarray The momentum flux at each node in the lattice.

Notes

This method computes the momentum flux by dotting the non-equilibrium distribution function with the lattice direction vectors.

Source code in src/boundary_conditions.py 
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
@partial(jit, static_argnums=(0,))
def momentum_flux(self, fneq):
    """
    Compute the momentum flux.

    Parameters
    ----------
    fneq : jax.numpy.ndarray
        The non-equilibrium distribution function at each node in the lattice.

    Returns
    -------
    jax.numpy.ndarray
        The momentum flux at each node in the lattice.

    Notes
    -----
    This method computes the momentum flux by dotting the non-equilibrium distribution function with the lattice
    direction vectors.
    """
    return jnp.dot(fneq, self.lattice.cc)

prepare_populations 

1
prepare_populations(fout, fin, implementation_step)

Prepares the distribution functions for the boundary condition.

Parameters

fout : jax.numpy.ndarray The incoming distribution functions. fin : jax.numpy.ndarray The outgoing distribution functions. implementation_step : str The step in the lattice Boltzmann method algorithm at which the preparation is applied.

Returns

jax.numpy.ndarray The prepared distribution functions.

Notes

This method should be overridden in subclasses if the boundary condition requires preparation of the distribution functions during post-collision or post-streaming. See ExtrapolationBoundaryCondition for an example.

Source code in src/boundary_conditions.py 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
@partial(jit, static_argnums=(0, 3), inline=True)
def prepare_populations(self, fout, fin, implementation_step):
    """
    Prepares the distribution functions for the boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The incoming distribution functions.
    fin : jax.numpy.ndarray
        The outgoing distribution functions.
    implementation_step : str
        The step in the lattice Boltzmann method algorithm at which the preparation is applied.

    Returns
    -------
    jax.numpy.ndarray
        The prepared distribution functions.

    Notes
    -----
    This method should be overridden in subclasses if the boundary condition requires preparation of the distribution functions during post-collision or post-streaming. See ExtrapolationBoundaryCondition for an example.
    """   
    return fout

Bases: BoundaryCondition

Bounce-back boundary condition for a lattice Boltzmann method simulation.

This class implements a full-way bounce-back boundary condition, where particles hitting the boundary are reflected back in the direction they came from. The boundary condition is applied after the collision step.

Attributes

name : str The name of the boundary condition. For this class, it is “BounceBackFullway”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostCollision”.

Source code in src/boundary_conditions.py 
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
class BounceBack(BoundaryCondition):
    """
    Bounce-back boundary condition for a lattice Boltzmann method simulation.

    This class implements a full-way bounce-back boundary condition, where particles hitting the boundary are reflected
    back in the direction they came from. The boundary condition is applied after the collision step.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "BounceBackFullway".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostCollision".
    """
    def __init__(self, indices, gridInfo, precision_policy):
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "BounceBackFullway"
        self.implementationStep = "PostCollision"

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the bounce-back boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.

        Notes
        -----
        This method applies the bounce-back boundary condition by reflecting the input distribution functions at the
        boundary nodes in the opposite direction.
        """

        return fin[self.indices][..., self.lattice.opp_indices]

apply 

1
apply(fout, fin)

Applies the bounce-back boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Notes

This method applies the bounce-back boundary condition by reflecting the input distribution functions at the boundary nodes in the opposite direction.

Source code in src/boundary_conditions.py 
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the bounce-back boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.

    Notes
    -----
    This method applies the bounce-back boundary condition by reflecting the input distribution functions at the
    boundary nodes in the opposite direction.
    """

    return fin[self.indices][..., self.lattice.opp_indices]

Bases: BoundaryCondition

Moving bounce-back boundary condition for a lattice Boltzmann method simulation.

This class implements a moving bounce-back boundary condition, where particles hitting the boundary are reflected back in the direction they came from, with an additional velocity due to the movement of the boundary. The boundary condition is applied after the collision step.

Attributes

name : str The name of the boundary condition. For this class, it is “BounceBackFullwayMoving”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostCollision”. isDynamic : bool Whether the boundary condition is dynamic (changes over time). For this class, it is True. update_function : function A function that updates the boundary condition. For this class, it is a function that updates the boundary condition based on the current time step. The signature of the function is update_function(time) -> (indices, vel),

Source code in src/boundary_conditions.py 
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
class BounceBackMoving(BoundaryCondition):
    """
    Moving bounce-back boundary condition for a lattice Boltzmann method simulation.

    This class implements a moving bounce-back boundary condition, where particles hitting the boundary are reflected
    back in the direction they came from, with an additional velocity due to the movement of the boundary. The boundary
    condition is applied after the collision step.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "BounceBackFullwayMoving".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostCollision".
    isDynamic : bool
        Whether the boundary condition is dynamic (changes over time). For this class, it is True.
    update_function : function
        A function that updates the boundary condition. For this class, it is a function that updates the boundary
        condition based on the current time step. The signature of the function is `update_function(time) -> (indices, vel)`,

    """
    def __init__(self, gridInfo, precision_policy, update_function=None):
        # We get the indices at time zero to pass to the parent class for initialization
        indices, _ = update_function(0)
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "BounceBackFullwayMoving"
        self.implementationStep = "PostCollision"
        self.isDynamic = True
        self.update_function = jit(update_function)

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin, time):
        """
        Applies the moving bounce-back boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.
        time : int
            The current time step.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.
        """
        indices, vel = self.update_function(time)
        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
        cu = 6.0 * self.lattice.w * jnp.dot(vel, c)
        return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu)

apply 

1
apply(fout, fin, time)

Applies the moving bounce-back boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions. time : int The current time step.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Source code in src/boundary_conditions.py 
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin, time):
    """
    Applies the moving bounce-back boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.
    time : int
        The current time step.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.
    """
    indices, vel = self.update_function(time)
    c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
    cu = 6.0 * self.lattice.w * jnp.dot(vel, c)
    return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu)

Bases: BoundaryCondition

Halfway bounce-back boundary condition for a lattice Boltzmann method simulation.

This class implements a halfway bounce-back boundary condition. The boundary condition is applied after the streaming step.

Attributes

name : str The name of the boundary condition. For this class, it is “BounceBackHalfway”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostStreaming”. needsExtraConfiguration : bool Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True. isSolid : bool Whether the boundary condition represents a solid boundary. For this class, it is True. vel : array-like The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default).

Source code in src/boundary_conditions.py 
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
class BounceBackHalfway(BoundaryCondition):
    """
    Halfway bounce-back boundary condition for a lattice Boltzmann method simulation.

    This class implements a halfway bounce-back boundary condition. The boundary condition is applied after
    the streaming step.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "BounceBackHalfway".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostStreaming".
    needsExtraConfiguration : bool
        Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
    isSolid : bool
        Whether the boundary condition represents a solid boundary. For this class, it is True.
    vel : array-like
        The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default).
    """
    def __init__(self, indices, gridInfo, precision_policy, vel=None):
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "BounceBackHalfway"
        self.implementationStep = "PostStreaming"
        self.needsExtraConfiguration = True
        self.isSolid = True
        self.vel = vel

    def configure(self, boundaryMask):
        """
        Configures the boundary condition.

        Parameters
        ----------
        boundaryMask : array-like
            The grid mask for the boundary voxels.

        Returns
        -------
        None

        Notes
        -----
        This method performs an index shift for the halfway bounce-back boundary condition. It updates the indices of
        the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes.
        """
        # Perform index shift for halfway BB.
        hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
        nbd_orig = len(self.indices[0])
        idx = np.array(self.indices).T
        idx_trg = []
        for i in range(self.lattice.q):
            idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
        indices_new = np.unique(np.vstack(idx_trg), axis=0)
        self.indices = tuple(indices_new.T)
        nbd_modified = len(self.indices[0])
        if (nbd_orig != nbd_modified) and self.vel is not None:
            vel_avg = np.mean(self.vel, axis=0)
            self.vel = jnp.zeros(indices_new.shape, dtype=self.precisionPolicy.compute_dtype) + vel_avg
            print("WARNING: assuming a constant averaged velocity vector is imposed at all BC cells!")

        return

    @partial(jit, static_argnums=(0,))
    def impose_boundary_vel(self, fbd, bindex):
        c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
        cu = 6.0 * self.lattice.w * jnp.dot(self.vel, c)
        fbd = fbd.at[bindex, self.imissing].add(-cu[bindex, self.iknown])
        return fbd

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the halfway bounce-back boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.
        """
        nbd = len(self.indices[0])
        bindex = np.arange(nbd)[:, None]
        fbd = fout[self.indices]

        fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
        if self.vel is not None:
            fbd = self.impose_boundary_vel(fbd, bindex)
        return fbd

apply 

1
apply(fout, fin)

Applies the halfway bounce-back boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Source code in src/boundary_conditions.py 
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the halfway bounce-back boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.
    """
    nbd = len(self.indices[0])
    bindex = np.arange(nbd)[:, None]
    fbd = fout[self.indices]

    fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
    if self.vel is not None:
        fbd = self.impose_boundary_vel(fbd, bindex)
    return fbd

configure 

1
configure(boundaryMask)

Configures the boundary condition.

Parameters

boundaryMask : array-like The grid mask for the boundary voxels.

Returns

None

Notes

This method performs an index shift for the halfway bounce-back boundary condition. It updates the indices of the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes.

Source code in src/boundary_conditions.py 
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
def configure(self, boundaryMask):
    """
    Configures the boundary condition.

    Parameters
    ----------
    boundaryMask : array-like
        The grid mask for the boundary voxels.

    Returns
    -------
    None

    Notes
    -----
    This method performs an index shift for the halfway bounce-back boundary condition. It updates the indices of
    the boundary nodes to be the indices of fluid nodes adjacent of the solid nodes.
    """
    # Perform index shift for halfway BB.
    hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
    nbd_orig = len(self.indices[0])
    idx = np.array(self.indices).T
    idx_trg = []
    for i in range(self.lattice.q):
        idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
    indices_new = np.unique(np.vstack(idx_trg), axis=0)
    self.indices = tuple(indices_new.T)
    nbd_modified = len(self.indices[0])
    if (nbd_orig != nbd_modified) and self.vel is not None:
        vel_avg = np.mean(self.vel, axis=0)
        self.vel = jnp.zeros(indices_new.shape, dtype=self.precisionPolicy.compute_dtype) + vel_avg
        print("WARNING: assuming a constant averaged velocity vector is imposed at all BC cells!")

    return

Bases: BoundaryCondition

Equilibrium boundary condition for a lattice Boltzmann method simulation.

This class implements an equilibrium boundary condition, where the distribution function at the boundary nodes is set to the equilibrium distribution function. The boundary condition is applied after the streaming step.

Attributes

name : str The name of the boundary condition. For this class, it is “EquilibriumBC”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostStreaming”. out : jax.numpy.ndarray The equilibrium distribution function at the boundary nodes.

Source code in src/boundary_conditions.py 
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
class EquilibriumBC(BoundaryCondition):
    """
    Equilibrium boundary condition for a lattice Boltzmann method simulation.

    This class implements an equilibrium boundary condition, where the distribution function at the boundary nodes is
    set to the equilibrium distribution function. The boundary condition is applied after the streaming step.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "EquilibriumBC".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostStreaming".
    out : jax.numpy.ndarray
        The equilibrium distribution function at the boundary nodes.
    """

    def __init__(self, indices, gridInfo, precision_policy, rho, u):
        super().__init__(indices, gridInfo, precision_policy)
        self.out = self.precisionPolicy.cast_to_output(self.equilibrium(rho, u))
        self.name = "EquilibriumBC"
        self.implementationStep = "PostStreaming"

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the equilibrium boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.

        Notes
        -----
        This method applies the equilibrium boundary condition by setting the output distribution functions at the
        boundary nodes to the equilibrium distribution function.
        """
        return self.out

apply 

1
apply(fout, fin)

Applies the equilibrium boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Notes

This method applies the equilibrium boundary condition by setting the output distribution functions at the boundary nodes to the equilibrium distribution function.

Source code in src/boundary_conditions.py 
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the equilibrium boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.

    Notes
    -----
    This method applies the equilibrium boundary condition by setting the output distribution functions at the
    boundary nodes to the equilibrium distribution function.
    """
    return self.out

Bases: BoundaryCondition

Source code in src/boundary_conditions.py 
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
class DoNothing(BoundaryCondition):
    def __init__(self, indices, gridInfo, precision_policy):
        """
        Do-nothing boundary condition for a lattice Boltzmann method simulation.

        This class implements a do-nothing boundary condition, where no action is taken at the boundary nodes. The boundary
        condition is applied after the streaming step.

        Attributes
        ----------
        name : str
            The name of the boundary condition. For this class, it is "DoNothing".
        implementationStep : str
            The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
            it is "PostStreaming".

        Notes
        -----
        This boundary condition enforces skipping of streaming altogether as it sets post-streaming equal to post-collision
        populations (so no streaming at this BC voxels). The problem with returning post-streaming values or "fout[self.indices]
        is that the information that exit the domain on the opposite side of this boundary, would "re-enter". This is because
        we roll the entire array and so the boundary condition acts like a one-way periodic BC. If EquilibriumBC is used as
        the BC for that opposite boundary, then the rolled-in values are taken from the initial condition at equilibrium.
        Otherwise if ZouHe is used for example the simulation looks like a run-down simulation at low-Re. The opposite boundary
        may be even a wall (consider pipebend example). If we correct imissing directions and assign "fin", this method becomes
        much less stable and also one needs to correctly take care of corner cases.
        """
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "DoNothing"
        self.implementationStep = "PostStreaming"


    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the do-nothing boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.

        Notes
        -----
        This method applies the do-nothing boundary condition by simply returning the input distribution functions at the
        boundary nodes.
        """
        return fin[self.indices]

__init__ 

1
__init__(indices, gridInfo, precision_policy)

Do-nothing boundary condition for a lattice Boltzmann method simulation.

This class implements a do-nothing boundary condition, where no action is taken at the boundary nodes. The boundary condition is applied after the streaming step.

Attributes

name : str The name of the boundary condition. For this class, it is “DoNothing”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostStreaming”.

Notes

This boundary condition enforces skipping of streaming altogether as it sets post-streaming equal to post-collision populations (so no streaming at this BC voxels). The problem with returning post-streaming values or “fout[self.indices] is that the information that exit the domain on the opposite side of this boundary, would “re-enter”. This is because we roll the entire array and so the boundary condition acts like a one-way periodic BC. If EquilibriumBC is used as the BC for that opposite boundary, then the rolled-in values are taken from the initial condition at equilibrium. Otherwise if ZouHe is used for example the simulation looks like a run-down simulation at low-Re. The opposite boundary may be even a wall (consider pipebend example). If we correct imissing directions and assign “fin”, this method becomes much less stable and also one needs to correctly take care of corner cases.

Source code in src/boundary_conditions.py 
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
def __init__(self, indices, gridInfo, precision_policy):
    """
    Do-nothing boundary condition for a lattice Boltzmann method simulation.

    This class implements a do-nothing boundary condition, where no action is taken at the boundary nodes. The boundary
    condition is applied after the streaming step.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "DoNothing".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostStreaming".

    Notes
    -----
    This boundary condition enforces skipping of streaming altogether as it sets post-streaming equal to post-collision
    populations (so no streaming at this BC voxels). The problem with returning post-streaming values or "fout[self.indices]
    is that the information that exit the domain on the opposite side of this boundary, would "re-enter". This is because
    we roll the entire array and so the boundary condition acts like a one-way periodic BC. If EquilibriumBC is used as
    the BC for that opposite boundary, then the rolled-in values are taken from the initial condition at equilibrium.
    Otherwise if ZouHe is used for example the simulation looks like a run-down simulation at low-Re. The opposite boundary
    may be even a wall (consider pipebend example). If we correct imissing directions and assign "fin", this method becomes
    much less stable and also one needs to correctly take care of corner cases.
    """
    super().__init__(indices, gridInfo, precision_policy)
    self.name = "DoNothing"
    self.implementationStep = "PostStreaming"

apply 

1
apply(fout, fin)

Applies the do-nothing boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Notes

This method applies the do-nothing boundary condition by simply returning the input distribution functions at the boundary nodes.

Source code in src/boundary_conditions.py 
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the do-nothing boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.

    Notes
    -----
    This method applies the do-nothing boundary condition by simply returning the input distribution functions at the
    boundary nodes.
    """
    return fin[self.indices]

Bases: BoundaryCondition

Zou-He boundary condition for a lattice Boltzmann method simulation.

This class implements the Zou-He boundary condition, which is a non-equilibrium bounce-back boundary condition. It can be used to set inflow and outflow boundary conditions with prescribed pressure or velocity.

Attributes

name : str The name of the boundary condition. For this class, it is “ZouHe”. implementationStep : str The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class, it is “PostStreaming”. type : str The type of the boundary condition. It can be either ‘velocity’ for a prescribed velocity boundary condition, or ‘pressure’ for a prescribed pressure boundary condition. prescribed : float or array-like The prescribed values for the boundary condition. It can be either the prescribed velocities for a ‘velocity’ boundary condition, or the prescribed pressures for a ‘pressure’ boundary condition.

References

Zou, Q., & He, X. (1997). On pressure and velocity boundary conditions for the lattice Boltzmann BGK model. Physics of Fluids, 9(6), 1591-1598. doi:10.1063/1.869307

Source code in src/boundary_conditions.py 
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
class ZouHe(BoundaryCondition):
    """
    Zou-He boundary condition for a lattice Boltzmann method simulation.

    This class implements the Zou-He boundary condition, which is a non-equilibrium bounce-back boundary condition.
    It can be used to set inflow and outflow boundary conditions with prescribed pressure or velocity.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "ZouHe".
    implementationStep : str
        The step in the lattice Boltzmann method algorithm at which the boundary condition is applied. For this class,
        it is "PostStreaming".
    type : str
        The type of the boundary condition. It can be either 'velocity' for a prescribed velocity boundary condition,
        or 'pressure' for a prescribed pressure boundary condition.
    prescribed : float or array-like
        The prescribed values for the boundary condition. It can be either the prescribed velocities for a 'velocity'
        boundary condition, or the prescribed pressures for a 'pressure' boundary condition.

    References
    ----------
    Zou, Q., & He, X. (1997). On pressure and velocity boundary conditions for the lattice Boltzmann BGK model.
    Physics of Fluids, 9(6), 1591-1598. doi:10.1063/1.869307
    """
    def __init__(self, indices, gridInfo, precision_policy, type, prescribed):
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "ZouHe"
        self.implementationStep = "PostStreaming"
        self.type = type
        self.prescribed = prescribed
        self.needsExtraConfiguration = True

    def configure(self, boundaryMask):
        """
        Correct boundary indices to ensure that only voxelized surfaces with normal vectors along main cartesian axes
        are assigned this type of BC.
        """
        nv = np.dot(self.lattice.c, ~boundaryMask.T)
        corner_voxels = np.count_nonzero(nv, axis=0) > 1
        # removed_voxels = np.array(self.indices)[:, corner_voxels]
        self.indices = tuple(np.array(self.indices)[:, ~corner_voxels])
        self.prescribed = self.prescribed[~corner_voxels]
        return

    @partial(jit, static_argnums=(0,), inline=True)
    def calculate_vel(self, fpop, rho):
        """
        Calculate velocity based on the prescribed pressure/density (Zou/He BC)
        """
        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))

        # Return the above unormal as a normal vector which sets the tangential velocities to zero
        vel = unormal * self.normals
        return vel

    @partial(jit, static_argnums=(0,), inline=True)
    def calculate_rho(self, fpop, vel):
        """
        Calculate density based on the prescribed velocity (Zou/He BC)
        """
        unormal = np.sum(self.normals*vel, axis=1)

        rho = (1.0/(1.0 + unormal))[..., None] * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
                                  2.*jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
        return rho

    @partial(jit, static_argnums=(0,), inline=True)
    def calculate_equilibrium(self, fpop):
        """
        This is the ZouHe method of calculating the missing macroscopic variables at the boundary.
        """
        if self.type == 'velocity':
            vel = self.prescribed
            rho = self.calculate_rho(fpop, vel)
        elif self.type == 'pressure':
            rho = self.prescribed
            vel = self.calculate_vel(fpop, rho)
        else:
            raise ValueError(f"type = {self.type} not supported! Use \'pressure\' or \'velocity\'.")

        # compute feq at the boundary
        feq = self.equilibrium(rho, vel)
        return feq

    @partial(jit, static_argnums=(0,), inline=True)
    def bounceback_nonequilibrium(self, fpop, feq):
        """
        Calculate unknown populations using bounce-back of non-equilibrium populations
        a la original Zou & He formulation
        """
        nbd = len(self.indices[0])
        bindex = np.arange(nbd)[:, None]
        fbd = fpop[self.indices]
        fknown = fpop[self.indices][bindex, self.iknown] + feq[bindex, self.imissing] - feq[bindex, self.iknown]
        fbd = fbd.at[bindex, self.imissing].set(fknown)
        return fbd

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, _):
        """
        Applies the Zou-He boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        _ : jax.numpy.ndarray
            The input distribution functions. This is not used in this method.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.

        Notes
        -----
        This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based
        on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions
        based on the non-equilibrium bounce-back method. 
        Tangential velocity is not ensured to be zero by adding transverse contributions based on
        Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher
        Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds.
        """
        # compute the equilibrium based on prescribed values and the type of BC
        feq = self.calculate_equilibrium(fout)

        # set the unknown f populations based on the non-equilibrium bounce-back method
        fbd = self.bounceback_nonequilibrium(fout, feq)


        return fbd

apply 

1
apply(fout, _)

Applies the Zou-He boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. _ : jax.numpy.ndarray The input distribution functions. This is not used in this method.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Notes

This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions based on the non-equilibrium bounce-back method. Tangential velocity is not ensured to be zero by adding transverse contributions based on Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher Reynolds numbers. One needs to use “Regularized” BC at higher Reynolds.

Source code in src/boundary_conditions.py 
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
@partial(jit, static_argnums=(0,))
def apply(self, fout, _):
    """
    Applies the Zou-He boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    _ : jax.numpy.ndarray
        The input distribution functions. This is not used in this method.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.

    Notes
    -----
    This method applies the Zou-He boundary condition by first computing the equilibrium distribution functions based
    on the prescribed values and the type of boundary condition, and then setting the unknown distribution functions
    based on the non-equilibrium bounce-back method. 
    Tangential velocity is not ensured to be zero by adding transverse contributions based on
    Hecth & Harting (2010) (doi:10.1088/1742-5468/2010/01/P01018) as it caused numerical instabilities at higher
    Reynolds numbers. One needs to use "Regularized" BC at higher Reynolds.
    """
    # compute the equilibrium based on prescribed values and the type of BC
    feq = self.calculate_equilibrium(fout)

    # set the unknown f populations based on the non-equilibrium bounce-back method
    fbd = self.bounceback_nonequilibrium(fout, feq)


    return fbd

bounceback_nonequilibrium 

1
bounceback_nonequilibrium(fpop, feq)

Calculate unknown populations using bounce-back of non-equilibrium populations a la original Zou & He formulation

Source code in src/boundary_conditions.py 
741
742
743
744
745
746
747
748
749
750
751
752
@partial(jit, static_argnums=(0,), inline=True)
def bounceback_nonequilibrium(self, fpop, feq):
    """
    Calculate unknown populations using bounce-back of non-equilibrium populations
    a la original Zou & He formulation
    """
    nbd = len(self.indices[0])
    bindex = np.arange(nbd)[:, None]
    fbd = fpop[self.indices]
    fknown = fpop[self.indices][bindex, self.iknown] + feq[bindex, self.imissing] - feq[bindex, self.iknown]
    fbd = fbd.at[bindex, self.imissing].set(fknown)
    return fbd

calculate_equilibrium 

1
calculate_equilibrium(fpop)

This is the ZouHe method of calculating the missing macroscopic variables at the boundary.

Source code in src/boundary_conditions.py 
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
@partial(jit, static_argnums=(0,), inline=True)
def calculate_equilibrium(self, fpop):
    """
    This is the ZouHe method of calculating the missing macroscopic variables at the boundary.
    """
    if self.type == 'velocity':
        vel = self.prescribed
        rho = self.calculate_rho(fpop, vel)
    elif self.type == 'pressure':
        rho = self.prescribed
        vel = self.calculate_vel(fpop, rho)
    else:
        raise ValueError(f"type = {self.type} not supported! Use \'pressure\' or \'velocity\'.")

    # compute feq at the boundary
    feq = self.equilibrium(rho, vel)
    return feq

calculate_rho 

1
calculate_rho(fpop, vel)

Calculate density based on the prescribed velocity (Zou/He BC)

Source code in src/boundary_conditions.py 
712
713
714
715
716
717
718
719
720
721
@partial(jit, static_argnums=(0,), inline=True)
def calculate_rho(self, fpop, vel):
    """
    Calculate density based on the prescribed velocity (Zou/He BC)
    """
    unormal = np.sum(self.normals*vel, axis=1)

    rho = (1.0/(1.0 + unormal))[..., None] * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
                              2.*jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
    return rho

calculate_vel 

1
calculate_vel(fpop, rho)

Calculate velocity based on the prescribed pressure/density (Zou/He BC)

Source code in src/boundary_conditions.py 
700
701
702
703
704
705
706
707
708
709
710
@partial(jit, static_argnums=(0,), inline=True)
def calculate_vel(self, fpop, rho):
    """
    Calculate velocity based on the prescribed pressure/density (Zou/He BC)
    """
    unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
                           2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))

    # Return the above unormal as a normal vector which sets the tangential velocities to zero
    vel = unormal * self.normals
    return vel

configure 

1
configure(boundaryMask)

Correct boundary indices to ensure that only voxelized surfaces with normal vectors along main cartesian axes are assigned this type of BC.

Source code in src/boundary_conditions.py 
688
689
690
691
692
693
694
695
696
697
698
def configure(self, boundaryMask):
    """
    Correct boundary indices to ensure that only voxelized surfaces with normal vectors along main cartesian axes
    are assigned this type of BC.
    """
    nv = np.dot(self.lattice.c, ~boundaryMask.T)
    corner_voxels = np.count_nonzero(nv, axis=0) > 1
    # removed_voxels = np.array(self.indices)[:, corner_voxels]
    self.indices = tuple(np.array(self.indices)[:, ~corner_voxels])
    self.prescribed = self.prescribed[~corner_voxels]
    return

Bases: ZouHe

Regularized boundary condition for a lattice Boltzmann method simulation.

This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure or velocity.

Attributes

name : str The name of the boundary condition. For this class, it is “Regularized”. Qi : numpy.ndarray The Qi tensor, which is used in the regularization of the distribution functions.

References

Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva. Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703

Source code in src/boundary_conditions.py 
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
class Regularized(ZouHe):
    """
    Regularized boundary condition for a lattice Boltzmann method simulation.

    This class implements the regularized boundary condition, which is a non-equilibrium bounce-back boundary condition
    with additional regularization. It can be used to set inflow and outflow boundary conditions with prescribed pressure
    or velocity.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "Regularized".
    Qi : numpy.ndarray
        The Qi tensor, which is used in the regularization of the distribution functions.

    References
    ----------
    Latt, J. (2007). Hydrodynamic limit of lattice Boltzmann equations. PhD thesis, University of Geneva.
    Latt, J., Chopard, B., Malaspinas, O., Deville, M., & Michler, A. (2008). Straight velocity boundaries in the
    lattice Boltzmann method. Physical Review E, 77(5), 056703. doi:10.1103/PhysRevE.77.056703
    """

    def __init__(self, indices, gridInfo, precision_policy, type, prescribed):
        super().__init__(indices, gridInfo, precision_policy, type, prescribed)
        self.name = "Regularized"
        #TODO for Hesam: check to understand why corner cases cause instability here.
        # self.needsExtraConfiguration = False
        self.construct_symmetric_lattice_moment()

    def construct_symmetric_lattice_moment(self):
        """
        Construct the symmetric lattice moment Qi.

        The Qi tensor is used in the regularization of the distribution functions. It is defined as Qi = cc - cs^2*I,
        where cc is the tensor of lattice velocities, cs is the speed of sound, and I is the identity tensor.
        """
        Qi = self.lattice.cc
        if self.dim == 3:
            diagonal = (0, 3, 5)
            offdiagonal = (1, 2, 4)
        elif self.dim == 2:
            diagonal = (0, 2)
            offdiagonal = (1,)
        else:
            raise ValueError(f"dim = {self.dim} not supported")

        # Qi = cc - cs^2*I
        Qi = Qi.at[:, diagonal].set(self.lattice.cc[:, diagonal] - 1./3.)

        # multiply off-diagonal elements by 2 because the Q tensor is symmetric
        Qi = Qi.at[:, offdiagonal].set(self.lattice.cc[:, offdiagonal] * 2.0)

        self.Qi = Qi.T
        return

    @partial(jit, static_argnums=(0,), inline=True)
    def regularize_fpop(self, fpop, feq):
        """
        Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.

        Parameters
        ----------
        fpop : jax.numpy.ndarray
            The distribution functions.
        feq : jax.numpy.ndarray
            The equilibrium distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The regularized distribution functions.
        """

        # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
        f_neq = fpop - feq
        PiNeq = self.momentum_flux(f_neq)
        # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq)

        # Compute double dot product Qi:Pi1
        # QiPi1 = np.zeros_like(fpop)
        # Pi1 = PiNeq
        # QiPi1 = jnp.dot(Qi, Pi1)
        QiPi1 = jnp.dot(PiNeq, self.Qi)

        # assign all populations based on eq 45 of Latt et al (2008)
        # fneq ~ f^1
        fpop1 = 9. / 2. * self.lattice.w[None, :] * QiPi1
        fpop_regularized = feq + fpop1

        return fpop_regularized

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, _):
        """
        Applies the regularized boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        _ : jax.numpy.ndarray
            The input distribution functions. This is not used in this method.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.

        Notes
        -----
        This method applies the regularized boundary condition by first computing the equilibrium distribution functions based
        on the prescribed values and the type of boundary condition, then setting the unknown distribution functions
        based on the non-equilibrium bounce-back method, and finally regularizing the distribution functions.
        """

        # compute the equilibrium based on prescribed values and the type of BC
        feq = self.calculate_equilibrium(fout)

        # set the unknown f populations based on the non-equilibrium bounce-back method
        fbd = self.bounceback_nonequilibrium(fout, feq)

        # Regularize the boundary fpop
        fbd = self.regularize_fpop(fbd, feq)
        return fbd

apply 

1
apply(fout, _)

Applies the regularized boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. _ : jax.numpy.ndarray The input distribution functions. This is not used in this method.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Notes

This method applies the regularized boundary condition by first computing the equilibrium distribution functions based on the prescribed values and the type of boundary condition, then setting the unknown distribution functions based on the non-equilibrium bounce-back method, and finally regularizing the distribution functions.

Source code in src/boundary_conditions.py 
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
@partial(jit, static_argnums=(0,))
def apply(self, fout, _):
    """
    Applies the regularized boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    _ : jax.numpy.ndarray
        The input distribution functions. This is not used in this method.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.

    Notes
    -----
    This method applies the regularized boundary condition by first computing the equilibrium distribution functions based
    on the prescribed values and the type of boundary condition, then setting the unknown distribution functions
    based on the non-equilibrium bounce-back method, and finally regularizing the distribution functions.
    """

    # compute the equilibrium based on prescribed values and the type of BC
    feq = self.calculate_equilibrium(fout)

    # set the unknown f populations based on the non-equilibrium bounce-back method
    fbd = self.bounceback_nonequilibrium(fout, feq)

    # Regularize the boundary fpop
    fbd = self.regularize_fpop(fbd, feq)
    return fbd

construct_symmetric_lattice_moment 

1
construct_symmetric_lattice_moment()

Construct the symmetric lattice moment Qi.

The Qi tensor is used in the regularization of the distribution functions. It is defined as Qi = cc - cs^2*I, where cc is the tensor of lattice velocities, cs is the speed of sound, and I is the identity tensor.

Source code in src/boundary_conditions.py 
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
def construct_symmetric_lattice_moment(self):
    """
    Construct the symmetric lattice moment Qi.

    The Qi tensor is used in the regularization of the distribution functions. It is defined as Qi = cc - cs^2*I,
    where cc is the tensor of lattice velocities, cs is the speed of sound, and I is the identity tensor.
    """
    Qi = self.lattice.cc
    if self.dim == 3:
        diagonal = (0, 3, 5)
        offdiagonal = (1, 2, 4)
    elif self.dim == 2:
        diagonal = (0, 2)
        offdiagonal = (1,)
    else:
        raise ValueError(f"dim = {self.dim} not supported")

    # Qi = cc - cs^2*I
    Qi = Qi.at[:, diagonal].set(self.lattice.cc[:, diagonal] - 1./3.)

    # multiply off-diagonal elements by 2 because the Q tensor is symmetric
    Qi = Qi.at[:, offdiagonal].set(self.lattice.cc[:, offdiagonal] * 2.0)

    self.Qi = Qi.T
    return

regularize_fpop 

1
regularize_fpop(fpop, feq)

Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.

Parameters

fpop : jax.numpy.ndarray The distribution functions. feq : jax.numpy.ndarray The equilibrium distribution functions.

Returns

jax.numpy.ndarray The regularized distribution functions.

Source code in src/boundary_conditions.py 
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
@partial(jit, static_argnums=(0,), inline=True)
def regularize_fpop(self, fpop, feq):
    """
    Regularizes the distribution functions by adding non-equilibrium contributions based on second moments of fpop.

    Parameters
    ----------
    fpop : jax.numpy.ndarray
        The distribution functions.
    feq : jax.numpy.ndarray
        The equilibrium distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The regularized distribution functions.
    """

    # Compute momentum flux of off-equilibrium populations for regularization: Pi^1 = Pi^{neq}
    f_neq = fpop - feq
    PiNeq = self.momentum_flux(f_neq)
    # PiNeq = self.momentum_flux(fpop) - self.momentum_flux(feq)

    # Compute double dot product Qi:Pi1
    # QiPi1 = np.zeros_like(fpop)
    # Pi1 = PiNeq
    # QiPi1 = jnp.dot(Qi, Pi1)
    QiPi1 = jnp.dot(PiNeq, self.Qi)

    # assign all populations based on eq 45 of Latt et al (2008)
    # fneq ~ f^1
    fpop1 = 9. / 2. * self.lattice.w[None, :] * QiPi1
    fpop_regularized = feq + fpop1

    return fpop_regularized

Bases: BoundaryCondition

Extrapolation outflow boundary condition for a lattice Boltzmann method simulation.

This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition that uses extrapolation to avoid strong wave reflections.

Attributes

name : str The name of the boundary condition. For this class, it is “ExtrapolationOutflow”. sound_speed : float The speed of sound in the simulation.

References

Geier, M., Schönherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547. doi:10.1016/j.camwa.2015.05.001.

Source code in src/boundary_conditions.py 
 915
 916
 917
 918
 919
 920
 921
 922
 923
 924
 925
 926
 927
 928
 929
 930
 931
 932
 933
 934
 935
 936
 937
 938
 939
 940
 941
 942
 943
 944
 945
 946
 947
 948
 949
 950
 951
 952
 953
 954
 955
 956
 957
 958
 959
 960
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
class ExtrapolationOutflow(BoundaryCondition):
    """
    Extrapolation outflow boundary condition for a lattice Boltzmann method simulation.

    This class implements the extrapolation outflow boundary condition, which is a type of outflow boundary condition
    that uses extrapolation to avoid strong wave reflections.

    Attributes
    ----------
    name : str
        The name of the boundary condition. For this class, it is "ExtrapolationOutflow".
    sound_speed : float
        The speed of sound in the simulation.

    References
    ----------
    Geier, M., Schönherr, M., Pasquali, A., & Krafczyk, M. (2015). The cumulant lattice Boltzmann equation in three
    dimensions: Theory and validation. Computers & Mathematics with Applications, 70(4), 507–547.
    doi:10.1016/j.camwa.2015.05.001.
    """

    def __init__(self, indices, gridInfo, precision_policy):
        super().__init__(indices, gridInfo, precision_policy)
        self.name = "ExtrapolationOutflow"
        self.needsExtraConfiguration = True
        self.sound_speed = 1./jnp.sqrt(3.)

    def configure(self, boundaryMask):
        """
        Configure the boundary condition by finding neighbouring voxel indices.

        Parameters
        ----------
        boundaryMask : np.ndarray
            The grid mask for the boundary voxels.
        """        
        hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
        idx = np.array(self.indices).T
        idx_trg = []
        for i in range(self.lattice.q):
            idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
        indices_nbr = np.unique(np.vstack(idx_trg), axis=0)
        self.indices_nbr = tuple(indices_nbr.T)

        return

    @partial(jit, static_argnums=(0, 3), inline=True)
    def prepare_populations(self, fout, fin, implementation_step):
        """
        Prepares the distribution functions for the boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The incoming distribution functions.
        fin : jax.numpy.ndarray
            The outgoing distribution functions.
        implementation_step : str
            The step in the lattice Boltzmann method algorithm at which the preparation is applied.

        Returns
        -------
        jax.numpy.ndarray
            The prepared distribution functions.

        Notes
        -----
        Because this function is called "PostCollision", f_poststreaming refers to previous time step or t-1
        """
        f_postcollision = fout
        f_poststreaming = fin
        if implementation_step == 'PostStreaming':
            return f_postcollision
        nbd = len(self.indices[0])
        bindex = np.arange(nbd)[:, None]
        fps_bdr = f_poststreaming[self.indices]
        fps_nbr = f_poststreaming[self.indices_nbr]
        fpc_bdr = f_postcollision[self.indices]
        fpop = fps_bdr[bindex, self.imissing]
        fpop_neighbour = fps_nbr[bindex, self.imissing]
        fpop_extrapolated = self.sound_speed * fpop_neighbour + (1. - self.sound_speed) * fpop

        # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data
        fpc_bdr = fpc_bdr.at[bindex, self.iknown].set(fpop_extrapolated)
        f_postcollision = f_postcollision.at[self.indices].set(fpc_bdr)
        return f_postcollision

    @partial(jit, static_argnums=(0,))
    def apply(self, fout, fin):
        """
        Applies the extrapolation outflow boundary condition.

        Parameters
        ----------
        fout : jax.numpy.ndarray
            The output distribution functions.
        fin : jax.numpy.ndarray
            The input distribution functions.

        Returns
        -------
        jax.numpy.ndarray
            The modified output distribution functions after applying the boundary condition.
        """
        nbd = len(self.indices[0])
        bindex = np.arange(nbd)[:, None]
        fbd = fout[self.indices]
        fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
        return fbd

apply 

1
apply(fout, fin)

Applies the extrapolation outflow boundary condition.

Parameters

fout : jax.numpy.ndarray The output distribution functions. fin : jax.numpy.ndarray The input distribution functions.

Returns

jax.numpy.ndarray The modified output distribution functions after applying the boundary condition.

Source code in src/boundary_conditions.py 
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
@partial(jit, static_argnums=(0,))
def apply(self, fout, fin):
    """
    Applies the extrapolation outflow boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The output distribution functions.
    fin : jax.numpy.ndarray
        The input distribution functions.

    Returns
    -------
    jax.numpy.ndarray
        The modified output distribution functions after applying the boundary condition.
    """
    nbd = len(self.indices[0])
    bindex = np.arange(nbd)[:, None]
    fbd = fout[self.indices]
    fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
    return fbd

configure 

1
configure(boundaryMask)

Configure the boundary condition by finding neighbouring voxel indices.

Parameters

boundaryMask : np.ndarray The grid mask for the boundary voxels.

Source code in src/boundary_conditions.py 
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
def configure(self, boundaryMask):
    """
    Configure the boundary condition by finding neighbouring voxel indices.

    Parameters
    ----------
    boundaryMask : np.ndarray
        The grid mask for the boundary voxels.
    """        
    hasFluidNeighbour = ~boundaryMask[:, self.lattice.opp_indices]
    idx = np.array(self.indices).T
    idx_trg = []
    for i in range(self.lattice.q):
        idx_trg.append(idx[hasFluidNeighbour[:, i], :] + self.lattice.c[:, i])
    indices_nbr = np.unique(np.vstack(idx_trg), axis=0)
    self.indices_nbr = tuple(indices_nbr.T)

    return

prepare_populations 

1
prepare_populations(fout, fin, implementation_step)

Prepares the distribution functions for the boundary condition.

Parameters

fout : jax.numpy.ndarray The incoming distribution functions. fin : jax.numpy.ndarray The outgoing distribution functions. implementation_step : str The step in the lattice Boltzmann method algorithm at which the preparation is applied.

Returns

jax.numpy.ndarray The prepared distribution functions.

Notes

Because this function is called “PostCollision”, f_poststreaming refers to previous time step or t-1

Source code in src/boundary_conditions.py 
 961
 962
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
@partial(jit, static_argnums=(0, 3), inline=True)
def prepare_populations(self, fout, fin, implementation_step):
    """
    Prepares the distribution functions for the boundary condition.

    Parameters
    ----------
    fout : jax.numpy.ndarray
        The incoming distribution functions.
    fin : jax.numpy.ndarray
        The outgoing distribution functions.
    implementation_step : str
        The step in the lattice Boltzmann method algorithm at which the preparation is applied.

    Returns
    -------
    jax.numpy.ndarray
        The prepared distribution functions.

    Notes
    -----
    Because this function is called "PostCollision", f_poststreaming refers to previous time step or t-1
    """
    f_postcollision = fout
    f_poststreaming = fin
    if implementation_step == 'PostStreaming':
        return f_postcollision
    nbd = len(self.indices[0])
    bindex = np.arange(nbd)[:, None]
    fps_bdr = f_poststreaming[self.indices]
    fps_nbr = f_poststreaming[self.indices_nbr]
    fpc_bdr = f_postcollision[self.indices]
    fpop = fps_bdr[bindex, self.imissing]
    fpop_neighbour = fps_nbr[bindex, self.imissing]
    fpop_extrapolated = self.sound_speed * fpop_neighbour + (1. - self.sound_speed) * fpop

    # Use the iknown directions of f_postcollision that leave the domain during streaming to store the BC data
    fpc_bdr = fpc_bdr.at[bindex, self.iknown].set(fpop_extrapolated)
    f_postcollision = f_postcollision.at[self.indices].set(fpc_bdr)
    return f_postcollision