class Multiphase(LBMBase):
"""
Multiphase model based on the Shan-Chen method.
The user supplies an equation of state (EOS). Pressure is evaluated from density and temperature before the effective mass. Both single-component
multiphase and multicomponent multiphase systems are supported.
Parameters
----------
k (list): Modification coefficient used to tune surface tension.
A (numpy.ndarray): Weighting factor for combining the Shan-Chen and Zhang-Chen forces.
g_kkprime (numpy.ndarray): Symmetric component-interaction matrix with shape (n_components, n_components).
wetting_formulation (str or None, optional): Contact-angle scheme. Select "geometric" or
"improved_virtual_density" when a boundary condition defines theta. Defaults to None.
geometric_preprocessing_backend (str, optional): "auto" uses GPU preprocessing when available, "gpu"
requires it, and "cpu" keeps the NumPy implementation. Used only by geometric wetting. Defaults to "auto".
geometric_preprocessing_batch_size (int, optional): Rays processed per GPU at once. Used only by geometric
wetting. Defaults to 65536.
References
----------
1. Shan, Xiaowen, and Hudong Chen. “Lattice Boltzmann Model for Simulating Flows with Multiple Phases and Components.”
Physical Review E 47, no. 3 (March 1, 1993): 1815-19. https://doi.org/10.1103/PhysRevE.47.1815.
2. Yuan, Peng, and Laura Schaefer. “Equations of State in a Lattice Boltzmann Model.”
Physics of Fluids 18, no. 4 (April 3, 2006): 042101. https://doi.org/10.1063/1.2187070.
Notes
-----
1. Boundary conditions are handled separately for each component. For example, define a wall condition once per component in a two-component system.
2. Pytrees contain one leaf per component in the order defined by ``initialize_macroscopic_fields``.
3. Component-specific lists and arrays must use the same ordering.
"""
def __init__(self, **kwargs):
self.n_components = kwargs.get("n_components")
super().__init__(**kwargs)
self.k = kwargs.get("k")
self.A = kwargs.get("A")
self.eos = kwargs.get("EOS", None)
self.g_kkprime = kwargs.get("g_kkprime") # Fluid-fluid interaction strength
self.body_force = kwargs.get("body_force", None)
self.wetting_formulation = kwargs.get("wetting_formulation")
self._uses_default_compute_force = type(self).compute_force is Multiphase.compute_force
self._uses_default_macroscopic_velocity = type(self).macroscopic_velocity is Multiphase.macroscopic_velocity
self._has_wetting_bc = tuple(
any(self._is_wetting_boundary_condition(bc) and bc.theta is not None for bc in component_bcs) for component_bcs in self.BCs
)
if self.wetting_formulation is None and any(self._has_wetting_bc):
raise ValueError(
"A wetting_formulation must be selected when a boundary condition defines theta. "
"Supported schemes: geometric and improved_virtual_density."
)
has_wetting_bc = any(self._has_wetting_bc)
uses_geometric_wetting = self.wetting_formulation == "geometric" and has_wetting_bc
uses_improved_wetting = self.wetting_formulation == "improved_virtual_density" and has_wetting_bc
if uses_geometric_wetting:
self.computed_nearest_next_nearest_nbr = False
self.geometric_preprocessing_backend = kwargs.get("geometric_preprocessing_backend", "auto")
if self.geometric_preprocessing_backend not in ("auto", "cpu", "gpu"):
raise ValueError("geometric_preprocessing_backend must be 'auto', 'cpu', or 'gpu'.")
self.geometric_preprocessing_batch_size = int(kwargs.get("geometric_preprocessing_batch_size", 65536))
if self.geometric_preprocessing_batch_size <= 0:
raise ValueError("geometric_preprocessing_batch_size must be positive.")
self.G_ff = self.compute_ff_greens_function()
self.g_kkprime = jnp.array(self.g_kkprime, dtype=self.precision_policy.compute_dtype)
A_host = np.asarray(self.A)
g_host = np.asarray(self.g_kkprime)
self._psi_interactions = tuple(
tuple((1.0 - A_host[output, source]) * g_host[output, source] != 0.0 for source in range(self.n_components))
for output in range(self.n_components)
)
self._U_interactions = tuple(
tuple(A_host[output, source] != 0.0 for source in range(self.n_components)) for output in range(self.n_components)
)
self._psi_stencil_components = tuple(
any(self._psi_interactions[output][source] for output in range(self.n_components)) for source in range(self.n_components)
)
self._U_stencil_components = tuple(
any(self._U_interactions[output][source] for output in range(self.n_components)) for source in range(self.n_components)
)
self._uses_psi_force = any(self._psi_stencil_components)
self._uses_U_force = any(self._U_stencil_components)
P = PartitionSpec
scalar_spec = P("x", None, None) if self.dim == 2 else P("x", None, None, None)
aux_spec = P("x", None, None)
self.local_improved_wetting = None
if uses_improved_wetting and self.n_devices > 1:
self.local_improved_wetting = jit(
shard_map(
self.local_improved_wetting_m,
mesh=self.mesh,
in_specs=(scalar_spec, scalar_spec, aux_spec, aux_spec, aux_spec, aux_spec),
out_specs=scalar_spec,
check_vma=False,
)
)
self.local_geometric_wetting = None
self._uses_local_neq_bc = self.n_devices > 1 and any(type(bc) in NEQ_BC_TYPES for component_bcs in self.BCs for bc in component_bcs)
if self._uses_local_neq_bc:
self.neq_bc_data = self._make_local_neq_bc_data()
variants = {data[:2] for component_data in self.neq_bc_data for data in component_data if data is not None}
exact_bc = next(
(bc for component_bcs in self.BCs for bc in component_bcs if type(bc) is ExactNonEquilibriumExtrapolation),
None,
)
correction_weights = tuple(np.asarray(exact_bc.w_NEQ).tolist()) if exact_bc is not None else None
self.local_neq_bc_kernels = self._build_local_neq_bc_kernels(scalar_spec, variants, correction_weights)
G_ff_host = np.array(self.G_ff)
# scalar_neighbor_sum: G_ff-weighted scalar neighbor sum, used by the wetting/average-density
# denominator. scalar_force_stencil: the same neighbor structure, weighted by G_ff*c (a per-direction
# vector instead of a scalar), used by the Shan-Chen/Zhang-Chen fluid-fluid force - both share
# _neighbor_stencil_m's one-x-halo-exchange machinery, bound to their own static weights.
self.scalar_neighbor_sum = (
jit(
shard_map(
partial(self._neighbor_stencil_m, weights=G_ff_host),
mesh=self.mesh,
in_specs=scalar_spec,
out_specs=scalar_spec,
check_vma=False,
)
)
if uses_improved_wetting
else None
)
self.scalar_force_stencil = jit(
shard_map(
partial(self._neighbor_stencil_m, weights=G_ff_host[None, :] * np.array(self.c)),
mesh=self.mesh,
in_specs=scalar_spec,
out_specs=scalar_spec,
check_vma=False,
)
)
self.solid_mask_streamed = None
self.average_density_denominator = None
if self.scalar_neighbor_sum is not None:
self.solid_mask_streamed = self.get_solid_mask_streamed()
self.average_density_denominator = [
self.scalar_neighbor_sum(1 - mask) if has_wetting_bc else None
for mask, has_wetting_bc in zip(self.solid_mask_streamed, self._has_wetting_bc, strict=True)
]
if self.average_density_denominator is not None:
for denominator in self.average_density_denominator:
if denominator is not None:
denominator.block_until_ready()
self.geometric_force_mask = None
self.geometric_wetting_data, self.geometric_fluid_mask = self._create_geometric_wetting_data() if uses_geometric_wetting else (None, None)
if uses_geometric_wetting and self.n_devices > 1:
self.local_geometric_wetting = self._build_local_geometric_wetting_kernels(scalar_spec)
self.local_improved_wetting_data = (
self._make_local_improved_wetting_data()
if uses_improved_wetting and self.n_devices > 1
else [[None] * len(component_bcs) for component_bcs in self.BCs]
)
@property
def omega(self):
return self._omega
@omega.setter
def omega(self, value):
if not isinstance(value, list):
raise ValueError("omega must be a list")
self._omega = value
@property
def n_components(self):
return self._n_components
@n_components.setter
def n_components(self, value):
if value is None:
raise ValueError("Number of components cannot be None")
if value <= 0:
raise ValueError("Number of components must be positive")
if not isinstance(value, int):
raise ValueError("Number of components must be an integer")
self._n_components = value
@property
def k(self):
return self._k
@k.setter
def k(self, value):
if value is None:
raise ValueError("Modification coefficient must be provided")
if isinstance(value, float) or isinstance(value, int):
if self.n_components != 1:
raise ValueError("The number of modification coefficients provided does not match the number of components in the system")
self._k = [value]
elif isinstance(value, list):
if len(value) != self.n_components:
raise ValueError("The number of modification coefficients provided does not match the number of components in the system")
self._k = value
@property
def A(self):
return self._A
@A.setter
def A(self, value):
if value is None:
raise ValueError("Weight coefficient value must be provided")
if isinstance(value, np.ndarray):
if value.shape != (self.n_components, self.n_components):
raise ValueError("The dimensions of A should match the number of components")
self._A = jnp.array(value, dtype=self.precision_policy.compute_dtype)
@property
def body_force(self):
return self._body_force
@body_force.setter
def body_force(self, value):
if value is None:
self._body_force = None
if isinstance(value, list):
self._body_force = jnp.array(np.array(value), dtype=self.precision_policy.compute_dtype)
if isinstance(value, np.ndarray):
self._body_force = jnp.array(value, dtype=self.precision_policy.compute_dtype)
@property
def g_kkprime(self):
return self._g_kkprime
@g_kkprime.setter
def g_kkprime(self, value):
if not isinstance(value, np.ndarray) and not isinstance(value, jax.numpy.ndarray):
raise ValueError("g_kkprime must be a numpy array or jax.numpy.ndarray")
if value.shape != (self.n_components, self.n_components):
raise ValueError("g_kkprime must be a matrix of size n_components x n_components")
if not np.allclose(value, np.transpose(value), atol=1e-6):
raise ValueError("g_kkprime must be a symmetric matrix")
self._g_kkprime = np.array(value)
@property
def wetting_formulation(self):
return self._wetting_formulation
@wetting_formulation.setter
def wetting_formulation(self, value):
if value is None or value in ["geometric", "improved_virtual_density"]:
self._wetting_formulation = value
else:
raise ValueError("Invalid wetting scheme type. Supported schemes: None, geometric, and improved_virtual_density.")
def _is_wetting_boundary_condition(self, bc):
"""
Check whether a boundary condition can carry wetting parameters.
Parameters
----------
bc (BoundaryCondition): Boundary condition object.
Returns
-------
(bool): True if the boundary condition supports contact angle data.
"""
return isinstance(bc, (BounceBackHalfway, BounceBack, BounceBackMoving, InterpolatedBounceBackBouzidi, InterpolatedBounceBackDifferentiable))
def _get_solid_indices(self, bc):
"""
Return the solid-node indices of a boundary condition.
BounceBackHalfway (and subclasses) shift bc.indices to the adjacent fluid nodes during configure and keep the
original solid nodes in bc.solid_indices. Wetting data must be built on the solid nodes.
Parameters
----------
bc (BoundaryCondition): Boundary condition object.
Returns
-------
(tuple): Solid-node index tuple.
"""
return getattr(bc, "solid_indices", bc.indices)
def _create_component_solid_mask(self, BC):
"""
Create a solid mask for computing wall normals in geometric wetting.
Parameters
----------
BC (list): Boundary conditions for one component.
Returns
-------
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
"""
shape = (self.nx, self.ny) if self.dim == 2 else (self.nx, self.ny, self.nz)
solid_mask = np.zeros(shape, dtype=bool)
for bc in BC:
if self._is_wetting_boundary_condition(bc):
indices = np.array(self._get_solid_indices(bc), dtype=np.int64)
if self.dim == 2:
bounds = [(self.nx, indices[0]), (self.ny, indices[1])]
else:
bounds = [(self.nx, indices[0]), (self.ny, indices[1]), (self.nz, indices[2])]
valid = np.ones((indices.shape[1],), dtype=bool)
for size, index in bounds:
valid &= (index >= 0) & (index < size)
solid_mask[tuple(indices[:, valid])] = True
return solid_mask
def _create_geometric_force_mask(self, fluid_mask):
"""Disable geometric force where no opposite D3Q lattice-neighbor pair resolves a pore passage."""
if self.dim != 3:
return fluid_mask
resolved = np.zeros_like(fluid_mask)
directions = np.asarray(self.lattice.c, dtype=np.int32).T
for direction in directions:
nonzero = np.flatnonzero(direction)
if len(nonzero) == 0 or direction[nonzero[0]] < 0:
continue
positive = np.roll(fluid_mask, tuple(direction), axis=tuple(range(self.dim)))
negative = np.roll(fluid_mask, tuple(-direction), axis=tuple(range(self.dim)))
resolved |= positive & negative
return fluid_mask & resolved
def _compute_geometric_normals(self, bc, solid_mask, indices=None, source_rows=None):
"""
Compute normals for geometric wetting using boundary data and solid mask.
Parameters
----------
bc (BoundaryCondition): Boundary condition with wettability data.
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
indices (numpy.ndarray, optional): Boundary-node subset. Defaults to all solid indices belonging to ``bc``.
source_rows (numpy.ndarray, optional): Rows of ``bc`` corresponding to ``indices``, used when boundary-provided
normals are available.
Returns
-------
normals (numpy.ndarray): Unit normals pointing from wall nodes toward fluid nodes.
"""
all_indices = np.array(self._get_solid_indices(bc), dtype=np.int64).T
if indices is None:
indices = all_indices
source_rows = np.arange(len(all_indices), dtype=np.int64)
else:
indices = np.asarray(indices, dtype=np.int64)
normals = np.zeros((indices.shape[0], self.dim), dtype=np.float64)
# bc.normals rows correspond to bc.indices; for halfway bounce-back those are the shifted
# fluid nodes, not the solid nodes used here, so fall back to the neighbor-based normals.
if bc.is_solid and hasattr(bc, "normals") and not hasattr(bc, "solid_indices"):
bc_normals = np.asarray(bc.normals, dtype=np.float64)
if bc_normals.shape == (len(all_indices), self.dim):
if source_rows is not None:
bc_normals = bc_normals[source_rows]
normal_norm = np.linalg.norm(bc_normals, axis=1, keepdims=True)
normals = np.divide(bc_normals, normal_norm, out=normals, where=normal_norm > 1e-12)
c = np.array(self.lattice.c).T
c = c[np.linalg.norm(c, axis=1) > 0]
missing_normal = np.linalg.norm(normals, axis=1) <= 1e-12
missing_rows = np.flatnonzero(missing_normal)
accumulated = np.zeros((len(missing_rows), self.dim), dtype=np.float64)
domain_shape = np.asarray(solid_mask.shape, dtype=np.int64)
# Keep lattice-direction accumulation order unchanged, but evaluate every boundary node together.
for direction in c:
neighbors = indices[missing_rows] + direction
in_bounds = np.all((neighbors >= 0) & (neighbors < domain_shape), axis=1)
in_bounds_rows = np.flatnonzero(in_bounds)
if len(in_bounds_rows) == 0:
continue
fluid = ~solid_mask[tuple(neighbors[in_bounds_rows].T)]
accumulated[in_bounds_rows[fluid]] += direction / np.linalg.norm(direction)
normal_norm = np.linalg.norm(accumulated, axis=1, keepdims=True)
normals[missing_rows] = np.divide(accumulated, normal_norm, out=np.zeros_like(accumulated), where=normal_norm > 1e-12)
return normals
def _solid_fluid_interface_mask(self, indices, solid_mask):
"""
Identify solid boundary nodes that touch at least one fluid node.
Parameters
----------
indices (numpy.ndarray): Boundary node coordinates with shape (n, dim).
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
Returns
-------
interface (numpy.ndarray): Boolean mask with True for solid-fluid interface nodes.
"""
directions = np.array(self.lattice.c, dtype=np.int64).T
directions = directions[np.linalg.norm(directions, axis=1) > 0]
interface = np.zeros((indices.shape[0],), dtype=bool)
domain_shape = np.asarray(solid_mask.shape, dtype=np.int64)
# Preserve direction order and early acceptance while replacing the per-node Python loop with bulk indexing.
for direction in directions:
unresolved_rows = np.flatnonzero(~interface)
if len(unresolved_rows) == 0:
break
neighbors = indices[unresolved_rows] + direction
in_bounds = np.all((neighbors >= 0) & (neighbors < domain_shape), axis=1)
in_bounds_rows = np.flatnonzero(in_bounds)
if len(in_bounds_rows) == 0:
continue
fluid = ~solid_mask[tuple(neighbors[in_bounds_rows].T)]
interface[unresolved_rows[in_bounds_rows[fluid]]] = True
return interface
def _first_mesh_intersection(self, indices, directions):
"""
Find first mesh-line intersections from boundary nodes.
Parameters
----------
indices (numpy.ndarray): Boundary node coordinates with shape (n, dim).
directions (numpy.ndarray): Characteristic directions with shape (n, dim).
Returns
-------
points (numpy.ndarray): First intersection points with mesh lines.
"""
eps = 1e-12
abs_dir = np.abs(directions)
t_axis = np.divide(1.0, abs_dir, out=np.full_like(abs_dir, np.inf, dtype=np.float64), where=abs_dir > eps)
t = np.min(t_axis, axis=1)
t = np.where(np.isfinite(t), t, 1.0)
points = indices + t[:, None] * directions
rounded = np.round(points)
return np.where(np.isclose(points, rounded, atol=eps), rounded, points)
def _uses_only_fluid_nodes(self, point, solid_mask):
"""
Check if a multilinear interpolation stencil contains only fluid nodes.
Parameters
----------
point (numpy.ndarray): Off-lattice or on-lattice interpolation point.
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
Returns
-------
(bool): True if every interpolation node with non-zero weight is inside the domain and fluid.
"""
eps = 1e-12
floor_point = np.floor(point)
lower = floor_point.astype(np.int64)
upper = lower + 1
frac = point - floor_point
stencil = []
for corner in np.ndindex(*(2 for _ in range(self.dim))):
index = np.where(corner, upper, lower)
weight = np.prod(np.where(corner, frac, 1.0 - frac))
stencil.append((index, weight))
for index, weight in stencil:
if weight <= eps:
continue
in_bounds = (0 <= index[0] < self.nx) and (0 <= index[1] < self.ny)
if self.dim == 3:
in_bounds = in_bounds and (0 <= index[2] < self.nz)
if not in_bounds or solid_mask[tuple(index)]:
return False
return True
def _uses_only_fluid_nodes_batch(self, points, solid_mask):
"""Vectorized equivalent of ``_uses_only_fluid_nodes`` for a batch of interpolation points."""
eps = 1e-12
floor_points = np.floor(points)
lower = floor_points.astype(np.int64)
upper = lower + 1
frac = points - floor_points
domain_shape = np.asarray(solid_mask.shape, dtype=np.int64)
valid = np.ones((len(points),), dtype=bool)
for corner in np.ndindex(*(2 for _ in range(self.dim))):
corner = np.asarray(corner, dtype=bool)
interpolation_indices = np.where(corner, upper, lower)
weights = np.prod(np.where(corner, frac, 1.0 - frac), axis=1)
active_rows = np.flatnonzero((weights > eps) & valid)
if len(active_rows) == 0:
continue
active_indices = interpolation_indices[active_rows]
in_bounds = np.all((active_indices >= 0) & (active_indices < domain_shape), axis=1)
active_valid = np.zeros((len(active_rows),), dtype=bool)
in_bounds_rows = np.flatnonzero(in_bounds)
if len(in_bounds_rows) > 0:
active_valid[in_bounds_rows] = ~solid_mask[tuple(active_indices[in_bounds_rows].T)]
valid[active_rows] &= active_valid
return valid
def _first_fluid_mesh_intersection(self, indices, directions, solid_mask, return_valid=False, max_intersections=None):
"""
Find first mesh-line intersections with fluid-only interpolation stencils.
Parameters
----------
indices (numpy.ndarray): Boundary node coordinates with shape (n, dim).
directions (numpy.ndarray): Characteristic directions with shape (n, dim).
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
return_valid (bool, optional): If True, return a boolean mask for nodes where a fluid-only stencil was found.
max_intersections (int, optional): Maximum number of candidate mesh intersections to test for each node.
Returns
-------
points (numpy.ndarray): First fluid-side mesh intersection points. If return_valid is True, returns
(points, valid), where valid is a boolean mask for accepted intersections.
"""
eps = 1e-12
indices = np.asarray(indices)
directions = np.asarray(directions)
points = self._first_mesh_intersection(indices, directions)
valid = np.zeros((indices.shape[0],), dtype=bool)
max_steps = self.nx + self.ny if self.dim == 2 else self.nx + self.ny + self.nz
if len(indices) == 0:
return (points, valid) if return_valid else points
# Bound temporary host memory while vectorizing rays. Each batch holds at most about two million float64
# candidates, independent of domain size or porous-interface area.
candidate_slots = self.dim * max_steps
batch_size = max(1, min(len(indices), 2_000_000 // candidate_slots))
steps = np.arange(1, max_steps + 1, dtype=np.float64)
for start in range(0, len(indices), batch_size):
stop = min(start + batch_size, len(indices))
batch_indices = indices[start:stop]
batch_directions = directions[start:stop]
absolute_directions = np.abs(batch_directions)
candidates = np.full((len(batch_indices), self.dim, max_steps), np.inf, dtype=np.float64)
np.divide(
steps[None, None, :],
absolute_directions[..., None],
out=candidates,
where=absolute_directions[..., None] > eps,
)
candidates = np.round(np.sort(candidates.reshape(len(batch_indices), -1), axis=1), decimals=12)
unresolved = np.any(absolute_directions > eps, axis=1)
candidate_counts = np.zeros((len(batch_indices),), dtype=np.int32)
batch_points = points[start:stop]
batch_valid = valid[start:stop]
for candidate_index in range(candidates.shape[1]):
values = candidates[:, candidate_index]
unique = np.isfinite(values)
if candidate_index > 0:
unique &= values != candidates[:, candidate_index - 1]
eligible = unresolved & unique
if max_intersections is not None:
eligible &= candidate_counts < max_intersections
eligible_rows = np.flatnonzero(eligible)
if len(eligible_rows) > 0:
candidate_points = batch_indices[eligible_rows] + values[eligible_rows, None] * batch_directions[eligible_rows]
rounded = np.round(candidate_points)
candidate_points = np.where(np.isclose(candidate_points, rounded, atol=eps), rounded, candidate_points)
accepted = self._uses_only_fluid_nodes_batch(candidate_points, solid_mask)
accepted_rows = eligible_rows[accepted]
batch_points[accepted_rows] = candidate_points[accepted]
batch_valid[accepted_rows] = True
unresolved[accepted_rows] = False
candidate_counts += unique
if max_intersections is not None:
unresolved &= candidate_counts < max_intersections
if not np.any(unresolved):
break
points[start:stop] = batch_points
valid[start:stop] = batch_valid
if return_valid:
return points, valid
return points
def _build_interpolation_data(self, points):
"""
Build multilinear interpolation data for density samples.
Parameters
----------
points (numpy.ndarray): Off-lattice or on-lattice sample points.
Returns
-------
data (tuple): Index arrays and weights for multilinear interpolation.
"""
floor_points = np.floor(points)
lower = floor_points.astype(np.int32)
upper = lower + 1
frac = points - floor_points
compute_dtype = np.dtype(self.precision_policy.compute_dtype)
if self.dim == 2:
x0 = np.clip(lower[:, 0], 0, self.nx - 1)
y0 = np.clip(lower[:, 1], 0, self.ny - 1)
x1 = np.clip(upper[:, 0], 0, self.nx - 1)
y1 = np.clip(upper[:, 1], 0, self.ny - 1)
wx = frac[:, 0]
wy = frac[:, 1]
return (
np.asarray(x0, dtype=np.int32),
np.asarray(y0, dtype=np.int32),
np.asarray(x1, dtype=np.int32),
np.asarray(y1, dtype=np.int32),
np.asarray((1.0 - wx) * (1.0 - wy), dtype=compute_dtype),
np.asarray(wx * (1.0 - wy), dtype=compute_dtype),
np.asarray((1.0 - wx) * wy, dtype=compute_dtype),
np.asarray(wx * wy, dtype=compute_dtype),
)
x0 = np.clip(lower[:, 0], 0, self.nx - 1)
y0 = np.clip(lower[:, 1], 0, self.ny - 1)
z0 = np.clip(lower[:, 2], 0, self.nz - 1)
x1 = np.clip(upper[:, 0], 0, self.nx - 1)
y1 = np.clip(upper[:, 1], 0, self.ny - 1)
z1 = np.clip(upper[:, 2], 0, self.nz - 1)
wx = frac[:, 0]
wy = frac[:, 1]
wz = frac[:, 2]
return (
np.asarray(x0, dtype=np.int32),
np.asarray(y0, dtype=np.int32),
np.asarray(z0, dtype=np.int32),
np.asarray(x1, dtype=np.int32),
np.asarray(y1, dtype=np.int32),
np.asarray(z1, dtype=np.int32),
np.asarray((1.0 - wx) * (1.0 - wy) * (1.0 - wz), dtype=compute_dtype),
np.asarray(wx * (1.0 - wy) * (1.0 - wz), dtype=compute_dtype),
np.asarray((1.0 - wx) * wy * (1.0 - wz), dtype=compute_dtype),
np.asarray(wx * wy * (1.0 - wz), dtype=compute_dtype),
np.asarray((1.0 - wx) * (1.0 - wy) * wz, dtype=compute_dtype),
np.asarray(wx * (1.0 - wy) * wz, dtype=compute_dtype),
np.asarray((1.0 - wx) * wy * wz, dtype=compute_dtype),
np.asarray(wx * wy * wz, dtype=compute_dtype),
)
def _build_geometric_3d_lattice_data(self, indices, normals, solid_mask):
"""
Build lattice-node stencil data for the 3D geometric wetting scheme.
Parameters
----------
indices (numpy.ndarray): Boundary node coordinates with shape (n, 3).
normals (numpy.ndarray): Unit normals pointing from wall nodes toward fluid nodes.
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
Returns
-------
active (numpy.ndarray): Boolean mask with True for nodes that have a valid normal stencil.
normal_2_indices (tuple): JAX index tuple for the second fluid node along the selected normal direction.
tangent_indices (tuple): JAX index tuples for up to two opposite tangent node pairs around the first normal node.
tangent_pair_valid (jax.numpy.ndarray): Boolean mask indicating which tangent pairs are valid for each active node.
Notes
-----
This helper constructs the older lattice-node 3D stencil. The active 3D geometric wetting path uses
_build_geometric_3d_characteristic_data to sample multiple off-lattice characteristic directions.
"""
lattice_directions = np.array(self.lattice.c, dtype=np.int64).T
nonzero_direction_indices = np.flatnonzero(np.linalg.norm(lattice_directions, axis=1) > 0)
nonzero_directions = lattice_directions[nonzero_direction_indices]
unit_directions = nonzero_directions / np.linalg.norm(nonzero_directions, axis=1, keepdims=True)
normal_2_indices = np.zeros_like(indices)
tangent_indices = [np.zeros_like(indices) for _ in range(4)]
active = np.zeros((indices.shape[0],), dtype=bool)
tangent_pair_valid = np.zeros((2, indices.shape[0]), dtype=bool)
for i, (index, normal) in enumerate(zip(indices, normals)):
fluid_direction = None
sorted_direction_indices = np.argsort(-(unit_directions @ normal))
for direction_index in sorted_direction_indices:
direction = nonzero_directions[direction_index]
unit_direction = unit_directions[direction_index]
if unit_direction @ normal <= 1e-12:
break
normal_1 = index + direction
normal_2 = index + 2 * direction
in_bounds_1 = (0 <= normal_1[0] < self.nx) and (0 <= normal_1[1] < self.ny) and (0 <= normal_1[2] < self.nz)
in_bounds_2 = (0 <= normal_2[0] < self.nx) and (0 <= normal_2[1] < self.ny) and (0 <= normal_2[2] < self.nz)
if in_bounds_1 and in_bounds_2 and not solid_mask[tuple(normal_1)] and not solid_mask[tuple(normal_2)]:
fluid_direction = direction
normal_2_indices[i] = normal_2
active[i] = True
break
if fluid_direction is None:
continue
normal_1 = index + fluid_direction
tangent_candidates = nonzero_directions[nonzero_directions @ fluid_direction == 0]
tangent_candidates = tangent_candidates[np.linalg.norm(tangent_candidates, axis=1) > 0]
tangent_scores = np.abs(tangent_candidates @ normal)
sorted_tangents = tangent_candidates[np.argsort(tangent_scores)]
selected_tangents = []
for tangent in sorted_tangents:
if any(np.all(tangent == selected) or np.all(tangent == -selected) for selected in selected_tangents):
continue
plus = normal_1 + tangent
minus = normal_1 - tangent
in_bounds_plus = (0 <= plus[0] < self.nx) and (0 <= plus[1] < self.ny) and (0 <= plus[2] < self.nz)
in_bounds_minus = (0 <= minus[0] < self.nx) and (0 <= minus[1] < self.ny) and (0 <= minus[2] < self.nz)
if in_bounds_plus and in_bounds_minus and not solid_mask[tuple(plus)] and not solid_mask[tuple(minus)]:
pair_index = len(selected_tangents)
selected_tangents.append(tangent)
tangent_indices[2 * pair_index][i] = plus
tangent_indices[2 * pair_index + 1][i] = minus
tangent_pair_valid[pair_index, i] = True
if len(selected_tangents) == 2:
break
for pair_index in range(len(selected_tangents), 2):
tangent_indices[2 * pair_index][i] = normal_1
tangent_indices[2 * pair_index + 1][i] = normal_1
if not np.any(active):
empty = tuple(jnp.array([], dtype=jnp.int32) for _ in range(3))
return active, empty, (), jnp.array([], dtype=jnp.bool_)
return (
active,
tuple(jnp.array(index, dtype=jnp.int32) for index in normal_2_indices[active].T),
tuple(tuple(jnp.array(index, dtype=jnp.int32) for index in point_indices[active].T) for point_indices in tangent_indices),
jnp.array(tangent_pair_valid[:, active], dtype=jnp.bool_),
)
def _build_geometric_gpu_ray_kernel(self, ray_count):
"""Build a Pallas kernel that traces one independent geometric ray per GPU program."""
from jax.experimental import pallas as pl
from jax.experimental.pallas import triton as pltriton
nx, ny, nz = self.nx, self.ny, self.nz
max_steps = nx + ny + nz
candidate_slots = self.dim * max_steps
def round_half_to_even(value, scale=1.0):
scaled = value * scale
lower = jnp.floor(scaled)
fraction = scaled - lower
upper = lower + 1.0
even_lower = jnp.floor(lower * 0.5) * 2.0 == lower
tie = jnp.where(even_lower, lower, upper)
return jnp.where(fraction < 0.5, lower, jnp.where(fraction > 0.5, upper, tie)) / scale
def ray_kernel(indices_ref, directions_ref, solid_ref, points_ref, valid_ref):
ray = pl.program_id(0)
index = tuple(indices_ref[ray, axis] for axis in range(3))
direction = tuple(directions_ref[ray, axis] for axis in range(3))
absolute = tuple(jnp.abs(value) for value in direction)
active = tuple(value > 1e-12 for value in absolute)
first_t = jnp.minimum(
jnp.minimum(
jnp.where(active[0], 1.0 / absolute[0], jnp.inf),
jnp.where(active[1], 1.0 / absolute[1], jnp.inf),
),
jnp.where(active[2], 1.0 / absolute[2], jnp.inf),
)
first_t = jnp.where(jnp.isfinite(first_t), first_t, 1.0)
initial_point = tuple(index[axis] + first_t * direction[axis] for axis in range(3))
initial_point = tuple(
jnp.where(
jnp.isclose(value, rounded := round_half_to_even(value), atol=1e-12),
rounded,
value,
)
for value in initial_point
)
def condition(state):
iteration, _, _, _, _, _, _, _, done = state
return (iteration < candidate_slots) & ~done
def body(state):
iteration, k0, k1, k2, out0, out1, out2, found, _ = state
t0 = jnp.where(active[0] & (k0 <= max_steps), round_half_to_even(k0 / absolute[0], 1e12), jnp.inf)
t1 = jnp.where(active[1] & (k1 <= max_steps), round_half_to_even(k1 / absolute[1], 1e12), jnp.inf)
t2 = jnp.where(active[2] & (k2 <= max_steps), round_half_to_even(k2 / absolute[2], 1e12), jnp.inf)
value = jnp.minimum(jnp.minimum(t0, t1), t2)
finite = jnp.isfinite(value)
safe_value = jnp.where(finite, value, 0.0)
point = tuple(index[axis] + safe_value * direction[axis] for axis in range(3))
point = tuple(
jnp.where(
jnp.isclose(component, rounded := round_half_to_even(component), atol=1e-12),
rounded,
component,
)
for component in point
)
lower = tuple(jnp.floor(component).astype(jnp.int32) for component in point)
upper = tuple(component + 1 for component in lower)
frac = tuple(point[axis] - lower[axis] for axis in range(3))
fluid = finite
for corner in np.ndindex(2, 2, 2):
coordinates = tuple(upper[axis] if corner[axis] else lower[axis] for axis in range(3))
weight = (
(frac[0] if corner[0] else 1.0 - frac[0])
* (frac[1] if corner[1] else 1.0 - frac[1])
* (frac[2] if corner[2] else 1.0 - frac[2])
)
in_bounds = (
(coordinates[0] >= 0)
& (coordinates[0] < nx)
& (coordinates[1] >= 0)
& (coordinates[1] < ny)
& (coordinates[2] >= 0)
& (coordinates[2] < nz)
)
safe = (
jnp.clip(coordinates[0], 0, nx - 1),
jnp.clip(coordinates[1], 0, ny - 1),
jnp.clip(coordinates[2], 0, nz - 1),
)
fluid &= (weight <= 1e-12) | (in_bounds & ~solid_ref[safe])
return (
iteration + 1,
k0 + (t0 == value),
k1 + (t1 == value),
k2 + (t2 == value),
jnp.where(fluid, point[0], out0),
jnp.where(fluid, point[1], out1),
jnp.where(fluid, point[2], out2),
found | fluid,
fluid | ~finite,
)
result = jax.lax.while_loop(condition, body, (0, 1, 1, 1, *initial_point, False, False))
points_ref[ray, 0] = result[4]
points_ref[ray, 1] = result[5]
points_ref[ray, 2] = result[6]
valid_ref[ray] = result[7]
return pl.pallas_call(
ray_kernel,
out_shape=(
jax.ShapeDtypeStruct((ray_count, 3), jnp.float64),
jax.ShapeDtypeStruct((ray_count,), jnp.bool_),
),
grid=(ray_count,),
compiler_params=pltriton.CompilerParams(num_warps=1),
name="geometric_fluid_intersections",
)
def _first_fluid_mesh_intersection_gpu(self, indices, directions, solid_mask):
"""Trace 3D characteristic rays on every local GPU with bounded temporary device memory."""
devices = tuple(device for device in jax.local_devices() if device.platform == "gpu")
if not devices:
raise RuntimeError("GPU geometric preprocessing requested, but no local GPU is available.")
indices = np.asarray(indices, dtype=np.int32)
directions = np.asarray(directions, dtype=np.float64)
points = np.empty_like(directions)
valid = np.empty((len(indices),), dtype=bool)
batch_size = self.geometric_preprocessing_batch_size
with jax.enable_x64():
kernel = getattr(self, "_geometric_gpu_ray_kernel", None)
if kernel is None:
kernel = self._build_geometric_gpu_ray_kernel(batch_size)
self._geometric_gpu_ray_kernel = kernel
device_masks = tuple(jax.device_put(jnp.asarray(solid_mask, dtype=jnp.bool_), device) for device in devices)
group_size = batch_size * len(devices)
for group_start in range(0, len(indices), group_size):
pending = []
for device_index, (device, device_mask) in enumerate(zip(devices, device_masks, strict=True)):
start = group_start + device_index * batch_size
stop = min(start + batch_size, len(indices))
if start >= stop:
continue
count = stop - start
batch_indices = np.zeros((batch_size, 3), dtype=np.int32)
batch_directions = np.zeros((batch_size, 3), dtype=np.float64)
batch_indices[:count] = indices[start:stop]
batch_directions[:count] = directions[start:stop]
with jax.default_device(device):
output = kernel(
jax.device_put(batch_indices, device),
jax.device_put(batch_directions, device),
device_mask,
)
pending.append((start, stop, output))
for start, stop, output in pending:
count = stop - start
points[start:stop] = np.asarray(output[0])[:count]
valid[start:stop] = np.asarray(output[1])[:count]
return points, valid
def _nearest_fluid_lattice_points(self, indices, normals, solid_mask):
"""Select the adjacent fluid lattice node most aligned with each wall normal."""
if len(indices) == 0:
return np.empty((0, self.dim), dtype=np.float64)
directions = np.asarray(self.lattice.c, dtype=np.int32).T
directions = directions[np.any(directions != 0, axis=1)]
candidates = np.asarray(indices, dtype=np.int32)[:, None, :] + directions[None, :, :]
domain_shape = np.asarray(solid_mask.shape, dtype=np.int32)
in_bounds = np.all((candidates >= 0) & (candidates < domain_shape), axis=-1)
safe_candidates = np.clip(candidates, 0, domain_shape - 1)
fluid = in_bounds & ~solid_mask[tuple(np.moveaxis(safe_candidates, -1, 0))]
if not np.all(np.any(fluid, axis=1)):
raise RuntimeError("Geometric wetting interface node has no adjacent fluid lattice node.")
unit_directions = directions / np.linalg.norm(directions, axis=1, keepdims=True)
alignment = np.asarray(normals, dtype=np.float64) @ unit_directions.T
choice = np.argmax(np.where(fluid, alignment, -np.inf), axis=1)
return np.asarray(candidates[np.arange(len(indices)), choice], dtype=np.float64)
def _repair_invalid_geometric_point_data(self, point_data, validity, indices, normals, solid_mask, fallback_points):
"""Replace unresolved characteristic samples with fluid-only samples during preprocessing."""
invalid_count = sum(np.count_nonzero(~valid) for valid in validity)
if invalid_count == 0:
return tuple(point_data)
has_valid = np.logical_or.reduce(validity)
all_invalid = ~has_valid
if np.any(all_invalid):
fallback_points[all_invalid] = self._nearest_fluid_lattice_points(
indices[all_invalid],
normals[all_invalid],
solid_mask,
)
repaired_data = []
for data, valid in zip(point_data, validity, strict=True):
invalid_rows = np.flatnonzero(~valid)
if len(invalid_rows):
replacements = self._build_interpolation_data(fallback_points[invalid_rows])
for values, replacement in zip(data, replacements, strict=True):
values[invalid_rows] = replacement
repaired_data.append(data)
logger.info(
"Replaced %d unresolved geometric rays at %d wall nodes (%d had no valid ray).",
invalid_count,
np.count_nonzero(~np.logical_and.reduce(validity)),
np.count_nonzero(all_invalid),
)
return tuple(repaired_data)
def _build_geometric_3d_characteristic_data(self, indices, normals, theta, solid_mask):
"""
Build cone-sampled interpolation data for 3D geometric wetting.
Parameters
----------
indices (numpy.ndarray): Boundary node coordinates with shape (n, 3).
normals (numpy.ndarray): Unit normals pointing from wall nodes toward fluid nodes.
theta (numpy.ndarray): Contact angle in radians for each boundary node.
solid_mask (numpy.ndarray): Boolean mask with True on boundary nodes.
Returns
-------
point_data (tuple): Multilinear interpolation data for characteristic directions sampled on the contact-angle
cone around the wall normal.
Notes
-----
Assumes theta is prescribed in radians. The 3D construction samples eight azimuthal directions on the cone
instead of selecting only two characteristic directions, then the wall density is selected from the extrema of
those samples in apply_contact_angle.
"""
sample_count = 8
azimuths = np.linspace(0.0, 2.0 * np.pi, sample_count, endpoint=False)
reference_x = np.array([1.0, 0.0, 0.0], dtype=np.float64)
reference_y = np.array([0.0, 1.0, 0.0], dtype=np.float64)
references = np.where((np.abs(normals[:, 0]) > 0.9)[:, None], reference_y, reference_x)
tangent_1 = np.cross(normals, references)
tangent_1_norm = np.linalg.norm(tangent_1, axis=1, keepdims=True)
tangent_1 = np.divide(tangent_1, tangent_1_norm, out=np.zeros_like(tangent_1), where=tangent_1_norm > 1e-12)
tangent_2 = np.cross(normals, tangent_1)
angle = np.pi / 2 - theta
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
gpu_devices = tuple(device for device in jax.local_devices() if device.platform == "gpu")
use_gpu = self.geometric_preprocessing_backend == "gpu" or (
self.geometric_preprocessing_backend == "auto" and bool(gpu_devices) and len(indices) * sample_count >= 262144
)
if self.geometric_preprocessing_backend == "gpu" and not gpu_devices:
raise RuntimeError("GPU geometric preprocessing requested, but no local GPU is available.")
point_data = []
validity = []
fallback_points = np.zeros_like(normals)
has_fallback = np.zeros((len(indices),), dtype=bool)
crosses_solid = np.zeros((len(indices),), dtype=bool)
for azimuth in azimuths:
tangent_direction = np.cos(azimuth) * tangent_1 + np.sin(azimuth) * tangent_2
directions = cos_angle[:, None] * normals + sin_angle[:, None] * tangent_direction
if use_gpu:
try:
points, valid = self._first_fluid_mesh_intersection_gpu(indices, directions, solid_mask)
except Exception as error:
if self.geometric_preprocessing_backend == "gpu":
raise
logger.warning("GPU geometric preprocessing failed; using CPU fallback: %s", error)
use_gpu = False
points, valid = self._first_fluid_mesh_intersection(indices, directions, solid_mask, return_valid=True)
else:
points, valid = self._first_fluid_mesh_intersection(indices, directions, solid_mask, return_valid=True)
# A characteristic belongs to this wall only when its first mesh crossing is fluid. Continuing through
# solid until a remote pore is found can import a density from the opposite side of a grain. In thin
# throats, the extrema selection then creates a liquid-density wall inside vapor and can drive the
# adjacent fluid density negative in one step. Treat later crossings like unresolved rays so they use a
# local same-wall sample (or the nearest adjacent fluid node when every ray is unresolved).
first_points = self._first_mesh_intersection(indices, directions)
first_distance = np.linalg.norm(first_points - indices, axis=1)
point_distance = np.linalg.norm(points - indices, axis=1)
local = point_distance <= first_distance + 1e-10
crosses_solid |= valid & ~local
valid &= local
first_valid = valid & ~has_fallback
fallback_points[first_valid] = points[first_valid]
has_fallback |= valid
validity.append(valid)
point_data.append(self._build_interpolation_data(points))
# A partial contact-angle cone has a directional bias. If any otherwise-valid ray first crossed solid,
# replace the whole cone at that wall by its nearest adjacent fluid node. Fully resolved cones—and thus
# ordinary flat/curved-wall behavior—remain unchanged.
if np.any(crosses_solid):
for valid in validity:
valid[crosses_solid] = False
return self._repair_invalid_geometric_point_data(
point_data,
validity,
indices,
normals,
solid_mask,
fallback_points,
)
def _create_geometric_wetting_data(self, localize=True):
"""
Precompute interpolation data for geometric wetting.
Returns
-------
geometric_wetting_data (list): Component-wise interpolation data for wetted boundary nodes.
geometric_fluid_mask (list): Component-wise boolean masks (jax.numpy.ndarray) with True on fluid nodes, used to clamp wall densities to the fluid density range.
Notes
-----
Assumes theta is prescribed in radians for each wetted boundary node. Boundary conditions without theta are included in the solid mask but skipped for contact-angle
interpolation. The 2D implementation keeps the original two-characteristic construction. The 3D implementation samples multiple characteristic directions on the
contact-angle cone around each wall normal and stores interpolation data for every sample.
References
----------
1. Fei, Linlin, Feifei Qin, Jianlin Zhao, Dominique Derome, and Jan Carmeliet.
“Lattice Boltzmann Modelling of Isothermal Two-Component Evaporation in Porous Media.”
Journal of Fluid Mechanics 955 (January 2023): A18. doi: 10.1017/jfm.2022.1048.
2. Wang, Lei, Hai-bo Huang, and Xi-Yun Lu. “Scheme for Contact Angle and Its Hysteresis in a Multiphase Lattice
Boltzmann Method.”
Physical Review E 87, no. 1 (2013): 013301. doi:10.1103/PhysRevE.87.013301.
"""
geometric_wetting_data = []
geometric_fluid_mask = []
geometric_force_mask = []
characteristics_time = 0.0
for BC in self.BCs:
solid_mask = self._create_component_solid_mask(BC)
fluid_mask_host = ~solid_mask[..., None]
geometric_fluid_mask.append(self.distributed_array_init(fluid_mask_host.shape, jnp.bool_, init_val=fluid_mask_host))
force_mask_host = self._create_geometric_force_mask(fluid_mask_host[..., 0])[..., None]
geometric_force_mask.append(self.distributed_array_init(force_mask_host.shape, jnp.bool_, init_val=force_mask_host))
component_data = []
for bc in BC:
if not self._is_wetting_boundary_condition(bc):
continue
if bc.theta is None:
continue
indices = np.array(self._get_solid_indices(bc), dtype=np.int32).T
theta = np.asarray(bc.theta, dtype=np.float64).reshape(-1)
if theta.size == 1:
theta = np.full((indices.shape[0],), theta.item(), dtype=np.float64)
if theta.shape[0] != indices.shape[0]:
raise ValueError("Geometric wetting theta must be scalar or match the number of boundary nodes.")
interface = self._solid_fluid_interface_mask(indices, solid_mask)
if not np.any(interface):
continue
source_rows = np.flatnonzero(interface)
indices = indices[interface]
theta = theta[interface]
normals = self._compute_geometric_normals(bc, solid_mask, indices=indices, source_rows=source_rows)
valid_normal = np.linalg.norm(normals, axis=1) > 1e-12
if not np.any(valid_normal):
continue
indices = indices[valid_normal]
theta = theta[valid_normal]
normals = normals[valid_normal]
if self.dim == 3:
characteristics_start = time.perf_counter()
points = self._build_geometric_3d_characteristic_data(indices, normals, theta, solid_mask)
characteristics_time += time.perf_counter() - characteristics_start
data = {
"indices": tuple(np.asarray(index, dtype=np.int32) for index in indices.T),
"theta": np.asarray(theta.reshape(-1, 1), dtype=np.dtype(self.precision_policy.compute_dtype)),
"points": points,
}
component_data.append(self._localize_geometric_wetting_data(data) if localize and self.n_devices > 1 else data)
continue
# Characteristics determination for density interpolation
# The density of nearest intersection to mesh is used for solid density.
# In most cases, intersection does not occur on a fluid point so the density is interpolated from neighboring fluid points.
# This ensures a local density value is used.
characteristics_start = time.perf_counter()
angle = np.pi / 2 - theta
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
direction_1 = np.column_stack((
normals[:, 0] * cos_angle - normals[:, 1] * sin_angle,
normals[:, 0] * sin_angle + normals[:, 1] * cos_angle,
))
direction_2 = np.column_stack((
normals[:, 0] * cos_angle + normals[:, 1] * sin_angle,
-normals[:, 0] * sin_angle + normals[:, 1] * cos_angle,
))
points_1, valid_1 = self._first_fluid_mesh_intersection(indices, direction_1, solid_mask, return_valid=True)
points_2, valid_2 = self._first_fluid_mesh_intersection(indices, direction_2, solid_mask, return_valid=True)
characteristics_time += time.perf_counter() - characteristics_start
fallback_points = np.zeros_like(normals)
fallback_points[valid_1] = points_1[valid_1]
fallback_points[~valid_1 & valid_2] = points_2[~valid_1 & valid_2]
point_1, point_2 = self._repair_invalid_geometric_point_data(
[self._build_interpolation_data(points_1), self._build_interpolation_data(points_2)],
[valid_1, valid_2],
indices,
normals,
solid_mask,
fallback_points,
)
data = {
"indices": tuple(np.asarray(index, dtype=np.int32) for index in indices.T),
"theta": np.asarray(theta.reshape(-1, 1), dtype=np.dtype(self.precision_policy.compute_dtype)),
"point_1": point_1,
"point_2": point_2,
}
component_data.append(self._localize_geometric_wetting_data(data) if localize and self.n_devices > 1 else data)
geometric_wetting_data.append(component_data)
logger.info(f"Time taken to determine geometric wetting characteristics: {characteristics_time:.6f} seconds")
self.geometric_force_mask = geometric_force_mask
return geometric_wetting_data, geometric_fluid_mask
def get_solid_mask_streamed(self):
"""
Define the solid mask used for the fluid-solid interaction (wetting) force. One flag per node, not per
lattice direction - neighbor-direction information is derived on demand by scalar_neighbor_sum instead of
being pre-streamed into a persistent per-direction mask. The boundary conditions must be passed separately.
Returns
-------
list of jax.Array or None: Component masks with shape (nx, ny, 1) for d == 2 or (nx, ny, nz, 1) for d == 3.
Components without a wetting boundary contain None.
"""
shape = (self.nx, self.ny, 1) if self.dim == 2 else (self.nx, self.ny, self.nz, 1)
solid_mask = []
for component_bcs, has_wetting_bc in zip(self.BCs, self._has_wetting_bc, strict=True):
if not has_wetting_bc:
solid_mask.append(None)
continue
solid_indices = [np.array(self._get_solid_indices(bc)).T for bc in component_bcs if self._is_wetting_boundary_condition(bc)]
mask_host = np.zeros(shape, dtype=np.int8)
if solid_indices:
index = np.vstack(solid_indices)
mask_host[tuple(index.T)] = 1
mask = self.distributed_array_init(shape, jnp.int8, init_val=mask_host)
solid_mask.append(mask)
return solid_mask
def _neighbor_stencil_m(self, field, weights):
"""
Sum a scalar (single-channel) field over its lattice neighbors, weighted per direction, using local
jnp.roll for every axis and exactly one x-halo exchange per distinct x-shift shared by every lattice
direction with that shift (-1, 0 or +1 for every lattice this library supports).
Reused for two purposes, bound to their own static weights via functools.partial when scalar_neighbor_sum
/ scalar_force_stencil are built (see __init__): the G_ff-weighted neighbor sum used by
compute_average_density's wetting denominator (scalar weights), and the G_ff*c weighted directional sum
used by compute_fluid_fluid_force's Shan-Chen/Zhang-Chen force (vector weights) - both without ever
streaming a q-channel array.
Parameters
----------
field (jax.numpy.ndarray): Local shard of a scalar field, shape (nx, ny, 1) or (nx, ny, nz, 1).
weights (numpy.ndarray): Per-direction weights, shape (q,) for a scalar-weighted neighbor sum, or
(dim, q) for a direction-vector weighted sum (one weight vector per lattice direction).
Returns
-------
(jax.numpy.ndarray): Local shard of the weighted neighbor sum, shape (..., 1) for scalar weights or
(..., dim) for vector weights.
"""
field = field.astype(self.precision_policy.compute_dtype)
x_shifted = {0: field}
for x_shift in (1, -1):
shifted = jnp.roll(field, x_shift, axis=0)
if x_shift == 1:
x_shifted[1] = shifted.at[:1].set(self.send_right(field[-1:], "x"))
else:
x_shifted[-1] = shifted.at[-1:].set(self.send_left(field[:1], "x"))
directions = np.array(self.lattice.c).T
is_vector = weights.ndim == 2
out_channels = weights.shape[0] if is_vector else 1
total = jnp.zeros((*field.shape[:-1], out_channels), dtype=self.precision_policy.compute_dtype)
for q_index, direction in enumerate(directions):
w = weights[:, q_index] if is_vector else weights[q_index]
if np.all(w == 0.0):
continue
base = x_shifted[int(direction[0])]
remaining_axes = tuple(int(component) for component in direction[1 : self.dim])
if any(remaining_axes):
base = jnp.roll(base, remaining_axes, axis=tuple(range(1, self.dim)))
total = total + base * jnp.asarray(w, dtype=self.precision_policy.compute_dtype)
return total
def _surface_stencil_m(self, psi, weights):
"""Compute first- and second-power scalar surface moments with one pair of x halo exchanges."""
psi = psi.astype(self.precision_policy.compute_dtype)
field = jnp.concatenate((psi, jnp.square(psi)), axis=-1)
x_shifted = {0: field}
for x_shift in (1, -1):
shifted = jnp.roll(field, x_shift, axis=0)
if x_shift == 1:
x_shifted[1] = shifted.at[:1].set(self.send_right(field[-1:], "x"))
else:
x_shifted[-1] = shifted.at[-1:].set(self.send_left(field[:1], "x"))
directions = np.asarray(self.lattice.c).T
n_moments = weights.shape[0]
total = jnp.zeros((*psi.shape[:-1], 2, n_moments), dtype=self.precision_policy.compute_dtype)
for q_index, direction in enumerate(directions):
weight = weights[:, q_index]
if np.all(weight == 0.0):
continue
base = x_shifted[int(direction[0])]
remaining_axes = tuple(int(component) for component in direction[1 : self.dim])
if any(remaining_axes):
base = jnp.roll(base, remaining_axes, axis=tuple(range(1, self.dim)))
total = total + base[..., :, None] * jnp.asarray(weight, dtype=self.precision_policy.compute_dtype)
return total.reshape((*psi.shape[:-1], 2 * n_moments))
def _create_boundary_data(self):
"""
Create boundary data for the Lattice Boltzmann simulation by setting boundary conditions,
creating grid mask, and preparing local masks and normal arrays.
"""
self.BCs = [[] for _ in range(self.n_components)]
self.set_boundary_conditions()
# Accumulate the indices of all BCs to create the grid mask with FALSE along directions that
# stream into a boundary voxel.
for i in range(self.n_components):
logger.info(f"Component: {i + 1}")
solid_halo_list = [np.array(bc.indices).T for bc in self.BCs[i] if bc.is_solid]
solid_halo_voxels = np.unique(np.vstack(solid_halo_list), axis=0) if solid_halo_list else None
# Create the grid mask on each process
start = time.time()
grid_mask = self.create_grid_mask(solid_halo_voxels)
logger.info("Time to create the grid mask: %.6f seconds", time.time() - start)
start = time.time()
for bc in self.BCs[i]:
assert bc.implementation_step in ["PostStreaming", "PostCollision"]
bc.create_local_mask_and_normal_arrays(grid_mask)
logger.info("Time to create the local masks and normal arrays: %.6f seconds", time.time() - start)
def _make_local_bounceback_indices(self):
"""
Distribute one padded local BounceBack index array per component (self.BCs is a list of per-component
boundary condition lists for Multiphase, unlike the flat list in LBMBase).
Returns
-------
(list): One distributed local index array (or None) per component, see LBMBase._make_local_bounceback_indices.
"""
sharding = NamedSharding(self.mesh, PartitionSpec("x", None, None))
local_bounceback_indices = []
for BCs in self.BCs:
local_indices = self._collect_bounceback_indices(BCs)
if local_indices is None:
local_bounceback_indices.append(None)
continue
indices = self.distributed_array_init(local_indices.shape, jnp.int32, init_val=local_indices, sharding=sharding)
indices.block_until_ready()
local_bounceback_indices.append(indices)
return local_bounceback_indices
def _make_local_wall_bc_data(self):
"""
Distribute, per component and per concrete wall boundary condition type in WALL_BC_TYPES, the padded
local fluid-node indices and auxiliary data, plus one merged padded local solid-node index array per
component (self.BCs is a list of per-component boundary condition lists for Multiphase, unlike the flat
list in LBMBase).
Returns
-------
(list, list): One wall_bc_data dict and one solid-pin index array (or None) per component, see
LBMBase._make_local_wall_bc_data.
"""
wall_bc_data_by_component = []
solid_pin_indices_by_component = []
for BCs in self.BCs:
wall_bc_data = {}
for bc_type, _, has_weights in WALL_BC_TYPES:
local_indices, (local_imissing, local_iknown, local_vel, local_weights) = self._collect_wall_bc_data(BCs, bc_type, has_weights)
wall_bc_data[bc_type] = (
self._distribute_local(local_indices, jnp.int32),
self._distribute_local(local_imissing, jnp.uint8),
self._distribute_local(local_iknown, jnp.uint8),
self._distribute_local(local_vel),
self._distribute_local(local_weights),
)
wall_bc_data_by_component.append(wall_bc_data)
solid_pin_indices_by_component.append(self._distribute_local(self._collect_solid_pin_indices(BCs), jnp.int32))
return wall_bc_data_by_component, solid_pin_indices_by_component
def _make_local_inlet_outlet_data(self):
"""
Distribute, per component and per (concrete type, prescribed-value type) pair in INLET_OUTLET_BC_TYPES,
the padded local fluid-node indices and auxiliary data (self.BCs is a list of per-component boundary
condition lists for Multiphase, unlike the flat list in LBMBase).
Returns
-------
(list): One inlet_outlet_bc_data dict per component, see LBMBase._make_local_inlet_outlet_data.
"""
inlet_outlet_bc_data_by_component = []
for BCs in self.BCs:
inlet_outlet_bc_data = {}
for bc_type, ttype, _ in INLET_OUTLET_BC_TYPES:
local_indices, (local_normals, local_imiddle_mask, local_iknown_mask, local_imissing, local_iknown, local_prescribed) = (
self._collect_inlet_outlet_bc_data(BCs, bc_type, ttype)
)
inlet_outlet_bc_data[(bc_type, ttype)] = (
self._distribute_local(local_indices, jnp.int32),
self._distribute_local(local_normals),
self._distribute_local(local_imiddle_mask, jnp.bool_),
self._distribute_local(local_iknown_mask, jnp.bool_),
self._distribute_local(local_imissing, jnp.uint8),
self._distribute_local(local_iknown, jnp.uint8),
self._distribute_local(local_prescribed),
)
inlet_outlet_bc_data_by_component.append(inlet_outlet_bc_data)
return inlet_outlet_bc_data_by_component
def _make_local_equilibrium_bc_data(self):
"""
Distribute, per component, the padded local EquilibriumBC indices and precomputed equilibrium values
(self.BCs is a list of per-component boundary condition lists for Multiphase, unlike the flat list in
LBMBase).
Returns
-------
(list, list): One distributed local index array and one distributed equilibrium-value array (or None)
per component, see LBMBase._make_local_equilibrium_bc_data.
"""
indices_by_component = []
out_by_component = []
for BCs in self.BCs:
local_indices, local_out = self._collect_equilibrium_bc_data(BCs)
indices_by_component.append(self._distribute_local(local_indices, jnp.int32))
out_by_component.append(self._distribute_local(local_out))
return indices_by_component, out_by_component
def local_neq_bc_m(
self,
fout,
local_indices,
local_neighbor_indices,
local_imissing,
local_prescribed,
*,
exact,
needs_halo,
correction_weights,
):
"""Apply one static non-equilibrium extrapolation boundary using shard-local indices."""
indices = local_indices[0]
neighbor_indices = local_neighbor_indices[0]
prescribed = local_prescribed[0]
if needs_halo:
left_halo = self.send_right(fout[-1:], "x")
right_halo = self.send_left(fout[:1], "x")
neighbor_field = jnp.concatenate((left_halo, fout, right_halo), axis=0)
neighbor_indices = neighbor_indices.at[:, 0].add(1)
else:
neighbor_field = fout
idx = tuple(indices[:, axis] for axis in range(self.dim))
neighbor_idx = tuple(neighbor_indices[:, axis] for axis in range(self.dim))
fbd = fout.at[idx].get(mode="fill", fill_value=0.0)
f_nbr = neighbor_field.at[neighbor_idx].get(mode="fill", fill_value=0.0)
c = jnp.asarray(self.lattice.c, dtype=self.precision_policy.compute_dtype)
imissing = local_imissing[0]
fbd = _neq_extrapolation_math(
fbd,
f_nbr,
prescribed,
imissing,
self.lattice.w,
c,
correction_weights if exact else None,
)
return fout.at[idx].set(fbd, mode="drop")
def _build_local_neq_bc_kernels(self, field_spec, variants, correction_weights):
"""Build only extrapolation-kernel variants used by configured multiphase boundaries."""
P = PartitionSpec
aux_spec = P("x", None, None)
return {
(exact, needs_halo): jit(
shard_map(
partial(
self.local_neq_bc_m,
exact=exact,
needs_halo=needs_halo,
correction_weights=correction_weights if exact else None,
),
mesh=self.mesh,
in_specs=(field_spec, aux_spec, aux_spec, aux_spec, aux_spec),
out_specs=field_spec,
check_vma=False,
)
)
for exact, needs_halo in variants
}
@staticmethod
def _expand_neq_prescribed(values, count):
"""Normalize prescribed density to one row per boundary node."""
values = np.asarray(values)
if values.ndim == 0:
return np.full((count, 1), values.item(), dtype=values.dtype)
if values.ndim == 1:
if values.shape[0] == count:
return values[:, None]
return np.broadcast_to(values, (count, values.shape[0])).copy()
if values.shape[0] == count:
return values
return np.broadcast_to(values, (count, *values.shape)).copy()
def _collect_neq_bc_data(self, bc):
"""Build local data for one multiphase non-equilibrium extrapolation boundary."""
if type(bc) not in NEQ_BC_TYPES:
return None
if not bc.neighbors_found:
bc.find_neighbors()
bc.neighbors_found = True
indices = np.asarray(bc.indices, dtype=np.int32).T
if len(indices) == 0:
return None
neighbor_indices = np.asarray(bc.indices_nbr, dtype=np.int32).T
prescribed = self._expand_neq_prescribed(bc.prescribed, len(indices))
imissing = np.asarray(bc.imissing)
local_indices, (local_neighbors, local_imissing, local_prescribed) = self._split_local_indices(
indices,
neighbor_indices,
imissing,
prescribed,
)
local_nx = self.nx // self.n_devices
owner = indices[:, 0] // local_nx
needs_halo = False
for device in range(self.n_devices):
count = int(np.count_nonzero(owner == device))
if count == 0:
continue
local_neighbors[device, :count, 0] -= device * local_nx
needs_halo = needs_halo or bool(np.any((local_neighbors[device, :count, 0] < 0) | (local_neighbors[device, :count, 0] >= local_nx)))
return (
type(bc) is ExactNonEquilibriumExtrapolation,
needs_halo,
self._distribute_local(local_indices, jnp.int32),
self._distribute_local(local_neighbors, jnp.int32),
self._distribute_local(local_imissing, jnp.uint8),
self._distribute_local(local_prescribed, self.precision_policy.compute_dtype),
)
def _make_local_neq_bc_data(self):
"""Build local extrapolation data per component while preserving BC ordering."""
return [[self._collect_neq_bc_data(bc) for bc in component_bcs] for component_bcs in self.BCs]
@partial(jit, static_argnums=(0, 3), inline=True)
def equilibrium(self, rho_tree, u_tree, cast_output=True):
"""
Compute the equilibrium distribution using density and velocity pytrees.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field
u_tree (pytree of jax.numpy.ndarray): Velocity field
cast_output (bool, optional): A flag to cast the density and velocity values to the compute and output
precision. Default: True
Returns
-------
feq_tree (pytree of jax.numpy.ndarray): Equilibrium distribution.
"""
if cast_output:
cast = lambda x: self.precision_policy.cast_to_compute(x)
rho_tree = tree_map(cast, rho_tree)
u_tree = tree_map(cast, u_tree)
c = jnp.array(self.c, dtype=self.precision_policy.compute_dtype)
cu_tree = tree_map(lambda u: 3.0 * jnp.dot(u, c), u_tree)
usqr_tree = tree_map(lambda u: 1.5 * jnp.sum(jnp.square(u), axis=-1, keepdims=True), u_tree)
feq_tree = tree_map(lambda rho, udote, udotu: rho * self.w * (1.0 + udote * (1.0 + 0.5 * udote) - udotu), rho_tree, cu_tree, usqr_tree)
if cast_output:
return tree_map(lambda f_eq: self.precision_policy.cast_to_output(f_eq), feq_tree)
else:
return feq_tree
def local_improved_wetting_m(self, rho, rho_ave, local_indices, local_theta, local_phi, local_delta_rho):
"""Apply improved-virtual-density wetting using local solid-node indices on each x shard."""
indices = local_indices[0]
idx = tuple(indices[:, axis] for axis in range(self.dim))
rho_ave_boundary = rho_ave.at[idx].get(mode="fill", fill_value=0.0)
theta = local_theta[0]
wall_density = (theta <= jnp.pi / 2) * (local_phi[0] * rho_ave_boundary) + (theta > jnp.pi / 2) * (rho_ave_boundary - local_delta_rho[0])
return rho.at[idx].set(wall_density, mode="drop")
def local_geometric_wetting_m(
self,
rho,
fluid_mask,
component_data,
):
"""Apply all geometric boundaries for one component inside one shard-local kernel."""
local_min = jnp.min(jnp.where(fluid_mask, rho, jnp.inf))
local_max = jnp.max(jnp.where(fluid_mask, rho, -jnp.inf))
rho_min = jax.lax.pmin(local_min, "x")
rho_max = jax.lax.pmax(local_max, "x")
for data in component_data:
indices = data["local_indices"][0]
interpolation_indices = data["request_indices"][0]
weights = data["request_weights"][0]
if self.dim == 2:
x0, y0, x1, y1 = (interpolation_indices[..., index] for index in range(4))
contributions = (
weights[..., 0, None] * rho.at[x0, y0].get(mode="fill", fill_value=0.0)
+ weights[..., 1, None] * rho.at[x1, y0].get(mode="fill", fill_value=0.0)
+ weights[..., 2, None] * rho.at[x0, y1].get(mode="fill", fill_value=0.0)
+ weights[..., 3, None] * rho.at[x1, y1].get(mode="fill", fill_value=0.0)
)
else:
x0, y0, z0, x1, y1, z1 = (interpolation_indices[..., index] for index in range(6))
contributions = (
weights[..., 0, None] * rho.at[x0, y0, z0].get(mode="fill", fill_value=0.0)
+ weights[..., 1, None] * rho.at[x1, y0, z0].get(mode="fill", fill_value=0.0)
+ weights[..., 2, None] * rho.at[x0, y1, z0].get(mode="fill", fill_value=0.0)
+ weights[..., 3, None] * rho.at[x1, y1, z0].get(mode="fill", fill_value=0.0)
+ weights[..., 4, None] * rho.at[x0, y0, z1].get(mode="fill", fill_value=0.0)
+ weights[..., 5, None] * rho.at[x1, y0, z1].get(mode="fill", fill_value=0.0)
+ weights[..., 6, None] * rho.at[x0, y1, z1].get(mode="fill", fill_value=0.0)
+ weights[..., 7, None] * rho.at[x1, y1, z1].get(mode="fill", fill_value=0.0)
)
if "request_slots" in data:
template = data["reduce_scatter_template"][0] if "reduce_scatter_template" in data else data["sample_template"][0]
samples = jnp.zeros((*template.shape, rho.shape[-1]), dtype=rho.dtype)
samples = samples.reshape((-1, rho.shape[-1])).at[data["request_slots"][0]].set(contributions, mode="drop")
samples = samples.reshape((*template.shape, rho.shape[-1]))
if "reduce_scatter_template" in data:
samples = jax.lax.psum_scatter(samples, "x", scatter_dimension=0)
else:
samples = jax.lax.psum(samples, "x")
samples = samples[data["local_sample_indices"][0]]
else:
samples = jax.lax.psum(contributions, "x")
samples = samples[data["local_sample_indices"][0]]
rho_wall = jnp.where(
data["local_select_max"][0],
jnp.max(samples, axis=-2),
jnp.min(samples, axis=-2),
)
idx = tuple(indices[:, axis] for axis in range(self.dim))
rho = rho.at[idx].set(jnp.clip(rho_wall, rho_min, rho_max), mode="drop")
return rho
def _localize_geometric_wetting_data(self, data):
"""Convert one geometric boundary's interpolation data to sharded int32 requests."""
wall_indices = np.stack(data["indices"], axis=-1).astype(np.int32, copy=False)
theta = np.asarray(data["theta"], dtype=np.dtype(self.precision_policy.compute_dtype))
select_max = theta <= np.asarray(np.pi / 2, dtype=theta.dtype)
point_data = data["points"] if self.dim == 3 else (data["point_1"], data["point_2"])
index_count = 6 if self.dim == 3 else 4
sample_indices = np.stack([np.stack(point[:index_count], axis=-1) for point in point_data], axis=1).astype(np.int32, copy=False)
sample_weights = np.stack([np.stack(point[index_count:], axis=-1) for point in point_data], axis=1).astype(
np.dtype(self.precision_policy.compute_dtype),
copy=False,
)
sample_rows = np.arange(len(wall_indices), dtype=np.int32)
local_indices, (local_select_max, local_sample_indices) = self._split_local_indices(wall_indices, select_max, sample_rows)
local_nx = self.nx // self.n_devices
sample_count = sample_indices.shape[1]
x_positions = (0, 3) if self.dim == 3 else (0, 2)
wall_owner = wall_indices[:, 0] // local_nx
max_local = local_indices.shape[1]
local_wall_slot = np.empty(len(wall_indices), dtype=np.int32)
for destination in range(self.n_devices):
destination_rows = wall_owner == destination
local_wall_slot[destination_rows] = np.arange(np.count_nonzero(destination_rows), dtype=np.int32)
# Route balanced wall rows directly to their destination shard. If destination padding exceeds 12.5%,
# keep the compact all-reduce layout to prevent an imbalanced wall from increasing temporary memory.
use_reduce_scatter = self.n_devices * max_local <= int(np.ceil(1.125 * len(wall_indices)))
x0_owner = sample_indices[..., x_positions[0]] // local_nx
x1_owner = sample_indices[..., x_positions[1]] // local_nx
active = [(x0_owner == source) | (x1_owner == source) for source in range(self.n_devices)]
max_requests = max(np.count_nonzero(source_active) for source_active in active)
dense_request_count = len(wall_indices) * sample_count
# Compaction needs a scatter into collective slots. Retain the direct dense kernel when fewer than
# 12.5% of gathers would be removed, avoiding a slowdown for walls whose samples all live on one shard.
if max_requests > 0.875 * dense_request_count:
source_ids = np.arange(self.n_devices, dtype=np.int32)[:, None, None]
request_indices = np.broadcast_to(sample_indices[None], (self.n_devices,) + sample_indices.shape).copy()
request_indices[..., x_positions[0]] -= source_ids * local_nx
request_indices[..., x_positions[1]] -= source_ids * local_nx
request_weights = np.broadcast_to(sample_weights[None], (self.n_devices,) + sample_weights.shape).copy()
request_weights[..., 0::2] *= (x0_owner[None] == source_ids)[..., None]
request_weights[..., 1::2] *= (x1_owner[None] == source_ids)[..., None]
return {
"local_indices": self._distribute_local(local_indices, jnp.int32),
"local_select_max": self._distribute_local(local_select_max, jnp.bool_),
"local_sample_indices": self._distribute_local(local_sample_indices, jnp.int32),
"request_indices": self._distribute_local(request_indices, jnp.int32),
"request_weights": self._distribute_local(request_weights, self.precision_policy.compute_dtype),
}
request_indices = np.zeros((self.n_devices, max_requests, index_count), dtype=np.int32)
request_indices[..., x_positions[0]] = local_nx
request_indices[..., x_positions[1]] = local_nx
request_weights = np.zeros((self.n_devices, max_requests, sample_weights.shape[-1]), dtype=sample_weights.dtype)
sample_buffer_size = (self.n_devices * max_local if use_reduce_scatter else len(wall_indices)) * sample_count
request_slots = np.full((self.n_devices, max_requests), sample_buffer_size, dtype=np.int32)
for source, source_active in enumerate(active):
rows, samples = np.nonzero(source_active)
count = len(rows)
localized = sample_indices[rows, samples].copy()
localized[..., x_positions[0]] -= source * local_nx
localized[..., x_positions[1]] -= source * local_nx
request_indices[source, :count] = localized
source_weights = sample_weights[rows, samples].copy()
source_weights[..., 0::2] *= (x0_owner[rows, samples] == source)[:, None]
source_weights[..., 1::2] *= (x1_owner[rows, samples] == source)[:, None]
request_weights[source, :count] = source_weights
if use_reduce_scatter:
request_slots[source, :count] = wall_owner[rows] * max_local * sample_count + local_wall_slot[rows] * sample_count + samples
else:
request_slots[source, :count] = rows * sample_count + samples
localized = {
"local_indices": self._distribute_local(local_indices, jnp.int32),
"local_select_max": self._distribute_local(local_select_max, jnp.bool_),
"request_indices": self._distribute_local(request_indices, jnp.int32),
"request_weights": self._distribute_local(request_weights, self.precision_policy.compute_dtype),
"request_slots": self._distribute_local(request_slots, jnp.int32),
}
if use_reduce_scatter:
localized["reduce_scatter_template"] = self._distribute_local(
np.zeros((self.n_devices, self.n_devices, max_local, sample_count), dtype=np.bool_), jnp.bool_
)
else:
localized["local_sample_indices"] = self._distribute_local(local_sample_indices, jnp.int32)
localized["sample_template"] = self._distribute_local(
np.zeros((self.n_devices, len(wall_indices), sample_count), dtype=np.bool_), jnp.bool_
)
return localized
def _build_local_geometric_wetting_kernels(self, scalar_spec):
"""Build one fused shard-local geometric-wetting kernel per active component."""
kernels = []
for component_data in self.geometric_wetting_data:
if not component_data:
kernels.append(None)
continue
data_specs = tree_map(lambda array: PartitionSpec("x", *([None] * (array.ndim - 1))), component_data)
kernels.append(
jit(
shard_map(
self.local_geometric_wetting_m,
mesh=self.mesh,
in_specs=(scalar_spec, scalar_spec, data_specs),
out_specs=scalar_spec,
check_vma=False,
)
)
)
return tuple(kernels)
def _wetting_parameter_at_indices(self, value, indices):
"""Convert scalar, per-node, or full-domain wetting data to one scalar row per solid node."""
value = np.asarray(value)
count = len(indices)
if value.size == 1:
return np.full((count, 1), value.reshape(()).item(), dtype=value.dtype)
spatial_shape = (self.nx, self.ny) if self.dim == 2 else (self.nx, self.ny, self.nz)
if value.shape[: self.dim] == spatial_shape:
selected = value[tuple(indices.T)]
return np.asarray(selected).reshape(count, -1)
if value.shape[0] == count:
return value.reshape(count, -1)
raise ValueError("Wetting parameters must be scalar, per-boundary-node arrays, or full-domain fields.")
def _make_local_improved_wetting_data(self):
"""Build static local wetting data per component and boundary condition."""
data_by_component = []
for component_bcs in self.BCs:
component_data = []
for bc in component_bcs:
if not self._is_wetting_boundary_condition(bc) or bc.theta is None or bc.is_dynamic:
component_data.append(None)
continue
indices = np.asarray(self._get_solid_indices(bc), dtype=np.int32).T
if len(indices) == 0:
component_data.append(None)
continue
theta = self._wetting_parameter_at_indices(bc.theta, indices)
phi = self._wetting_parameter_at_indices(bc.phi, indices)
delta_rho = self._wetting_parameter_at_indices(bc.delta_rho, indices)
local_indices, (local_theta, local_phi, local_delta_rho) = self._split_local_indices(indices, theta, phi, delta_rho)
component_data.append((
self._distribute_local(local_indices, jnp.int32),
self._distribute_local(local_theta, self.precision_policy.compute_dtype),
self._distribute_local(local_phi, self.precision_policy.compute_dtype),
self._distribute_local(local_delta_rho, self.precision_policy.compute_dtype),
))
data_by_component.append(component_data)
return data_by_component
@partial(jit, static_argnums=(0,))
def compute_average_density(self, rho_tree):
"""
Compute component densities averaged over neighboring fluid nodes, using the scalar solid mask and the
denominator cached once at construction time (self.average_density_denominator), instead of streaming a
q-channel mask and dividing every timestep.
Parameters
----------
rho_tree (pytree of jax.Array): Component density fields with shape ``(nx, ny, 1)`` in 2D or ``(nx, ny, nz, 1)``
in 3D.
Returns
-------
pytree of jax.Array
Averaged component density fields with the same shapes as the inputs.
"""
if self.scalar_neighbor_sum is None or self.average_density_denominator is None:
return rho_tree
return [
self.scalar_neighbor_sum(rho * (1 - solid_mask)) / denominator if denominator is not None else rho
for rho, solid_mask, denominator in zip(
rho_tree,
self.solid_mask_streamed,
self.average_density_denominator,
strict=True,
)
]
@partial(jit, static_argnums=(0,))
def apply_contact_angle(self, rho_tree):
"""
Apply prescribed contact angles to wall-node densities.
For the geometric scheme, only theta is used. The 2D path interpolates two characteristic samples and chooses the appropriate extrema.
The 3D path interpolates multiple samples on the contact-angle cone and chooses the maximum density for theta <= pi / 2 or the minimum
density for theta > pi / 2. For improved virtual density, theta is used with phi and delta_rho according to the selected wettability branch.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field.
Returns
-------
(pytree of jax.numpy.ndarray) Density field with adjusted contact angle values at the boundary nodes.
References
----------
1. Li, Q., Yu, Y. & Luo, K. H. "Implementation of contact angles in pseudopotential lattice Boltzmann simulations with
curved boundaries." Phys. Rev. E 100, 053313 (2019).
2. Fei, Linlin, Feifei Qin, Jianlin Zhao, Dominique Derome, and Jan Carmeliet. “Lattice Boltzmann Modelling of
Isothermal Two-Component Evaporation in Porous Media.”
Journal of Fluid Mechanics 955 (January 2023): A18.
3. Wang, Lei, Hai-bo Huang, and Xi-Yun Lu. “Scheme for Contact Angle and Its Hysteresis in a Multiphase Lattice
Boltzmann Method.” Physical Review E 87, no. 1 (2013): 013301.
"""
if self.wetting_formulation is None or not any(self._has_wetting_bc):
return rho_tree
if self.wetting_formulation == "improved_virtual_density":
rho_ave_tree = self.compute_average_density(rho_tree)
def set_contact_angle(rho, rho_ave, BC, local_component_data):
rho_min = jnp.min(rho)
rho_max = jnp.max(rho)
for bc_index, bc in enumerate(BC):
if isinstance(
bc, (BounceBackHalfway, BounceBack, BounceBackMoving, InterpolatedBounceBackBouzidi, InterpolatedBounceBackDifferentiable)
):
if bc.theta is not None:
local_data = local_component_data[bc_index]
if local_data is not None:
local_indices, local_theta, local_phi, local_delta_rho = local_data
rho = self.local_improved_wetting(rho, rho_ave, local_indices, local_theta, local_phi, local_delta_rho)
else:
indices = self._get_solid_indices(bc)
rho = rho.at[indices].set(
(bc.theta <= jnp.pi / 2) * (bc.phi * rho_ave[indices])
+ (bc.theta > jnp.pi / 2) * (rho_ave[indices] - bc.delta_rho)
)
rho = jnp.clip(rho, min=rho_min, max=rho_max)
return rho
return [
set_contact_angle(rho, rho_ave, BC, local_component_data) if has_wetting_bc else rho
for rho, rho_ave, BC, local_component_data, has_wetting_bc in zip(
rho_tree,
rho_ave_tree,
self.BCs,
self.local_improved_wetting_data,
self._has_wetting_bc,
strict=True,
)
]
elif self.wetting_formulation == "geometric":
if self.local_geometric_wetting is not None:
return [
kernel(rho, fluid_mask, component_data) if kernel is not None else rho
for rho, fluid_mask, component_data, kernel in zip(
rho_tree,
self.geometric_fluid_mask,
self.geometric_wetting_data,
self.local_geometric_wetting,
strict=True,
)
]
def interpolate_density(rho, interpolation_data):
"""
Interpolate density at precomputed geometric wetting sample points.
Parameters
----------
rho (jax.numpy.ndarray): Density field for one component.
interpolation_data (tuple): Index arrays and weights generated by _build_interpolation_data.
Returns
-------
(jax.numpy.ndarray): Interpolated density values at the sample points.
"""
if self.dim == 2:
x0, y0, x1, y1, w00, w10, w01, w11 = interpolation_data
return w00[:, None] * rho[x0, y0] + w10[:, None] * rho[x1, y0] + w01[:, None] * rho[x0, y1] + w11[:, None] * rho[x1, y1]
x0, y0, z0, x1, y1, z1, w000, w100, w010, w110, w001, w101, w011, w111 = interpolation_data
return (
w000[:, None] * rho[x0, y0, z0]
+ w100[:, None] * rho[x1, y0, z0]
+ w010[:, None] * rho[x0, y1, z0]
+ w110[:, None] * rho[x1, y1, z0]
+ w001[:, None] * rho[x0, y0, z1]
+ w101[:, None] * rho[x1, y0, z1]
+ w011[:, None] * rho[x0, y1, z1]
+ w111[:, None] * rho[x1, y1, z1]
)
def set_geometric_contact_angle(rho, component_data, fluid_mask):
"""
Set wall density values for one component using precomputed geometric wetting data.
Parameters
----------
rho (jax.numpy.ndarray): Density field for one component.
component_data (list): Boundary-wise geometric wetting data for one component.
fluid_mask (jax.numpy.ndarray): Boolean mask with True on fluid nodes.
Returns
-------
rho (jax.numpy.ndarray): Density field with wall values updated at wetted boundary nodes.
"""
# Bound wall densities by the fluid density range so the wall can never introduce
# a density outside what exists in the fluid. Using the global field range instead
# would include the unphysical densities stored at solid nodes by bounce-back and
# let the wall values ratchet the fluid range upward.
rho_min = jnp.min(jnp.where(fluid_mask, rho, jnp.inf))
rho_max = jnp.max(jnp.where(fluid_mask, rho, -jnp.inf))
for data in component_data:
if self.dim == 2:
rho_1 = interpolate_density(rho, data["point_1"])
rho_2 = interpolate_density(rho, data["point_2"])
rho_wall = jnp.where(data["theta"] <= jnp.pi / 2, jnp.maximum(rho_1, rho_2), jnp.minimum(rho_1, rho_2))
rho = rho.at[data["indices"]].set(jnp.clip(rho_wall, rho_min, rho_max))
else:
rho_samples = [interpolate_density(rho, point_data) for point_data in data["points"]]
rho_sample_min = rho_samples[0]
rho_sample_max = rho_samples[0]
for rho_sample in rho_samples[1:]:
rho_sample_min = jnp.minimum(rho_sample_min, rho_sample)
rho_sample_max = jnp.maximum(rho_sample_max, rho_sample)
rho_wall = jnp.where(data["theta"] <= jnp.pi / 2, rho_sample_max, rho_sample_min)
rho = rho.at[data["indices"]].set(jnp.clip(rho_wall, rho_min, rho_max))
return rho
return tree_map(
lambda rho, component_data, fluid_mask: set_geometric_contact_angle(rho, component_data, fluid_mask),
rho_tree,
self.geometric_wetting_data,
self.geometric_fluid_mask,
)
return rho_tree
@partial(jit, static_argnums=(0,))
def collision(self, fin_tree, T=None):
"""
Apply collision step of LBM. The optional temperature field T is used
by thermal EOS variants (see compute_pressure).
"""
pass
def compute_ff_greens_function(self):
"""
Define the fluid-fluid interaction force Green's function used to compute interaction phase-phase interaction forces.
The interaction coefficient between k^th and kprime^th component: self.gkkprime[k, kprime]
During computation, this value is multiplied with corresponding g_kkprime value to get the Green's function:
G_kkprime = self.g_kk[k, k_prime] * self.G_ff
G_kkprime(x, x') = g1 * g_kkprime, if |x - x'| = 1
= g2 * g_kkprime, if |x - x'| = sqrt(2)
= 0, otherwise
Here d is the dimension of problem and x' are the neighboring points.
Some examples values could be:
For D2Q9:
g1 = 1/3 and g2 = 1/12
For D3Q19
g1 = 1/6 and g2 = 1/12
Returns
-------
G_ff (jax.numpy.ndarray): Dimension: (q, )
"""
c = np.array(self.lattice.c).T
G_ff = np.zeros((self.q,), dtype=np.float64)
cl = np.linalg.norm(c, axis=-1)
if isinstance(self.lattice, LatticeD2Q9):
g1 = 1 / 3
g2 = 1 / 12
G_ff[np.isclose(cl, 1.0, atol=1e-6)] = g1
G_ff[np.isclose(cl, jnp.sqrt(2.0), atol=1e-6)] = g2
elif isinstance(self.lattice, LatticeD3Q19):
g1 = 1 / 6
g2 = 1 / 12
G_ff[np.isclose(cl, 1.0, atol=1e-6)] = g1
G_ff[np.isclose(cl, jnp.sqrt(2.0), atol=1e-6)] = g2
else:
raise NotImplementedError("Please define Green's function for D3Q27 lattice by modifying compute_ff_greens_function.")
return jnp.array(G_ff, dtype=self.precision_policy.compute_dtype)
def assign_fields_sharded(self):
"""
This function is used to initialize pytree of the distribution arrays using the initial velocities and velocity defined in self.initialize_macroscopic_fields function.
To do this, function first uses the initialize_macroscopic_fields function to get the initial values of rho (rho0) and velocity (u0).
If this function is not modified then, the distribution pytree is initialized with density value of 1.0 everywhere and velocity of 0.0 everywhere
The distribution is initialized with rho0 and u0 values, using the self.equilibrium function.
Returns
-------
f: pytree of distributed JAX array of shape: (self.nx, self.ny, self.q) for 2D and (self.nx, self.ny, self.nz, self.q) for 3D.
"""
rho0_tree, u0_tree = self.initialize_macroscopic_fields()
if self.dim == 2:
shape = (self.nx, self.ny, self.q)
if self.dim == 3:
shape = (self.nx, self.ny, self.nz, self.q)
f_tree = []
if rho0_tree is not None and u0_tree is not None:
assert len(rho0_tree) == self.n_components, "The initial density values for all components must be provided"
assert len(u0_tree) == self.n_components, "The initial velocity values for all components must be provided."
for i in range(self.n_components):
rho0, u0 = rho0_tree[i], u0_tree[i]
rho0 = self.precision_policy.cast_to_compute(rho0)
u0 = self.precision_policy.cast_to_compute(u0)
f_tree.append(self.initialize_populations(rho0, u0))
else:
for i in range(self.n_components):
f_tree.append(self.distributed_array_init(shape, self.precision_policy.output_dtype, init_val=self.w))
return f_tree
@partial(jit, static_argnums=(0,), inline=True)
def update_macroscopic(self, f_tree):
"""
update_macroscopic from base.py extended to pytrees.
Parameters
----------
f_tree (pytree of jax.numpy.ndarray): Distribution field.
Returns
-------
rho_tree (pytree of jax.numpy.ndarray): Density field.
u_tree (pytree of jax.numpy.ndarray): Velocity field.
"""
f_tree = tree_map(lambda f: self.precision_policy.cast_to_compute(f), f_tree)
rho_tree = tree_map(lambda f: jnp.sum(f, axis=-1, keepdims=True), f_tree)
c = jnp.array(self.c, dtype=self.precision_policy.compute_dtype).T
u_tree = tree_map(lambda f, rho: jnp.dot(f, c) / rho, f_tree, rho_tree) # Component velocity
return rho_tree, u_tree
@partial(jit, static_argnums=(0,), inline=True)
def macroscopic_velocity(self, f_tree, rho_tree, T=None):
"""
macroscopic_velocity computes the velocity and incorporates forces into velocity for Exact Difference Method (EDM) (used for SRT and MRT collision) models
and the consistent forcing scheme developed by LinLin Fei et. al (for Cascaded LBM). This is used for post-processing only and not for equilibrium distribution computation.
Parameters
----------
f_tree (pytree of jax.numpy.ndarray): Distribution field.
rho_tree (pytree of jax.numpy.ndarray): Density field.
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal.
Returns
-------
u_tree (pytree of jax.numpy.ndarray): Velocity field.
"""
# rho_tree = tree_map(lambda f: jnp.sum(f, axis=-1, keepdims=True), f_tree)
c = jnp.array(self.c, dtype=self.precision_policy.compute_dtype).T
u_tree = tree_map(lambda f, rho: jnp.dot(f, c) / rho, f_tree, rho_tree)
F_tree = self.compute_force(rho_tree, T=T)
return tree_map(lambda rho, u, F: u + 0.5 * F / rho, rho_tree, u_tree, F_tree)
@partial(jit, static_argnums=(0,))
def compute_total_density(self, rho_tree):
"""
Compute the total density using component velocity and density values.
Parameters
----------
rho_tree (Pytree of jax.numpy.ndarray): Density field.
Returns
-------
(jax.numpy.ndarray): Total density field.
"""
return reduce(operator.add, rho_tree)
@partial(jit, static_argnums=(0,))
def compute_total_velocity(self, rho_tree, u_tree):
"""
Compute the total velocity using component velocity and density values.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field
u_tree (pytree of jax.numpy.ndarray): Velocity field
Returns
-------
(jax.numpy.ndarray): Total velocity field.
"""
n = reduce(operator.add, tree_map(lambda rho, u: rho * u, rho_tree, u_tree))
d = reduce(operator.add, rho_tree)
return n / d
@partial(jit, static_argnums=(0,))
def compute_pressure(self, rho_tree, psi_tree=None, T=None):
"""
Generalized function for computing pressure. By default it uses equation
of state but it can be modified if the pseudopotential is computed using
a different method.
For a thermal EOS (temperature_field_type == "thermal") the local
temperature field T must be provided and the pressure is evaluated with
EOS_thermal, coupling the flow to the temperature solver in thermal.py.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field.
psi_tree (pytree of jax.numpy.ndarray): Pseudopotential field.
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal.
Returns
-------
(pytree of jax.numpy.ndarray): Pressure field.
"""
if self.eos.temperature_field_type == "thermal":
if T is None:
raise ValueError("Temperature field T must be passed through step/collision when using a thermal EOS.")
return self.eos.EOS_thermal(rho_tree, T)
return self.eos.EOS(rho_tree)
@partial(jit, static_argnums=(0,))
def compute_total_pressure(self, p_tree, rho_tree=None):
"""
Compute the total combined pressure from all components.
Parameters
----------
p_tree (pytree of jax.numpy.ndarray): Pressure field.
rho_tree (pytree of jax.numpy.ndarray, default=None): Density field.
Returns
-------
(jax.numpy.ndarray): Total pressure field.
"""
return reduce(operator.add, p_tree)
@partial(jit, static_argnums=(0,))
def compute_potential(self, rho_tree, T=None):
"""
Compute the potential (psi and U) which is required for computing interaction forces.
The psi values are obtained using the corresponding EOS. This function can be overloaded to handle cases where one or more component does not
have EOS.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field.
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal.
Returns
-------
psi_tree (pytree of jax.numpy.ndarray): Pseudopotential field.
"""
rho_tree = tree_map(lambda rho: self.precision_policy.cast_to_compute(rho), rho_tree)
p_tree = self.compute_pressure(rho_tree, T=T)
# Shan-Chen potential using modified pressure
psi_tree = tree_map(
lambda k, p, rho, G: jnp.sqrt(2 * (k * p - self.lattice.cs2 * rho) / G), self.k, p_tree, rho_tree, self.g_kkprime.diagonal().tolist()
)
# Zhang-Chen potential
U_tree = tree_map(lambda k, p, rho: k * p - self.lattice.cs2 * rho, self.k, p_tree, rho_tree)
return psi_tree, U_tree
# Compute the force using the effective mass (psi) and the interaction potential (phi)
@partial(jit, static_argnums=(0,))
def _compute_force_fields(self, rho_tree, T=None):
"""Return force plus the contact-angle-adjusted potentials used to construct it."""
rho_tree = self.apply_contact_angle(rho_tree)
psi_tree, U_tree = self.compute_potential(rho_tree, T=T)
fluid_fluid_force = self.compute_fluid_fluid_force(psi_tree, U_tree)
# fluid_solid_force = self.compute_fluid_solid_force(rho_tree)
if self.body_force is not None:
force_tree = tree_map(lambda ff, rho: ff + self.body_force * rho, fluid_fluid_force, rho_tree)
else:
force_tree = fluid_fluid_force
if self.wetting_formulation == "geometric" and any(self._has_wetting_bc):
force_tree = tree_map(lambda force, force_mask: force * force_mask, force_tree, self.geometric_force_mask)
return force_tree, psi_tree, U_tree
# Compute the force using the effective mass (psi) and the interaction potential (phi)
@partial(jit, static_argnums=(0,))
def compute_force(self, rho_tree, T=None):
"""
Compute the force acting on each component(fluid). This includes fluid-fluid, fluid-solid, and body forces.
Parameters
----------
rho_tree (pytree of jax.numpy.ndarray): Density field.
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal.
Returns
-------
fluid_fluid_force (pytree of jax.numpy.ndarray): Total force field.
"""
force_tree, _, _ = self._compute_force_fields(rho_tree, T=T)
return force_tree
@partial(jit, static_argnums=(0,))
def compute_fluid_fluid_force(self, psi_tree, U_tree):
"""
Compute the fluid-fluid interaction force using the effective mass (psi).
The force calculation is based on the Shan-Chen method using the weighted sum
of Shan-Chen and Zhang-Chen potential where modified pressure is used:
modified pressure = k * pressure;
k is defined by user. Set k=1 for default pseudopotential formulation.
Parameters
----------
psi_tree (pytree of jax.numpy.ndarray): Pseudo-potential field (Yuan-Schaefer, with modification)
U_tree (pytree of jax.numpy.ndarray): Pseudo-potential field (Zhang-Chen, with modification)
Returns
-------
(pytree of jax.numpy.ndarray): Fluid-fluid interaction forces.
Notes
-----
jnp.dot(G_ff * field_s, c), with field_s the q-channel streamed field, is a G_ff*c weighted directional
sum of neighbor values - exactly what scalar_force_stencil computes directly from the unstreamed scalar
field (see _neighbor_stencil_m). Computed once per component here, outside the per-output-component
vmap below, matching the original psi_s_tree/U_s_tree precompute (scalar_force_stencil is itself a
shard_map'd call and must not be invoked from inside vmap).
"""
psi_stencil_tree = [
self.scalar_force_stencil(psi) if used else None for psi, used in zip(psi_tree, self._psi_stencil_components, strict=True)
]
U_stencil_tree = [self.scalar_force_stencil(U) if used else None for U, used in zip(U_tree, self._U_stencil_components, strict=True)]
force_tree = []
for output in range(self.n_components):
psi_terms = [
(1.0 - self.A[output, source]) * self.g_kkprime[output, source] * psi_stencil_tree[source]
for source in range(self.n_components)
if self._psi_interactions[output][source]
]
U_terms = [self.A[output, source] * U_stencil_tree[source] for source in range(self.n_components) if self._U_interactions[output][source]]
force = psi_tree[output] * reduce(operator.add, psi_terms) if psi_terms else None
if U_terms:
U_force = reduce(operator.add, U_terms)
force = U_force if force is None else force + U_force
if force is None:
force = jnp.zeros((*psi_tree[output].shape[:-1], self.dim), dtype=self.precision_policy.compute_dtype)
force_tree.append(force)
return force_tree
@partial(jit, static_argnums=(0,), inline=True)
def apply_force(self, f_postcollision_tree, feq_tree, rho_tree, u_tree, T=None):
"""
Modified version of the apply_force defined in LBMBase to account for modified force.
Adds the force contribution using the exact-difference method (Kupershtokh), computing
feq(rho, u + F/rho) - feq(rho, u) directly from cu, dcu and delta_usqr instead of building a second full
equilibrium distribution and subtracting feq_tree from it.
Parameters
----------
f_postcollision_tree (pytree of jax.numpy.ndarray): Post-collision distribution field.
feq_tree (pytree of jax.numpy.ndarray): Equilibrium distribution functions. Unused - kept for interface
compatibility with existing callers, since the compact difference formula only needs rho, u and F.
rho_tree (pytree of jax.numpy.ndarray): Density field.
u_tree (pytree of jax.numpy.ndarray): Velocity field.
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal.
Returns
-------
f_postcollision_tree (pytree of jax.numpy.ndarray): The post-collision distribution field with the force applied.
References
----------
1. Kupershtokh, A. (2004). New method of incorporating a body force term into the lattice Boltzmann
equation. In Proceedings of the 5th International EHD Workshop (pp. 241-246). University of Poitiers.
"""
F_tree = self.compute_force(rho_tree, T=T)
du_tree = tree_map(lambda F, rho: F / rho, F_tree, rho_tree)
c = jnp.array(self.c, dtype=self.precision_policy.compute_dtype)
cu_tree = tree_map(lambda u: 3.0 * jnp.dot(u, c), u_tree)
dcu_tree = tree_map(lambda du: 3.0 * jnp.dot(du, c), du_tree)
delta_usqr_tree = tree_map(
lambda u, du: 1.5 * (2.0 * jnp.sum(u * du, axis=-1, keepdims=True) + jnp.sum(jnp.square(du), axis=-1, keepdims=True)),
u_tree,
du_tree,
)
delta_feq_tree = tree_map(
lambda rho, cu, dcu, delta_usqr: rho * self.w * (dcu * (1.0 + cu + 0.5 * dcu) - delta_usqr),
rho_tree,
cu_tree,
dcu_tree,
delta_usqr_tree,
)
return tree_map(lambda f_postcollision, delta_feq: f_postcollision + delta_feq, f_postcollision_tree, delta_feq_tree)
@partial(jit, static_argnums=(0, 4), donate_argnums=(1,))
def apply_bc(self, fout_tree, fin_tree, timestep, implementation_step):
"""
This function extends apply_bc to pytrees.
Full-way BounceBack, every wall boundary condition in WALL_BC_TYPES (BounceBackHalfway,
InterpolatedBounceBackBouzidi, InterpolatedBounceBackDifferentiable), EquilibriumBC, and every inlet/
outlet boundary condition in INLET_OUTLET_BC_TYPES (ZouHe, Regularized) are handled separately per
component, in batched calls using local (per-shard, int32) indices, instead of the generic per-BC
global-index loop.
Parameters
----------
fout_tree (pytree of jax.numpy.ndarray): The post-collision or post-streaming distribution functions where bc
needs to be applied.
fin_tree (pytree of jax.numpy.ndarray): The pre-collision or post-collision distribution functions.
timestep (int): Current simulation timestep, used by dynamic boundary conditions.
implementation_step (str): The implementation step at which the boundary conditions should be applied.
Returns
-------
(pytree of jax.numpy.ndarray): The output distribution functions after applying the boundary conditions.
"""
def _apply_bc_(fin, fout, bc):
if isinstance(bc, (BounceBack, BounceBackHalfway, EquilibriumBC, ZouHe)):
return fout
fout = bc.prepare_populations(fout, fin, implementation_step)
if bc.implementation_step == implementation_step:
if bc.is_dynamic:
fout = bc.apply(fout, fin, timestep)
else:
fout = fout.at[bc.indices].set(bc.apply(fout, fin))
return fout
def __apply_bc__(fout, fin, BCs, component_neq_data=None):
for bc_index, bc in enumerate(BCs):
local_neq_data = component_neq_data[bc_index] if component_neq_data is not None else None
if local_neq_data is not None:
exact, needs_halo, local_indices, local_neighbors, local_imissing, local_prescribed = local_neq_data
if bc.implementation_step == implementation_step:
fout = self.local_neq_bc_kernels[(exact, needs_halo)](
fout,
local_indices,
local_neighbors,
local_imissing,
local_prescribed,
)
continue
fout = _apply_bc_(fin, fout, bc)
return fout
if self._uses_local_neq_bc:
fout_tree = [
__apply_bc__(fout, fin, BCs, component_neq_data)
for fout, fin, BCs, component_neq_data in zip(fout_tree, fin_tree, self.BCs, self.neq_bc_data, strict=True)
]
else:
fout_tree = [__apply_bc__(fout, fin, BCs) for fout, fin, BCs in zip(fout_tree, fin_tree, self.BCs, strict=True)]
if implementation_step == "PostCollision":
fout_tree = [
self.local_bounceback(fout, fin, local_indices) if local_indices is not None else fout
for fout, fin, local_indices in zip(fout_tree, fin_tree, self.local_bounceback_indices, strict=True)
]
if implementation_step == "PostStreaming":
new_fout_tree = []
for fout, fin, wall_bc_data, solid_pin_indices, inlet_outlet_bc_data, equilibrium_bc_indices, equilibrium_bc_values in zip(
fout_tree,
fin_tree,
self.wall_bc_data,
self.solid_pin_indices,
self.inlet_outlet_bc_data,
self.local_equilibrium_bc_indices,
self.local_equilibrium_bc_values,
strict=True,
):
if solid_pin_indices is not None:
fout = self.local_solid_pin(fout, solid_pin_indices)
for bc_type, _, _ in WALL_BC_TYPES:
local_indices, local_imissing, local_iknown, local_vel, local_weights = wall_bc_data[bc_type]
if local_indices is not None:
fout = self.local_wall_bc_kernels[bc_type](fout, fin, local_indices, local_imissing, local_iknown, local_vel, local_weights)
if equilibrium_bc_indices is not None:
fout = self.local_equilibrium_bc(fout, equilibrium_bc_indices, equilibrium_bc_values)
for bc_type, ttype, _ in INLET_OUTLET_BC_TYPES:
local_indices, local_normals, local_imiddle_mask, local_iknown_mask, local_imissing, local_iknown, local_prescribed = (
inlet_outlet_bc_data[(bc_type, ttype)]
)
if local_indices is not None:
fout = self.local_inlet_outlet_kernels[(bc_type, ttype)](
fout, local_indices, local_normals, local_imiddle_mask, local_iknown_mask, local_imissing, local_iknown, local_prescribed
)
new_fout_tree.append(fout)
fout_tree = new_fout_tree
return fout_tree
@partial(jit, static_argnums=(0, 3), donate_argnums=(1,))
def step(self, f_poststreaming_tree, timestep, return_fpost=False, T=None):
"""
This function performs a single step of the LBM simulation.
It first performs the collision step, which is the relaxation of the distribution functions
towards their equilibrium values. It then applies the respective boundary conditions to the
post-collision distribution functions.
The function then performs the streaming step, which is the propagation of the distribution
functions in the lattice. It then applies the respective boundary conditions to the post-streaming
distribution functions.
Parameters
----------
f_poststreaming_tree (pytree of jax.numpy.ndarray): Post-streaming distribution function.
timestep (int): Current timestep
return_fpost (bool): Return post-collision distribution function (pytree).
T (jax.numpy.ndarray, optional): Temperature field, required when the EOS is thermal
(see the hybrid thermal solver in thermal.py).
Returns
-------
f_poststreaming_tree (pytree of jax.numpy.ndarray): Post-streamed distribution function.
f_collision_tree (pytree of jax.numpy.ndarray {Optional}): Post-collision distribution function.
"""
f_postcollision_tree = self.collision(f_poststreaming_tree, T=T)
f_postcollision_tree = self.apply_bc(f_postcollision_tree, f_poststreaming_tree, timestep, "PostCollision")
f_poststreaming_tree = tree_map(lambda f_postcollision: self.streaming(f_postcollision), f_postcollision_tree)
f_poststreaming_tree = self.apply_bc(f_poststreaming_tree, f_postcollision_tree, timestep, "PostStreaming")
if return_fpost:
return f_poststreaming_tree, f_postcollision_tree
else:
return f_poststreaming_tree, None
def run(self, t_max):
"""
This function runs the LBM simulation for a specified number of time steps.
It first initializes the distribution functions and then enters a loop where it performs the
simulation steps (collision, streaming, and boundary conditions) for each time step.
The function can also print the progress of the simulation, save the simulation data, and
compute the performance of the simulation in million lattice updates per second (MLUPS).
Parameters
----------
t_max (int): The total number of time steps to run the simulation.
Returns
-------
f_tree (pytree of jax.numpy.ndarray): Distribution function after t_max timesteps.
"""
f_tree = self.assign_fields_sharded()
start_step = 0
if self.restore_checkpoint:
latest_step = self.mngr.latest_step()
if latest_step is not None: # existing checkpoint present
# Assert that the checkpoint manager is not None
assert self.mngr is not None, "Checkpoint manager does not exist."
c_name = lambda i: f"component_{i}"
restore_target = lambda value: jax.ShapeDtypeStruct(value.shape, value.dtype, sharding=self.sharding)
state = jax.tree.map(restore_target, {c_name(i): f_tree[i] for i in range(self.n_components)})
try:
restored_state = self.mngr.restore(latest_step, args=orb.args.StandardRestore(state))
f_tree = [restored_state[c_name(i)] for i in range(self.n_components)]
logger.info(f"Restored checkpoint at step {latest_step}.")
except ValueError:
raise ValueError(f"Failed to restore checkpoint at step {latest_step}.")
start_step = latest_step + 1
if not (t_max > start_step):
raise ValueError(f"Simulation already exceeded maximum allowable steps (t_max = {t_max}). Consider increasing t_max.")
if self.compute_MLUPS:
start = time.time()
# Loop over all time steps
for timestep in range(start_step, t_max + 1):
io_flag = self.io_rate > 0 and (timestep % self.io_rate == 0 or timestep == t_max)
print_iter_flag = self.print_info_rate > 0 and timestep % self.print_info_rate == 0
checkpoint_flag = self.checkpoint_rate > 0 and timestep % self.checkpoint_rate == 0
# if io_flag:
# # Update the macroscopic variables and save the previous values (for error computation)
# rho_prev_tree, _ = self.update_macroscopic(f_tree)
# # update_macroscopic sums f_tree directly, so rho_prev_tree inherits f_tree's storage precision.
# # macroscopic_velocity -> compute_force -> apply_contact_angle scatters into rho at its own dtype
# # using values derived from G_ff (permanently fixed at compute precision), so under mixed
# # precision (storage narrower than compute) that scatter's source and target dtypes mismatch.
# u_prev_tree = self.macroscopic_velocity(f_tree, rho_prev_tree)
# rho_prev_tree = tree_map(
# lambda rho_prev: downsample_field(rho_prev, self.downsampling_factor),
# rho_prev_tree,
# )
# psi_prev_tree, _ = self.compute_potential(rho_prev_tree)
# p_prev_tree = self.compute_pressure(rho_prev_tree, psi_prev_tree)
# p_prev_total = self.compute_total_pressure(p_prev_tree, rho_prev_tree)
# p_prev_total = downsample_field(p_prev_total, self.downsampling_factor)
# u_prev_tree = tree_map(lambda u_prev: downsample_field(u_prev, self.downsampling_factor), u_prev_tree)
# rho_total_prev = self.compute_total_density(rho_prev_tree)
# u_total_prev = self.compute_total_velocity(rho_prev_tree, u_prev_tree)
# # Gather the data from all processes and convert it to numpy arrays (move to host memory)
# p_prev_total = process_allgather(p_prev_total)
# rho_prev_tree = tree_map(lambda rho_prev: process_allgather(rho_prev), rho_prev_tree)
# u_prev_tree = tree_map(lambda u_prev: process_allgather(u_prev), u_prev_tree)
# rho_total_prev = process_allgather(rho_total_prev)
# u_total_prev = process_allgather(u_total_prev)
# Perform one time-step (collision, streaming, and boundary conditions)
f_tree, fstar_tree = self.step(f_tree, timestep)
# Print the progress of the simulation
if print_iter_flag:
logger.info(
colored("Timestep ", "blue")
+ colored(f"{timestep}", "green")
+ colored(" of ", "blue")
+ colored(f"{t_max}", "green")
+ colored(" completed", "blue")
)
if io_flag:
# Save the simulation data
logger.info(f"Saving data at timestep {timestep}/{t_max}")
rho_tree, _ = self.update_macroscopic(f_tree)
u_tree = self.macroscopic_velocity(f_tree, rho_tree)
psi_tree, _ = self.compute_potential(rho_tree)
p_tree = self.compute_pressure(rho_tree, psi_tree)
p_total = self.compute_total_pressure(p_tree, rho_tree)
p_total = downsample_field(p_total, self.downsampling_factor)
rho_tree = tree_map(
lambda rho: downsample_field(rho, self.downsampling_factor),
rho_tree,
)
u_tree = tree_map(lambda u: downsample_field(u, self.downsampling_factor), u_tree)
rho_total = self.compute_total_density(rho_tree)
u_total = self.compute_total_velocity(rho_tree, u_tree)
# Gather the data from all processes and convert it to numpy arrays (move to host memory)
p_total = process_allgather(p_total)
rho_tree = tree_map(lambda rho: process_allgather(rho), rho_tree)
u_tree = tree_map(lambda u: process_allgather(u), u_tree)
rho_total = process_allgather(rho_total)
u_total = process_allgather(u_total)
# Save the data
self.handle_io_timestep(
timestep,
f_tree,
fstar_tree,
p_tree,
p_total,
u_tree,
u_total,
rho_total,
rho_tree,
# p_prev_tree,
# p_prev_total,
# u_total_prev,
# u_prev_tree,
# rho_total_prev,
# rho_prev_tree,
)
if checkpoint_flag:
# Save the checkpoint
logger.info(f"Saving checkpoint at timestep {timestep}/{t_max}")
state = {}
c_name = lambda i: f"component_{i}"
for i in range(self.n_components):
state[c_name(i)] = f_tree[i]
self.mngr.save(timestep, args=orb.args.StandardSave(state))
# Start the timer for the MLUPS computation after the first timestep (to remove compilation overhead)
if self.compute_MLUPS and timestep == 1:
jax.block_until_ready(f_tree)
start = time.time()
if self.compute_MLUPS:
# Compute and print the performance of the simulation in MLUPS
jax.block_until_ready(f_tree)
end = time.time()
if self.dim == 2:
logger.info(
colored("Domain: ", "blue") + colored(f"{self.nx} x {self.ny}", "green")
if self.dim == 2
else colored(f"{self.nx} x {self.ny} x {self.nz}", "green")
)
logger.info(
colored("Number of voxels: ", "blue") + colored(f"{self.nx * self.ny}", "green")
if self.dim == 2
else colored(f"{self.nx * self.ny * self.nz}", "green")
)
logger.info(
colored("MLUPS: ", "blue")
+ colored(
f"{self.n_components * self.nx * self.ny * t_max / (end - start) / 1e6}",
"red",
)
)
elif self.dim == 3:
logger.info(colored("Domain: ", "blue") + colored(f"{self.nx} x {self.ny} x {self.nz}", "green"))
logger.info(colored("Number of voxels: ", "blue") + colored(f"{self.nx * self.ny * self.nz}", "green"))
logger.info(
colored("MLUPS: ", "blue")
+ colored(
f"{self.n_components * self.nx * self.ny * self.nz * t_max / (end - start) / 1e6}",
"red",
)
)
if self.mngr is not None:
self.mngr.wait_until_finished()
return f_tree
def handle_io_timestep(
self,
timestep,
f_tree,
fstar_tree,
p_tree,
p_total,
u_tree,
u_total,
rho_total,
rho_tree,
# p_prev_tree,
# p_prev_total,
# u_total_prev,
# u_prev_tree,
# rho_total_prev,
# rho_prev_tree,
):
"""
This function handles the input/output (I/O) operations at each time step of the simulation.
It prepares the data to be saved and calls the output_data function, which can be overwritten
by the user to customize the I/O operations.
Parameters
----------
timestep (int): The current time step of the simulation.
f_tree (pytree of jax.numpy.ndarray): Post-streaming distribution functions at the current time step.
fstar_tree (pytree of jax.numpy.ndarray): Post-collision distribution functions at the current time step.
p_tree (pytree of jax.numpy.ndarray): Pressure field at the current time step.
p_total (jax.numpy.ndarray): Total pressure field at the current time step.
u_total (jax.numpy.ndarray): Total velocity field at the current time step.
u_tree (pytree of jax.numpy.ndarray): Velocity field at the current time step.
rho_total (jax.numpy.ndarray): Total density field at the current time step.
rho_tree (pytree of jax.numpy.ndarray): Density field at the current time step.
# p_prev_tree (pytree of jax.numpy.ndarray): Pressure field at the previous time step.
# p_prev_total (jax.numpy.ndarray): Total pressure field at the previous time step.
# u_total_prev (jax.numpy.ndarray): Total velocity field at the previous time step.
# u_prev_tree (pytree of jax.numpy.ndarray): Velocity field at the previous time step.
# rho_total_prev (jax.numpy.ndarray): Total density field at the previous time step.
# rho_prev_tree (pytree of jax.numpy.ndarray): Density field at the previous time step.
Returns
-------
None
"""
kwargs = {
"n_components": self.n_components,
"timestep": timestep,
"rho_total": rho_total,
"rho_tree": rho_tree,
"p_tree": p_tree,
"p": p_total,
"u_total": u_total,
"u_tree": u_tree,
# "rho_total_prev": rho_total_prev,
# "rho_prev_tree": rho_prev_tree,
# "p_prev_tree": p_prev_tree,
# "p_prev": p_prev_total,
# "u_total_prev": u_total_prev,
# "u_prev_tree": u_prev_tree,
"f_poststreaming_tree": f_tree,
"f_postcollision_tree": fstar_tree,
}
self.output_data(**kwargs)