Harmony managed empty placeholder bugfix

docs/release-notes/0.17.0.md

@@ -3,6 +3,7 @@
 ```{rubric} Bug fixes
 ```
 * Make the ``lanczos`` SVD breakdown check scale-aware; the previous fixed threshold never triggered in float32. {pr}`755` {smaller}`S Dicks`
+* Fix {func}`~rapids_singlecell.pp.harmony_integrate` with RMM managed memory by making multi-key clustering arrays optional and allocating them only when needed. {pr}`766` {smaller}`A Holly & S Dicks`
 ```{rubric} Features
 ```
 * Split {func}`~rapids_singlecell.gr.calculate_niche` into {func}`~rapids_singlecell.gr.calculate_niche_neighborhood`, {func}`~rapids_singlecell.gr.calculate_niche_utag` and {func}`~rapids_singlecell.gr.calculate_niche_cellcharter`, with ``mask``, ``library_key`` and cross-flavor ``min_niche_size``, following {mod}`squidpy` {pr}`758` {smaller}`S Dicks`

src/rapids_singlecell/_cuda/harmony/clustering/clustering.cu

@@ -1,9 +1,11 @@
 #include <cub/device/device_radix_sort.cuh>
 #include <cuda_runtime.h>
+#include <nanobind/stl/optional.h>
 
 #include <algorithm>
 #include <climits>
 #include <cmath>
+#include <optional>
 #include <stdexcept>
 #include <vector>
 
@@ -325,9 +327,15 @@ static void clustering_loop_impl(const ClusteringArgs<T>& a) {
     if (a.n_covariates < 1)
         throw std::invalid_argument(
             "clustering_loop requires at least one covariate");
-    if (a.use_joint_scatter && a.n_joint_categories < 1)
-        throw std::invalid_argument(
-            "joint scatter requires at least one joint category");
+    if (a.use_joint_scatter) {
+        if (a.n_joint_categories < 1)
+            throw std::invalid_argument(
+                "joint scatter requires at least one joint category");
+        if (!a.joint_codes || !a.marginal_joint_offsets ||
+            !a.marginal_joint_indices || !a.O_joint || !a.joint_codes_in)
+            throw std::invalid_argument(
+                "joint scatter requires all joint input and workspace arrays");
+    }
 
     size_t cub_temp_bytes = get_cub_sort_temp_bytes(a.n_cells);
 
@@ -538,19 +546,20 @@ static void register_clustering_loop(nb::module_& m) {
            gpu_array_c<const T, Device> Pr_b,
            gpu_array_c<const int, Device> cats,
            gpu_array_c<const T, Device> theta,
-           gpu_array_c<const int, Device> joint_codes,
-           gpu_array_c<const int, Device> marginal_joint_offsets,
-           gpu_array_c<const int, Device> marginal_joint_indices,
-           gpu_array_c<T, Device> O_joint, gpu_array_c<T, Device> Y,
-           gpu_array_c<T, Device> Y_norm, gpu_array_c<T, Device> similarities,
+           std::optional<gpu_array_c<const int, Device>> joint_codes,
+           std::optional<gpu_array_c<const int, Device>> marginal_joint_offsets,
+           std::optional<gpu_array_c<const int, Device>> marginal_joint_indices,
+           std::optional<gpu_array_c<T, Device>> O_joint,
+           gpu_array_c<T, Device> Y, gpu_array_c<T, Device> Y_norm,
+           gpu_array_c<T, Device> similarities,
            gpu_array_c<int, Device> idx_list,
            gpu_array_c<int, Device> idx_list_alt,
            gpu_array_c<unsigned int, Device> sort_keys,
            gpu_array_c<unsigned int, Device> sort_keys_alt,
            gpu_array_c<uint8_t, Device> cub_temp,
            gpu_array_c<T, Device> R_out_buffer,
            gpu_array_c<int, Device> cats_in,
-           gpu_array_c<int, Device> joint_codes_in,
+           std::optional<gpu_array_c<int, Device>> joint_codes_in,
            gpu_array_c<T, Device> R_in_sum, gpu_array_c<T, Device> R_out_sum,
            gpu_array_c<T, Device> penalty_buf,
            gpu_array_c<T, Device> obj_scalar,
@@ -569,10 +578,12 @@ static void register_clustering_loop(nb::module_& m) {
                 Pr_b.data(),
                 cats.data(),
                 theta.data(),
-                joint_codes.data(),
-                marginal_joint_offsets.data(),
-                marginal_joint_indices.data(),
-                O_joint.data(),
+                joint_codes ? joint_codes->data() : nullptr,
+                marginal_joint_offsets ? marginal_joint_offsets->data()
+                                       : nullptr,
+                marginal_joint_indices ? marginal_joint_indices->data()
+                                       : nullptr,
+                O_joint ? O_joint->data() : nullptr,
                 Y.data(),
                 Y_norm.data(),
                 similarities.data(),
@@ -583,7 +594,7 @@ static void register_clustering_loop(nb::module_& m) {
                 cub_temp.data(),
                 R_out_buffer.data(),
                 cats_in.data(),
-                joint_codes_in.data(),
+                joint_codes_in ? joint_codes_in->data() : nullptr,
                 R_in_sum.data(),
                 R_out_sum.data(),
                 penalty_buf.data(),
@@ -610,16 +621,17 @@ static void register_clustering_loop(nb::module_& m) {
             clustering_loop_impl(a);
         },
         "Z_norm"_a, nb::kw_only(), "R"_a, "E"_a, "O"_a, "Pr_b"_a, "cats"_a,
-        "theta"_a, "joint_codes"_a, "marginal_joint_offsets"_a,
-        "marginal_joint_indices"_a, "O_joint"_a, "Y"_a, "Y_norm"_a,
-        "similarities"_a, "idx_list"_a, "idx_list_alt"_a, "sort_keys"_a,
-        "sort_keys_alt"_a, "cub_temp"_a, "R_out_buffer"_a, "cats_in"_a,
-        "joint_codes_in"_a, "R_in_sum"_a, "R_out_sum"_a, "penalty"_a,
-        "obj_scalar"_a, "ones_vec"_a, "last_obj"_a, "n_cells"_a, "n_pcs"_a,
-        "n_clusters"_a, "n_batches"_a, "n_covariates"_a, "n_joint_categories"_a,
-        "block_size"_a, "colsum_algo"_a, "sigma"_a, "tol"_a, "max_iter"_a,
-        "seed"_a, "stabilized"_a, "use_joint_scatter"_a, "stream"_a = 0,
-        "handle"_a);
+        "theta"_a, "joint_codes"_a = nb::none(),
+        "marginal_joint_offsets"_a = nb::none(),
+        "marginal_joint_indices"_a = nb::none(), "O_joint"_a = nb::none(),
+        "Y"_a, "Y_norm"_a, "similarities"_a, "idx_list"_a, "idx_list_alt"_a,
+        "sort_keys"_a, "sort_keys_alt"_a, "cub_temp"_a, "R_out_buffer"_a,
+        "cats_in"_a, "joint_codes_in"_a = nb::none(), "R_in_sum"_a,
+        "R_out_sum"_a, "penalty"_a, "obj_scalar"_a, "ones_vec"_a, "last_obj"_a,
+        "n_cells"_a, "n_pcs"_a, "n_clusters"_a, "n_batches"_a, "n_covariates"_a,
+        "n_joint_categories"_a, "block_size"_a, "colsum_algo"_a, "sigma"_a,
+        "tol"_a, "max_iter"_a, "seed"_a, "stabilized"_a, "use_joint_scatter"_a,
+        "stream"_a = 0, "handle"_a);
 }
 
 template <typename T, typename Device>

src/rapids_singlecell/preprocessing/_harmony/__init__.py

@@ -350,8 +350,6 @@ def harmonize(
             n_batches=n_joint_categories,
         )
 
-    empty_int = cp.empty(0, dtype=cp.int32)
-
     # Main harmony iterations
     is_converged = False
 
@@ -373,17 +371,9 @@ def harmonize(
             colsum_func=colsum_func_small,
             n_batches=n_batches,
             n_covariates=n_covariates,
-            joint_codes=joint_codes if joint_codes is not None else empty_int,
-            marginal_joint_offsets=(
-                marginal_joint_offsets
-                if marginal_joint_offsets is not None
-                else empty_int
-            ),
-            marginal_joint_indices=(
-                marginal_joint_indices
-                if marginal_joint_indices is not None
-                else empty_int
-            ),
+            joint_codes=joint_codes,
+            marginal_joint_offsets=marginal_joint_offsets,
+            marginal_joint_indices=marginal_joint_indices,
             n_joint_categories=n_joint_categories,
             use_joint_scatter=use_joint_scatter,
             random_state=random_state + i * 1000003,
@@ -546,7 +536,7 @@ def _allocate_clustering_workspace(
 ) -> dict:
     """Pre-allocate workspace buffers for the C++ clustering loop."""
     cub_temp_bytes = _hc_cl.get_cub_sort_temp_bytes(n_cells=n_cells)
-    return {
+    workspace = {
         "Y": cp.empty((n_clusters, n_pcs), dtype=dtype),
         "Y_norm": cp.empty((n_clusters, n_pcs), dtype=dtype),
         "similarities": cp.empty((n_cells, n_clusters), dtype=dtype),
@@ -557,19 +547,21 @@ def _allocate_clustering_workspace(
         "cub_temp": cp.empty(cub_temp_bytes, dtype=cp.uint8),
         "R_out_buffer": cp.empty((block_size, n_clusters), dtype=dtype),
         "cats_in": cp.empty(block_size * n_covariates, dtype=cp.int32),
-        "O_joint": cp.zeros(
-            (n_joint_categories if use_joint_scatter else 1, n_clusters), dtype=dtype
-        ),
-        "joint_codes_in": cp.empty(
-            max(1, block_size) if use_joint_scatter else 1, dtype=cp.int32
-        ),
         "R_in_sum": cp.empty(n_clusters, dtype=dtype),
         "R_out_sum": cp.empty(n_clusters, dtype=dtype),
         "penalty": cp.empty((n_batches, n_clusters), dtype=dtype),
         "obj_scalar": cp.empty(1, dtype=dtype),
         "ones_vec": cp.ones(block_size, dtype=dtype),
         "last_obj": cp.zeros(1, dtype=dtype),
     }
+    if use_joint_scatter:
+        workspace.update(
+            {
+                "O_joint": cp.zeros((n_joint_categories, n_clusters), dtype=dtype),
+                "joint_codes_in": cp.empty(block_size, dtype=cp.int32),
+            }
+        )
+    return workspace
 
 
 # Map colsum function to C++ enum: 0=columns, 1=atomics, 2=gemm
@@ -597,9 +589,9 @@ def _clustering(
     colsum_func: callable = None,
     n_batches: int = 0,
     n_covariates: int = 1,
-    joint_codes: cp.ndarray,
-    marginal_joint_offsets: cp.ndarray,
-    marginal_joint_indices: cp.ndarray,
+    joint_codes: cp.ndarray | None,
+    marginal_joint_offsets: cp.ndarray | None,
+    marginal_joint_indices: cp.ndarray | None,
     n_joint_categories: int,
     use_joint_scatter: bool,
     random_state: int = 0,
@@ -619,17 +611,29 @@ def _clustering(
     block_size = int(n_cells * block_proportion)
     colsum_algo_int = _COLSUM_MAP.get(colsum_func, 2)
 
+    joint_args = {}
+    if use_joint_scatter:
+        if (
+            joint_codes is None
+            or marginal_joint_offsets is None
+            or marginal_joint_indices is None
+        ):
+            raise ValueError("Joint scatter requires all joint category arrays.")
+        joint_args = {
+            "joint_codes": joint_codes,
+            "marginal_joint_offsets": marginal_joint_offsets,
+            "marginal_joint_indices": marginal_joint_indices,
+        }
+
     _hc_cl.clustering_loop(
         Z_norm,
         R=R,
         E=E,
         O=O,
         Pr_b=Pr_b.ravel(),
         cats=cats,
-        joint_codes=joint_codes,
-        marginal_joint_offsets=marginal_joint_offsets,
-        marginal_joint_indices=marginal_joint_indices,
         theta=theta,
+        **joint_args,
         **cpp_workspace,
         n_cells=n_cells,
         n_pcs=Z_norm.shape[1],

tests/test_managed_memory.py

@@ -7,6 +7,7 @@
 
 import cupy as cp
 import numpy as np
+import pandas as pd
 import pytest
 import rmm
 from anndata import AnnData
@@ -74,3 +75,15 @@ def test_qc_metrics(managed_memory):
     rsc.pp.calculate_qc_metrics(adata, log1p=False)
     assert "total_counts" in adata.obs.columns
     assert "n_genes_by_counts" in adata.obs.columns
+
+
+@pytest.mark.parametrize("dtype", [np.float32, np.float64])
+def test_harmony_integrate(managed_memory, dtype):
+    """Regression test for GH-763: harmony_integrate with managed memory."""
+    rng = np.random.default_rng(0)
+    adata = AnnData(rng.standard_normal((300, 5)).astype(dtype))
+    adata.obs["batch"] = pd.Categorical(rng.choice(["a", "b"], 300))
+    adata.obsm["X_pca"] = rng.standard_normal((300, 10)).astype(dtype)
+    rsc.pp.harmony_integrate(adata, key="batch", dtype=dtype)
+    assert adata.obsm["X_pca_harmony"].shape == (300, 10)
+    assert adata.obsm["X_pca_harmony"].dtype == dtype