add pseudobulk distances for pertpy GPU

[Intron7]

This PR adds the pertpy's Pseudobulk distances to RSC

[Intron7]

@coderabbitai review

[coderabbitai]

✅ Actions performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

Release Notes

• New Features

  • Added GPU-accelerated pseudobulk distance metrics: Euclidean, Mean Squared Error, Mean Absolute Error, Pearson, Cosine, and R2 Score—providing expanded distance calculation options.

• Documentation

  • Updated release notes documenting new pseudobulk distance metric additions.

• Tests

  • Added comprehensive test coverage for pseudobulk distance metrics across input formats and configurations.

Walkthrough

This PR adds six GPU-accelerated pseudobulk distance metrics (Euclidean, MSE, MAE, Pearson, Cosine, R²) for comparing group-level mean vectors in AnnData objects. It introduces paired/pairwise CUDA kernels, CuPy wrappers, metric implementations inheriting a shared workflow, and integration into the Distance API dispatcher.

Changes

Pseudobulk GPU metrics

Layer / File(s)	Summary
CUDA kernel implementation and bindings src/rapids_singlecell/_cuda/pseudobulk/kernels_pseudobulk.cuh, src/rapids_singlecell/_cuda/pseudobulk/pseudobulk.cu, CMakeLists.txt	Two template kernels (paired_kernel, pairwise_kernel) reduce per-feature differences using warp shuffles and shared memory, parameterized by PseudobulkOp (Squared, AbsMean). Nanobind bindings expose four Python functions that compute grid/block sizes and launch kernels with CUDA error checking.
CuPy kernel wrappers and stream management src/rapids_singlecell/pertpy_gpu/_metrics/_kernels/__init__.py, src/rapids_singlecell/pertpy_gpu/_metrics/_kernels/_pseudobulk.py	Four exported functions (paired_squared, paired_abs_mean, pairwise_squared, pairwise_abs_mean) validate 2D shapes, convert inputs to float64 contiguous arrays, allocate output buffers, and dispatch to CUDA kernels using CuPy's current stream.
BaseMetric helper methods for validation src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py	_parse_contrasts validates and splits contrast DataFrames into groupby and stratification columns. _resolve_onesided_inputs validates categorical groupby, normalizes selected groups, and computes the union of needed groups for subsetting.
PseudobulkMetric base class and implementations src/rapids_singlecell/pertpy_gpu/_metrics/_pseudobulk.py	PseudobulkMetric orchestrates embedding extraction, per-group mean aggregation, GPU coercion, and three public APIs (pairwise, onesided_distances, contrast_distances). Six concrete subclasses implement metric-specific distance formulas: Euclidean, MeanSquared (MSE), MeanAbsolute, PearsonDistance (with zero-variance NaN handling), CosineDistance (with SciPy-like clipping), and R2ScoreDistance (with upper-triangle canonicalization).
Distance API metric dispatch src/rapids_singlecell/pertpy_gpu/_distance.py, src/rapids_singlecell/pertpy_gpu/_metrics/__init__.py, src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py	Distance now accepts a Metric union type covering all supported metrics; _initialize_metric dispatches to PSEUDOBULK_METRICS registry. EDistanceMetric refactored to use _parse_contrasts and _resolve_onesided_inputs helpers instead of inline logic.
Tests and documentation tests/pertpy/test_distance_pseudobulk.py, docs/release-notes/0.15.2.md	Comprehensive pytest suite validates single-pair, pairwise, onesided, and contrast distances against reference implementations; unit tests for GPU kernels; edge-case coverage for zero-variance/zero-norm scenarios; multi-GPU fallback; cross-validation with pertpy when available. Release notes document the new metrics.

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 19.72% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'add pseudobulk distances for pertpy GPU' accurately and specifically describes the primary change: adding pseudobulk distance metrics to the pertpy GPU module.
Description check	✅ Passed	The description 'This PR adds the pertpy's Pseudobulk distances to RSC' is directly related to the changeset and clearly indicates the main objective of adding pseudobulk distance functionality.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch add-pseudobulk-metrics

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

src/rapids_singlecell/_cuda/pseudobulk/kernels_pseudobulk.cuh (1)

48-48: 💤 Low value

Use named constant for shared memory array size.

The shared memory array size 32 should be a named constant for consistency with the coding guidelines and to match WARP_SIZE. This appears in both kernels.

Suggested fix

 static constexpr int WARP_SIZE = 32;
+static constexpr int MAX_WARPS_PER_BLOCK = 32;  // supports up to 1024 threads/block
 
 // ... in paired_kernel and pairwise_kernel:
-    __shared__ double s[32];
+    __shared__ double s[MAX_WARPS_PER_BLOCK];

Also applies to: 83-83

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/rapids_singlecell/_cuda/pseudobulk/kernels_pseudobulk.cuh` at line 48,
Replace the hard-coded shared-memory size literal 32 with the named constant
used elsewhere (WARP_SIZE) in both shared array declarations (the lines
declaring "__shared__ double s[32];" in each kernel) so the shared array is
defined as "__shared__ double s[WARP_SIZE];" to follow the coding guideline and
keep consistency with the existing WARP_SIZE constant.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/pertpy/test_distance_pseudobulk.py`:
- Line 397: The docstring text contains a unicode multiplication symbol "×"
which triggers Ruff RUF002; update the docstring in
tests/pertpy/test_distance_pseudobulk.py (the string: """K=1 pairwise returns a
1×1 zero matrix; onesided returns a scalar zero.""") to use a plain "x" instead
("""K=1 pairwise returns a 1x1 zero matrix; onesided returns a scalar zero.""")
so the linter no longer flags it.
- Around line 30-31: The helper incorrectly treats metric ==
"root_mean_squared_error" as the Euclidean norm; change the branch handling
metric to compute RMSE as np.sqrt(np.mean((mu_X - mu_Y) ** 2)) using the
existing mu_X and mu_Y instead of np.linalg.norm so the test compares against
the true root-mean-square error.

---

Nitpick comments:
In `@src/rapids_singlecell/_cuda/pseudobulk/kernels_pseudobulk.cuh`:
- Line 48: Replace the hard-coded shared-memory size literal 32 with the named
constant used elsewhere (WARP_SIZE) in both shared array declarations (the lines
declaring "__shared__ double s[32];" in each kernel) so the shared array is
defined as "__shared__ double s[WARP_SIZE];" to follow the coding guideline and
keep consistency with the existing WARP_SIZE constant.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: a9be5f11-7432-4dcb-90b3-7b992ccd355a

📥 Commits

Reviewing files that changed from the base of the PR and between 16c369d and 7307a8a.

📒 Files selected for processing (12)

• CMakeLists.txt
• docs/release-notes/0.15.2.md
• src/rapids_singlecell/_cuda/pseudobulk/kernels_pseudobulk.cuh
• src/rapids_singlecell/_cuda/pseudobulk/pseudobulk.cu
• src/rapids_singlecell/pertpy_gpu/_distance.py
• src/rapids_singlecell/pertpy_gpu/_metrics/__init__.py
• src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py
• src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py
• src/rapids_singlecell/pertpy_gpu/_metrics/_kernels/__init__.py
• src/rapids_singlecell/pertpy_gpu/_metrics/_kernels/_pseudobulk.py
• src/rapids_singlecell/pertpy_gpu/_metrics/_pseudobulk.py
• tests/pertpy/test_distance_pseudobulk.py

[Zethson]

Thank you very much! Awesome stuff 🥇

[Zethson]

Oh and @Intron7 it'd be awesome if you also showed this off super briefly somewhere in the tutorials, please. Like just 1-2 cells and then some intersphinx link to show what's available. Claude is also very good with tutorials and can surely do that for you super quickly.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 93.05019% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 88.83%. Comparing base (16c369d) to head (e52a0d2).

Files with missing lines	Patch %	Lines
...pids_singlecell/pertpy_gpu/_metrics/_pseudobulk.py	93.61%	12 Missing ⚠️
...nglecell/pertpy_gpu/_metrics/_utils/_pseudobulk.py	86.66%	6 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #676      +/-   ##
==========================================
+ Coverage   88.65%   88.83%   +0.18%     
==========================================
  Files          98      100       +2     
  Lines        7375     7615     +240     
==========================================
+ Hits         6538     6765     +227     
- Misses        837      850      +13     

Files with missing lines	Coverage Δ
src/rapids_singlecell/pertpy_gpu/_distance.py	87.31% <100.00%> (+3.59%)	⬆️
.../rapids_singlecell/pertpy_gpu/_metrics/__init__.py	100.00% <ø> (ø)
...ids_singlecell/pertpy_gpu/_metrics/_base_metric.py	89.58% <100.00%> (+6.82%)	⬆️
...apids_singlecell/pertpy_gpu/_metrics/_edistance.py	95.21% <100.00%> (-0.17%)	⬇️
...nglecell/pertpy_gpu/_metrics/_utils/_pseudobulk.py	86.66% <86.66%> (ø)
...pids_singlecell/pertpy_gpu/_metrics/_pseudobulk.py	93.61% <93.61%> (ø)

... and 1 file with indirect coverage changes