sparse-ot 0.1.0__tar.gz

sparse_ot-0.1.0/.clangd

@@ -0,0 +1,6 @@
+CompileFlags:
+  Add:
+    - "-I/Users/jonatanbobrutsky-haim/Documents/Code/sparse-optimal-transport/.venv/lib/python3.12/site-packages/pybind11/include"
+    - "-I/Users/jonatanbobrutsky-haim/.local/share/uv/python/cpython-3.12.8-macos-aarch64-none/include/python3.12"
+    - "-Isrc/cpp"
+    - "-std=c++17"

sparse_ot-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,63 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+concurrency:
+  group: ci-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test:
+    name: ${{ matrix.os }} / py${{ matrix.python-version }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+
+      - name: Install build deps (Linux)
+        if: runner.os == 'Linux'
+        run: sudo apt-get update && sudo apt-get install -y build-essential cmake
+
+      - name: Install build deps (macOS)
+        if: runner.os == 'macOS'
+        run: brew install cmake
+
+      - name: Install package + test extras
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+
+      - name: Run tests (skip slow)
+        run: pytest tests/ -v -m "not slow" --timeout=300
+
+  # The slow tests (memory smoke test on a 10k×10k problem) run only on Linux
+  # against the latest Python, to keep CI time reasonable.
+  slow-tests:
+    name: slow tests (ubuntu / py3.12)
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+          cache: pip
+      - run: sudo apt-get update && sudo apt-get install -y build-essential cmake
+      - run: |
+          python -m pip install --upgrade pip
+          pip install -e ".[dev]"
+      - run: pytest tests/ -v -m "slow" --timeout=600

sparse_ot-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,107 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+  # Manual trigger for dry runs against TestPyPI.
+  workflow_dispatch:
+    inputs:
+      target:
+        description: "Publish target"
+        required: true
+        default: "testpypi"
+        type: choice
+        options:
+          - testpypi
+          - pypi
+
+permissions:
+  contents: read
+
+jobs:
+  build_wheels:
+    name: Wheels on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        # macos-13 (Intel) dropped: the GitHub-hosted runner pool is
+        # chronically oversubscribed and stalls publish jobs by 30+ min.
+        # Apple Silicon dominates the macOS Python user base on 3.10+;
+        # users on Intel Macs can build from the sdist.
+        os: [ubuntu-latest, ubuntu-24.04-arm, macos-14]
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: pypa/cibuildwheel@v2.21
+        env:
+          # PyPI-supported Python versions; skip 32-bit Linux, PyPy, musllinux
+          # (scipy wheels not universally available on musl).
+          CIBW_BUILD: "cp310-* cp311-* cp312-* cp313-*"
+          CIBW_SKIP: "*-musllinux_* *-win32 pp*"
+          # macOS targets: arm64 on macos-14, x86_64 on macos-13.
+          CIBW_ARCHS_MACOS: "auto64"
+          CIBW_ARCHS_LINUX: "auto64"
+          # Recent scipy releases publish only manylinux_2_28 wheels for x86_64
+          # and aarch64; the default manylinux2014 (glibc 2.17) container falls
+          # back to source-building scipy, which needs OpenBLAS and fails.
+          CIBW_MANYLINUX_X86_64_IMAGE: manylinux_2_28
+          CIBW_MANYLINUX_AARCH64_IMAGE: manylinux_2_28
+          # Minimal smoke test: import the package and run one tiny emd() call.
+          # The full pytest suite already runs on every supported (os, py) pair
+          # in ci.yml — re-running it inside each wheel container would force
+          # scipy + pot + matplotlib installs we don't need just to verify the
+          # binary loads.
+          CIBW_TEST_REQUIRES: "numpy scipy"
+          CIBW_TEST_COMMAND: >-
+            python -c "import numpy as np, sparse_ot;
+            G = sparse_ot.emd(np.array([0.4, 0.6]), np.array([0.5, 0.5]), np.array([[0.0, 1.0], [1.0, 0.0]]));
+            assert G.shape == (2, 2);
+            print('sparse_ot wheel smoke test ok')"
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: wheels-${{ matrix.os }}
+          path: ./wheelhouse/*.whl
+
+  build_sdist:
+    name: sdist
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+      - run: python -m pip install --upgrade build
+      - run: python -m build --sdist
+      - uses: actions/upload-artifact@v4
+        with:
+          name: sdist
+          path: dist/*.tar.gz
+
+  publish:
+    name: Publish to ${{ (github.event_name == 'workflow_dispatch' && inputs.target) || 'pypi' }}
+    needs: [build_wheels, build_sdist]
+    runs-on: ubuntu-latest
+    # Trusted-publishing OIDC requires an environment + id-token write.
+    environment:
+      name: ${{ (github.event_name == 'workflow_dispatch' && inputs.target) || 'pypi' }}
+      url: https://${{ ((github.event_name == 'workflow_dispatch' && inputs.target) || 'pypi') == 'testpypi' && 'test.' || '' }}pypi.org/p/sparse-ot
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          path: dist
+          merge-multiple: true
+
+      - name: Publish to PyPI
+        if: (github.event_name == 'workflow_dispatch' && inputs.target == 'pypi') || github.event_name == 'release'
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+      - name: Publish to TestPyPI
+        if: github.event_name == 'workflow_dispatch' && inputs.target == 'testpypi'
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://test.pypi.org/legacy/

sparse_ot-0.1.0/.gitignore

@@ -0,0 +1,25 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+dist/
+build/
+
+# Virtual environment
+.venv/
+
+# scikit-build-core / CMake build artifacts
+_skbuild/
+CMakeFiles/
+*.cmake
+Makefile
+
+# Compiled extensions
+*.so
+*.pyd
+
+# pytest
+.pytest_cache/
+
+# macOS
+.DS_Store

sparse_ot-0.1.0/CMakeLists.txt

@@ -0,0 +1,16 @@
+cmake_minimum_required(VERSION 3.18)
+project(sparse_ot_ext LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+find_package(pybind11 CONFIG REQUIRED)
+
+pybind11_add_module(_bonneel src/cpp/bonneel_solver.cpp)
+target_include_directories(_bonneel PRIVATE src/cpp)
+target_compile_options(_bonneel PRIVATE -O3)
+# Disable OpenMP for the Bonneel network-simplex wrapper: OpenMP thread overhead
+# causes 30-40x slowdown on the problem sizes we handle (n ≤ ~5 000).  POT's
+# emd_wrap is also built without OpenMP for the same reason.
+target_compile_definitions(_bonneel PRIVATE NOOMP)
+install(TARGETS _bonneel DESTINATION sparse_ot/_ext)

sparse_ot-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Jonatan Bobrutsky-Haim
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sparse_ot-0.1.0/PKG-INFO

@@ -0,0 +1,17 @@
+Metadata-Version: 2.2
+Name: sparse-ot
+Version: 0.1.0
+Summary: Sparse optimal transport with drop-in POT API
+License: MIT
+Requires-Python: >=3.10
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.10
+Provides-Extra: torch
+Requires-Dist: torch>=2.0; extra == "torch"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: pytest-timeout>=2.0; extra == "dev"
+Requires-Dist: pot>=0.9; extra == "dev"
+Requires-Dist: pybind11>=2.12; extra == "dev"
+Requires-Dist: matplotlib>=3.7; extra == "dev"
+

sparse_ot-0.1.0/README.md

@@ -0,0 +1,256 @@
+# sparse-ot
+
+[![CI](https://github.com/JBobrutsky/sparse-optimal-transport/actions/workflows/ci.yml/badge.svg)](https://github.com/JBobrutsky/sparse-optimal-transport/actions/workflows/ci.yml)
+
+Drop-in replacement for [POT](https://github.com/PythonOT/POT)'s `emd` /
+`emd2`, with native support for **sparse cost matrices**. One solver
+([Bonneel's network simplex](https://github.com/nbonneel/network_simplex))
+covers both regimes:
+
+- **Dense** `numpy.ndarray` cost matrix → dense plan.
+- **`scipy.sparse` CSR** cost matrix → sparse plan, with memory and per-pivot
+  work both scaling in the number of candidate edges `k` rather than `n × m`.
+
+## Quickstart
+
+```python
+import numpy as np, scipy.sparse, sparse_ot as sot
+
+# Dense: identical interface to ot.emd.
+G = sot.emd(a, b, M)
+cost = sot.emd2(a, b, M)
+
+# Sparse: pass a CSR cost matrix. Absent entries are forbidden edges, not
+# zero-cost shortcuts.
+M_csr = scipy.sparse.csr_matrix(...)
+G_csr = sot.emd(a, b, M_csr)
+
+# POT-compatible log dict (cost, u, v, warning, result_code).
+G, info = sot.emd(a, b, M, log=True)
+```
+
+The `u` and `v` returned in the log dict are the dual potentials with
+POT's sign convention (`u[i] + v[j] ≤ M[i,j]` at the optimum). With
+`center_dual=True` (default) `u` is shifted to zero mean, preserving
+`u[i] + v[j]` on every edge.
+
+## Why sparse?
+
+A 10 000 × 10 000 problem with 10 candidate edges per row (k = 100 000):
+
+| Solver path        | Memory        | Wall time |
+|--------------------|---------------|-----------|
+| Bonneel-dense      | ≈ 800 MB (cost matrix) | (does not run; OOM at this scale on small machines) |
+| **Bonneel-sparse** | **≈ 6 MB**    | seconds   |
+
+Most real OT problems (k-NN, transformer attention masks, point-cloud
+matching) are intrinsically sparse. Materialising them as dense costs
+matrices is wasteful and can be infeasible. This package gives you
+Bonneel's tight constants without the O(n·m) memory penalty.
+
+## Feasibility on sparse supports
+
+When you pass a sparse `M`, the transport plan is restricted to the
+edges you provide. The package checks that the support is connected
+and that supply totals match (`check_feasibility`), but **this does
+not guarantee an LP-feasible plan exists**.
+
+A small support can fail [Hall's condition](https://en.wikipedia.org/wiki/Hall%27s_marriage_theorem):
+some local block of rows `S` may collectively need to move more mass
+than the columns they reach can absorb. For example, a band-7 support
+(each row connects only to its 7 nearest columns) cannot route generic
+Dirichlet marginals at `n = 1000` — the corner rows have nowhere to
+shed their excess.
+
+When that happens we don't lie. The solver returns its best-effort
+flow, `info["result_code"] == 0`, and a `RuntimeWarning` fires
+explaining that the marginals weren't met. Compare to POT, which
+silently routes mass through any zero-cost or penalty edge in the
+densified representation and reports `success` with an arbitrary
+cost.
+
+In practice: build supports that are slightly denser than your
+marginals strictly require (k-NN with k chosen by validation, plus a
+small slack), or run with very dense support whenever you don't know
+the marginal distribution ahead of time.
+
+## Convergence and the `numItermax` knob
+
+Bonneel's network simplex stops at `numItermax` pivots without raising.
+If the iteration cap is hit before convergence the returned flow can
+violate marginals by orders of magnitude more than machine epsilon. We
+guard against this in two ways:
+
+1. The default `numItermax` is **problem-size-aware**:
+   `min(50M, max(100k, 100·(n + m + k)))`.
+2. After every solve we re-check the marginals. If `max(|G.sum(1) - a|,
+   |G.sum(0) - b|) > 1e-6`, we emit a `RuntimeWarning` and report
+   `result_code = 0` with a diagnostic in `info["warning"]`. No
+   exception is raised, matching POT's behavior.
+
+You can pass `numItermax=…` to override.
+
+## Build and install
+
+```bash
+pip install -e . --no-build-isolation
+```
+
+`pyproject.toml` sets `editable.rebuild = true`, so the pybind11
+extension is rebuilt automatically the next time `sparse_ot` is imported
+after a `src/cpp/` edit.
+
+## Benchmarks
+
+Two independent suites:
+
+```bash
+python benchmarks/bench_solvers.py --mid     # ~15 minutes
+python benchmarks/bench_solvers.py --quick   # ~30 seconds (used by CI)
+python benchmarks/bench_solvers.py           # full sweep, hours
+```
+
+- **Dense suite** — fully-random `n × n` cost matrix. Compares
+  `bonneel_dense` against `pot_reference` (`ot.emd`).
+- **Sparse suite** — feasible-by-construction kNN-grid (`benchmarks/problems.py`).
+  Runs `bonneel_sparse` only; there is no honest dense representation of
+  a kNN cost (absent edges must be +∞, which neither POT nor Bonneel's
+  dense path can express).
+
+Results are written to `benchmarks/results/{efficiency,accuracy}_{tag}.json`
+with the structure:
+
+```jsonc
+{
+  "dense":  { "<n>": { "bonneel_dense": {...}, "pot_reference": {...} } },
+  "sparse": { "<n>": { "<k>": { "bonneel_sparse": {...} } } }
+}
+```
+
+See "Benchmark results" below for the current numbers on the maintainer's
+laptop.
+
+## Benchmark results
+
+Numbers below are from `python benchmarks/bench_solvers.py --mid` on an
+Apple-Silicon laptop (Sonoma, 64 GB). Wall times are median of 1 run; the
+sparse-suite peak memory is `tracemalloc` peak during `emd()`.
+
+### Dense suite (fully random `n × n` cost)
+
+| n     | `bonneel_dense` | `pot_reference` | dense / pot |
+|------:|----------------:|----------------:|------------:|
+|   200 |          0.003s |          0.002s |       0.66× |
+|   500 |          0.019s |          0.015s |       0.80× |
+|  1000 |          0.085s |          0.072s |       0.85× |
+|  2000 |          0.406s |          0.346s |       0.85× |
+|  4000 |          2.166s |          1.414s |       0.65× |
+
+POT and `bonneel_dense` share the same C++ engine (POT vendors Bonneel's
+network simplex), so the wall-time ratio reflects pure wrapping overhead;
+POT's Cython wrapper is marginally tighter than our pybind11 wrapper.
+Costs agree to machine precision for n ≤ 2000. At n = 4000, POT's cost
+is 1.1 % higher than ours — POT's default `numItermax = 100 000`
+truncates before convergence, while our problem-size-aware default
+finishes the pivots.
+
+### Sparse suite (knn-grid)
+
+The same knn problem is run through both Bonneel paths so the
+speed/memory tradeoff is directly comparable. `bonneel_dense` runs only
+where `n ≤ MAX_DENSE_N`; above that the cost matrix doesn't fit and the
+cell is sparse-only. For the dense path we densify with a finite
+penalty (`max(M.data) · (n·m + 1)`) on absent edges — with the
+problem-size-aware `numItermax` the optimal basis never lands on a
+penalty edge.
+
+| n      | k    | nnz        | sparse wall | sparse peak | dense wall | dense peak | winner       |
+|-------:|-----:|-----------:|------------:|------------:|-----------:|-----------:|:-------------|
+|    200 |    2 |        400 |      0.006s |    0.04 MB |     0.003s |    0.32 MB |  dense 2.4×  |
+|    200 |   32 |      6 400 |      0.051s |    0.14 MB |     0.004s |    0.32 MB |  dense 14×   |
+|    200 |  128 |     25 600 |      0.204s |    0.52 MB |     0.003s |    0.32 MB |  dense 71×   |
+|  1 000 |    2 |      2 000 |      0.033s |    0.15 MB |     0.075s |    7.69 MB |  **sparse 2.3×** |
+|  1 000 |    8 |      8 000 |      0.100s |    0.22 MB |     0.178s |    7.69 MB |  **sparse 1.8×** |
+|  1 000 |   32 |     32 000 |      0.315s |    0.67 MB |     0.182s |    7.69 MB |  dense 1.7×  |
+|  1 000 |  128 |    128 000 |      1.103s |    2.59 MB |     0.256s |    7.69 MB |  dense 4.3×  |
+|  4 000 |    2 |      8 000 |      0.134s |    0.60 MB |     1.652s |  122.3  MB |  **sparse 12×**  |
+|  4 000 |    8 |     32 000 |      0.577s |    0.87 MB |     8.918s |  122.3  MB |  **sparse 15×**  |
+|  4 000 |   32 |    128 000 |      1.572s |    2.66 MB |     8.574s |  122.3  MB |  **sparse 5.5×** |
+|  4 000 |  128 |    512 000 |      4.540s |   10.35 MB |    10.298s |  122.3  MB |  **sparse 2.3×** |
+|  4 000 |  400 |  1 600 000 |     14.473s |   32.14 MB |     8.393s |  122.3  MB |  dense 1.7×  |
+|  4 000 |  512 |  2 048 000 |     17.956s |   41.11 MB |     9.535s |  122.3  MB |  dense 1.9×  |
+| 16 000 |    2 |     32 000 |      0.531s |    2.39 MB |       —    |      —     |  sparse only |
+| 16 000 |   32 |    512 000 |     12.066s |   10.62 MB |       —    |      —     |  sparse only |
+| 16 000 |  128 |  2 048 000 |     32.955s |   41.38 MB |       —    |      —     |  sparse only |
+| 16 000 |  512 |  8 192 000 |     75.201s |  164.4  MB |       —    |      —     |  sparse only |
+| 16 000 | 1600 | 25 600 000 |    269.488s |  513.1  MB |       —    |      —     |  sparse only |
+
+Reading the table:
+
+- At **n = 200** the dense path always wins because the n² cost matrix is
+  tiny and Bonneel's flat-array constants dominate over the sparse
+  digraph's per-arc indirection.
+- At **n = 1 000** the crossover is around k ≈ n / 50: below that, sparse
+  wins; above, dense wins.
+- At **n = 4 000** sparse wins by 2–15× up to k ≈ n / 20. Above that
+  density, dense again wins on time but its memory cost is fixed at
+  122 MB regardless of k.
+- At **n = 16 000** the dense path is out of reach (cost matrix ≈ 2 GB);
+  only sparse runs.
+
+### Accuracy
+
+Marginals stay at machine precision (worst case `2.5 × 10⁻¹⁶`) across
+every cell of both suites. On the knn problems, `bonneel_sparse` and
+`bonneel_dense` agree on cost to ≈ machine precision in 22 of 24 cells.
+Two exceptions:
+
+| n     | k    | sparse cost | dense cost | relative diff |
+|------:|-----:|------------:|-----------:|--------------:|
+| 4 000 |  400 |  271.7297   | 272.2578   |    1.9 × 10⁻³ |
+| 4 000 |  512 |  557.3260   | 557.3270   |    1.7 × 10⁻⁶ |
+
+In both cases `bonneel_sparse` finds a strictly lower-cost plan. The
+densified-with-penalty input introduces costs on the order of
+`max(M) · n² ≈ 10¹²`, and floating-point reduced-cost computations on
+that scale accumulate enough rounding noise to push the pivot rule off
+the true optimum. The sparse path never sees those large numbers and is
+the more accurate of the two when both can run.
+
+## Memory cutoffs
+
+`bench_solvers.py` skips cells beyond these defaults (16 GB target):
+
+| Constant         | Default       | Effect                                     |
+|------------------|---------------|--------------------------------------------|
+| `MAX_DENSE_N`    | 8 192         | Dense suite skipped above this             |
+| `MAX_SPARSE_NNZ` | 200 000 000   | Sparse cell skipped above this nnz         |
+
+Raise the constants for larger hardware.
+
+## Releasing
+
+PyPI uploads are automated via GitHub Actions and PyPI's
+[trusted-publishing OIDC](https://docs.pypi.org/trusted-publishers/). To
+cut a release:
+
+1. Bump `project.version` in `pyproject.toml`, commit, tag (`git tag vX.Y.Z`),
+   push (`git push --tags`).
+2. Create a GitHub Release pointing at the tag.
+
+The `.github/workflows/publish.yml` workflow then builds wheels via
+`cibuildwheel` for Linux (x86_64, arm64) and macOS (x86_64, arm64) across
+Python 3.10–3.13, builds an sdist, and uploads everything to PyPI.
+
+First-time setup (one-time, requires owner action on pypi.org):
+
+- Add a trusted publisher for **sparse-ot** with owner = `JBobrutsky`,
+  repository = `sparse-optimal-transport`, workflow = `publish.yml`,
+  environment = `pypi`.
+- For TestPyPI dry runs, register the same on test.pypi.org with
+  environment = `testpypi`. Then trigger `Publish to PyPI` via the
+  Actions UI (workflow_dispatch) with target = `testpypi`.
+
+## License
+
+[MIT](LICENSE).