harmony hotfix (#351)

* harmony hotfix

* update

Author: Intron7

docs/release-notes/0.12.3.md

@@ -1,13 +1,5 @@
-### 0.12.3 {small}`the-future`
-
-```{rubric} Features
-```
-
-```{rubric} Performance
-```
+### 0.12.3 {small}`2025-04-11`
 
 ```{rubric} Bug fixes
 ```
-
-```{rubric} Misc
-```
+* Fixed `harmony_integrate` breakage caused by undocumented changes to fused kernel float handling in` CuPy 13.4.1` {pr}`351` {smaller}`S Dicks`

docs/release-notes/0.12.4.md

@@ -0,0 +1,13 @@
+### 0.12.4 {small}`the-future`
+
+```{rubric} Features
+```
+
+```{rubric} Performance
+```
+
+```{rubric} Bug fixes
+```
+
+```{rubric} Misc
+```

docs/release-notes/index.md

@@ -3,6 +3,8 @@
 # Release notes
 
 ## Version 0.12.0
+```{include} /release-notes/0.12.3.md
+```
 ```{include} /release-notes/0.12.2.md
 ```
 ```{include} /release-notes/0.12.1.md

src/rapids_singlecell/preprocessing/_harmony/__init__.py

@@ -92,7 +92,7 @@ def harmonize(
     batch_mat: pd.DataFrame,
     batch_key: str | list[str],
     *,
-    n_clusters: int = None,
+    n_clusters: int | None = None,
     max_iter_harmony: int = 10,
     max_iter_clustering: int = 200,
     tol_harmony: float = 1e-4,
@@ -259,14 +259,13 @@ def _initialize_centroids(
     kmeans.fit(Z_norm)
     Y = kmeans.cluster_centers_.astype(Z_norm.dtype)
     Y_norm = _normalize_cp(Y, p=2)
-
     # Initialize R
-    R = _calc_R(-2 / sigma, cp.dot(Z_norm, Y_norm.T))
+    term = cp.float64(-2 / sigma).astype(Z_norm.dtype)
+    R = _calc_R(term, cp.dot(Z_norm, Y_norm.T))
     R = _normalize_cp(R, p=1)
 
     E = cp.dot(Pr_b, cp.sum(R, axis=0, keepdims=True))
     O = cp.dot(Phi.T, R)
-
     objectives_harmony = []
     _compute_objective(
         Y_norm,
@@ -278,7 +277,6 @@ def _initialize_centroids(
         E=E,
         objective_arr=objectives_harmony,
     )
-
     return R, E, O, objectives_harmony
 
 
@@ -308,7 +306,7 @@ def _clustering(
     n_cells = Z_norm.shape[0]
     objectives_clustering = []
     block_size = int(n_cells * block_proportion)
-    term = -2 / sigma
+    term = cp.float64(-2 / sigma).astype(Z_norm.dtype)
     for _ in range(max_iter):
         # Compute Cluster Centroids
         Y = cp.dot(R.T, Z_norm)  # Compute centroids

src/rapids_singlecell/preprocessing/_harmony/_fuses.py

@@ -14,7 +14,7 @@ def _get_pen(e: cp.ndarray, o: cp.ndarray, theta: cp.ndarray) -> cp.ndarray:
 
 
 @cp.fuse
-def _calc_R(term: cp.ndarray, dotproduct: cp.ndarray) -> cp.ndarray:
+def _calc_R(term: float, dotproduct: cp.ndarray) -> cp.ndarray:
     return cp.exp(term * (1 - dotproduct))
 
 

tests/test_harmony.py

@@ -1,23 +1,24 @@
 from __future__ import annotations
 
+import cupy as cp
 import pytest
 import scanpy as sc
 
 import rapids_singlecell as rsc
 
 
 @pytest.mark.parametrize("correction_method", ["fast", "original"])
-def test_harmony_integrate(correction_method):
+@pytest.mark.parametrize("dtype", [cp.float32, cp.float64])
+def test_harmony_integrate(correction_method, dtype):
     """
     Test that Harmony integrate works.
 
     This is a very simple test that just checks to see if the Harmony
     integrate wrapper successfully added a new field to ``adata.obsm``
     and makes sure it has the same dimensions as the original PCA table.
     """
-    adata = sc.datasets.pbmc3k()
-    sc.pp.recipe_zheng17(adata)
-    sc.tl.pca(adata)
-    adata.obs["batch"] = 1350 * ["a"] + 1350 * ["b"]
-    rsc.pp.harmony_integrate(adata, "batch", correction_method=correction_method)
+    adata = sc.datasets.pbmc68k_reduced()
+    rsc.pp.harmony_integrate(
+        adata, "bulk_labels", correction_method=correction_method, dtype=dtype
+    )
     assert adata.obsm["X_pca_harmony"].shape == adata.obsm["X_pca"].shape