Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions AGENTS.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down
76 changes: 4 additions & 72 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
Expand All @@ -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`.
Expand Down Expand Up @@ -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/
4 changes: 2 additions & 2 deletions docs/content/docs/contributing/dependencies.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 6 additions & 2 deletions docs/content/docs/contributing/index.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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.
4 changes: 2 additions & 2 deletions docs/content/docs/getting-started.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down
47 changes: 32 additions & 15 deletions src/src_method/_group.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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)
13 changes: 7 additions & 6 deletions src/src_method/_kernels.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, ...]


Expand Down Expand Up @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand All @@ -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:
Expand Down
Loading
Loading