Adding xlb backend option

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 *

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

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")