[main] Integration/peerlab (#38)

Major:

CUDA/Xenium install #21: CUDA 12.1 support; Xenium v1 + binary column datatype support; pixi/conda install updates
Reference gene correlation #32: Add custom reference gene correlation support via `--gene-corr-reference-path`
Memory optimisations #36: GPU/memory optimisations for large datasets; chunked subgraphing fallback
Large dataset fixes #44: Fix torch array overflow and related scaling issues
Tiling fixes #42: Fix tiling/bincount errors and transcript thresholding

Minor:

Logging improvements #31: Improve debug logging; support `SEGGER_LOG_LEVEL`
Margin handling #34: Warn instead of fail when tiling margins are too large
Misc cleanup: Refactors; assertions/debugging; remove unused code and comments

---------

Co-authored-by: Nandula
Co-authored-by: Kalfus

Author: Tobiaspk

.gitignore

@@ -205,4 +205,5 @@ __marimo__/
 # Custom
 .dev
 .dev/*
-*.pyc
+*.pyc
+*memray*

README.md

@@ -1,8 +1,30 @@
 # Installation
 
-## pip
+We recommend CUDA 12.1 with `cu*` packages version ≥24.2 and <26.0. Ensure your CUDA driver version matches or exceeds your toolkit version (≥12.1 for CUDA 12.1). 
+Adjust package versions in the environment files below if your system requires a different package versions.
 
-Before installing **segger**, please install GPU-accelerated versions of PyTorch, RAPIDS, and related packages compatible with your system. *Please ensure all CUDA-enabled packages are compiled for the same CUDA version.*
+## Clone the repository
+```bash
+git clone https://github.com/dpeerlab/segger.git segger && cd segger
+```
+
+## Using `conda`
+```bash
+conda env create -n segger -f environment_cuda121.yml
+```
+
+Adjust `environment_cuda121.yml` for other CUDA versions (e.g., `environment_cuda118.yml` for CUDA 11.8).
+
+## Using `pixi`
+```bash
+pixi install -e cuda121
+```
+
+Adjust the environment name in `pixi.toml` as needed for other CUDA versions.
+
+## `pip`
+
+Install GPU-accelerated PyTorch and RAPIDS compatible with your CUDA version before installing **segger**. All CUDA-enabled packages must be compiled for the same CUDA version.
 
 - **PyTorch & torchvision:** [Installation guide](https://pytorch.org/get-started/locally/)
 - **torch_scatter:** [Installation guide](https://github.com/rusty1s/pytorch_scatter#installation)
@@ -28,7 +50,6 @@ pip install cupy-cuda12x
 
 ```bash
 # Clone segger repo and install locally
-git clone https://github.com/dpeerlab/segger.git segger && cd segger
 pip install -e .
 ```
 

environment_cuda121.yml

@@ -0,0 +1,29 @@
+channels:
+  - pytorch
+  - nvidia
+  - rapidsai
+  - conda-forge
+
+dependencies:
+  - python>=3.11,<3.12
+
+  # RAPIDS 24.10 — last release compatible with CUDA 12.1 runtime
+  - cuspatial=24.10
+  - cudf=24.10
+  - cuml=24.10
+  - cugraph=24.10
+  - cupy
+
+  - pip
+  - pip:
+      - --extra-index-url https://pypi.nvidia.com
+      - --extra-index-url https://download.pytorch.org/whl/cu121
+      - --find-links https://data.pyg.org/whl/torch-2.5.0+cu121.html
+      - torch==2.5.*
+      - torchvision==0.20.*
+      - lightning
+      - torch-geometric
+      - torch-scatter
+      - torch-sparse
+      - cupy-cuda12x>=12.2,<13.0
+      - --editable .

pixi.lock

@@ -0,0 +1,45 @@
+[workspace]
+channels = ["pytorch", "nvidia", "rapidsai", "conda-forge"]
+platforms = ["linux-64"]
+
+[dependencies]
+python = ">=3.11"
+
+[pypi-dependencies]
+segger = { path = ".", editable = true }
+
+[feature.cuda121.dependencies]
+python = ">=3.11,<3.12"
+
+[feature.cuda121.system-requirements]
+cuda = "12"
+
+[feature.cuda121.pypi-options]
+index-strategy = "unsafe-best-match"
+index-url = "https://pypi.org/simple"
+extra-index-urls = [
+    "https://pypi.nvidia.com",
+    "https://download.pytorch.org/whl/cu121"
+]
+find-links = [
+    { url = "https://data.pyg.org/whl/torch-2.5.0+cu121.html" }
+]
+
+[feature.cuda121.pypi-dependencies]
+torch = "==2.5.*"
+torchvision = "==0.20.*"
+lightning = "*"
+torch_geometric = "*"
+torch-scatter = "*"
+torch-sparse = "*"
+cuspatial-cu12 = "==24.10.*"
+cudf-cu12 = "==24.10.*"
+cuml-cu12 = "==24.10.*"
+cugraph-cu12 = "==24.10.*"
+cupy-cuda12x = ">=12.2,<13.0"
+
+[tool.uv.extra-build-dependencies]
+torch-scatter = ["torch"]
+
+[environments]
+cuda121 = ["cuda121"]

pixi.toml

@@ -15,7 +15,6 @@ dependencies = [
     "anndata",
     "cyclopts",
     "geopandas",
-    "lightning",
     "numba",
     "numpy",
     "opencv-python",
@@ -27,8 +26,7 @@ dependencies = [
     "shapely",
     "scikit-image",
     "scikit-learn",
-    "tifffile",
-    "torch_geometric",
+    "tifffile"
 ]
 
 [build-system]

pyproject.toml

@@ -0,0 +1,30 @@
+from pathlib import Path
+import cupy as cp
+import torch
+import rmm
+from rmm.allocators.cupy import rmm_cupy_allocator
+from rmm.allocators.torch import rmm_torch_allocator
+from rmm.statistics import enable_statistics, get_statistics
+
+# Single RMM pool shared by CuPy/cuDF/cuSpatial AND PyTorch. Must be set before
+# any CUDA tensor is created.
+rmm.reinitialize(pool_allocator=True, managed_memory=True)
+cp.cuda.set_allocator(rmm_cupy_allocator)
+torch.cuda.memory.change_current_allocator(rmm_torch_allocator)
+enable_statistics()
+
+# Apply pytorch patches for issue pytorch/pytorch#51871 (CUDA nonzero INT_MAX limit).
+# Must run BEFORE any segger module imports HeteroData / bipartite_subgraph.
+from ._patches import apply as _apply_patches
+_apply_patches()
+
+def free_mem_str() -> str:
+    stats = get_statistics()
+    return (
+        f"GPU: {stats.current_bytes / 1e9:.2f} GB "
+        f"(peak {stats.peak_bytes / 1e9:.2f} GB)"
+    )
+
+
+def print_free_mem():
+    print(free_mem_str())

src/segger/__init__.py

@@ -0,0 +1,74 @@
+"""Workaround for pytorch/pytorch#51871 (CUDA `nonzero` INT_MAX limit).
+
+Patches `torch_geometric.utils.bipartite_subgraph` and the references already
+imported by `torch_geometric.data.hetero_data` / `._subgraph` so that
+`HeteroData.subgraph` falls back to a chunked-nonzero path when the edge
+tensor on CUDA exceeds INT_MAX (~2.15B) elements.
+
+See: https://github.com/dpeerlab/segger/issues/44
+"""
+import torch
+import torch_geometric.utils._subgraph as _sg
+import torch_geometric.utils as _tgu
+import torch_geometric.data.hetero_data as _hd
+from torch_geometric.utils import index_to_mask
+from torch_geometric.utils.map import map_index
+
+_INT_MAX = 2**31 - 1
+_pyg_bipartite = _sg.bipartite_subgraph
+
+
+def chunked_nonzero(mask: torch.Tensor, chunk: int = 2**30) -> torch.Tensor:
+    """Chunked version of `mask.nonzero()` that works on CUDA tensors with > INT_MAX elements."""
+    if mask.numel() <= _INT_MAX or mask.device.type != "cuda":
+        return mask.nonzero(as_tuple=False).flatten()
+    parts = []
+    for i, m in enumerate(mask.split(chunk)):
+        idx = m.nonzero(as_tuple=False).flatten()
+        if idx.numel():
+            parts.append(idx + i * chunk)
+    return torch.cat(parts)
+
+
+def bipartite_safe(subset, edge_index, edge_attr=None, relabel_nodes=False,
+                   size=None, return_edge_mask=False):
+    """
+    Replacement for `torch_geometric.utils.bipartite_subgraph`.
+    Falls back to a chunked subgraph version when the edge_index is too large for CUDA.
+    """
+    # original
+    if edge_index.numel() <= _INT_MAX or edge_index.device.type != "cuda":
+        return _pyg_bipartite(subset, edge_index, edge_attr, relabel_nodes,
+                              size, return_edge_mask)
+
+    # same as source
+    src_sub, dst_sub = subset
+    src_mask = index_to_mask(src_sub, size=size[0])
+    dst_mask = index_to_mask(dst_sub, size=size[1])
+    edge_mask = src_mask[edge_index[0]] & dst_mask[edge_index[1]]
+
+    # replaced this
+    idx = chunked_nonzero(edge_mask)
+
+    # same as source (but indices instead of mask)
+    edge_index = edge_index[:, idx]
+    edge_attr = edge_attr[edge_mask] if edge_attr is not None else None
+    if relabel_nodes:
+        src_index, _ = map_index(edge_index[0], src_sub, max_index=size[0], inclusive=True)
+        dst_index, _ = map_index(edge_index[1], dst_sub, max_index=size[1], inclusive=True)
+        edge_index = torch.stack([src_index, dst_index], dim=0)
+    return (edge_index, edge_attr, edge_mask) if return_edge_mask else (edge_index, edge_attr)
+
+
+_patches_applied = False
+
+
+def apply():
+    """Apply the patches."""
+    global _patches_applied
+    if _patches_applied:
+        return
+    _sg.bipartite_subgraph = bipartite_safe
+    _tgu.bipartite_subgraph = bipartite_safe
+    _hd.bipartite_subgraph = bipartite_safe
+    _patches_applied = True

src/segger/_patches.py

@@ -10,4 +10,3 @@
 
 # Debugging utilities
 app.command(debug)
-

src/segger/cli/main.py

@@ -1,7 +1,6 @@
 import os
 import logging
 from ..utils import setup_logging
-setup_logging(level=os.environ.get("LOG_LEVEL", "WARNING"))
 
 from cyclopts import App, Parameter, Group, validators
 from typing import Annotated, Literal
@@ -78,7 +77,6 @@ def segment(
         validator=validators.Path(exists=True, dir_okay=True),
     )] = registry.get_default("output_directory"),
 
-
     # Cell Representation
     node_representation_dim: Annotated[int, Parameter(
         help="Number of dimensions used to represent each node type.",
@@ -124,6 +122,20 @@ def segment(
         group=group_nodes,
     )] = registry.get_default("genes_clusters_resolution"),
 
+    gene_corr_reference_path: Annotated[Path | None, Parameter(
+        help=(
+            "Path to a reference AnnData .h5ad file used to compute a shared "
+            "gene-gene correlation matrix."
+        ),
+        group=group_nodes,
+    )] = None,
+
+
+    gene_missing_strategy: Annotated[Literal["error", "remove", "fill"], registry.get_parameter(
+        "gene_missing_strategy",
+        group=group_nodes,
+    )] = registry.get_default("gene_missing_strategy"),
+
 
     # Transcript-Transcript Graph
     transcripts_max_k: Annotated[int, registry.get_parameter(
@@ -154,7 +166,7 @@ def segment(
         group=group_prediction,
     )] = registry.get_default("prediction_graph_max_k"),
 
-    prediction_expansion_ratio: Annotated[float | None, registry.get_parameter(
+    prediction_graph_buffer_ratio: Annotated[float | None, registry.get_parameter(
         "prediction_graph_buffer_ratio",
         validator=validators.Number(gt=0),
         group=group_prediction,
@@ -289,6 +301,7 @@ def segment(
         group=group_loss,
     )] = registry.get_default("sg_weight_end"),
 
+    # Reference
     save_anndata: Annotated[bool, registry.get_parameter(
         "save_anndata",
         group=group_io,
@@ -301,8 +314,21 @@ def segment(
     """Run cell segmentation on spatial transcriptomics data."""
 
     # Setup logger and debug directory
+    setup_logging(level=os.environ.get("LOG_LEVEL", "WARNING"), debug=debug)
     logger = logging.getLogger(__name__)
 
+    debug_dir = None
+    if debug:
+        import json
+        debug_dir = Path(output_directory) / "debug"
+        debug_dir.mkdir(exist_ok=True, parents=True)
+        params = {k: (str(v) if not isinstance(v, (str, int, float, bool, type(None))) else v)
+                  for k, v in locals().items()
+                  if k not in {"logger", "debug_dir", "json"}}
+        with open(debug_dir / "params.json", "w") as f:
+            json.dump(params, f, indent=2, default=str)
+        logger.info(f"Saved run params to {debug_dir / 'params.json'}")
+
     # Remove SLURM environment autodetect
     from lightning.pytorch.plugins.environments import SLURMEnvironment
     SLURMEnvironment.detect = lambda: False
@@ -323,11 +349,14 @@ def segment(
         transcripts_graph_max_dist=transcripts_max_dist,
         prediction_graph_mode=prediction_mode,
         prediction_graph_max_k=prediction_max_k,
-        prediction_graph_buffer_ratio=prediction_expansion_ratio,
+        prediction_graph_buffer_ratio=prediction_graph_buffer_ratio,
         tiling_margin_training=tiling_margin_training,
         tiling_margin_prediction=tiling_margin_prediction,
         tiling_nodes_per_tile=max_nodes_per_tile,
         edges_per_batch=max_edges_per_batch,
+        gene_corr_reference_path=gene_corr_reference_path,
+        gene_missing_strategy=gene_missing_strategy,
+        debug_dir=debug_dir,
     )
 
     # Setup Lightning Model