Add GPU backend for the searcher-findee / cophenetic pipeline - #3
Open
hutaobo wants to merge 5 commits into
Open
Add GPU backend for the searcher-findee / cophenetic pipeline#3hutaobo wants to merge 5 commits into
hutaobo wants to merge 5 commits into
Conversation
The cophenetic / StructureMap pipeline is built on one heavy computation: the
directed inter-group average nearest-neighbour distance ("searcher-findee").
This adds a PyTorch GPU implementation of that kernel and threads a `backend`
switch ("cpu" default, "gpu"/"cuda", "auto") through the main README entry
points: compute_searcher_findee_distance_matrix_from_df,
compute_cophenetic_distances_from_df, and compute_cophenetic_distances_from_adata.
The GPU path is numerically equivalent to the scikit-learn path: with the
default gpu_dtype="float64" it reproduces the CPU result exactly (max abs diff
= 0.0 across 2D / 3D / integer-coordinate cases, k = 5..20); gpu_dtype="float32"
is the fast mode at ~1e-5 relative deviation. backend="cpu" stays the default,
so existing behaviour is unchanged. Only the O(N^2) nearest-neighbour step moves
to the GPU; the small k x k linkage/cophenet tail stays on CPU.
New:
- src/sfplot/analysis/searcher_findee_score_gpu.py: GPU kernel + shared NN primitive
- tests/test_searcher_findee_gpu_equivalence.py: CPU/GPU equivalence (runs on CPU in CI)
- examples/gpu_searcher_findee_benchmark.py: verify + benchmark on real CUDA hardware
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
审阅者指南为 searcher→findee 最近邻内核新增一个基于 PyTorch 的 GPU 后端，并将后端切换选项贯穿接线到 cophenetic 流水线中，同时确保数值等价并保持现有的 CPU 行为不变。 cophenetic 流水线中 CPU 与 GPU 后端选择的流程图flowchart LR
User["User code
compute_cophenetic_distances_from_df / _from_adata"]
Backend["_resolve_backend(backend)"]
CPUPath["CPU path
sklearn NearestNeighbors
per cluster"]
GPUPath["GPU path
compute_searcher_findee_distance_matrix_from_df_gpu"]
Primitive["nearest_cluster_distance_columns_gpu
(torch.cdist)"]
DistMat["searcher_findee distance matrix"]
Coph["cophenetic distances
(row, col)"]
User --> Backend
Backend -->|cpu| CPUPath
Backend -->|gpu or auto and cuda| GPUPath
CPUPath --> DistMat
GPUPath --> Primitive --> DistMat
DistMat --> Coph
文件级变更
提示和命令与 Sourcery 交互
自定义你的体验访问你的 dashboard 以：
获取帮助Original review guide in EnglishReviewer's GuideAdds a PyTorch-based GPU backend for the searcher→findee nearest-neighbor kernel and wires a backend switch through the cophenetic pipeline, while ensuring numerical equivalence and preserving existing CPU behavior. Flow diagram for CPU vs GPU backend selection in cophenetic pipelineflowchart LR
User["User code
compute_cophenetic_distances_from_df / _from_adata"]
Backend["_resolve_backend(backend)"]
CPUPath["CPU path
sklearn NearestNeighbors
per cluster"]
GPUPath["GPU path
compute_searcher_findee_distance_matrix_from_df_gpu"]
Primitive["nearest_cluster_distance_columns_gpu
(torch.cdist)"]
DistMat["searcher_findee distance matrix"]
Coph["cophenetic distances
(row, col)"]
User --> Backend
Backend -->|cpu| CPUPath
Backend -->|gpu or auto and cuda| GPUPath
CPUPath --> DistMat
GPUPath --> Primitive --> DistMat
DistMat --> Coph
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
There was a problem hiding this comment.
Pull request overview
Adds an optional PyTorch-based backend for the core “searcher→findee” nearest-neighbor kernel and threads a backend selector through the public cophenetic/StructureMap entry points so the documented pipeline can run on GPU while preserving the existing CPU-default behavior.
Changes:
- Introduces
src/sfplot/analysis/searcher_findee_score_gpu.pyimplementing a batchedtorch.cdistnearest-neighbor primitive and a GPU version of the distance-matrix API. - Adds
backend/GPU control parameters to the main computation entry points and wires in"cpu" | "gpu" | "auto"backend resolution. - Adds GPU/CPU equivalence tests (running the GPU code path on
device="cpu"), a benchmark example script, and README documentation for GPU usage.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
src/sfplot/analysis/searcher_findee_score.py |
Adds backend selection and routes computation to the GPU implementation when requested. |
src/sfplot/analysis/searcher_findee_score_gpu.py |
New PyTorch implementation of the nearest-neighbor kernel and GPU distance-matrix entry point. |
src/sfplot/analysis/__init__.py |
Exposes the new GPU entry point via lazy exports. |
src/sfplot/__init__.py |
Exposes the new GPU entry point at the legacy top-level namespace. |
tests/test_searcher_findee_gpu_equivalence.py |
New equivalence tests for CPU vs GPU backend behavior. |
examples/gpu_searcher_findee_benchmark.py |
New benchmark/equivalence script for real CUDA hardware. |
README.md |
Documents the new backend switch and GPU dtype tradeoffs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
There was a problem hiding this comment.
Hey - 我发现了 1 个问题
给 AI 代理的提示
请根据这次代码审查中的评论进行修改：
## 单独评论
### 评论 1
<location path="examples/gpu_searcher_findee_benchmark.py" line_range="79" />
<code_context>
+ return (time.perf_counter() - t0) / reps
+
+
+def main() -> int:
+ p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+ p.add_argument("--max-n", type=int, default=100_000, help="largest cell count to benchmark")
</code_context>
<issue_to_address>
**issue (complexity):** 建议通过将等价性检查/基准测试逻辑以及通用辅助函数抽取到更专注的函数，或共享的工具模块中，来重构体积较大的 main() 函数，从而让脚本更易于理解和维护。
你可以在不改变任何行为的前提下，通过把 `main()` 中的逻辑拆分成更聚焦的辅助函数，并将“工具”部分隔离出来，来降低这个示例的认知负担。这样既能保持示例的可读性，又能保留所有 CLI 选项和 JSON 输出。
**1. 将等价性检查和基准测试提取为单独函数**
这会让 `main()` 的执行流程更容易跟踪，并避免变成一个 140 行的大块函数：
```python
EQUIV_N = 20_000
BENCHMARK_SIZES = (1_000, 2_000, 5_000, 10_000, 20_000, 50_000,
100_000, 200_000, 500_000, 1_000_000)
def run_equivalence(args, cpu_kw, device: str) -> tuple[float, float, bool]:
n = min(EQUIV_N, args.max_n)
df = make_df(n, args.k, args.dims, seed=12345)
cpu = sf(df, **cpu_kw)
d64 = max_abs_diff(cpu, sf(df, z_col=cpu_kw["z_col"], backend="gpu",
device=device, gpu_dtype="float64",
gpu_max_memory_gb=args.max_memory_gb))
d32 = max_abs_diff(cpu, sf(df, z_col=cpu_kw["z_col"], backend="gpu",
device=device, gpu_dtype="float32",
gpu_max_memory_gb=args.max_memory_gb))
ok = d64 < 1e-9
print(f"[1] EQUIVALENCE (n={n})")
print(f" float64 max|delta| = {d64:.3e} -> {'OK (GPU == CPU)' if ok else 'MISMATCH'}")
print(f" float32 max|delta| = {d32:.3e} (fast mode)\n")
return d64, d32, ok
def run_benchmark(args, cpu_kw, gpu_kw) -> list[dict]:
sizes = [n for n in BENCHMARK_SIZES if n <= args.max_n]
print("[2] BENCHMARK (seconds/call; speedup = CPU/GPU)")
print(f" {'N':>9} {'CPU(s)':>10} {'GPU(s)':>10} {'speedup':>9}")
rows: list[dict] = []
for n in sizes:
df = make_df(n, args.k, args.dims, seed=n)
reps = max(1, int(200_000 // n))
t_cpu = time_call(lambda: sf(df, **cpu_kw), reps, "cpu")
t_gpu = time_call(lambda: sf(df, **gpu_kw), reps, args.device)
speedup = t_cpu / t_gpu if t_gpu > 0 else float("nan")
print(f" {n:>9} {t_cpu:>10.4f} {t_gpu:>10.4f} {speedup:>8.2f}x")
rows.append(dict(n=n, cpu_s=t_cpu, gpu_s=t_gpu, speedup=speedup, reps=reps))
return rows
```
然后 `main()` 主要就变成了一个调度函数：
```python
def main() -> int:
args = parse_args() # 一个小的辅助函数，只负责构建并返回 argparse.Namespace
z_col = "z" if args.dims == 3 else None
cuda_ok = torch.cuda.is_available()
print(f"torch {torch.__version__} device={args.device} cuda_available={cuda_ok}")
if args.device.startswith("cuda") and not cuda_ok:
print("WARNING: CUDA not available; GPU timings will fall back to CPU and be meaningless.")
print(f"params: k={args.k} dims={args.dims} gpu_dtype={args.dtype} max_n={args.max_n}\n")
cpu_kw = dict(z_col=z_col, backend="cpu")
gpu_kw = dict(z_col=z_col, backend="gpu", device=args.device,
gpu_dtype=args.dtype, gpu_max_memory_gb=args.max_memory_gb)
d64, d32, ok = run_equivalence(args, cpu_kw, args.device)
rows = run_benchmark(args, cpu_kw, gpu_kw)
if args.json:
write_json(args, cuda_ok, d64, d32, ok, rows) # 小的辅助函数
return 0 if ok else 1
```
**2. 将通用辅助函数移动到共享工具模块（可选但成本较低）**
`max_abs_diff` 和 `time_call` 足够通用，可以放到一个内部的小模块中（例如 `examples/_utils.py` 或 `cellgps/_bench_utils.py`），以便在其他示例/测试中复用：
```python
# examples/_utils.py
import numpy as np
import time
import torch
def max_abs_diff(a, b) -> float:
a = a.reindex(index=b.index, columns=b.columns)
av, bv = a.values.astype(float), b.values.astype(float)
both = ~(np.isnan(av) | np.isnan(bv))
return float(np.abs(av[both] - bv[both]).max()) if both.any() else 0.0
def time_call(fn, reps: int, device: str) -> float:
fn()
if device.startswith("cuda"):
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(reps):
fn()
if device.startswith("cuda"):
torch.cuda.synchronize()
return (time.perf_counter() - t0) / reps
```
然后在这个示例中：
```python
from examples._utils import max_abs_diff, time_call
```
这样可以让示例专注于“如何使用 GPU 后端”，同时保留当前的所有功能（CLI 接口、JSON、基准测试、等价性检查）。
</issue_to_address>帮我变得更有用！请在每条评论上点击 👍 或 👎，我会根据你的反馈改进之后的审查。
Original comment in English
Hey - I've found 1 issue
Prompt for AI Agents
Please address the comments from this code review:
## Individual Comments
### Comment 1
<location path="examples/gpu_searcher_findee_benchmark.py" line_range="79" />
<code_context>
+ return (time.perf_counter() - t0) / reps
+
+
+def main() -> int:
+ p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+ p.add_argument("--max-n", type=int, default=100_000, help="largest cell count to benchmark")
</code_context>
<issue_to_address>
**issue (complexity):** Consider refactoring the large main() function by extracting equivalence/benchmark logic and generic helpers into focused functions or a shared utility module to make the script easier to understand and maintain.
You can reduce the cognitive load of this example without changing any behavior by splitting the logic in `main()` into focused helpers and isolating the “utility” pieces. That keeps the example readable while preserving all CLI options and JSON output.
**1. Extract equivalence + benchmark into separate functions**
This makes the flow of `main()` much easier to follow and avoids the 140-line monolith:
```python
EQUIV_N = 20_000
BENCHMARK_SIZES = (1_000, 2_000, 5_000, 10_000, 20_000, 50_000,
100_000, 200_000, 500_000, 1_000_000)
def run_equivalence(args, cpu_kw, device: str) -> tuple[float, float, bool]:
n = min(EQUIV_N, args.max_n)
df = make_df(n, args.k, args.dims, seed=12345)
cpu = sf(df, **cpu_kw)
d64 = max_abs_diff(cpu, sf(df, z_col=cpu_kw["z_col"], backend="gpu",
device=device, gpu_dtype="float64",
gpu_max_memory_gb=args.max_memory_gb))
d32 = max_abs_diff(cpu, sf(df, z_col=cpu_kw["z_col"], backend="gpu",
device=device, gpu_dtype="float32",
gpu_max_memory_gb=args.max_memory_gb))
ok = d64 < 1e-9
print(f"[1] EQUIVALENCE (n={n})")
print(f" float64 max|delta| = {d64:.3e} -> {'OK (GPU == CPU)' if ok else 'MISMATCH'}")
print(f" float32 max|delta| = {d32:.3e} (fast mode)\n")
return d64, d32, ok
def run_benchmark(args, cpu_kw, gpu_kw) -> list[dict]:
sizes = [n for n in BENCHMARK_SIZES if n <= args.max_n]
print("[2] BENCHMARK (seconds/call; speedup = CPU/GPU)")
print(f" {'N':>9} {'CPU(s)':>10} {'GPU(s)':>10} {'speedup':>9}")
rows: list[dict] = []
for n in sizes:
df = make_df(n, args.k, args.dims, seed=n)
reps = max(1, int(200_000 // n))
t_cpu = time_call(lambda: sf(df, **cpu_kw), reps, "cpu")
t_gpu = time_call(lambda: sf(df, **gpu_kw), reps, args.device)
speedup = t_cpu / t_gpu if t_gpu > 0 else float("nan")
print(f" {n:>9} {t_cpu:>10.4f} {t_gpu:>10.4f} {speedup:>8.2f}x")
rows.append(dict(n=n, cpu_s=t_cpu, gpu_s=t_gpu, speedup=speedup, reps=reps))
return rows
```
Then `main()` becomes mostly orchestration:
```python
def main() -> int:
args = parse_args() # small helper that just builds/returns argparse.Namespace
z_col = "z" if args.dims == 3 else None
cuda_ok = torch.cuda.is_available()
print(f"torch {torch.__version__} device={args.device} cuda_available={cuda_ok}")
if args.device.startswith("cuda") and not cuda_ok:
print("WARNING: CUDA not available; GPU timings will fall back to CPU and be meaningless.")
print(f"params: k={args.k} dims={args.dims} gpu_dtype={args.dtype} max_n={args.max_n}\n")
cpu_kw = dict(z_col=z_col, backend="cpu")
gpu_kw = dict(z_col=z_col, backend="gpu", device=args.device,
gpu_dtype=args.dtype, gpu_max_memory_gb=args.max_memory_gb)
d64, d32, ok = run_equivalence(args, cpu_kw, args.device)
rows = run_benchmark(args, cpu_kw, gpu_kw)
if args.json:
write_json(args, cuda_ok, d64, d32, ok, rows) # small helper
return 0 if ok else 1
```
**2. Move generic helpers into a shared utility (optional but low-cost)**
`max_abs_diff` and `time_call` are generic enough to live in a small internal module (e.g. `examples/_utils.py` or `cellgps/_bench_utils.py`) and be reused by other examples/tests:
```python
# examples/_utils.py
import numpy as np
import time
import torch
def max_abs_diff(a, b) -> float:
a = a.reindex(index=b.index, columns=b.columns)
av, bv = a.values.astype(float), b.values.astype(float)
both = ~(np.isnan(av) | np.isnan(bv))
return float(np.abs(av[both] - bv[both]).max()) if both.any() else 0.0
def time_call(fn, reps: int, device: str) -> float:
fn()
if device.startswith("cuda"):
torch.cuda.synchronize()
t0 = time.perf_counter()
for _ in range(reps):
fn()
if device.startswith("cuda"):
torch.cuda.synchronize()
return (time.perf_counter() - t0) / reps
```
Then in the example:
```python
from examples._utils import max_abs_diff, time_call
```
This keeps the example focused on “how to use the GPU backend” while preserving all current functionality (CLI surface, JSON, benchmarking, equivalence check).
</issue_to_address>Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The whole Cell-GPS cophenetic / StructureMap pipeline (the README's main architecture) rests on one heavy computation — the directed inter-group average nearest-neighbour distance ("searcher→findee"), one
sklearn.neighbors.NearestNeighbors1-NN model per cell type. For large cell counts that O(N²) step dominates; everything downstream (linkage/copheneton the small k×k cluster matrix) is negligible.This PR adds a PyTorch GPU implementation of that kernel and threads a
backendswitch through the main README entry points, so the documented pipeline can run end-to-end on a GPU.What changed
src/sfplot/analysis/searcher_findee_score_gpu.py— batchedtorch.cdistGPU kernel + a sharednearest_cluster_distance_columns_gpuprimitive.backendswitch ("cpu"default,"gpu"/"cuda","auto") added to:compute_searcher_findee_distance_matrix_from_dfcompute_cophenetic_distances_from_dfcompute_cophenetic_distances_from_adatacompute_searcher_findee_distance_matrix_from_df_gpu(viacellgpsandsfplot).tests/test_searcher_findee_gpu_equivalence.py— run the GPU code path ondevice="cpu"so they execute in CI without a GPU.examples/gpu_searcher_findee_benchmark.py— verify equivalence + measure speedup on real CUDA hardware (warm-up +cuda.synchronize).Equivalence (the contract)
The GPU path is numerically equivalent to the scikit-learn path. With the default
gpu_dtype="float64"and the exact Euclidean formula (compute_mode="donot_use_mm_for_euclid_dist"), it reproduces the CPU result exactly:verified across 2D / 3D / integer-coordinate inputs and k = 5..20.
gpu_dtype="float32"is the fast mode at ~1e-5 relative deviation.Compatibility
backend="cpu"remains the default — existing behaviour is unchanged.torch(already declared as theCell-GPS[gpu]extra). The GPU module is imported lazily, only when a GPU backend is requested.🤖 Generated with Claude Code
Summary by Sourcery
添加一个基于 GPU 的 PyTorch 实现，用于 searcher–findee 最近邻核函数，并在 cophenetic/StructureMap 流水线中贯穿后端切换开关，从而在保持 CPU 默认行为的同时支持在 GPU 上运行。
New Features:
Enhancements:
observed=False，以使 groupby 行为与现有语义保持一致。Tests:
cellgps/sfplot命名空间对外导出。Original summary in English
Summary by Sourcery
Add a GPU-backed PyTorch implementation of the searcher–findee nearest-neighbor kernel and thread a backend switch through the cophenetic/StructureMap pipeline so it can run on GPU while preserving CPU-default behavior.
New Features:
Enhancements:
Tests: