return series for single control

Author: Intron7

src/rapids_singlecell/pertpy_gpu/_distance.py

@@ -239,7 +239,12 @@ def onesided_distances(
         n_bootstrap: int = 100,
         random_state: int = 0,
         multi_gpu: bool | list[int] | str | None = None,
-    ) -> pd.Series | tuple[pd.Series, pd.Series]:
+    ) -> (
+        pd.Series
+        | pd.DataFrame
+        | tuple[pd.Series, pd.Series]
+        | tuple[pd.DataFrame, pd.DataFrame]
+    ):
         """
         Compute distances from one selected group to all other groups.
 

src/rapids_singlecell/pertpy_gpu/_metrics/_edistance.py

@@ -162,7 +162,12 @@ def onesided_distances(
         n_bootstrap: int = 100,
         random_state: int = 0,
         multi_gpu: bool | list[int] | str | None = None,
-    ) -> pd.DataFrame | tuple[pd.DataFrame, pd.DataFrame]:
+    ) -> (
+        pd.Series
+        | pd.DataFrame
+        | tuple[pd.Series, pd.Series]
+        | tuple[pd.DataFrame, pd.DataFrame]
+    ):
         """
         Compute energy distances from selected reference group(s) to all other groups.
 
@@ -175,6 +180,8 @@ def onesided_distances(
         selected_group
             Reference group(s) to compute distances from. Can be a single
             group name or a sequence of group names for multiple controls.
+            When a single string is passed, returns a Series. When a sequence
+            is passed, returns a DataFrame with one column per control.
         groups
             Specific groups to compute distances to (if None, use all)
         bootstrap
@@ -194,13 +201,14 @@ def onesided_distances(
         Returns
         -------
         distances
-            DataFrame with groups as index and selected_group(s) as columns.
+            Series (single control) or DataFrame (multiple controls).
             If bootstrap=True, returns tuple of (distances, distances_var).
         """
         _assert_categorical_obs(adata, key=groupby)
 
-        # Normalize selected_group to a list
-        if isinstance(selected_group, str):
+        # Normalize selected_group to a list, track if input was a string
+        single_control = isinstance(selected_group, str)
+        if single_control:
             selected_groups = [selected_group]
         else:
             selected_groups = list(selected_group)
@@ -265,6 +273,9 @@ def onesided_distances(
                 distances = distances.loc[groups_list]
                 variances = variances.loc[groups_list]
 
+            if single_control:
+                sg = selected_groups[0]
+                return distances[sg], variances[sg]
             return distances, variances
 
         # Non-bootstrap path
@@ -293,6 +304,8 @@ def onesided_distances(
         if groups_list != all_groups:
             df = df.loc[groups_list]
 
+        if single_control:
+            return df[selected_groups[0]]
         return df
 
     def bootstrap(

tests/pertpy/test_distances.py

@@ -74,12 +74,11 @@ def test_distance_class_onesided_distances(small_adata: AnnData) -> None:
         small_adata, groupby="group", selected_group="g0"
     )
 
-    assert isinstance(result, pd.DataFrame)
+    assert isinstance(result, pd.Series)
     assert len(result) == 3  # 3 groups total
     assert "g0" in result.index
     assert "g1" in result.index
     assert "g2" in result.index
-    assert list(result.columns) == ["g0"]
 
 
 def test_distance_class_onesided_matches_pairwise(small_adata: AnnData) -> None:
@@ -96,10 +95,10 @@ def test_distance_class_onesided_matches_pairwise(small_adata: AnnData) -> None:
         )
         # Should match the row from pairwise matrix
         np.testing.assert_allclose(
-            onesided[group].values, pairwise_df.loc[group].values, atol=1e-5
+            onesided.values, pairwise_df.loc[group].values, atol=1e-5
         )
         # Self-distance should be 0
-        assert onesided.loc[group, group] == pytest.approx(0.0, abs=1e-6)
+        assert onesided.loc[group] == pytest.approx(0.0, abs=1e-6)
 
 
 def test_distance_class_onesided_multiple_controls(small_adata: AnnData) -> None:
@@ -149,18 +148,18 @@ def test_distance_class_onesided_bootstrap(small_adata: AnnData) -> None:
     assert len(result) == 2
     distances, distances_var = result
 
-    assert isinstance(distances, pd.DataFrame)
-    assert isinstance(distances_var, pd.DataFrame)
+    assert isinstance(distances, pd.Series)
+    assert isinstance(distances_var, pd.Series)
     assert len(distances) == 3
     assert len(distances_var) == 3
 
     # Self-distance variance should be 0
-    assert distances.loc["g0", "g0"] == pytest.approx(0.0, abs=1e-6)
-    assert distances_var.loc["g0", "g0"] == pytest.approx(0.0, abs=1e-6)
+    assert distances.loc["g0"] == pytest.approx(0.0, abs=1e-6)
+    assert distances_var.loc["g0"] == pytest.approx(0.0, abs=1e-6)
 
     # Non-self variances should be positive
-    assert distances_var.loc["g1", "g0"] > 0
-    assert distances_var.loc["g2", "g0"] > 0
+    assert distances_var.loc["g1"] > 0
+    assert distances_var.loc["g2"] > 0
 
 
 def test_distance_class_onesided_bootstrap_matches_pairwise(
@@ -189,11 +188,9 @@ def test_distance_class_onesided_bootstrap_matches_pairwise(
     )
 
     # Should match the corresponding row from pairwise
+    np.testing.assert_allclose(onesided.values, pairwise_df.loc["g0"].values, atol=1e-6)
     np.testing.assert_allclose(
-        onesided["g0"].values, pairwise_df.loc["g0"].values, atol=1e-6
-    )
-    np.testing.assert_allclose(
-        onesided_var["g0"].values, pairwise_var_df.loc["g0"].values, atol=1e-6
+        onesided_var.values, pairwise_var_df.loc["g0"].values, atol=1e-6
     )
 
 
@@ -329,7 +326,7 @@ def test_onesided_distances_correctness_vs_cpu(small_adata: AnnData) -> None:
             else:
                 expected = _compute_energy_distance_cpu(X, Y)
 
-            actual = onesided.loc[target_group, selected_group]
+            actual = onesided.loc[target_group]
             np.testing.assert_allclose(
                 actual,
                 expected,
@@ -1106,12 +1103,10 @@ def test_onesided_output_format(small_adata: AnnData) -> None:
         small_adata, groupby="group", selected_group="g0"
     )
 
-    assert isinstance(result, pd.DataFrame), (
-        "onesided_distances should return DataFrame"
+    assert isinstance(result, pd.Series), (
+        "onesided_distances with single control should return Series"
     )
     assert result.index.name == "group"
-    assert result.columns.name == "selected_group"
-    assert list(result.columns) == ["g0"]
 
 
 # ============================================================================
@@ -1414,7 +1409,7 @@ def test_single_cell_mixed_groups() -> None:
     # Test onesided_distances with single-cell selected group
     onesided = distance.onesided_distances(adata, groupby="group", selected_group="g0")
     assert np.all(np.isfinite(onesided.values)), "Onesided distances should be finite"
-    assert onesided.loc["g0", "g0"] == 0.0, "Self-distance should be 0"
+    assert onesided.loc["g0"] == 0.0, "Self-distance should be 0"
 
 
 def test_high_dimensional_features() -> None: