Make the XLB backend an optional extra

.github/workflows/tests.yml

@@ -8,7 +8,21 @@ on:
 
 jobs:
   test:
+    name: test (${{ matrix.install }})
     runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          # The full install, with the optional XLB backend, runs every test.
+          - install: with xlb extra
+            sync: uv sync --group dev --extra xlb
+            without_xlb: "0"
+          # The default install leaves XLB, JAX and Warp out; the XLB tests skip
+          # and the missing-extra error is checked against a real environment.
+          - install: without xlb extra
+            sync: uv sync --group dev
+            without_xlb: "1"
     steps:
       - uses: actions/checkout@v6
 
@@ -24,12 +38,15 @@ jobs:
           python-version: "3.12"
 
       - name: Install dependencies
-        run: uv sync --group dev
+        run: ${{ matrix.sync }}
 
       - name: Run tests
-        run: uv run pytest tests/ -v --cov=kabs --cov-report=xml
+        env:
+          KABS_TEST_WITHOUT_XLB: ${{ matrix.without_xlb }}
+        run: uv run --no-sync pytest tests/ -v --cov=kabs --cov-report=xml
 
       - name: Upload coverage to Codecov
+        if: matrix.without_xlb == '0'
         uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}

README.md

@@ -22,11 +22,22 @@ cd kabs
 uv sync
 ```
 
-Alternatively, install with `pip install -e .`. This installs both solvers:
-Taichi and XLB, with `warp-lang==1.10.0` pinned for compatibility with XLB.
+Alternatively, install with `pip install -e .`. This installs the default
+Taichi solver. The XLB solver, with its JAX and Warp compute backends, is the
+optional `xlb` extra, which adds XLB, JAX and `warp-lang==1.10.0` (pinned for
+compatibility with XLB):
 
-Key dependencies: `taichi` (GPU/CPU acceleration), `xlb`, `numpy`, `pyevtk`.
-Optional: `porespy` (used in the examples below to generate synthetic images).
+```bash
+uv sync --extra xlb         # in a clone
+pip install -e ".[xlb]"     # or with pip
+```
+
+Selecting `backend="xlb"` without the extra raises an `ImportError` that names
+it.
+
+Key dependencies: `taichi` (GPU/CPU acceleration), `numpy`, `pyevtk`.
+Optional: `xlb` and `warp-lang` (the `xlb` extra), and `porespy` (used in the
+examples below to generate synthetic images).
 
 ---
 
@@ -71,8 +82,9 @@ ti.init(arch=ti.metal) # Apple Silicon GPU
 ### Solver backends
 
 `solve_flow` uses the Taichi implementation by default. Select XLB through the
-same API; its default JAX compute backend works on CPU and supports multi-GPU
-execution on supported accelerator platforms:
+same API once the `xlb` extra is installed; its default JAX compute backend
+works on CPU and supports multi-GPU execution on supported accelerator
+platforms:
 
 ```python
 # Default: solve_flow(im, backend="taichi")

docs/warp-backend-readiness.md

@@ -222,7 +222,7 @@ fresh environment:
 ```powershell
 git fetch origin
 git switch finalizing-warp-option
-uv sync --group dev
+uv sync --group dev --extra xlb
 ```
 
 Verify that Warp detects CUDA:

kabs/_solve_flow_xlb.py

@@ -224,10 +224,13 @@ def solve_flow_xlb(
     call it with custom settings before invoking this function, set
     ``compute_backend`` to match your prior ``xlb.init()`` call.
 
-    **Installation**::
+    **Installation**: XLB and Warp are not installed with kabs by default.
+    They come with the optional ``xlb`` extra::
 
-        pip install "xlb>=0.3.1" "warp-lang==1.10.0"  # CPU/JAX
-        pip install "xlb[cuda]>=0.3.1" "warp-lang==1.10.0"  # NVIDIA GPU via JAX
+        pip install "kabs[xlb]"  # CPU/JAX and Warp
+        pip install "kabs[xlb]" "xlb[cuda]>=0.3.1"  # NVIDIA GPU via JAX
+
+    Without it, this function raises ``ImportError`` naming the extra.
     """
     if not isinstance(direction, str) or direction.lower() not in _FACE_NAMES:
         raise ValueError(f"direction must be 'x', 'y', or 'z', got {direction!r}")
@@ -271,8 +274,11 @@ def solve_flow_xlb(
         from xlb.operator.macroscopic import Macroscopic
     except ImportError as exc:
         raise ImportError(
-            "XLB is required for solve_flow_xlb. Install compatible dependencies with:\n"
-            "    pip install 'xlb>=0.3.1' 'warp-lang==1.10.0'\n"
+            "The XLB backend needs kabs's optional 'xlb' extra, which installs "
+            "XLB and Warp. Install it with:\n"
+            "    pip install 'kabs[xlb]'\n"
+            "or, in a clone of the kabs repository:\n"
+            "    uv sync --extra xlb\n"
         ) from exc
 
     _backend_map = {"jax": ComputeBackend.JAX, "warp": ComputeBackend.WARP}

pyproject.toml

@@ -17,6 +17,10 @@ dependencies = [
     "trame>=3.12.0",
     "trame-vtk>=2.11.3",
     "trame-vuetify>=3.2.1",
+]
+
+[project.optional-dependencies]
+xlb = [
     "xlb>=0.3.1",
     "warp-lang==1.10.0",
 ]

tests/test_backend_import_isolation.py

@@ -33,6 +33,21 @@ def test_import_kabs_does_not_import_taichi():
     assert completed.returncode == 0, completed.stdout + completed.stderr
 
 
+def test_import_kabs_does_not_import_the_xlb_stack():
+    completed = _run_fresh_python(
+        """
+        import sys
+        import kabs
+
+        loaded = {"xlb", "jax", "warp"} & {name.split(".")[0] for name in sys.modules}
+        assert not loaded, sorted(loaded)
+        assert "kabs._warp_convergence" not in sys.modules
+        """
+    )
+
+    assert completed.returncode == 0, completed.stdout + completed.stderr
+
+
 def test_public_single_phase_solver_export_remains_lazy():
     completed = _run_fresh_python(
         """

tests/test_convergence_xlb.py

@@ -3,14 +3,17 @@
 import importlib
 import os
 
-import jax
-import jax.numpy as jnp
 import numpy as np
 import pytest
 
-from kabs import compute_permeability, solve_flow_xlb
-from kabs._convergence import ConvergenceConfig
-from kabs._solve_flow_xlb import (
+pytest.importorskip("xlb", reason="xlb is not installed; skipping XLB tests")
+
+import jax  # noqa: E402
+import jax.numpy as jnp  # noqa: E402
+
+from kabs import compute_permeability, solve_flow_xlb  # noqa: E402
+from kabs._convergence import ConvergenceConfig  # noqa: E402
+from kabs._solve_flow_xlb import (  # noqa: E402
     _convergence_sums_to_host,
     _get_jax_convergence_reducer,
 )

tests/test_solve_flow_xlb_validation.py

@@ -1,11 +1,13 @@
 """Public validation for the optional XLB solver."""
 
 import builtins
+import importlib.util
+import os
 
 import numpy as np
 import pytest
 
-from kabs import solve_flow_xlb
+from kabs import solve_flow, solve_flow_xlb
 
 
 _IMAGE = np.ones((3, 3, 3), dtype=np.int8)
@@ -22,6 +24,18 @@ def guarded_import(name, *args, **kwargs):
     monkeypatch.setattr(builtins, "__import__", guarded_import)
 
 
+def _hide_xlb(monkeypatch):
+    """Make ``import xlb`` fail as it does when the ``xlb`` extra is not installed."""
+    original_import = builtins.__import__
+
+    def import_without_xlb(name, *args, **kwargs):
+        if name == "xlb" or name.startswith("xlb."):
+            raise ModuleNotFoundError(f"No module named {name!r}", name=name)
+        return original_import(name, *args, **kwargs)
+
+    monkeypatch.setattr(builtins, "__import__", import_without_xlb)
+
+
 @pytest.mark.parametrize("direction", ["invalid", None])
 def test_invalid_direction_is_reported_without_importing_xlb(direction, monkeypatch):
     _block_xlb_import(monkeypatch)
@@ -36,3 +50,22 @@ def test_invalid_compute_backend_is_reported_without_importing_xlb(
     _block_xlb_import(monkeypatch)
     with pytest.raises(ValueError, match="compute_backend"):
         solve_flow_xlb(_IMAGE, compute_backend=compute_backend)
+
+
+@pytest.mark.parametrize("compute_backend", ["jax", "warp"])
+def test_missing_xlb_extra_is_named_in_the_error(compute_backend, monkeypatch):
+    _hide_xlb(monkeypatch)
+    with pytest.raises(ImportError, match=r"kabs\[xlb\]") as excinfo:
+        solve_flow(_IMAGE, backend="xlb", compute_backend=compute_backend)
+    assert isinstance(excinfo.value.__cause__, ModuleNotFoundError)
+
+
+@pytest.mark.skipif(
+    os.environ.get("KABS_TEST_WITHOUT_XLB") != "1",
+    reason="set KABS_TEST_WITHOUT_XLB=1 in an environment installed without the xlb extra",
+)
+def test_install_without_the_extra_omits_the_xlb_stack():
+    for module in ("xlb", "warp", "jax"):
+        assert importlib.util.find_spec(module) is None, module
+    with pytest.raises(ImportError, match=r"kabs\[xlb\]"):
+        solve_flow(_IMAGE, backend="xlb")