refactor: ♻️ adopt the array API standard for backend dispatch and array typing (#69)

* refactor: ♻️ adopt the array API standard for backend dispatch and array typing

The algorithms take the array namespace from their arrays through
array-api-compat instead of receiving an `xp` module. The namespace is
picked once from `device=` by `DeviceGroup`, which now also resolves the
device ids; `EnvironmentStore` follows the arrays it is given, and
`truncated_qr` and the splitting kernels follow their inputs.

Arrays of either backend are annotated with one alias,
`src_method.utils.Array`; host-only arrays (the results, the exact path,
the sketches) are `np.ndarray`. Array creation, `linalg.qr`, `conj`,
`reshape`, `moveaxis`, `concat` and `asarray` use the standard names.
Streams, pinned memory and CuPy's pool stay behind `utils._backend`, and
the Gaussian sketches are always drawn on the host.

Closes #45

Assisted-by: GitHub Copilot:Claude Opus 5.5

* chore: clean-up readme.

Author: Panadestein

AGENTS.md

@@ -20,9 +20,12 @@ tensor, so there is no wrapper type.
   `A`); that licence does not extend to anything else.
 - Comment invariants, contracts and non-obvious numerical choices only. Never
   narrate the code.
-- Backend-agnostic code: go through `src_method.utils._backend` instead of
-  importing `numpy` or `cupy` directly in the algorithms, so CPU and GPU paths
-  stay in sync.
+- Backend-agnostic code: use the array API standard, with the namespace taken
+  from the arrays (`array_api_compat.array_namespace`) or from the
+  `DeviceGroup`, never an `xp` parameter; annotate arrays as
+  `src_method.utils.Array`. What the standard lacks (streams, pinned memory,
+  CuPy's pool) goes through `src_method.utils._backend`. Never import `cupy`
+  in the algorithms, so CPU and GPU paths stay in sync.
 - Log with the standard library: `logger = logging.getLogger(__name__)` per
   module, lazy `%`-style arguments, never `print`. Never add handlers, levels
   or `basicConfig` in the package; progress messages go at `DEBUG`.

README.md

@@ -2,6 +2,8 @@
 
 [![Documentation](https://github.com/Algorithmiq/src-method/actions/workflows/docpages.yml/badge.svg)](https://algorithmiq.github.io/src-method/)
 [![Test src_method](https://github.com/Algorithmiq/src-method/actions/workflows/test.yml/badge.svg)](https://github.com/Algorithmiq/src-method/actions/workflows/test.yml)
+[![Coverage](https://sonarcloud.io/api/project_badges/measure?project=Algorithmiq_src-method&metric=coverage&token=ab544eca4948ef2f80dffa566c31a9c00026a825)](https://sonarcloud.io/summary/new_code?id=Algorithmiq_src-method)
+[![Quality gate status](https://sonarcloud.io/api/project_badges/measure?project=Algorithmiq_src-method&metric=alert_status&token=ab544eca4948ef2f80dffa566c31a9c00026a825)](https://sonarcloud.io/summary/new_code?id=Algorithmiq_src-method)
 [![License](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
 
 An implementation of the SRC algorithm introduced in [Camaño, Epperly and Tropp, *Quantum* **10**, 2022 (2026)](https://doi.org/10.22331/q-2026-03-10-2022), extending the idea to other kinds of tensor networks.
@@ -109,21 +111,6 @@ On a node with several GPUs, `Resources(devices=4)` splits the sweep among four
 
 See [Large problems](https://algorithmiq.github.io/src-method/features/large-problems) for the budgets, the scratch directory, several GPUs and how to read the logged plan.
 
-### Logging
-
-`src_method` logs through the standard library `logging` module, under the
-`src_method` logger, and stays silent by default: it adds only a `NullHandler`
-and never touches the root logger or its handlers. Progress and timing messages
-are emitted at `DEBUG`, the small-network fallback at `WARNING`. To see them,
-configure logging in your application:
-
-```python
-import logging
-
-logging.basicConfig(level=logging.INFO)
-logging.getLogger("src_method").setLevel(logging.DEBUG)
-```
-
 ## Installation
 
 ```bash
@@ -150,45 +137,9 @@ dependencies = ["src_method[gpu-nvidia]"]
 dependencies = ["src_method[gpu-rocm]"]
 ```
 
-## Setting up the development environment
-
-The code has a [DevContainer] configuration that will get you up and running
-with all dependencies installed and configured, including sane defaults for the
-editor.
-
-You will need:
-
-1. A working [Docker] installation:
-   - For macOS and Windows, install [Docker Desktop](https://docs.docker.com/get-docker/)
-   - For Linux, install [Docker Engine](https://docs.docker.com/engine/install/#server) following the instructions for your specific distro.
-2. The [Visual Studio Code] editor. A recent version is recommended, _e.g._ >=1.78
-3. The VSCode [DevContainers extension](https://marketplace.visualstudio.com/items?itemName=ms-vscode-remote.remote-containers).
-4. The [GitHub CLI] tool.
-
-You can clone the repository with:
-
-```
-gh repo clone Algorithmiq/src-method
-```
-
-We recommend using a Git credential manager, such as [GitHub CLI], configured to
-use HTTPS as protocol for Git operations.
-
-Once the code is locally available, you can open its containing folder in
-[Visual Studio Code]. The editor will then set up the [DevContainer] for you.
-The first time you open the folder the startup will take a few minutes. Once the
-process is done, you will have _all_ project dependencies installed, including
-the git hooks.
-[Visual Studio Code] will be already configured with all the extensions helpful for Python development.
-
-> [!TIP]
-> The order in which Visual Studio Code loads the extensions in the
-> DevContainer is non-deterministic.  You might have to execute the *Reload Window*
-> command to get everything to work as expected after a fresh build of the container.
-
-### Alternative: Nix flake
+## Nix flake
 
-If you prefer [Nix] over Docker, the repository ships a `flake.nix` that provides
+The repository ships a `flake.nix` that provides
 a development shell with [uv], [Git] and the [GitHub CLI], plus the native
 libraries the binary wheels need at runtime. Python itself and all project
 dependencies remain managed by `uv`.
@@ -219,28 +170,9 @@ The documentation is a [Fumadocs] site under `docs/`, published [at this link].
 It covers concepts, features, tutorials, the API reference generated from the
 docstrings, and the developer guide.
 
-To preview it locally you need [Node.js] 22 or later:
-
-```bash
-uv sync --group docs
-cd docs
-npm ci
-uv run python scripts/gen_api_dump.py src_method -d .
-node scripts/generate-api.mjs
-uv run python scripts/notebooks_to_mdx.py
-npm run dev
-```
-
-and open http://localhost:3000. See [`docs/README.md`](docs/README.md) for details.
-
-[DevContainer]: https://containers.dev/
-[Docker]: https://docs.docker.com/get-docker/
-[Visual Studio Code]: https://code.visualstudio.com/
 [GitHub CLI]: https://cli.github.com/
 [Git]: https://git-scm.com/
-[Nix]: https://nixos.org/download/
 [direnv]: https://direnv.net/
 [uv]: https://docs.astral.sh/uv/
 [Fumadocs]: https://fumadocs.dev
-[Node.js]: https://nodejs.org
 [at this link]: https://algorithmiq.github.io/src-method/

docs/content/docs/contributing/dependencies.mdx

@@ -7,8 +7,8 @@ Dependencies are managed with [uv](https://docs.astral.sh/uv/) and declared in
 `pyproject.toml`:
 
 - **Runtime dependencies** (`[project] dependencies`) are what users install:
-  NumPy and `opt_einsum`. Keep this list short; every entry is imposed on every
-  downstream project.
+  NumPy, `opt_einsum` and `array-api-compat`. Keep this list short; every entry is
+  imposed on every downstream project.
 - **Extras** (`[project.optional-dependencies]`) enable optional features in
   production: `gpu-nvidia` and `gpu-rocm` add CuPy.
 - **Dependency groups** (`[dependency-groups]`) are for development only and are

docs/content/docs/contributing/index.mdx

@@ -74,5 +74,9 @@ not just the shapes.
 The project targets Python 3.11+, is checked with `ruff` under a strict rule set
 and type-checked with [`ty`](https://github.com/astral-sh/ty). Public functions
 carry type hints and Google-style docstrings without types in the argument list.
-The algorithms go through `src_method.utils._backend` rather than importing NumPy
-or CuPy directly, so the CPU and GPU paths stay in sync.
+The algorithms speak the [array API standard](https://data-apis.org/array-api/)
+through `array-api-compat`: they take the namespace from their arrays
+(`array_namespace(x)`) instead of importing NumPy or CuPy, and annotate them as
+`src_method.utils.Array`. What the standard leaves out (streams, page-locked
+memory, CuPy's memory pool, host-side random draws) lives in
+`src_method.utils._backend`, so the CPU and GPU paths stay in sync.

docs/content/docs/getting-started.mdx

@@ -16,8 +16,8 @@ uv pip install "src_method[gpu-nvidia]"
 uv pip install "src_method[gpu-rocm]"
 ```
 
-The runtime dependencies are NumPy and `opt_einsum`. The GPU extras add CuPy;
-see [GPU execution](/features/gpu).
+The runtime dependencies are NumPy, `opt_einsum` and `array-api-compat`. The GPU
+extras add CuPy; see [GPU execution](/features/gpu).
 
 ## The three entry points
 

pyproject.toml

@@ -32,7 +32,7 @@ classifiers = [
   "Typing :: Typed",
 ]
 dynamic = ["version"]
-dependencies = ["numpy>=2.0,<3", "opt_einsum>=3.4.0"]
+dependencies = ["array-api-compat>=1.12", "numpy>=2.0,<3", "opt_einsum>=3.4.0"]
 
 
 [project.optional-dependencies]

src/src_method/_group.py

@@ -19,13 +19,21 @@
 from concurrent.futures import ThreadPoolExecutor
 from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
 
-from .utils import copy_into, current_stream, enable_peer_access, use_device
+from ._plan import device_ids
+from .utils import (
+    copy_into,
+    current_device,
+    current_stream,
+    device_count,
+    enable_peer_access,
+    get_xp,
+    use_device,
+)
 
 if TYPE_CHECKING:
     from collections.abc import Callable, Sequence
-    from types import ModuleType
 
-    from numpy.typing import NDArray
+    from .utils import Array
 
 T = TypeVar("T")
 
@@ -60,17 +68,26 @@ class DeviceGroup:
     """The devices of one sweep and the collectives between them.
 
     Args:
-        xp: Array module (``numpy`` or ``cupy``).
-        devices: The device ids, in rank order; rank 0 is the root.
+        device: ``"cpu"`` or ``"gpu"``, as in `src_method.stack.src`.
+        devices: A count ``n`` for the first ``n`` visible devices, or the device
+            ids in rank order; rank 0 is the root. Defaults to the current device.
+
+    Raises:
+        ValueError: If a device does not exist.
     """
 
-    def __init__(self, xp: ModuleType, devices: Sequence[int]) -> None:
-        self.xp = xp
-        self.devices = tuple(devices)
+    def __init__(self, device: str, devices: int | Sequence[int] | None = None) -> None:
+        self.xp = get_xp(device)
+        ids = device_ids(devices) or (current_device(self.xp),)
+        count = device_count(self.xp)
+        if count is not None and max(ids) >= count:
+            msg = f"Requested devices {list(ids)}, but only {count} are visible."
+            raise ValueError(msg)
+        self.devices = ids
         self.size = len(self.devices)
         # Uploading 1/size of each core and gathering the rest pays off only when
         # the devices copy directly between them.
-        self.split_uploads = self.size > 1 and enable_peer_access(xp, self.devices)
+        self.split_uploads = self.size > 1 and enable_peer_access(self.xp, ids)
         self._barrier = threading.Barrier(self.size)
         # Two sets of slots, alternating between consecutive collectives: a rank may
         # publish into the next collective while slower ranks still read this one.
@@ -171,12 +188,12 @@ def _settle(self, rank: int, parity: int) -> None:
             stream.wait_event(event)
         self._parts[parity][rank] = None
 
-    def _pull(self, part: _Part, dst: NDArray) -> None:
+    def _pull(self, part: _Part, dst: Array) -> None:
         stream = current_stream(self.xp)
         stream.wait_event(part.ready)
-        copy_into(dst, part.array, self.xp, stream)
+        copy_into(dst, part.array, stream)
 
-    def allgather(self, rank: int, local: NDArray) -> NDArray:
+    def allgather(self, rank: int, local: Array) -> Array:
         """Concatenate the blocks of every rank along the first axis, on every rank.
 
         Args:
@@ -199,7 +216,7 @@ def allgather(self, rank: int, local: NDArray) -> NDArray:
         self._settle(rank, parity)
         return full
 
-    def gather(self, rank: int, local: NDArray, root: int = 0) -> NDArray | None:
+    def gather(self, rank: int, local: Array, root: int = 0) -> Array | None:
         """Concatenate the blocks of every rank along the first axis, on ``root``.
 
         Args:
@@ -224,7 +241,7 @@ def gather(self, rank: int, local: NDArray, root: int = 0) -> NDArray | None:
         self._settle(rank, parity)
         return full
 
-    def scatter(self, rank: int, full: NDArray | None, root: int = 0) -> NDArray:
+    def scatter(self, rank: int, full: Array | None, root: int = 0) -> Array:
         """Give every rank its block (see `blocks`) of the first axis of ``full``.
 
         Args:
@@ -246,7 +263,7 @@ def scatter(self, rank: int, full: NDArray | None, root: int = 0) -> NDArray:
         self._settle(rank, parity)
         return local
 
-    def _empty_like_parts(self, parts: Sequence[_Part]) -> NDArray:
+    def _empty_like_parts(self, parts: Sequence[_Part]) -> Array:
         first = parts[0].array
         rows = sum(part.array.shape[0] for part in parts)
         return self.xp.empty((rows, *first.shape[1:]), dtype=first.dtype)

src/src_method/_kernels.py

@@ -20,9 +20,10 @@
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
-    from numpy.typing import NDArray
     from opt_einsum.contract import ContractExpression
 
+    from .utils import Array
+
 Shape = tuple[int, ...]
 
 
@@ -71,7 +72,7 @@ class _Contractions:
     def __init__(self) -> None:
         self._compiled: dict[tuple[str, tuple[Shape, ...]], ContractExpression] = {}
 
-    def __call__(self, eq: str, *operands: NDArray) -> NDArray:
+    def __call__(self, eq: str, *operands: Array) -> Array:
         shapes = tuple(op.shape for op in operands)
         expr = self._compiled.get((eq, shapes))
         if expr is None:
@@ -90,7 +91,7 @@ def __init__(self, depth: int) -> None:
         self.eqs = equations(depth)
         self._contract = _Contractions()
 
-    def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray:
+    def env(self, env: Array, omega: Array, cores: Sequence[Array]) -> Array:
         """Advance a batch of sketch columns of the environment by one site.
 
         Args:
@@ -103,7 +104,7 @@ def env(self, env: NDArray, omega: NDArray, cores: Sequence[NDArray]) -> NDArray
         """
         return self._contract(self.eqs.ltr, env, omega, *cores)
 
-    def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray:
+    def sketch(self, env: Array, cores: Sequence[Array], proj: Array) -> Array:
         """Sketch a batch of columns of the site's running core.
 
         Args:
@@ -116,7 +117,7 @@ def sketch(self, env: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArr
         """
         return self._contract(self.eqs.rtl_m, env, *cores, proj)
 
-    def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDArray:
+    def project(self, eta: Array, cores: Sequence[Array], proj: Array) -> Array:
         """Project a batch of rows of the new projected environment.
 
         Args:
@@ -129,7 +130,7 @@ def project(self, eta: NDArray, cores: Sequence[NDArray], proj: NDArray) -> NDAr
         """
         return self._contract(self.eqs.rtl_s, eta, *cores, proj)
 
-    def first(self, cores: Sequence[NDArray], proj: NDArray) -> NDArray:
+    def first(self, cores: Sequence[Array], proj: Array) -> Array:
         """Contract the first site with a batch of rows of ``S``.
 
         Args: