diff --git a/kabs/__init__.py b/kabs/__init__.py index 5365362..82f00fd 100644 --- a/kabs/__init__.py +++ b/kabs/__init__.py @@ -13,6 +13,7 @@ from ._single_phase_solver import * from ._solve_flow import * +from ._solve_flow_xlb import * from ._compute_permeability import * from ._compute_hydraulic_conductance import * from .plots import * diff --git a/kabs/_solve_flow_xlb.py b/kabs/_solve_flow_xlb.py new file mode 100644 index 0000000..2efaeda --- /dev/null +++ b/kabs/_solve_flow_xlb.py @@ -0,0 +1,281 @@ +"""High-level entry point: run a single-phase LBM flow simulation using XLB.""" + +import time + +import numpy as np + +from ._solve_flow import FlowResult, _RHO_IN, _RHO_OUT + + +__all__ = ["solve_flow_xlb"] + +# XLB 3D face names for each flow direction. +# In XLB's grid convention: left=x0, right=x1, front=y0, back=y1, bottom=z0, top=z1 +_FACE_NAMES = { + "x": ("left", "right"), + "y": ("front", "back"), + "z": ("bottom", "top"), +} + + +def solve_flow_xlb( + im, + direction="x", + n_steps=15000, + nu=1.0 / 6.0, + log_every=500, + verbose=True, + tol=1e-3, + compute_backend="jax", +): + """ + Run a pressure-driven single-phase LBM simulation using the XLB library. + + This function is the XLB-based equivalent of ``solve_flow()``. It returns + the same ``FlowResult`` object and is fully compatible with the downstream + ``compute_permeability()`` and ``compute_hydraulic_conductance()`` functions. + + Parameters + ---------- + im : np.ndarray, shape (nx, ny, nz) + Binary image of the pore space. 1 (or True) = pore, 0 (or False) = solid. + This matches the PoreSpy convention. + direction : {'x', 'y', 'z'} + Axis along which the pressure gradient is applied. Default ``'x'``. + n_steps : int + Maximum number of LBM time steps to run. Default 15000. + nu : float + Kinematic viscosity in lattice units. Default 1/6. + log_every : int + Print a progress line and check convergence every this many steps. + Default 500. + verbose : bool + Print progress to stdout. Default True. + tol : float or None + Convergence tolerance. The simulation stops early when the relative + change in the total velocity magnitude between log intervals falls + below this value: ``delta|v| / |v| < tol``. Set to ``None`` to + always run the full ``n_steps``. Default 1e-3. + compute_backend : {'jax', 'warp'} + XLB compute backend. ``'jax'`` enables multi-GPU via JAX (CPU also + works). ``'warp'`` uses NVIDIA Warp for single-GPU runs. + Default ``'jax'``. + + Returns + ------- + result : FlowResult + Result object containing ``solid``, ``rho``, ``velocity``, + ``direction``, and ``nu`` as numpy arrays/values. Pass directly to + ``compute_permeability()`` or ``compute_hydraulic_conductance()``, + or call ``result.export_to_vtk(prefix)`` to save a VTR file. + + Notes + ----- + **Collision model**: this solver uses the BGK (single relaxation time) + operator. BGK is appropriate for Stokes-regime (low Reynolds number) + porous-media flow and produces physically equivalent results to the MRT + operator used by ``solve_flow()`` in that regime. If you suspect + inertial effects (higher flow rates, large pores), prefer ``solve_flow()`` + which uses MRT. + + **Boundary conditions**: inlet and outlet faces receive an equilibrium + pressure (density) BC with ρ_in = 1.00 and ρ_out = 0.99, matching + ``solve_flow()``. Unassigned transverse faces are automatically periodic + (XLB's streaming operator wraps around by default). + + **JIT compilation**: JAX will compile the stepper on the first call. + Expect an additional overhead of 10–60 s on the first invocation. + + **Initialization**: ``xlb.init()`` is called internally. If you need to + call it with custom settings before invoking this function, set + ``compute_backend`` to match your prior ``xlb.init()`` call. + + **Installation**:: + + pip install xlb # CPU-only + pip install "xlb[cuda]" # NVIDIA GPU via JAX + pip install "xlb[warp]" # NVIDIA Warp (single GPU) + """ + try: + import xlb + from xlb.compute_backend import ComputeBackend + from xlb.precision_policy import PrecisionPolicy + from xlb.grid import grid_factory + from xlb.operator.stepper import IncompressibleNavierStokesStepper + from xlb.operator.boundary_condition import ZouHeBC, FullwayBounceBackBC + from xlb.operator.macroscopic import Macroscopic + except ImportError as exc: + raise ImportError( + "XLB is required for solve_flow_xlb. Install it with:\n" + " pip install xlb # CPU-only\n" + " pip install 'xlb[cuda]' # NVIDIA GPU via JAX\n" + " pip install 'xlb[warp]' # NVIDIA Warp (single GPU)\n" + ) from exc + + direction = direction.lower() + if direction not in _FACE_NAMES: + raise ValueError(f"direction must be 'x', 'y', or 'z', got {direction!r}") + + _backend_map = {"jax": ComputeBackend.JAX, "warp": ComputeBackend.WARP} + backend_key = compute_backend.lower() + if backend_key not in _backend_map: + raise ValueError( + f"compute_backend must be 'jax' or 'warp', got {compute_backend!r}" + ) + _backend = _backend_map[backend_key] + + # --- Initialize XLB (safe to call multiple times) --- + precision_policy = PrecisionPolicy.FP32FP32 + velocity_set = xlb.velocity_set.D3Q19( + precision_policy=precision_policy, + compute_backend=_backend, + ) + xlb.init( + velocity_set=velocity_set, + default_backend=_backend, + default_precision_policy=precision_policy, + ) + + # --- Prepare image --- + # Public (PoreSpy) convention: 1 = pore, 0 = solid. + # Internal (FlowResult) convention: 1 = solid, 0 = pore. + solid_im = (im == 0).astype(np.int8) + solid_mask = solid_im.astype(bool) # True where solid + nx, ny, nz = im.shape + + # BGK relaxation parameter from viscosity: tau = 3*nu + 0.5, omega = 1/tau + omega = 1.0 / (3.0 * nu + 0.5) + + # --- Build XLB grid and resolve boundary indices --- + grid = grid_factory((nx, ny, nz), compute_backend=_backend) + + # remove_edges=True avoids assigning corner/edge voxels to both the + # pressure BC and any transverse-face BC, consistent with XLB conventions. + box_no_edge = grid.bounding_box_indices(remove_edges=True) + + inlet_name, outlet_name = _FACE_NAMES[direction] + + def _pore_face_indices(face_key): + """Return face voxel indices restricted to pore voxels.""" + idx = np.array(box_no_edge[face_key]) # shape (3, n) + xs, ys, zs = idx[0], idx[1], idx[2] + is_pore = ~solid_mask[xs, ys, zs] + return [idx[d][is_pore].tolist() for d in range(3)] + + inlet_indices = _pore_face_indices(inlet_name) + outlet_indices = _pore_face_indices(outlet_name) + + # All solid voxels throughout the domain + solid_coords = np.where(solid_mask) + solid_indices = [solid_coords[d].tolist() for d in range(3)] + + # --- Boundary conditions --- + # Pressure BCs: Zou-He pressure BC at inlet/outlet pore voxels. + # This sets the density to the prescribed value and determines velocity + # from mass conservation — analogous to the Taichi pressure BC. + # Solid BCs: full-way bounce-back throughout the domain. + # Transverse faces have no explicit BC → streaming wraps around (periodic). + bc_inlet = ZouHeBC(bc_type="pressure", prescribed_value=float(_RHO_IN), indices=inlet_indices) + bc_outlet = ZouHeBC(bc_type="pressure", prescribed_value=float(_RHO_OUT), indices=outlet_indices) + boundary_conditions = [bc_inlet, bc_outlet] + + if solid_coords[0].size > 0: + bc_solid = FullwayBounceBackBC(indices=solid_indices) + boundary_conditions.append(bc_solid) + + # --- Macroscopic operator (always JAX for numpy extraction) --- + # For the Warp backend, f arrays are converted to JAX before calling macro. + macro = Macroscopic( + compute_backend=ComputeBackend.JAX, + precision_policy=precision_policy, + velocity_set=xlb.velocity_set.D3Q19( + precision_policy=precision_policy, + compute_backend=ComputeBackend.JAX, + ), + ) + + def _to_jax_array(f): + """Convert f to a JAX array if it is a Warp array.""" + import jax.numpy as jnp + + if isinstance(f, jnp.ndarray): + return f + from xlb.utils import warp_array_to_jax + + return warp_array_to_jax(f) + + # --- Stepper --- + stepper = IncompressibleNavierStokesStepper( + grid=grid, + boundary_conditions=boundary_conditions, + collision_type="BGK", + ) + f_0, f_1, bc_mask, missing_mask = stepper.prepare_fields() + + # --- Time loop --- + time_init = time.time() + time_pre = time_init + v_prev = None + final_step = n_steps + final_criterion = None + + for i in range(n_steps + 1): + f_0, f_1 = stepper(f_0, f_1, bc_mask, missing_mask, omega, i) + f_0, f_1 = f_1, f_0 # swap buffers; f_0 now holds the updated state + + if i % log_every == 0: + time_now = time.time() + diff = int(time_now - time_pre) + elap = int(time_now - time_init) + m_d, s_d = divmod(diff, 60) + h_d, m_d = divmod(m_d, 60) + m_e, s_e = divmod(elap, 60) + h_e, m_e = divmod(m_e, 60) + + if verbose: + print( + f"Step {i:6d}/{n_steps} " + f"interval {h_d:02d}h{m_d:02d}m{s_d:02d}s " + f"elapsed {h_e:02d}h{m_e:02d}m{s_e:02d}s" + ) + + # Convergence check: extract velocity on CPU. + # rho shape: (1, nx, ny, nz); u shape: (3, nx, ny, nz) + _, u_jax = macro(_to_jax_array(f_0)) + v_now = np.array(u_jax) # (3, nx, ny, nz) + + if v_prev is not None: + v_total = np.sum(np.abs(v_now)) + v_change = np.sum(np.abs(v_now - v_prev)) + if v_total > 0: + final_criterion = v_change / v_total + if verbose: + print(f" |v|={v_total:.4e} delta|v|={v_change:.4e}") + if tol is not None and v_total > 0 and final_criterion < tol: + if verbose: + print( + f"Converged at step {i} " + f"(delta|v|/|v| = {final_criterion:.2e} < tol={tol:.2e})" + ) + final_step = i + break + v_prev = v_now + time_pre = time_now + + # --- Extract final fields --- + rho_jax, u_jax = macro(_to_jax_array(f_0)) + # rho: (1, nx, ny, nz) → (nx, ny, nz); u: (3, nx, ny, nz) → (nx, ny, nz, 3) + rho_np = np.array(rho_jax[0]).astype(np.float32) + vel_np = np.moveaxis(np.array(u_jax), 0, -1).astype(np.float32) + + result = FlowResult.from_arrays( + solid=solid_im, + rho=rho_np, + velocity=vel_np, + direction=direction, + nu=nu, + ) + result.n_iterations = final_step + result.convergence_criterion = final_criterion + + return result diff --git a/tests/benchmark_solvers.py b/tests/benchmark_solvers.py new file mode 100644 index 0000000..8c74a42 --- /dev/null +++ b/tests/benchmark_solvers.py @@ -0,0 +1,92 @@ +"""Compare wall-clock time for Taichi vs XLB solvers. + +Run with: + python benchmark_solvers.py + +Geometry: 4-cylinder bundle (same as test_solve_flow_xlb.py / test_permeability.py). +N_STEPS is kept the same for both so times are directly comparable. +""" + +import time +import numpy as np +import taichi as ti + +ti.init(arch=ti.cpu) + +import kabs + +# --------------------------------------------------------------------------- +# Geometry +# --------------------------------------------------------------------------- +R = 10 +L = 30 +NY = NZ = 50 +NU = 1.0 / 6.0 +CENTRES = [(13, 13), (13, 37), (37, 13), (37, 37)] +N_STEPS = 4000 + +a, b = np.mgrid[0:NY, 0:NZ] +cross = np.zeros((NY, NZ), dtype=bool) +for ca, cb in CENTRES: + cross |= (a - ca) ** 2 + (b - cb) ** 2 <= R**2 +im = np.broadcast_to(cross[np.newaxis], (L, NY, NZ)).copy().astype(int) + +print(f"Domain: {im.shape} pore fraction: {im.mean():.3f} n_steps: {N_STEPS}\n") + +# --------------------------------------------------------------------------- +# Taichi solve (includes JIT warm-up on first call) +# --------------------------------------------------------------------------- +print("=== Taichi (MRT) ===") +t0 = time.perf_counter() +result_ti = kabs.solve_flow(im, direction="x", n_steps=N_STEPS, nu=NU, + log_every=N_STEPS, verbose=False, tol=None) +t_taichi = time.perf_counter() - t0 + +res_ti = kabs.compute_permeability(result_ti, verbose=False) +k_ti = res_ti["k_lu"] +print(f" Time : {t_taichi:.2f} s") +print(f" k : {k_ti:.4f} (LU)") + +# --------------------------------------------------------------------------- +# XLB solve (includes JAX JIT warm-up on first call) +# --------------------------------------------------------------------------- +print("\n=== XLB (BGK / JAX) ===") +t0 = time.perf_counter() +result_xlb = kabs.solve_flow_xlb(im, direction="x", n_steps=N_STEPS, nu=NU, + log_every=N_STEPS, verbose=False, tol=None) +t_xlb = time.perf_counter() - t0 + +res_xlb = kabs.compute_permeability(result_xlb, verbose=False) +k_xlb = res_xlb["k_lu"] +print(f" Time : {t_xlb:.2f} s") +print(f" k : {k_xlb:.4f} (LU)") + +# --------------------------------------------------------------------------- +# Summary +# --------------------------------------------------------------------------- +print("\n=== Summary ===") +print(f" Taichi : {t_taichi:.2f} s k = {k_ti:.4f}") +print(f" XLB : {t_xlb:.2f} s k = {k_xlb:.4f}") +ratio = t_taichi / t_xlb if t_xlb > 0 else float("nan") +print(f" Taichi / XLB speedup: {ratio:.2f}x (>1 means XLB is faster)") +print(f" |k_ti - k_xlb| / k_ti: {abs(k_ti - k_xlb) / k_ti * 100:.2f} %") + +# --------------------------------------------------------------------------- +# Optional: second run to separate JIT from steady-state throughput +# --------------------------------------------------------------------------- +print("\n=== Second run (JIT already warm) ===") + +t0 = time.perf_counter() +kabs.solve_flow(im, direction="x", n_steps=N_STEPS, nu=NU, + log_every=N_STEPS, verbose=False, tol=None) +t_taichi2 = time.perf_counter() - t0 + +t0 = time.perf_counter() +kabs.solve_flow_xlb(im, direction="x", n_steps=N_STEPS, nu=NU, + log_every=N_STEPS, verbose=False, tol=None) +t_xlb2 = time.perf_counter() - t0 + +ratio2 = t_taichi2 / t_xlb2 if t_xlb2 > 0 else float("nan") +print(f" Taichi : {t_taichi2:.2f} s") +print(f" XLB : {t_xlb2:.2f} s") +print(f" Taichi / XLB speedup: {ratio2:.2f}x") diff --git a/tests/test_solve_flow_xlb.py b/tests/test_solve_flow_xlb.py new file mode 100644 index 0000000..7198aa7 --- /dev/null +++ b/tests/test_solve_flow_xlb.py @@ -0,0 +1,227 @@ +"""Tests for solve_flow_xlb — XLB-based LBM solver. + +Two suites: + +1. ``TestSolveFlowXlbSmoke`` — fast structural checks (field shapes, types, + FlowResult attributes) on a tiny domain. + +2. ``TestXlbVsTaichi`` — numerical agreement between the XLB (BGK) and Taichi + (MRT) solvers on the same 4-cylinder bundle-of-tubes geometry used in + ``test_permeability.py``. Both solvers should agree with each other and + with the analytical Hagen-Poiseuille result. + +Notes on tolerance +------------------ +BGK and MRT are different collision operators but produce identical steady-state +solutions in the Stokes (low-Re) regime. The permeability results are expected +to agree within 5 % — in practice they are much closer, but a generous bound +accounts for minor floating-point differences due to the differing collision +paths. +""" + +import numpy as np +import pytest + +pytest.importorskip("xlb", reason="xlb is not installed; skipping XLB tests") + +from kabs import solve_flow, solve_flow_xlb, compute_permeability # noqa: E402 + + +# --------------------------------------------------------------------------- +# Shared geometry helpers (mirrors test_permeability.py) +# --------------------------------------------------------------------------- + +_R = 10 +_L = 30 +_NY = _NZ = 50 +_NU = 1.0 / 6.0 +_CENTRES = [(13, 13), (13, 37), (37, 13), (37, 37)] +_N_TUBES = len(_CENTRES) + + +def _bundle_image(direction="x"): + """Binary image (1=pore) of 4 parallel cylinders along *direction*.""" + a, b = np.mgrid[0:_NY, 0:_NZ] + cross = np.zeros((_NY, _NZ), dtype=bool) + for ca, cb in _CENTRES: + cross |= (a - ca) ** 2 + (b - cb) ** 2 <= _R**2 + if direction == "x": + return np.broadcast_to(cross[np.newaxis], (_L, _NY, _NZ)).copy().astype(int) + elif direction == "y": + return np.broadcast_to(cross[:, np.newaxis, :], (_NY, _L, _NZ)).copy().astype(int) + else: + return np.broadcast_to(cross[:, :, np.newaxis], (_NY, _NZ, _L)).copy().astype(int) + + +def _k_analytical(): + return _N_TUBES * np.pi * _R**4 / (8 * _NY * _NZ) + + +def _tiny_image(): + """Single cylinder r=4 in a 12×12 cross-section, length 10 — for fast smoke tests.""" + a, b = np.mgrid[0:12, 0:12] + cross = (a - 6) ** 2 + (b - 6) ** 2 <= 4**2 + return np.broadcast_to(cross[np.newaxis], (10, 12, 12)).copy().astype(int) + + +# --------------------------------------------------------------------------- +# Suite 1: Structural / smoke tests +# --------------------------------------------------------------------------- + +_TINY_KW = dict(direction="x", nu=_NU, n_steps=500, tol=1e-3, log_every=100, verbose=False) + + +class TestSolveFlowXlbSmoke: + @classmethod + def setup_class(cls): + cls.im = _tiny_image() + cls.result = solve_flow_xlb(cls.im, **_TINY_KW) + + def test_returns_flow_result(self): + from kabs._solve_flow import FlowResult + + assert isinstance(self.result, FlowResult) + + def test_velocity_shape(self): + assert self.result.velocity.shape == (10, 12, 12, 3) + + def test_rho_shape(self): + assert self.result.rho.shape == (10, 12, 12) + + def test_solid_shape(self): + assert self.result.solid.shape == (10, 12, 12) + + def test_velocity_dtype(self): + assert self.result.velocity.dtype == np.float32 + + def test_rho_dtype(self): + assert self.result.rho.dtype == np.float32 + + def test_solid_is_flipped(self): + """Internal convention: 1=solid, 0=pore (opposite of input).""" + assert self.result.solid.dtype == np.int8 + # Input im has 1=pore; solid should have 0 where im==1 + assert np.all(self.result.solid[self.im == 1] == 0) + assert np.all(self.result.solid[self.im == 0] == 1) + + def test_direction_stored(self): + assert self.result.direction == "x" + + def test_nu_stored(self): + assert self.result.nu == pytest.approx(_NU) + + def test_n_iterations_set(self): + assert self.result.n_iterations is not None + assert self.result.n_iterations <= 500 + + def test_rho_in_range(self): + """Density should stay close to 1 everywhere.""" + assert float(self.result.rho.min()) > 0.95 + assert float(self.result.rho.max()) < 1.05 + + def test_flow_in_x_direction(self): + """Pressure gradient along x should produce positive x-velocity.""" + pore = self.result.solid == 0 + assert float(self.result.velocity[..., 0][pore].mean()) > 0.0 + + def test_transverse_velocity_small(self): + """y and z velocities should be negligibly small for axial flow.""" + pore = self.result.solid == 0 + vy_mean = float(np.abs(self.result.velocity[..., 1][pore]).mean()) + vz_mean = float(np.abs(self.result.velocity[..., 2][pore]).mean()) + vx_mean = float(np.abs(self.result.velocity[..., 0][pore]).mean()) + assert vy_mean < 0.05 * vx_mean + assert vz_mean < 0.05 * vx_mean + + def test_compatible_with_compute_permeability(self): + """FlowResult must be accepted by compute_permeability without error.""" + out = compute_permeability(self.result, verbose=False) + assert out["k_lu"] > 0.0 + + def test_invalid_direction_raises(self): + with pytest.raises(ValueError, match="direction"): + solve_flow_xlb(self.im, direction="w") + + def test_invalid_backend_raises(self): + with pytest.raises(ValueError, match="compute_backend"): + solve_flow_xlb(self.im, compute_backend="cuda_nope") + + +# --------------------------------------------------------------------------- +# Suite 2: XLB vs Taichi numerical agreement +# --------------------------------------------------------------------------- + +_SOLVE_KW = dict(nu=_NU, n_steps=4000, tol=1e-3, log_every=200, verbose=False) + + +class TestXlbVsTaichi: + """Compare XLB (BGK) against Taichi (MRT) on the 4-cylinder bundle.""" + + @classmethod + def setup_class(cls): + im = _bundle_image(direction="x") + cls.result_ti = solve_flow(im, direction="x", **_SOLVE_KW) + cls.result_xlb = solve_flow_xlb(im, direction="x", **_SOLVE_KW) + cls.k_ana = _k_analytical() + + # --- Both solvers vs analytical --- + + def test_taichi_vs_analytical(self): + """Taichi result must be within 7 % of Hagen-Poiseuille (reference check).""" + out = compute_permeability(self.result_ti, verbose=False) + assert out["k_lu"] == pytest.approx(self.k_ana, rel=0.07) + + def test_xlb_vs_analytical(self): + """XLB result must be within 7 % of Hagen-Poiseuille.""" + out = compute_permeability(self.result_xlb, verbose=False) + assert out["k_lu"] == pytest.approx(self.k_ana, rel=0.07) + + # --- Mutual agreement --- + + def test_permeability_agreement(self): + """XLB and Taichi permeabilities must agree within 5 %.""" + k_ti = compute_permeability(self.result_ti, verbose=False)["k_lu"] + k_xlb = compute_permeability(self.result_xlb, verbose=False)["k_lu"] + assert k_xlb == pytest.approx(k_ti, rel=0.05) + + def test_darcy_velocity_agreement(self): + """XLB and Taichi Darcy velocities must agree within 5 %.""" + u_ti = compute_permeability(self.result_ti, verbose=False)["u_darcy"] + u_xlb = compute_permeability(self.result_xlb, verbose=False)["u_darcy"] + assert u_xlb == pytest.approx(u_ti, rel=0.05) + + def test_mean_rho_agreement(self): + """Volume-averaged density should be close between solvers.""" + rho_ti = float(self.result_ti.rho.mean()) + rho_xlb = float(self.result_xlb.rho.mean()) + assert rho_xlb == pytest.approx(rho_ti, rel=0.01) + + +# --------------------------------------------------------------------------- +# Allow direct execution: python tests/test_solve_flow_xlb.py +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + import taichi as ti + + ti.init(arch=ti.cpu) + + obj_smoke = TestSolveFlowXlbSmoke() + TestSolveFlowXlbSmoke.setup_class() + obj_smoke = TestSolveFlowXlbSmoke() + + obj_cross = TestXlbVsTaichi() + TestXlbVsTaichi.setup_class() + + for cls, obj in [(TestSolveFlowXlbSmoke, obj_smoke), (TestXlbVsTaichi, obj_cross)]: + print(f"\n{'=' * 60}") + print(f" {cls.__name__}") + print(f"{'=' * 60}") + for name in sorted(dir(obj)): + if name.startswith("test_"): + method = getattr(obj, name) + try: + method() + print(f" PASS {name}") + except Exception as exc: + print(f" FAIL {name}: {exc}")