update kernel and tests

Author: Intron7

src/rapids_singlecell/pertpy_gpu/_distance.py

@@ -3,6 +3,8 @@
 from typing import TYPE_CHECKING, Literal, NamedTuple
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     import cupy as cp
     import numpy as np
     import pandas as pd
@@ -20,50 +22,33 @@ class Distance:
     """
     GPU-accelerated distance computation between groups of cells.
 
-    This class provides an extensible framework for computing various distance metrics
-    between cell groups in single-cell data, with GPU acceleration via CuPy.
-
-    The API is designed to be compatible with pertpy's Distance class.
+    API compatible with pertpy's Distance class.
 
     Parameters
     ----------
     metric : str
-        Distance metric to use. Currently supported:
-        - 'edistance': Energy distance (GPU-accelerated)
+        Distance metric. Currently supported: 'edistance' (energy distance).
     layer_key : str | None
-        Name of the counts layer containing raw counts to calculate distances for.
-        Mutually exclusive with 'obsm_key'. If None, the parameter is ignored.
+        Key in adata.layers for cell data. Mutually exclusive with 'obsm_key'.
     obsm_key : str | None
         Key in adata.obsm for embeddings. Mutually exclusive with 'layer_key'.
-        Defaults to 'X_pca' if neither layer_key nor obsm_key is specified.
-    kernel : str
-        Kernel strategy: 'auto' or 'manual'.
-        - 'auto': Dynamically choose optimal blocks_per_pair (default)
-        - 'manual': Use the specified blocks_per_pair directly
-    blocks_per_pair : int
-        Number of blocks per pair (default: 32). For 'auto', this is the maximum.
-        Higher values increase parallelism but add atomic overhead.
+        Defaults to 'X_pca' if neither is specified.
 
     Examples
     --------
     >>> import rapids_singlecell as rsc
     >>> distance = rsc.ptg.Distance(metric='edistance')
     >>> result = distance.pairwise(adata, groupby='perturbation')
-    >>> print(result.distances)
 
-    # Direct computation on arrays (pertpy-compatible API)
-    >>> X = adata.obsm["X_pca"][adata.obs["group"] == "A"]
-    >>> Y = adata.obsm["X_pca"][adata.obs["group"] == "B"]
-    >>> d = distance(X, Y)  # Returns energy distance as float
+    >>> # Direct computation on arrays
+    >>> d = distance(X, Y)
     """
 
     def __init__(
         self,
         metric: Literal["edistance"] = "edistance",
         layer_key: str | None = None,
         obsm_key: str | None = None,
-        kernel: Literal["auto", "manual"] = "auto",
-        blocks_per_pair: int = 32,
     ):
         """Initialize Distance calculator with specified metric."""
         if layer_key and obsm_key:
@@ -77,8 +62,6 @@ def __init__(
         self.metric = metric
         self.layer_key = layer_key
         self.obsm_key = obsm_key
-        self.kernel = kernel
-        self.blocks_per_pair = blocks_per_pair
         self._metric_impl = None
         self._initialize_metric()
 
@@ -92,8 +75,6 @@ def _initialize_metric(self):
             self._metric_impl = EDistanceMetric(
                 layer_key=self.layer_key,
                 obsm_key=self.obsm_key,
-                kernel=self.kernel,
-                blocks_per_pair=self.blocks_per_pair,
             )
         else:
             raise ValueError(
@@ -140,11 +121,10 @@ def pairwise(
         adata: AnnData,
         groupby: str,
         *,
-        groups: list[str] | None = None,
+        groups: Sequence[str] | None = None,
         bootstrap: bool = False,
         n_bootstrap: int = 100,
         random_state: int = 0,
-        inplace: bool = False,
     ):
         """
         Compute pairwise distances between all cell groups.
@@ -155,28 +135,25 @@ def pairwise(
             Annotated data matrix
         groupby : str
             Key in adata.obs for grouping cells
-        groups : list[str] | None
+        groups : Sequence[str] | None
             Specific groups to compute (if None, use all)
         bootstrap : bool
             Whether to compute bootstrap variance estimates
         n_bootstrap : int
             Number of bootstrap iterations (if bootstrap=True)
         random_state : int
             Random seed for reproducibility
-        inplace : bool
-            Whether to store results in adata.uns
 
         Returns
         -------
         result
-            Result object containing distances and optional variance DataFrames.
-            The exact type depends on the metric used.
+            DataFrame with pairwise distances. If bootstrap=True, returns
+            tuple of (distances, distances_var) DataFrames.
 
         Examples
         --------
         >>> distance = Distance(metric='edistance')
         >>> result = distance.pairwise(adata, groupby='condition')
-        >>> print(result.distances)
         """
         return self._metric_impl.pairwise(
             adata=adata,
@@ -185,7 +162,6 @@ def pairwise(
             bootstrap=bootstrap,
             n_bootstrap=n_bootstrap,
             random_state=random_state,
-            inplace=inplace,
         )
 
     def onesided_distances(
@@ -194,11 +170,11 @@ def onesided_distances(
         groupby: str,
         selected_group: str,
         *,
-        groups: list[str] | None = None,
+        groups: Sequence[str] | None = None,
         bootstrap: bool = False,
         n_bootstrap: int = 100,
         random_state: int = 0,
-    ) -> pd.Series:
+    ) -> pd.Series | tuple[pd.Series, pd.Series]:
         """
         Compute distances from one selected group to all other groups.
 
@@ -210,7 +186,7 @@ def onesided_distances(
             Key in adata.obs for grouping cells
         selected_group : str
             Reference group to compute distances from
-        groups : list[str] | None
+        groups : Sequence[str] | None
             Specific groups to compute distances to (if None, use all)
         bootstrap : bool
             Whether to compute bootstrap variance estimates
@@ -221,16 +197,16 @@ def onesided_distances(
 
         Returns
         -------
-        distances : pd.Series
-            Series containing distances from selected_group to all other groups
+        distances : pd.Series | tuple[pd.Series, pd.Series]
+            Series containing distances from selected_group to all other groups.
+            If bootstrap=True, returns tuple of (distances, distances_var).
 
         Examples
         --------
         >>> distance = Distance(metric='edistance')
         >>> distances = distance.onesided_distances(
         ...     adata, groupby='condition', selected_group='control'
         ... )
-        >>> print(distances)
         """
         if not hasattr(self._metric_impl, "onesided_distances"):
             raise NotImplementedError(

src/rapids_singlecell/pertpy_gpu/_metrics/_base_metric.py

@@ -4,6 +4,8 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
+    from collections.abc import Sequence
+
     from anndata import AnnData
 
 
@@ -30,58 +32,50 @@ def pairwise(
         adata: AnnData,
         groupby: str,
         *,
-        groups: list[str] | None = None,
+        groups: Sequence[str] | None = None,
         bootstrap: bool = False,
         n_bootstrap: int = 100,
         random_state: int = 0,
-        inplace: bool = False,
     ):
         """
         Compute pairwise distances between all cell groups.
 
-        This method must be implemented by all metric subclasses.
-
         Parameters
         ----------
         adata : AnnData
             Annotated data matrix
         groupby : str
             Key in adata.obs for grouping cells
-        groups : list[str] | None
+        groups : Sequence[str] | None
             Specific groups to compute (if None, use all)
         bootstrap : bool
             Whether to compute bootstrap variance estimates
         n_bootstrap : int
             Number of bootstrap iterations (if bootstrap=True)
         random_state : int
             Random seed for reproducibility
-        inplace : bool
-            Whether to store results in adata.uns
 
         Returns
         -------
         result
             Result object containing distances and optional variance information.
-            The exact type depends on the specific metric implementation.
         """
-        pass
+        ...
 
     def onesided_distances(
         self,
         adata: AnnData,
         groupby: str,
         selected_group: str,
         *,
-        groups: list[str] | None = None,
+        groups: Sequence[str] | None = None,
         bootstrap: bool = False,
         n_bootstrap: int = 100,
         random_state: int = 0,
     ):
         """
         Compute distances from one selected group to all other groups.
 
-        This method is optional and may not be implemented by all metrics.
-
         Parameters
         ----------
         adata : AnnData
@@ -90,7 +84,7 @@ def onesided_distances(
             Key in adata.obs for grouping cells
         selected_group : str
             Reference group to compute distances from
-        groups : list[str] | None
+        groups : Sequence[str] | None
             Specific groups to compute distances to (if None, use all)
         bootstrap : bool
             Whether to compute bootstrap variance estimates
@@ -121,8 +115,6 @@ def bootstrap(
         """
         Compute bootstrap mean and variance for distance between two specific groups.
 
-        This method is optional and may not be implemented by all metrics.
-
         Parameters
         ----------
         adata : AnnData