GPU-accelerated SEACells (opt-in, end-to-end on GPU)

.gitignore

@@ -1,8 +1,11 @@
 build
 dist
-pyproject.toml
 SEACells.egg-info
 **/__pycache__
 **/.ipynb_checkpoints/
 **/.DS_Store
 **/.idea
+data/
+docs/
+s41587-023-01716-9.pdf
+uv.lock

.python-version

@@ -0,0 +1 @@
+3.11

README.md

@@ -2,60 +2,72 @@
 
 **S**ingle-c**E**ll **A**ggregation for High Resolution **Cell S**tates
 
-#### Installation and dependencies
+SEACells identifies **metacells**: groups of cells in the same biological state, found by
+archetypal analysis on a nearest-neighbor kernel built from a low-dimensional embedding
+(`X_pca` for scRNA-seq, `X_svd` for scATAC-seq). This denoises the data while preserving
+heterogeneity, giving a high-resolution set of cell states for downstream analysis. See the
+[paper](https://www.nature.com/articles/s41587-023-01716-9) for the method.
 
-1.  SEACells has been implemented in Python3.8 can be installed via pip:
-    $> pip install cmake
-    $> pip install SEACells
-    It can also be installed directly from source.
+#### Installation
 
-        $> git clone https://github.com/dpeerlab/SEACells.git
-        $> cd SEACells
-        $> python setup.py install
-
-2.  If you are using `conda`, you can use the `environment.yaml` to create a new environment and install SEACells.
+Uses [**uv**](https://docs.astral.sh/uv/). Install uv once with
+`curl -LsSf https://astral.sh/uv/install.sh | sh`, then:
 
 ```
-conda env create -n seacells --file environment.yaml
-conda activate seacells
+git clone https://github.com/dpeerlab/SEACells.git && cd SEACells
+uv sync                 # CPU — works anywhere
+uv sync --extra gpu     # + RAPIDS/CuPy/FAISS (NVIDIA GPU, CUDA 13, Linux/x86_64)
+uv sync --extra dev     # + linting/pre-commit hooks
 ```
 
-3. You can also use `pip` to install the requirements
-
-```
-pip install -r requirements.txt
-```
+This builds a `.venv` with SEACells installed editable. Run code with
+`uv run python ...` or `source .venv/bin/activate`. The GPU wheels come from the
+NVIDIA pip index (preconfigured in `pyproject.toml`); validated on A100 80GB.
 
-And then follow step (1)
+#### Running SEACells (CPU & GPU)
 
-4. MulticoreTSNE issues can be solved using
+The core API is unchanged. A minimal run:
 
-```
-conda create --name seacells -c conda-forge -c bioconda cython python=3.8
-conda activate seacells
-pip install git+https://github.com/settylab/Palantir@removeTSNE
-git clone https://github.com/dpeerlab/SEACells.git
-cd SEACells
-python setup.py install
-```
+```python
+import SEACells
 
-4. SEACells depends on a number of `python3` packages available on pypi and these dependencies are listed in `setup.py`.
-
-    All the dependencies will be automatically installed using the above commands
-
-5. To uninstall:
-   $> pip uninstall SEACells
-
-6. To install the developer installation of SEACells, run
+# ad: AnnData with a low-dim embedding in ad.obsm ('X_pca' for RNA, 'X_svd' for ATAC)
+model = SEACells.core.SEACells(
+    ad,
+    build_kernel_on='X_pca',   # 'X_svd' for scATAC
+    n_SEACells=90,             # number of metacells (heuristic: ~1 per 75 cells)
+)
+model.construct_kernel_matrix()
+model.fit(min_iter=10, max_iter=100)   # converges in ~15-50 iterations
 
+# metacell assignments are written to ad.obs['SEACell']; aggregate raw counts:
+meta_ad = SEACells.core.summarize_by_SEACell(ad, SEACells_label='SEACell', summarize_layer='raw')
 ```
-git clone https://github.com/dpeerlab/SEACells.git
-cd SEACells.git
 
-pip install -e ".[dev]"
-pre-commit install
+**GPU acceleration (optimized).** Pass `use_gpu=True, use_unified=True` to run the
+end-to-end GPU implementation (`SEACells.model.SEACellsModel`) — everything else is
+identical:
+
+```python
+model = SEACells.core.SEACells(
+    ad, build_kernel_on='X_pca', n_SEACells=90,
+    use_gpu=True,        # run on GPU (needs cupy + cuML / RAPIDS)
+    use_unified=True,    # use the optimized unified backend
+)
+model.construct_kernel_matrix()
+model.fit(min_iter=10, max_iter=100)
 ```
 
+It keeps the kernel and weight matrices resident on the GPU, uses exact GPU kNN
+(cuML) and a memory-scalable reconstruction error, and scales to ~100k cells in a few GB
+(a full 100k-cell fit runs in ~10 min on one A100; see
+[`docs/gpu_speed_and_scale.md`](docs/gpu_speed_and_scale.md)).
+`use_unified=True` also works with `use_gpu=False` (an optimized, single-source CPU path).
+
+**Backward compatible.** `use_unified` defaults to `False`, so existing code is unchanged:
+the default CPU path (`use_gpu=False`) and the legacy `use_gpu=True` / `use_sparse=True`
+backends all behave exactly as before. `use_unified` is strictly opt-in.
+
 #### Usage
 
 1. <b>ATAC preprocessing</b>:

SEACells/build_graph.py

@@ -5,7 +5,7 @@
 import numpy as np
 from joblib import Parallel, delayed
 from scipy.sparse import lil_matrix
-from tqdm.notebook import tqdm
+from tqdm.auto import tqdm
 
 # get number of cores for multiprocessing
 NUM_CORES = cpu_count()

SEACells/core.py

@@ -22,6 +22,7 @@ def SEACells(
     l2_penalty: float = 0,
     max_franke_wolfe_iters: int = 50,
     use_sparse: bool = False,
+    use_unified: bool = False,
 ):
     """Core SEACells class.
 
@@ -37,9 +38,34 @@ def SEACells(
     :param l2_penalty: (float) L2 penalty for Franke-Wolfe algorithm
     :param max_franke_wolfe_iters: (int) maximum number of iterations for Franke-Wolfe algorithm
     :param use_sparse: (bool) whether to use sparse matrix operations. Currently only supported for CPU implementation.
+    :param use_unified: (bool) use the optimized unified CPU/GPU implementation (``model.SEACellsModel``)
+                        instead of the legacy ``cpu``/``cpu_dense``/``gpu`` backends. Honors ``use_gpu``.
+                        Recommended for GPU runs and large datasets; the legacy backends are retained
+                        for backward compatibility. Not compatible with ``use_sparse``.
 
-    See cpu.py or gpu.py for descriptions of model attributes and methods.
+    See model.py (unified) or cpu.py/gpu.py (legacy) for descriptions of model attributes and methods.
     """
+    if use_unified:
+        assert (
+            not use_sparse
+        ), "use_sparse is a legacy CPU-only option; the unified backend does not support it."
+        try:
+            from . import model as _model
+        except ImportError:
+            import model as _model
+        return _model.SEACellsModel(
+            ad,
+            build_kernel_on,
+            n_SEACells,
+            use_gpu=use_gpu,
+            verbose=verbose,
+            n_waypoint_eigs=n_waypoint_eigs,
+            n_neighbors=n_neighbors,
+            convergence_epsilon=convergence_epsilon,
+            l2_penalty=l2_penalty,
+            max_franke_wolfe_iters=max_franke_wolfe_iters,
+        )
+
     if use_sparse:
         assert (
             not use_gpu

SEACells/gpu.py

@@ -423,7 +423,7 @@ def _updateA(self, B, A_prev):
 
         Ag = cp.array(A)
         Bg = cp.array(B)
-        Kg = cupyx.scipy.sparse.csc_matrix(self.K)
+        Kg = cupyx.scipy.sparse.csc_matrix(self.K)  # self.K sits on CPU, so it is re-uploaded to the GPU every call. The upload itself is cheap (<1% of an iteration); the real cost of keeping K on the host is that compute_RSS then builds the n x n reconstruction on the CPU (see model.py for the resident, reduced-form version)
 
         # precompute some gradient terms
         t2g = Kg.dot(Bg).T