diff --git a/docs/release-notes/0.17.0.md b/docs/release-notes/0.17.0.md index eaacf8b7c..c6910eb46 100644 --- a/docs/release-notes/0.17.0.md +++ b/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` diff --git a/src/rapids_singlecell/_cuda/harmony/clustering/clustering.cu b/src/rapids_singlecell/_cuda/harmony/clustering/clustering.cu index eddd2b30e..50ec53b0e 100644 --- a/src/rapids_singlecell/_cuda/harmony/clustering/clustering.cu +++ b/src/rapids_singlecell/_cuda/harmony/clustering/clustering.cu @@ -1,9 +1,11 @@ #include #include +#include #include #include #include +#include #include #include @@ -325,9 +327,15 @@ static void clustering_loop_impl(const ClusteringArgs& 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,11 +546,12 @@ static void register_clustering_loop(nb::module_& m) { gpu_array_c Pr_b, gpu_array_c cats, gpu_array_c theta, - gpu_array_c joint_codes, - gpu_array_c marginal_joint_offsets, - gpu_array_c marginal_joint_indices, - gpu_array_c O_joint, gpu_array_c Y, - gpu_array_c Y_norm, gpu_array_c similarities, + std::optional> joint_codes, + std::optional> marginal_joint_offsets, + std::optional> marginal_joint_indices, + std::optional> O_joint, + gpu_array_c Y, gpu_array_c Y_norm, + gpu_array_c similarities, gpu_array_c idx_list, gpu_array_c idx_list_alt, gpu_array_c sort_keys, @@ -550,7 +559,7 @@ static void register_clustering_loop(nb::module_& m) { gpu_array_c cub_temp, gpu_array_c R_out_buffer, gpu_array_c cats_in, - gpu_array_c joint_codes_in, + std::optional> joint_codes_in, gpu_array_c R_in_sum, gpu_array_c R_out_sum, gpu_array_c penalty_buf, gpu_array_c 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 diff --git a/src/rapids_singlecell/preprocessing/_harmony/__init__.py b/src/rapids_singlecell/preprocessing/_harmony/__init__.py index 5e2958e3b..ca04a3216 100644 --- a/src/rapids_singlecell/preprocessing/_harmony/__init__.py +++ b/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,12 +547,6 @@ 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), @@ -570,6 +554,14 @@ def _allocate_clustering_workspace( "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,6 +611,20 @@ 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, @@ -626,10 +632,8 @@ def _clustering( 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], diff --git a/tests/test_managed_memory.py b/tests/test_managed_memory.py index 253e11792..9c78df348 100644 --- a/tests/test_managed_memory.py +++ b/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