From 17df5712beb7d52ae93e568870083bd8040681dd Mon Sep 17 00:00:00 2001 From: dicks1 Date: Thu, 25 Apr 2024 15:04:15 +0200 Subject: [PATCH 01/69] add first functions --- src/rapids_singlecell/_compat.py | 35 ++ .../preprocessing/__init__.py | 2 +- .../preprocessing/_normalize.py | 189 +++++++-- src/rapids_singlecell/preprocessing/_qc.py | 398 ++++++++++++++++++ .../preprocessing/_simple.py | 197 +-------- src/rapids_singlecell/preprocessing/_utils.py | 11 +- 6 files changed, 597 insertions(+), 235 deletions(-) create mode 100644 src/rapids_singlecell/_compat.py create mode 100644 src/rapids_singlecell/preprocessing/_qc.py diff --git a/src/rapids_singlecell/_compat.py b/src/rapids_singlecell/_compat.py new file mode 100644 index 00000000..ecb20551 --- /dev/null +++ b/src/rapids_singlecell/_compat.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +import cupy as cp +from cupyx.scipy.sparse import csr_matrix + +try: + from dask.array import Array as DaskArray +except ImportError: + + class DaskArray: + pass + + +try: + from dask.distributed import Client as DaskClient +except ImportError: + + class DaskClient: + pass + + +def _get_dask_client(client=None): + from dask.distributed import default_client + + if client is None or not isinstance(client, DaskClient): + return default_client() + return client + + +def _meta_dense(dtype): + return cp.zeros([0], dtype=dtype) + + +def _meta_sparse(dtype): + return csr_matrix(cp.array((1.0,), dtype=dtype)) diff --git a/src/rapids_singlecell/preprocessing/__init__.py b/src/rapids_singlecell/preprocessing/__init__.py index 796f9eed..aad8b79a 100644 --- a/src/rapids_singlecell/preprocessing/__init__.py +++ b/src/rapids_singlecell/preprocessing/__init__.py @@ -5,11 +5,11 @@ from ._neighbors import neighbors from ._normalize import log1p, normalize_pearson_residuals, normalize_total from ._pca import pca +from ._qc import calculate_qc_metrics from ._regress_out import regress_out from ._scale import scale from ._scrublet import scrublet, scrublet_simulate_doublets from ._simple import ( - calculate_qc_metrics, filter_cells, filter_genes, filter_highly_variable, diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 2430d3f4..8e92e9db 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -2,12 +2,22 @@ import math import warnings +from functools import singledispatch from typing import TYPE_CHECKING import cupy as cp +from cuml.dask.common.part_utils import _extract_partitions from cupyx.scipy import sparse from scanpy.get import _get_obs_rep, _set_obs_rep +from rapids_singlecell._compat import ( + DaskArray, + DaskClient, + _get_dask_client, + _meta_dense, + _meta_sparse, +) + from ._utils import _check_gpu_X, _check_nonnegative_integers if TYPE_CHECKING: @@ -21,6 +31,7 @@ def normalize_total( layer: int | str = None, inplace: bool = True, copy: bool = False, + client: DaskClient | None = None, ) -> AnnData | sparse.csr_matrix | cp.ndarray | None: """ Normalizes rows in matrix so they sum to `target_sum` @@ -42,6 +53,9 @@ def normalize_total( copy Whether to return a copy or update `adata`. Not compatible with inplace=False. + client + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. + Returns ------- Returns a normalized copy or updates `adata` with a normalized version of \ @@ -54,7 +68,7 @@ def normalize_total( adata = adata.copy() X = _get_obs_rep(adata, layer=layer) - _check_gpu_X(X) + _check_gpu_X(X, allow_dask=True) if not inplace: X = X.copy() @@ -62,49 +76,145 @@ def normalize_total( if sparse.isspmatrix_csc(X): X = X.tocsr() if not target_sum: - if sparse.issparse(X): - from ._kernels._norm_kernel import _get_sparse_sum_major + target_sum = _get_target_sum(X, client=client) + + X = _normalize_total(X, target_sum, client=client) + + if inplace: + _set_obs_rep(adata, X, layer=layer) + + if copy: + return adata + elif not inplace: + return X - counts_per_cell = cp.zeros(X.shape[0], dtype=X.dtype) - sum_kernel = _get_sparse_sum_major(X.dtype) - sum_kernel( - (X.shape[0],), - (64,), - (X.indptr, X.data, counts_per_cell, X.shape[0]), - ) - else: - counts_per_cell = X.sum(axis=1) - target_sum = cp.median(counts_per_cell) - if sparse.isspmatrix_csr(X): +@singledispatch +def _normalize_total(X: cp.ndarray, target_sum: int, client=None) -> cp.ndarray: + from ._kernels._norm_kernel import _mul_dense + + if not X.flags.c_contiguous: + X = cp.asarray(X, order="C") + mul_kernel = _mul_dense(X.dtype) + mul_kernel( + (math.ceil(X.shape[0] / 128),), + (128,), + (X, X.shape[0], X.shape[1], int(target_sum)), + ) + return X + + +@_normalize_total.register(sparse.csr_matrix) +def _(X: sparse.csr_matrix, target_sum: int, client=None) -> sparse.csr_matrix: + from ._kernels._norm_kernel import _mul_csr + + mul_kernel = _mul_csr(X.dtype) + mul_kernel( + (math.ceil(X.shape[0] / 128),), + (128,), + (X.indptr, X.data, X.shape[0], int(target_sum)), + ) + return X + + +@_normalize_total.register(DaskArray) +def _(X: DaskArray, target_sum: int, client=None) -> DaskArray: + client = _get_dask_client(client) + if isinstance(X._meta, sparse.csr_matrix): from ._kernels._norm_kernel import _mul_csr mul_kernel = _mul_csr(X.dtype) - mul_kernel( - (math.ceil(X.shape[0] / 128),), - (128,), - (X.indptr, X.data, X.shape[0], int(target_sum)), - ) + mul_kernel.compile() - else: + def __mul(X_part): + mul_kernel( + (math.ceil(X.shape[0] / 128),), + (128,), + (X_part.indptr, X_part.data, X_part.shape[0], int(target_sum)), + ) + return X_part + + X = X.map_blocks(lambda X: __mul(X), meta=_meta_sparse(X.dtype)) + elif isinstance(X.meta, cp.ndarray): from ._kernels._norm_kernel import _mul_dense - if not X.flags.c_contiguous: - X = cp.asarray(X, order="C") mul_kernel = _mul_dense(X.dtype) - mul_kernel( - (math.ceil(X.shape[0] / 128),), - (128,), - (X, X.shape[0], X.shape[1], int(target_sum)), - ) + mul_kernel.compile() - if inplace: - _set_obs_rep(adata, X, layer=layer) + def __mul(X_part): + mul_kernel( + (math.ceil(X_part.shape[0] / 128),), + (128,), + (X_part, X_part.shape[0], X_part.shape[1], int(target_sum)), + ) + return X_part - if copy: - return adata - elif not inplace: - return X + X = X.map_blocks(lambda X: __mul(X), meta=_meta_dense(X.dtype)) + else: + raise ValueError(f"Cannot normalize {type(X)}") + return X + + +@singledispatch +def _get_target_sum(X: cp.ndarray, client=None) -> int: + return cp.median(X.sum(axis=1)) + + +@_get_target_sum.register(sparse.csr_matrix) +def _(X: sparse.csr_matrix, client=None) -> int: + from ._kernels._norm_kernel import _get_sparse_sum_major + + counts_per_cell = cp.zeros(X.shape[0], dtype=X.dtype) + sum_kernel = _get_sparse_sum_major(X.dtype) + sum_kernel( + (X.shape[0],), + (64,), + (X.indptr, X.data, counts_per_cell, X.shape[0]), + ) + + target_sum = cp.median(counts_per_cell) + return target_sum + + +@_get_target_sum.register(DaskArray) +def _(X: DaskArray, client=None) -> int: + import dask.array as da + + client = _get_dask_client(client) + + if isinstance(X._meta, sparse.csr_matrix): + from ._kernels._norm_kernel import _get_sparse_sum_major + + sum_kernel = _get_sparse_sum_major(X.dtype) + sum_kernel.compile() + + def __sum(X_part): + counts_per_cell = cp.zeros(X_part.shape[0], dtype=X_part.dtype) + sum_kernel( + (X.shape[0],), + (64,), + (X_part.indptr, X_part.data, counts_per_cell, X_part.shape[0]), + ) + return counts_per_cell + + elif isinstance(X._meta, cp.ndarray): + + def __sum(X_part): + return X_part.sum(axis=1) + else: + raise ValueError(f"Cannot compute target sum for {type(X)}") + + parts = client.sync(_extract_partitions, X) + futures = [client.submit(__sum, part, workers=[w]) for w, part in parts] + # Gather results from futures + futures = client.gather(futures) + objs = [] + for i in futures: + objs.append(da.from_array(i, chunks=i.shape)) + + counts_per_cell = da.concatenate(objs).compute() + target_sum = cp.median(counts_per_cell) + return target_sum def log1p( @@ -148,13 +258,20 @@ def log1p( adata = adata.copy() X = _get_obs_rep(adata, layer=layer, obsm=obsm) - _check_gpu_X(X) + _check_gpu_X(X, allow_dask=True) + + if not inplace: + X = X.copy() if isinstance(X, cp.ndarray): X = cp.log1p(X) - else: + elif sparse.issparse(X): X = X.log1p() - + elif isinstance(X, DaskArray): + if isinstance(X._meta, cp.ndarray): + X = X.map_blocks(cp.log1p, meta=_meta_dense(X.dtype)) + elif isinstance(X._meta, sparse.csr_matrix): + X = X.map_blocks(lambda X: X.log1p(), meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} if inplace: _set_obs_rep(adata, X, layer=layer, obsm=obsm) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py new file mode 100644 index 00000000..d87d8614 --- /dev/null +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -0,0 +1,398 @@ +from __future__ import annotations + +import math +from typing import TYPE_CHECKING + +import cupy as cp +from cuml.dask.common.part_utils import _extract_partitions +from cuml.internals.memory_utils import with_cupy_rmm +from cupyx.scipy import sparse +from scanpy.get import _get_obs_rep + +from rapids_singlecell._compat import DaskArray, DaskClient, _get_dask_client + +from ._utils import _check_gpu_X + +if TYPE_CHECKING: + from anndata import AnnData + + +def calculate_qc_metrics( + adata: AnnData, + *, + expr_type: str = "counts", + var_type: str = "genes", + qc_vars: str | list = None, + log1p: bool = True, + layer: str = None, + client: DaskClient | None = None, +) -> None: + """\ + Calculates basic qc Parameters. Calculates number of genes per cell (n_genes) and number of counts per cell (n_counts). + Loosely based on calculate_qc_metrics from scanpy [Wolf et al. 2018]. Updates :attr:`~anndata.AnnData.obs` and :attr:`~anndata.AnnData.var` with columns with qc data. + + Parameters + ---------- + adata + AnnData object + expr_type + Name of kind of values in X. + var_type + The kind of thing the variables are. + qc_vars + Keys for boolean columns of :attr:`~anndata.AnnData.var` which identify variables you could want to control for (e.g. Mito). + Run flag_gene_family first + log1p + Set to `False` to skip computing `log1p` transformed annotations. + layer + If provided, use :attr:`~anndata.AnnData.layers` for expression values instead of :attr:`~anndata.AnnData.X`. + client + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. + + Returns + ------- + adds the following columns in :attr:`~anndata.AnnData.obs` : + `total_{var_type}_by_{expr_type}` + E.g. 'total_genes_by_counts'. Number of genes with positive counts in a cell. + `total_{expr_type}` + E.g. 'total_counts'. Total number of counts for a cell. + for `qc_var` in `qc_vars` + `total_{expr_type}_{qc_var}` + number of counts per qc_var (e.g total counts mitochondrial genes) + `pct_{expr_type}_{qc_var}` + Proportion of counts of qc_var (percent of counts mitochondrial genes) + + adds the following columns in :attr:`~anndata.AnnData.var` : + `total_{expr_type}` + E.g. 'total_counts'. Sum of counts for a gene. + `n_genes_by_{expr_type}` + E.g. 'n_cells_by_counts'. Number of cells this expression is measured in. + `mean_{expr_type}` + E.g. "mean_counts". Mean expression over all cells. + `pct_dropout_by_{expr_type}` + E.g. 'pct_dropout_by_counts'. Percentage of cells this feature does not appear in. + + """ + + X = _get_obs_rep(adata, layer=layer) + + _check_gpu_X(X, allow_dask=True) + + sums_cells, sums_genes, cell_ex, gene_ex = _first_pass_qc(X, client=client) + # .var + adata.var[f"n_cells_by_{expr_type}"] = cp.asnumpy(gene_ex) + adata.var[f"total_{expr_type}"] = cp.asnumpy(sums_genes) + mean_array = sums_genes / adata.n_obs + adata.var[f"mean_{expr_type}"] = cp.asnumpy(mean_array) + adata.var[f"pct_dropout_by_{expr_type}"] = cp.asnumpy( + (1 - gene_ex / adata.n_obs) * 100 + ) + if log1p: + adata.var[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_genes)) + adata.var[f"log1p_mean_{expr_type}"] = cp.asnumpy(cp.log1p(mean_array)) + # .obs + adata.obs[f"n_{var_type}_by_{expr_type}"] = cp.asnumpy(cell_ex) + adata.obs[f"total_{expr_type}"] = cp.asnumpy(sums_cells) + if log1p: + adata.obs[f"log1p_n_{var_type}_by_{expr_type}"] = cp.asnumpy(cp.log1p(cell_ex)) + adata.obs[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_cells)) + + if qc_vars: + if isinstance(qc_vars, str): + qc_vars = [qc_vars] + for qc_var in qc_vars: + mask = cp.array(adata.var[qc_var], dtype=cp.bool_) + sums_cells_sub = _second_pass_qc(X, mask, client=client) + + adata.obs[f"total_{expr_type}_{qc_var}"] = cp.asnumpy(sums_cells_sub) + adata.obs[f"pct_{expr_type}_{qc_var}"] = cp.asnumpy( + sums_cells_sub / sums_cells * 100 + ) + if log1p: + adata.obs[f"log1p_total_{expr_type}_{qc_var}"] = cp.asnumpy( + cp.log1p(sums_cells_sub) + ) + + +def _first_pass_qc(X, client=None): + if isinstance(X, DaskArray): + return _first_pass_qc_dask(X, client=client) + + sums_cells = cp.zeros(X.shape[0], dtype=X.dtype) + sums_genes = cp.zeros(X.shape[1], dtype=X.dtype) + cell_ex = cp.zeros(X.shape[0], dtype=cp.int32) + gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) + if sparse.issparse(X): + if sparse.isspmatrix_csr(X): + from ._kernels._qc_kernels import _sparse_qc_csr + + block = (32,) + grid = (int(math.ceil(X.shape[0] / block[0])),) + sparse_qc_csr = _sparse_qc_csr(X.data.dtype) + sparse_qc_csr( + grid, + block, + ( + X.indptr, + X.indices, + X.data, + sums_cells, + sums_genes, + cell_ex, + gene_ex, + X.shape[0], + ), + ) + elif sparse.isspmatrix_csc(X): + from ._kernels._qc_kernels import _sparse_qc_csc + + block = (32,) + grid = (int(math.ceil(X.shape[1] / block[0])),) + sparse_qc_csc = _sparse_qc_csc(X.data.dtype) + sparse_qc_csc( + grid, + block, + ( + X.indptr, + X.indices, + X.data, + sums_cells, + sums_genes, + cell_ex, + gene_ex, + X.shape[1], + ), + ) + else: + raise ValueError("Please use a csr or csc matrix") + else: + from ._kernels._qc_kernels import _sparse_qc_dense + + if not X.flags.c_contiguous: + X = cp.asarray(X, order="C") + block = (16, 16) + grid = ( + int(math.ceil(X.shape[0] / block[0])), + int(math.ceil(X.shape[1] / block[1])), + ) + sparse_qc_dense = _sparse_qc_dense(X.dtype) + sparse_qc_dense( + grid, + block, + (X, sums_cells, sums_genes, cell_ex, gene_ex, X.shape[0], X.shape[1]), + ) + return sums_cells, sums_genes, cell_ex, gene_ex + + +@with_cupy_rmm +def _first_pass_qc_dask(X, client=None): + import dask + from cuml.dask.common.part_utils import _extract_partitions + + client = _get_dask_client(client) + + if isinstance(X._meta, sparse.csr_matrix): + from ._kernels._qc_kernels import _sparse_qc_csr + + sparse_qc_csr = _sparse_qc_csr(X.dtype) + sparse_qc_csr.compile() + + def __qc_calc(X_part): + sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) + sums_genes = cp.zeros((X_part.shape[1], 1), dtype=X_part.dtype) + cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) + gene_ex = cp.zeros((X_part.shape[1], 1), dtype=cp.int32) + block = (32,) + grid = (int(math.ceil(X_part.shape[0] / block[0])),) + sparse_qc_csr( + grid, + block, + ( + X_part.indptr, + X_part.indices, + X_part.data, + sums_cells, + sums_genes, + cell_ex, + gene_ex, + X_part.shape[0], + ), + ) + return sums_cells, sums_genes, cell_ex, gene_ex + + elif isinstance(X._meta, cp.ndarray): + from ._kernels._qc_kernels import _sparse_qc_dense + + _sparse_qc_dense = _sparse_qc_dense(X.dtype) + _sparse_qc_dense.compile() + + def __qc_calc(X_part): + sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) + sums_genes = cp.zeros((X_part.shape[1], 1), dtype=X_part.dtype) + cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) + gene_ex = cp.zeros((X_part.shape[1], 1), dtype=cp.int32) + if not X_part.flags.c_contiguous: + X_part = cp.asarray(X_part, order="C") + block = (16, 16) + grid = ( + int(math.ceil(X_part.shape[0] / block[0])), + int(math.ceil(X_part.shape[1] / block[1])), + ) + sparse_qc_dense = _sparse_qc_dense(X.dtype) + sparse_qc_dense( + grid, + block, + ( + X_part, + sums_cells, + sums_genes, + cell_ex, + gene_ex, + X_part.shape[0], + X_part.shape[1], + ), + ) + return sums_cells, sums_genes, cell_ex, gene_ex + else: + raise ValueError( + "Please use a cupy csr_matrix or cp.ndarray. csc_matrix are not supported with dask." + ) + + parts = client.sync(_extract_partitions, X) + futures = [client.submit(__qc_calc, part, workers=[w]) for w, part in parts] + # Gather results from futures + results = client.gather(futures) + + # Initialize lists to hold the Dask arrays + sums_cells_objs = [] + sums_genes_objs = [] + cell_ex_objs = [] + gene_ex_objs = [] + + # Process each result + for sums_cells, sums_genes, cell_ex, gene_ex in results: + # Append the arrays to their respective lists as Dask arrays + sums_cells_objs.append( + dask.array.from_array(sums_cells, chunks=sums_cells.shape) + ) + sums_genes_objs.append( + dask.array.from_array(sums_genes, chunks=sums_genes.shape) + ) + cell_ex_objs.append(dask.array.from_array(cell_ex, chunks=cell_ex.shape)) + gene_ex_objs.append(dask.array.from_array(gene_ex, chunks=gene_ex.shape)) + sums_cells = dask.array.concatenate(sums_cells_objs) + sums_genes = dask.array.concatenate(sums_genes_objs, axis=1).sum(axis=1) + cell_ex = dask.array.concatenate(cell_ex_objs) + gene_ex = dask.array.concatenate(gene_ex_objs, axis=1).sum(axis=1) + sums_cells, sums_genes, cell_ex, gene_ex = dask.compute( + sums_cells, sums_genes, cell_ex, gene_ex + ) + return sums_cells.ravel(), sums_genes.ravel(), cell_ex.ravel(), gene_ex.ravel() + + +def _second_pass_qc(X, mask, client=None): + if isinstance(X, DaskArray): + return _second_pass_qc_dask(X, mask, client=client) + sums_cells_sub = cp.zeros(X.shape[0], dtype=X.dtype) + if sparse.issparse(X): + if sparse.isspmatrix_csr(X): + from ._kernels._qc_kernels import _sparse_qc_csr_sub + + block = (32,) + grid = (int(math.ceil(X.shape[0] / block[0])),) + sparse_qc_csr_sub = _sparse_qc_csr_sub(X.data.dtype) + sparse_qc_csr_sub( + grid, + block, + (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[0]), + ) + elif sparse.isspmatrix_csc(X): + from ._kernels._qc_kernels import _sparse_qc_csc_sub + + block = (32,) + grid = (int(math.ceil(X.shape[1] / block[0])),) + sparse_qc_csc_sub = _sparse_qc_csc_sub(X.data.dtype) + sparse_qc_csc_sub( + grid, + block, + (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[1]), + ) + + else: + from ._kernels._qc_kernels import _sparse_qc_dense_sub + + block = (16, 16) + grid = ( + int(math.ceil(X.shape[0] / block[0])), + int(math.ceil(X.shape[1] / block[1])), + ) + sparse_qc_dense_sub = _sparse_qc_dense_sub(X.dtype) + sparse_qc_dense_sub( + grid, block, (X, sums_cells_sub, mask, X.shape[0], X.shape[1]) + ) + return sums_cells_sub + + +@with_cupy_rmm +def _second_pass_qc_dask(X, mask, client=None): + import dask + + client = _get_dask_client(client) + + if isinstance(X._meta, sparse.csr_matrix): + from ._kernels._qc_kernels import _sparse_qc_csr_sub + + sparse_qc_csr = _sparse_qc_csr_sub(X.dtype) + sparse_qc_csr.compile() + + def __qc_calc(X_part): + sums_cells_sub = cp.zeros((X_part.shape[0]), dtype=X_part.dtype) + block = (32,) + grid = (int(math.ceil(X_part.shape[0] / block[0])),) + sparse_qc_csr( + grid, + block, + ( + X_part.indptr, + X_part.indices, + X_part.data, + sums_cells_sub, + mask, + X_part.shape[0], + ), + ) + return sums_cells_sub + + elif isinstance(X._meta, cp.ndarray): + from ._kernels._qc_kernels import _sparse_qc_dense + + _sparse_qc_dense = _sparse_qc_dense(X.dtype) + _sparse_qc_dense.compile() + + def __qc_calc(X_part): + sums_cells_sub = cp.zeros((X_part.shape[0]), dtype=X_part.dtype) + if not X_part.flags.c_contiguous: + X_part = cp.asarray(X_part, order="C") + block = (16, 16) + grid = ( + int(math.ceil(X_part.shape[0] / block[0])), + int(math.ceil(X_part.shape[1] / block[1])), + ) + sparse_qc_dense = _sparse_qc_dense(X.dtype) + sparse_qc_dense( + grid, + block, + (X_part, sums_cells_sub, mask, X_part.shape[0], X_part.shape[1]), + ) + return sums_cells_sub + + parts = client.sync(_extract_partitions, X) + futures = [client.submit(__qc_calc, part, workers=[w]) for w, part in parts] + # Gather results from futures + futures = client.gather(futures) + objs = [] + for i in futures: + objs.append(dask.array.from_array(i, chunks=i.shape)) + + sums_cells_sub = dask.array.concatenate(objs).compute() + return sums_cells_sub.ravel() diff --git a/src/rapids_singlecell/preprocessing/_simple.py b/src/rapids_singlecell/preprocessing/_simple.py index e0db508f..3e63445f 100644 --- a/src/rapids_singlecell/preprocessing/_simple.py +++ b/src/rapids_singlecell/preprocessing/_simple.py @@ -1,211 +1,16 @@ from __future__ import annotations -import math from typing import TYPE_CHECKING import cupy as cp import numpy as np -from cupyx.scipy import sparse -from scanpy.get import _get_obs_rep -from ._utils import _check_gpu_X +from ._qc import calculate_qc_metrics if TYPE_CHECKING: from anndata import AnnData -def calculate_qc_metrics( - adata: AnnData, - *, - expr_type: str = "counts", - var_type: str = "genes", - qc_vars: str | list = None, - log1p: bool = True, - layer: str = None, -) -> None: - """\ - Calculates basic qc Parameters. Calculates number of genes per cell (n_genes) and number of counts per cell (n_counts). - Loosely based on calculate_qc_metrics from scanpy [Wolf et al. 2018]. Updates :attr:`~anndata.AnnData.obs` and :attr:`~anndata.AnnData.var` with columns with qc data. - - Parameters - ---------- - adata - AnnData object - expr_type - Name of kind of values in X. - var_type - The kind of thing the variables are. - qc_vars - Keys for boolean columns of :attr:`~anndata.AnnData.var` which identify variables you could want to control for (e.g. Mito). - Run flag_gene_family first - log1p - Set to `False` to skip computing `log1p` transformed annotations. - layer - If provided, use :attr:`~anndata.AnnData.layers` for expression values instead of :attr:`~anndata.AnnData.X`. - - Returns - ------- - adds the following columns in :attr:`~anndata.AnnData.obs` : - `total_{var_type}_by_{expr_type}` - E.g. 'total_genes_by_counts'. Number of genes with positive counts in a cell. - `total_{expr_type}` - E.g. 'total_counts'. Total number of counts for a cell. - for `qc_var` in `qc_vars` - `total_{expr_type}_{qc_var}` - number of counts per qc_var (e.g total counts mitochondrial genes) - `pct_{expr_type}_{qc_var}` - Proportion of counts of qc_var (percent of counts mitochondrial genes) - - adds the following columns in :attr:`~anndata.AnnData.var` : - `total_{expr_type}` - E.g. 'total_counts'. Sum of counts for a gene. - `n_genes_by_{expr_type}` - E.g. 'n_cells_by_counts'. Number of cells this expression is measured in. - `mean_{expr_type}` - E.g. "mean_counts". Mean expression over all cells. - `pct_dropout_by_{expr_type}` - E.g. 'pct_dropout_by_counts'. Percentage of cells this feature does not appear in. - - """ - - X = _get_obs_rep(adata, layer=layer) - - _check_gpu_X(X) - - sums_cells = cp.zeros(X.shape[0], dtype=X.dtype) - sums_genes = cp.zeros(X.shape[1], dtype=X.dtype) - cell_ex = cp.zeros(X.shape[0], dtype=cp.int32) - gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) - if sparse.issparse(X): - if sparse.isspmatrix_csr(X): - from ._kernels._qc_kernels import _sparse_qc_csr - - block = (32,) - grid = (int(math.ceil(X.shape[0] / block[0])),) - sparse_qc_csr = _sparse_qc_csr(X.data.dtype) - sparse_qc_csr( - grid, - block, - ( - X.indptr, - X.indices, - X.data, - sums_cells, - sums_genes, - cell_ex, - gene_ex, - X.shape[0], - ), - ) - elif sparse.isspmatrix_csc(X): - from ._kernels._qc_kernels import _sparse_qc_csc - - block = (32,) - grid = (int(math.ceil(X.shape[1] / block[0])),) - sparse_qc_csc = _sparse_qc_csc(X.data.dtype) - sparse_qc_csc( - grid, - block, - ( - X.indptr, - X.indices, - X.data, - sums_cells, - sums_genes, - cell_ex, - gene_ex, - X.shape[1], - ), - ) - else: - raise ValueError("Please use a csr or csc matrix") - else: - from ._kernels._qc_kernels import _sparse_qc_dense - - if not X.flags.c_contiguous: - X = cp.asarray(X, order="C") - block = (16, 16) - grid = ( - int(math.ceil(X.shape[0] / block[0])), - int(math.ceil(X.shape[1] / block[1])), - ) - sparse_qc_dense = _sparse_qc_dense(X.dtype) - sparse_qc_dense( - grid, - block, - (X, sums_cells, sums_genes, cell_ex, gene_ex, X.shape[0], X.shape[1]), - ) - - # .var - adata.var[f"n_cells_by_{expr_type}"] = cp.asnumpy(gene_ex) - adata.var[f"total_{expr_type}"] = cp.asnumpy(sums_genes) - mean_array = sums_genes / adata.n_obs - adata.var[f"mean_{expr_type}"] = cp.asnumpy(mean_array) - adata.var[f"pct_dropout_by_{expr_type}"] = cp.asnumpy( - (1 - gene_ex / adata.n_obs) * 100 - ) - if log1p: - adata.var[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_genes)) - adata.var[f"log1p_mean_{expr_type}"] = cp.asnumpy(cp.log1p(mean_array)) - # .obs - adata.obs[f"n_{var_type}_by_{expr_type}"] = cp.asnumpy(cell_ex) - adata.obs[f"total_{expr_type}"] = cp.asnumpy(sums_cells) - if log1p: - adata.obs[f"log1p_n_{var_type}_by_{expr_type}"] = cp.asnumpy(cp.log1p(cell_ex)) - adata.obs[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_cells)) - - if qc_vars: - if isinstance(qc_vars, str): - qc_vars = [qc_vars] - for qc_var in qc_vars: - sums_cells_sub = cp.zeros(X.shape[0], dtype=X.dtype) - mask = cp.array(adata.var[qc_var], dtype=cp.bool_) - if sparse.issparse(X): - if sparse.isspmatrix_csr(X): - from ._kernels._qc_kernels import _sparse_qc_csr_sub - - block = (32,) - grid = (int(math.ceil(X.shape[0] / block[0])),) - sparse_qc_csr_sub = _sparse_qc_csr_sub(X.data.dtype) - sparse_qc_csr_sub( - grid, - block, - (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[0]), - ) - elif sparse.isspmatrix_csc(X): - from ._kernels._qc_kernels import _sparse_qc_csc_sub - - block = (32,) - grid = (int(math.ceil(X.shape[1] / block[0])),) - sparse_qc_csc_sub = _sparse_qc_csc_sub(X.data.dtype) - sparse_qc_csc_sub( - grid, - block, - (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[1]), - ) - - else: - from ._kernels._qc_kernels import _sparse_qc_dense_sub - - block = (16, 16) - grid = ( - int(math.ceil(X.shape[0] / block[0])), - int(math.ceil(X.shape[1] / block[1])), - ) - sparse_qc_dense_sub = _sparse_qc_dense_sub(X.dtype) - sparse_qc_dense_sub( - grid, block, (X, sums_cells_sub, mask, X.shape[0], X.shape[1]) - ) - adata.obs[f"total_{expr_type}_{qc_var}"] = cp.asnumpy(sums_cells_sub) - adata.obs[f"pct_{expr_type}_{qc_var}"] = cp.asnumpy( - sums_cells_sub / sums_cells * 100 - ) - if log1p: - adata.obs[f"log1p_total_{expr_type}_{qc_var}"] = cp.asnumpy( - cp.log1p(sums_cells_sub) - ) - - def flag_gene_family( adata: AnnData, *, diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index f68de056..e3d9d684 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -5,6 +5,8 @@ import cupy as cp from cupyx.scipy.sparse import issparse, isspmatrix_csc, isspmatrix_csr +from rapids_singlecell._compat import DaskArray + def _mean_var_major(X, major, minor): from ._kernels._mean_var_kernel import _get_mean_var_major @@ -85,15 +87,20 @@ def _check_nonnegative_integers(X): return True -def _check_gpu_X(X, require_cf=False): +def _check_gpu_X(X, require_cf=False, allow_dask=False): if isinstance(X, cp.ndarray): return True elif issparse(X): - if X.has_canonical_format or not require_cf: + if not require_cf: + return True + elif X.has_canonical_format: return True else: X.sort_indices() X.sum_duplicates() + elif allow_dask: + if isinstance(X, DaskArray): + return _check_gpu_X(X._meta) else: raise TypeError( "The input is not a CuPy ndarray or CuPy sparse matrix. " From 40167cae98bf348eef7114b53c5e5a4777418f67 Mon Sep 17 00:00:00 2001 From: dicks1 Date: Thu, 25 Apr 2024 17:06:34 +0200 Subject: [PATCH 02/69] add hvg part1 --- src/rapids_singlecell/preprocessing/_hvg.py | 23 ++- src/rapids_singlecell/preprocessing/_utils.py | 131 +++++++++++++++++- 2 files changed, 150 insertions(+), 4 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 7c82d71b..5ed12d60 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -7,9 +7,11 @@ import cupy as cp import numpy as np import pandas as pd -from cupyx.scipy.sparse import issparse +from cupyx.scipy.sparse import csr_matrix, issparse from scanpy.get import _get_obs_rep +from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse + from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var if TYPE_CHECKING: @@ -43,6 +45,7 @@ def highly_variable_genes( chunksize: int = 1000, n_samples: int = 10000, batch_key: str = None, + client: DaskClient | None = None, ) -> None: """\ Annotate highly variable genes. @@ -185,6 +188,7 @@ def highly_variable_genes( n_top_genes=n_top_genes, n_bins=n_bins, flavor=flavor, + client=client, ) else: adata.obs[batch_key] = adata.obs[batch_key].astype("category") @@ -197,6 +201,13 @@ def highly_variable_genes( if not isinstance(X, cp.ndarray): X_batch = X[adata.obs[batch_key] == batch,].tocsc() nnz_per_gene = cp.diff(X_batch.indptr).ravel() + elif isinstance(X, DaskArray): + X_batch = X[adata.obs[batch_key] == batch,] + if isinstance(X._meta, cp.ndarray): + nnz_per_gene = (X > 0).sum(axis=0).compute().ravel() + elif isinstance(X._meta, csr_matrix): + pass + # to do implement this else: X_batch = X[adata.obs[batch_key] == batch,].copy() nnz_per_gene = cp.sum(X_batch > 0, axis=0).ravel() @@ -216,6 +227,7 @@ def highly_variable_genes( n_top_genes=n_top_genes, n_bins=n_bins, flavor=flavor, + client=client, ) hvg_inter["gene"] = inter_genes missing_hvg = pd.DataFrame( @@ -299,6 +311,7 @@ def _highly_variable_genes_single_batch( n_top_genes=None, flavor="seurat", n_bins=20, + client=None, ): """\ See `highly_variable_genes`. @@ -311,9 +324,15 @@ def _highly_variable_genes_single_batch( if flavor == "seurat": if issparse(X): X = X.expm1() + elif isinstance(X, DaskArray): + if isinstance(X._meta, cp.ndarray): + X = X.map_blocks(cp.expm1, meta=_meta_dense(X.dtype)) + elif isinstance(X._meta, csr_matrix): + X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) else: X = cp.expm1(X) - mean, var = _get_mean_var(X, axis=0) + + mean, var = _get_mean_var(X, axis=0, client=client) mean[mean == 0] = 1e-12 disp = var / mean if flavor == "seurat": # logarithmized mean as in Seurat diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index e3d9d684..a854a4f8 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -3,9 +3,11 @@ import math import cupy as cp +from cuml.dask.common.part_utils import _extract_partitions +from cuml.internals.memory_utils import with_cupy_rmm from cupyx.scipy.sparse import issparse, isspmatrix_csc, isspmatrix_csr -from rapids_singlecell._compat import DaskArray +from rapids_singlecell._compat import DaskArray, _get_dask_client def _mean_var_major(X, major, minor): @@ -40,7 +42,121 @@ def _mean_var_minor(X, major, minor): return mean, var -def _get_mean_var(X, axis=0): +@with_cupy_rmm +def _mean_var_minor_dask(X, major, minor, client=None): + """ + Implements sum operation for dask array when the backend is cupy sparse csr matrix + """ + import dask.array as da + + from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( + _get_mean_var_minor, + ) + + client = _get_dask_client(client) + + get_mean_var_minor = _get_mean_var_minor(X.dtype) + get_mean_var_minor.compile() + + def __mean_var(X_part, minor, major): + mean = cp.zeros((minor, 1), dtype=cp.float64) + var = cp.zeros((minor, 1), dtype=cp.float64) + block = (32,) + grid = (int(math.ceil(X_part.nnz / block[0])),) + get_mean_var_minor( + grid, block, (X_part.indices, X_part.data, mean, var, major, X_part.nnz) + ) + return mean, var + + parts = client.sync(_extract_partitions, X) + futures = [ + client.submit(__mean_var, part, minor, major, workers=[w]) for w, part in parts + ] + # Gather results from futures + results = client.gather(futures) + + # Initialize lists to hold the Dask arrays + means_objs = [] + var_objs = [] + + # Process each result + for means, vars in results: + # Append the arrays to their respective lists as Dask arrays + means_objs.append(da.from_array(means, chunks=means.shape)) + var_objs.append(da.from_array(vars, chunks=vars.shape)) + mean = da.concatenate(means_objs, axis=1).sum(axis=1) + var = da.concatenate(var_objs, axis=1).sum(axis=1) + mean, var = da.compute(mean, var) + mean, var = mean.ravel(), var.ravel() + var = (var - mean**2) * (major / (major - 1)) + return mean, var + + +# todo: Implement this dynamically for csc matrix as well +@with_cupy_rmm +def _mean_var_major_dask(X, major, minor, client=None): + """ + Implements sum operation for dask array when the backend is cupy sparse csr matrix + """ + import dask.array as da + + from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( + _get_mean_var_major, + ) + + client = _get_dask_client(client) + + get_mean_var_major = _get_mean_var_major(X.dtype) + get_mean_var_major.compile() + + def __mean_var(X_part, minor, major): + mean = cp.zeros(X_part.shape[0], dtype=cp.float64) + var = cp.zeros(X_part.shape[0], dtype=cp.float64) + block = (64,) + grid = (X_part.shape[0],) + get_mean_var_major( + grid, + block, + ( + X_part.indptr, + X_part.indices, + X_part.data, + mean, + var, + X_part.shape[0], + minor, + ), + ) + return mean, var + + parts = client.sync(_extract_partitions, X) + futures = [ + client.submit(__mean_var, part, minor, major, workers=[w]) for w, part in parts + ] + # Gather results from futures + results = client.gather(futures) + + # Initialize lists to hold the Dask arrays + means_objs = [] + var_objs = [] + + # Process each result + for means, vars_ in results: + # Append the arrays to their respective lists as Dask arrays + means_objs.append(da.from_array(means, chunks=means.shape)) + var_objs.append(da.from_array(vars_, chunks=vars_.shape)) + mean = da.concatenate(means_objs) + var = da.concatenate(var_objs) + mean, var = da.compute(mean, var) + mean, var = mean.ravel(), var.ravel() + mean = mean / minor + var = var / minor + var -= cp.power(mean, 2) + var *= minor / (minor - 1) + return mean, var + + +def _get_mean_var(X, axis=0, client=None): if issparse(X): if axis == 0: if isspmatrix_csr(X): @@ -65,6 +181,17 @@ def _get_mean_var(X, axis=0): major = X.shape[1] minor = X.shape[0] mean, var = _mean_var_minor(X, major, minor) + elif isinstance(X, DaskArray): + if isspmatrix_csr(X._meta): + if axis == 0: + major = X.shape[0] + minor = X.shape[1] + mean, var = _mean_var_minor_dask(X, major, minor, client) + if axis == 1: + major = X.shape[0] + minor = X.shape[1] + mean, var = _mean_var_major_dask(X, major, minor, client) + else: mean = X.mean(axis=axis) var = X.var(axis=axis) From 0cdb85d412062d1cbf67eb43defba1c1b48947f7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 2 May 2024 11:58:43 +0000 Subject: [PATCH 03/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_hvg.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 525e616f..9ae8356e 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -12,8 +12,8 @@ from cupyx.scipy.sparse import csr_matrix, issparse from scanpy.get import _get_obs_rep - from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse + from ._simple import calculate_qc_metrics from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var @@ -271,7 +271,6 @@ def _highly_variable_genes_single_batch( n_bins: int = 20, flavor: Literal["seurat", "cell_ranger"] = "seurat", ) -> pd.DataFrame: - """\ See `highly_variable_genes`. From 48b68f6ed13e441be96fc362847dc7e029dbb5a0 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 2 May 2024 15:21:01 +0200 Subject: [PATCH 04/69] reset to main for hvg --- src/rapids_singlecell/preprocessing/_hvg.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 9ae8356e..b6a9f80f 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -9,12 +9,10 @@ import cupy as cp import numpy as np import pandas as pd -from cupyx.scipy.sparse import csr_matrix, issparse +from cupyx.scipy.sparse import issparse from scanpy.get import _get_obs_rep -from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse - -from ._simple import calculate_qc_metrics +from ._qc import calculate_qc_metrics from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var if TYPE_CHECKING: @@ -49,7 +47,6 @@ def highly_variable_genes( chunksize: int = 1000, n_samples: int = 10000, batch_key: str = None, - client: DaskClient | None = None, ) -> None: """\ Annotate highly variable genes. @@ -289,15 +286,9 @@ def _highly_variable_genes_single_batch( X = X.copy() if issparse(X): X = X.expm1() - elif isinstance(X, DaskArray): - if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(cp.expm1, meta=_meta_dense(X.dtype)) - elif isinstance(X._meta, csr_matrix): - X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) else: X = cp.expm1(X) - - mean, var = _get_mean_var(X, axis=0, client=client) + mean, var = _get_mean_var(X, axis=0) mean[mean == 0] = 1e-12 disp = var / mean if flavor == "seurat": # logarithmized mean as in Seurat From 886cafa54ccb90bfb84fe221d25f343a59dd66fb Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 2 May 2024 16:31:29 +0200 Subject: [PATCH 05/69] add support for hvg --- src/rapids_singlecell/preprocessing/_hvg.py | 45 ++++++++++++++++----- 1 file changed, 35 insertions(+), 10 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index b6a9f80f..2e97332c 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -9,9 +9,11 @@ import cupy as cp import numpy as np import pandas as pd -from cupyx.scipy.sparse import issparse +from cupyx.scipy.sparse import csr_matrix, issparse from scanpy.get import _get_obs_rep +from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse + from ._qc import calculate_qc_metrics from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var @@ -47,6 +49,7 @@ def highly_variable_genes( chunksize: int = 1000, n_samples: int = 10000, batch_key: str = None, + client: DaskClient | None = None, ) -> None: """\ Annotate highly variable genes. @@ -116,6 +119,8 @@ def highly_variable_genes( of enrichment of zeros for each gene (only for `flavor='poisson_gene_selection'`). batch_key If specified, highly-variable genes are selected within each batch separately and merged. + client + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. Returns ------- @@ -188,7 +193,12 @@ def highly_variable_genes( if batch_key is None: df = _highly_variable_genes_single_batch( - adata, layer=layer, cutoff=cutoff, n_bins=n_bins, flavor=flavor + adata, + layer=layer, + cutoff=cutoff, + n_bins=n_bins, + flavor=flavor, + client=client, ) else: df = _highly_variable_genes_batched( @@ -198,6 +208,7 @@ def highly_variable_genes( cutoff=cutoff, n_bins=n_bins, flavor=flavor, + client=client, ) adata.uns["hvg"] = {"flavor": flavor} @@ -267,6 +278,7 @@ def _highly_variable_genes_single_batch( cutoff: _Cutoffs | int, n_bins: int = 20, flavor: Literal["seurat", "cell_ranger"] = "seurat", + client: DaskClient | None = None, ) -> pd.DataFrame: """\ See `highly_variable_genes`. @@ -277,18 +289,25 @@ def _highly_variable_genes_single_batch( `highly_variable`, `means`, `dispersions`, and `dispersions_norm`. """ X = _get_obs_rep(adata, layer=layer) - + _check_gpu_X(X, allow_dask=True) if hasattr(X, "_view_args"): # AnnData array view # For compatibility with anndata<0.9 X = X.copy() # Doesn't actually copy memory, just removes View class wrapper if flavor == "seurat": - X = X.copy() - if issparse(X): - X = X.expm1() + if isinstance(X, DaskArray): + if isinstance(X._meta, cp.ndarray): + X = X.map_blocks(cp.expm1, meta=_meta_dense(X.dtype)) + elif isinstance(X._meta, csr_matrix): + X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) else: - X = cp.expm1(X) - mean, var = _get_mean_var(X, axis=0) + X = X.copy() + if issparse(X): + X = X.expm1() + else: + X = cp.expm1(X) + + mean, var = _get_mean_var(X, axis=0, client=client) mean[mean == 0] = 1e-12 disp = var / mean if flavor == "seurat": # logarithmized mean as in Seurat @@ -407,6 +426,7 @@ def _highly_variable_genes_batched( n_bins: int, flavor: Literal["seurat", "cell_ranger"], cutoff: _Cutoffs | int, + client: DaskClient | None = None, ) -> pd.DataFrame: adata._sanitize() batches = adata.obs[batch_key].cat.categories @@ -415,12 +435,17 @@ def _highly_variable_genes_batched( for batch in batches: adata_subset = adata[adata.obs[batch_key] == batch] - calculate_qc_metrics(adata_subset, layer=layer) + calculate_qc_metrics(adata_subset, layer=layer, client=client) filt = adata_subset.var["n_cells_by_counts"].to_numpy() > 0 adata_subset = adata_subset[:, filt] hvg = _highly_variable_genes_single_batch( - adata_subset, layer=layer, cutoff=cutoff, n_bins=n_bins, flavor=flavor + adata_subset, + layer=layer, + cutoff=cutoff, + n_bins=n_bins, + flavor=flavor, + client=client, ) hvg.reset_index(drop=False, inplace=True, names=["gene"]) From d7bf01ea60096e6893fa905523295ab5556aabef Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 2 May 2024 16:46:08 +0200 Subject: [PATCH 06/69] first pass pca --- src/rapids_singlecell/preprocessing/_pca.py | 283 ++++-------------- .../_sparse_pca/_dask_sparse_pca.py | 203 +++++++++++++ .../_kernels/_pca_sparse_kernel.py | 0 .../preprocessing/_sparse_pca/_sparse_pca.py | 190 ++++++++++++ 4 files changed, 455 insertions(+), 221 deletions(-) create mode 100644 src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py rename src/rapids_singlecell/preprocessing/{ => _sparse_pca}/_kernels/_pca_sparse_kernel.py (100%) create mode 100644 src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 1db6355c..61928831 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -5,14 +5,15 @@ import cupy as cp import numpy as np -from cuml.decomposition import PCA, TruncatedSVD +from cuml.decomposition import TruncatedSVD from cuml.internals.input_utils import sparse_scipy_to_cp -from cupyx.scipy.sparse import csr_matrix, isspmatrix_csr +from cupyx.scipy.sparse import csr_matrix from cupyx.scipy.sparse import issparse as cpissparse from scanpy._utils import Empty, _empty from scanpy.preprocessing._pca import _handle_mask_var from scipy.sparse import issparse +from rapids_singlecell._compat import DaskArray, DaskClient from rapids_singlecell.get import _get_obs_rep if TYPE_CHECKING: @@ -34,6 +35,7 @@ def pca( copy: bool = False, chunked: bool = False, chunk_size: int = None, + client: DaskClient | None = None, ) -> None | AnnData: """ Performs PCA using the cuml decomposition function. @@ -88,6 +90,9 @@ def pca( Number of observations to include in each chunk. \ Required if `chunked=True` was passed. + client + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. + Returns ------- adds fields to `adata`: @@ -126,51 +131,71 @@ def pca( n_comps = min_dim - 1 else: n_comps = 50 + if isinstance(X, DaskArray): + if chunked: + raise ValueError( + "Dask arrays are not supported for chunked PCA computation." + ) + if isinstance(X._meta, cp.ndarray): + from cuml.dask.decomposition import PCA - if chunked: - from cuml.decomposition import IncrementalPCA + if svd_solver == "auto": + svd_solver = "jacobi" + pca_func = PCA(n_components=n_comps, svd_solver=svd_solver) + X_pca = pca_func.fit_transform(X) + elif isinstance(X._meta, csr_matrix): + from cuml.dask.decomposition import PCA_sparse - X_pca = np.zeros((X.shape[0], n_comps), X.dtype) + pca_func = PCA_sparse(n_components=n_comps) + X_pca = pca_func.fit_transform(X) - pca_func = IncrementalPCA( - n_components=n_comps, output_type="numpy", batch_size=chunk_size - ) - pca_func.fit(X) - - n_batches = math.ceil(X.shape[0] / chunk_size) - for batch in range(n_batches): - start_idx = batch * chunk_size - stop_idx = min(batch * chunk_size + chunk_size, X.shape[0]) - chunk = X[start_idx:stop_idx, :] - if issparse(chunk) or cpissparse(chunk): - chunk = chunk.toarray() - X_pca[start_idx:stop_idx] = pca_func.transform(chunk) else: - if zero_center: - if cpissparse(X) or issparse(X): - if issparse(X): - X = sparse_scipy_to_cp(X, dtype=X.dtype) - X = csr_matrix(X) - pca_func = PCA_sparse(n_components=n_comps) - X_pca = pca_func.fit_transform(X) - else: - pca_func = PCA( + if chunked: + from cuml.decomposition import IncrementalPCA + + X_pca = np.zeros((X.shape[0], n_comps), X.dtype) + + pca_func = IncrementalPCA( + n_components=n_comps, output_type="numpy", batch_size=chunk_size + ) + pca_func.fit(X) + + n_batches = math.ceil(X.shape[0] / chunk_size) + for batch in range(n_batches): + start_idx = batch * chunk_size + stop_idx = min(batch * chunk_size + chunk_size, X.shape[0]) + chunk = X[start_idx:stop_idx, :] + if issparse(chunk) or cpissparse(chunk): + chunk = chunk.toarray() + X_pca[start_idx:stop_idx] = pca_func.transform(chunk) + else: + if zero_center: + if cpissparse(X) or issparse(X): + if issparse(X): + X = sparse_scipy_to_cp(X, dtype=X.dtype) + X = csr_matrix(X) + from ._sparse_pca._sparse_pca import PCA_sparse + + pca_func = PCA_sparse(n_components=n_comps) + X_pca = pca_func.fit_transform(X) + else: + pca_func = PCA( + n_components=n_comps, + svd_solver=svd_solver, + random_state=random_state, + output_type="numpy", + ) + X_pca = pca_func.fit_transform(X) + + elif not zero_center: + pca_func = TruncatedSVD( n_components=n_comps, - svd_solver=svd_solver, random_state=random_state, + algorithm=svd_solver, output_type="numpy", ) X_pca = pca_func.fit_transform(X) - elif not zero_center: - pca_func = TruncatedSVD( - n_components=n_comps, - random_state=random_state, - algorithm=svd_solver, - output_type="numpy", - ) - X_pca = pca_func.fit_transform(X) - if X_pca.dtype.descr != np.dtype(dtype).descr: X_pca = X_pca.astype(dtype) @@ -194,187 +219,3 @@ def pca( adata.varm["PCs"] = pca_func.components_.T if copy: return adata - - -class PCA_sparse: - def __init__(self, n_components) -> None: - self.n_components = n_components - - def fit(self, x): - if self.n_components is None: - n_rows = x.shape[0] - n_cols = x.shape[1] - self.n_components_ = min(n_rows, n_cols) - else: - self.n_components_ = self.n_components - - if not isspmatrix_csr(x): - x = x.tocsr() - _check_matrix_for_zero_genes(x) - self.n_samples_ = x.shape[0] - self.n_features_in_ = x.shape[1] if x.ndim == 2 else 1 - self.dtype = x.data.dtype - - covariance, self.mean_, _ = _cov_sparse(x=x, return_mean=True) - - self.explained_variance_, self.components_ = cp.linalg.eigh( - covariance, UPLO="U" - ) - - # NOTE: We reverse the eigen vector and eigen values here - # because cupy provides them in ascending order. Make a copy otherwise - # it is not C_CONTIGUOUS anymore and would error when converting to - # CumlArray - self.explained_variance_ = self.explained_variance_[::-1] - - self.components_ = cp.flip(self.components_, axis=1) - - self.components_ = self.components_.T[: self.n_components_, :] - - self.explained_variance_ratio_ = self.explained_variance_ / cp.sum( - self.explained_variance_ - ) - - self.explained_variance_ = self.explained_variance_[: self.n_components_] - - self.explained_variance_ratio_ = self.explained_variance_ratio_[ - : self.n_components_ - ] - - return self - - def transform(self, X): - X = X - self.mean_ - X_transformed = X.dot(self.components_.T) - self.components_ = self.components_.get() - self.explained_variance_ = self.explained_variance_.get() - self.explained_variance_ratio_ = self.explained_variance_ratio_.get() - return X_transformed.get() - - def fit_transform(self, X, y=None): - return self.fit(X).transform(X) - - -def _cov_sparse(x, return_gram=False, return_mean=False): - """ - Computes the mean and the covariance of matrix X of - the form Cov(X, X) = E(XX) - E(X)E(X) - - This is a temporary fix for - cuml issue #5475 and cupy issue #7699, - where the operation `x.T.dot(x)` did not work for - larger sparse matrices. - - Parameters - ---------- - - x : cupyx.scipy.sparse of size (m, n) - return_gram : boolean (default = False) - If True, gram matrix of the form (1 / n) * X.T.dot(X) - will be returned. - When True, a copy will be created - to store the results of the covariance. - When False, the local gram matrix result - will be overwritten - return_mean: boolean (default = False) - If True, the Maximum Likelihood Estimate used to - calculate the mean of X and X will be returned, - of the form (1 / n) * mean(X) and (1 / n) * mean(X) - - Returns - ------- - - result : cov(X, X) when return_gram and return_mean are False - cov(X, X), gram(X, X) when return_gram is True, - return_mean is False - cov(X, X), mean(X), mean(X) when return_gram is False, - return_mean is True - cov(X, X), gram(X, X), mean(X), mean(X) - when return_gram is True and return_mean is True - """ - - from ._kernels._pca_sparse_kernel import ( - _copy_kernel, - _cov_kernel, - _gramm_kernel_csr, - ) - - gram_matrix = cp.zeros((x.shape[1], x.shape[1]), dtype=x.data.dtype) - - block = (128,) - grid = (x.shape[0],) - compute_mean_cov = _gramm_kernel_csr(x.data.dtype) - compute_mean_cov( - grid, - block, - ( - x.indptr, - x.indices, - x.data, - x.shape[0], - x.shape[1], - gram_matrix, - ), - ) - - copy_gram = _copy_kernel(x.data.dtype) - block = (32, 32) - grid = (math.ceil(x.shape[1] / block[0]), math.ceil(x.shape[1] / block[1])) - copy_gram( - grid, - block, - (gram_matrix, x.shape[1]), - ) - - mean_x = x.sum(axis=0) * (1 / x.shape[0]) - gram_matrix *= 1 / x.shape[0] - - if return_gram: - cov_result = cp.zeros( - (gram_matrix.shape[0], gram_matrix.shape[0]), - dtype=gram_matrix.dtype, - ) - else: - cov_result = gram_matrix - - compute_cov = _cov_kernel(x.dtype) - - block_size = (32, 32) - grid_size = (math.ceil(gram_matrix.shape[0] / 8),) * 2 - compute_cov( - grid_size, - block_size, - (cov_result, gram_matrix, mean_x, mean_x, gram_matrix.shape[0]), - ) - - if not return_gram and not return_mean: - return cov_result - elif return_gram and not return_mean: - return cov_result, gram_matrix - elif not return_gram and return_mean: - return cov_result, mean_x, mean_x - elif return_gram and return_mean: - return cov_result, gram_matrix, mean_x, mean_x - - -def _check_matrix_for_zero_genes(X): - gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) - - from ._kernels._pca_sparse_kernel import _zero_genes_kernel - - block = (32,) - grid = (int(math.ceil(X.nnz / block[0])),) - _zero_genes_kernel( - grid, - block, - ( - X.indices, - gene_ex, - X.nnz, - ), - ) - if cp.any(gene_ex == 0): - raise ValueError( - "There are genes with zero expression. " - "Please remove them before running PCA." - ) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py new file mode 100644 index 00000000..ab2483e6 --- /dev/null +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import math + +import cupy as cp +import dask +from cuml.dask.common.part_utils import _extract_partitions +from cuml.internals.memory_utils import with_cupy_rmm +from cupyx import cusparse + +from rapids_singlecell._compat import ( + _get_dask_client, + _meta_dense, +) +from rapids_singlecell.preprocessing._utils import _get_mean_var + + +class PCA_sparse_dask: + def __init__(self, n_components, client) -> None: + self.n_components = n_components + self.client = _get_dask_client(client) + + def fit(self, x): + if self.n_components is None: + n_rows = x.shape[0] + n_cols = x.shape[1] + self.n_components_ = min(n_rows, n_cols) + else: + self.n_components_ = self.n_components + + self.n_samples_ = x.shape[0] + self.n_features_in_ = x.shape[1] if x.ndim == 2 else 1 + self.dtype = x.dtype + covariance, self.mean_, _ = _cov_sparse_dask(self.client, x=x, return_mean=True) + self.explained_variance_, self.components_ = cp.linalg.eigh( + covariance, UPLO="U" + ) + # NOTE: We reverse the eigen vector and eigen values here + # because cupy provides them in ascending order. Make a copy otherwise + # it is not C_CONTIGUOUS anymore and would error when converting to + # CumlArray + self.explained_variance_ = self.explained_variance_[::-1] + + self.components_ = cp.flip(self.components_, axis=1) + + self.components_ = self.components_.T[: self.n_components_, :] + + self.explained_variance_ratio_ = self.explained_variance_ / cp.sum( + self.explained_variance_ + ) + if self.n_components_ < min(self.n_samples_, self.n_features_in_): + self.noise_variance_ = self.explained_variance_[self.n_components_ :].mean() + else: + self.noise_variance_ = cp.array([0.0]) + self.explained_variance_ = self.explained_variance_[: self.n_components_] + + self.explained_variance_ratio_ = self.explained_variance_ratio_[ + : self.n_components_ + ] + return self + + def transform(self, X): + def _transform(X_part, mean_, components_): + dense = cusparse.csr2dense(X_part) + dense = dense - mean_ + X_pca = dense.dot(components_.T) + return X_pca + + X_pca = X.map_blocks( + _transform, + mean_=self.mean_, + components_=self.components_, + dtype=X.dtype, + meta=_meta_dense(X.dtype), + ) + + self.components_ = self.components_.get() + self.explained_variance_ = self.explained_variance_.get() + self.explained_variance_ratio_ = self.explained_variance_ratio_.get() + return X_pca.persist() + + def fit_transform(self, X, y=None): + return self.fit(X).transform(X) + + +@with_cupy_rmm +def _cov_sparse_dask(client, x, return_gram=False, return_mean=False): + """ + Computes the mean and the covariance of matrix X of + the form Cov(X, X) = E(XX) - E(X)E(X) + + This is a temporary fix for + cuml issue #5475 and cupy issue #7699, + where the operation `x.T.dot(x)` did not work for + larger sparse matrices. + + Parameters + ---------- + + x : cupyx.scipy.sparse of size (m, n) + return_gram : boolean (default = False) + If True, gram matrix of the form (1 / n) * X.T.dot(X) + will be returned. + When True, a copy will be created + to store the results of the covariance. + When False, the local gram matrix result + will be overwritten + return_mean: boolean (default = False) + If True, the Maximum Likelihood Estimate used to + calculate the mean of X and X will be returned, + of the form (1 / n) * mean(X) and (1 / n) * mean(X) + + Returns + ------- + + result : cov(X, X) when return_gram and return_mean are False + cov(X, X), gram(X, X) when return_gram is True, + return_mean is False + cov(X, X), mean(X), mean(X) when return_gram is False, + return_mean is True + cov(X, X), gram(X, X), mean(X), mean(X) + when return_gram is True and return_mean is True + """ + + from rapids_singlecell.preprocessing._kernels._pca_sparse_kernel import ( + _copy_kernel, + _cov_kernel, + _gramm_kernel_csr, + ) + + compute_mean_cov = _gramm_kernel_csr(x.dtype) + compute_mean_cov.compile() + + def __gram_block(x_part, n_cols): + gram_matrix = cp.zeros((n_cols, n_cols), dtype=x.dtype) + + block = (128,) + grid = (x_part.shape[0],) + compute_mean_cov( + grid, + block, + ( + x_part.indptr, + x_part.indices, + x_part.data, + x_part.shape[0], + n_cols, + gram_matrix, + ), + ) + return gram_matrix + + parts = client.sync(_extract_partitions, x) + futures = [ + client.submit(__gram_block, part, x.shape[1], workers=[w]) for w, part in parts + ] + # Gather results from futures + objs = [] + for i in range(len(futures)): + obj = dask.array.from_delayed( + futures[i], shape=(x.shape[1], x.shape[1]), dtype=x.dtype + ) + objs.append(obj) + gram_matrix = dask.array.stack(objs).sum(axis=0).compute() + mean_x, _ = _get_mean_var(x, client=client) + mean_x = mean_x.astype(x.dtype) + copy_gram = _copy_kernel(x.dtype) + block = (32, 32) + grid = (math.ceil(x.shape[1] / block[0]), math.ceil(x.shape[1] / block[1])) + copy_gram( + grid, + block, + (gram_matrix, x.shape[1]), + ) + + gram_matrix *= 1 / x.shape[0] + + if return_gram: + cov_result = cp.zeros( + (gram_matrix.shape[0], gram_matrix.shape[0]), + dtype=gram_matrix.dtype, + ) + else: + cov_result = gram_matrix + + compute_cov = _cov_kernel(gram_matrix.dtype) + + block_size = (32, 32) + grid_size = (math.ceil(gram_matrix.shape[0] / 8),) * 2 + compute_cov( + grid_size, + block_size, + (cov_result, gram_matrix, mean_x, mean_x, gram_matrix.shape[0]), + ) + + if not return_gram and not return_mean: + return cov_result + elif return_gram and not return_mean: + return cov_result, gram_matrix + elif not return_gram and return_mean: + return cov_result, mean_x, mean_x + elif return_gram and return_mean: + return cov_result, gram_matrix, mean_x, mean_x diff --git a/src/rapids_singlecell/preprocessing/_kernels/_pca_sparse_kernel.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py similarity index 100% rename from src/rapids_singlecell/preprocessing/_kernels/_pca_sparse_kernel.py rename to src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py new file mode 100644 index 00000000..27ef71b6 --- /dev/null +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -0,0 +1,190 @@ +from __future__ import annotations + +import math + +import cupy as cp +from cupyx.scipy.sparse import isspmatrix_csr + + +class PCA_sparse: + def __init__(self, n_components) -> None: + self.n_components = n_components + + def fit(self, x): + if self.n_components is None: + n_rows = x.shape[0] + n_cols = x.shape[1] + self.n_components_ = min(n_rows, n_cols) + else: + self.n_components_ = self.n_components + + if not isspmatrix_csr(x): + x = x.tocsr() + _check_matrix_for_zero_genes(x) + self.n_samples_ = x.shape[0] + self.n_features_in_ = x.shape[1] if x.ndim == 2 else 1 + self.dtype = x.data.dtype + + covariance, self.mean_, _ = _cov_sparse(x=x, return_mean=True) + + self.explained_variance_, self.components_ = cp.linalg.eigh( + covariance, UPLO="U" + ) + + # NOTE: We reverse the eigen vector and eigen values here + # because cupy provides them in ascending order. Make a copy otherwise + # it is not C_CONTIGUOUS anymore and would error when converting to + # CumlArray + self.explained_variance_ = self.explained_variance_[::-1] + + self.components_ = cp.flip(self.components_, axis=1) + + self.components_ = self.components_.T[: self.n_components_, :] + + self.explained_variance_ratio_ = self.explained_variance_ / cp.sum( + self.explained_variance_ + ) + + self.explained_variance_ = self.explained_variance_[: self.n_components_] + + self.explained_variance_ratio_ = self.explained_variance_ratio_[ + : self.n_components_ + ] + + return self + + def transform(self, X): + X = X - self.mean_ + X_transformed = X.dot(self.components_.T) + self.components_ = self.components_.get() + self.explained_variance_ = self.explained_variance_.get() + self.explained_variance_ratio_ = self.explained_variance_ratio_.get() + return X_transformed.get() + + def fit_transform(self, X, y=None): + return self.fit(X).transform(X) + + +def _cov_sparse(x, return_gram=False, return_mean=False): + """ + Computes the mean and the covariance of matrix X of + the form Cov(X, X) = E(XX) - E(X)E(X) + + This is a temporary fix for + cuml issue #5475 and cupy issue #7699, + where the operation `x.T.dot(x)` did not work for + larger sparse matrices. + + Parameters + ---------- + + x : cupyx.scipy.sparse of size (m, n) + return_gram : boolean (default = False) + If True, gram matrix of the form (1 / n) * X.T.dot(X) + will be returned. + When True, a copy will be created + to store the results of the covariance. + When False, the local gram matrix result + will be overwritten + return_mean: boolean (default = False) + If True, the Maximum Likelihood Estimate used to + calculate the mean of X and X will be returned, + of the form (1 / n) * mean(X) and (1 / n) * mean(X) + + Returns + ------- + + result : cov(X, X) when return_gram and return_mean are False + cov(X, X), gram(X, X) when return_gram is True, + return_mean is False + cov(X, X), mean(X), mean(X) when return_gram is False, + return_mean is True + cov(X, X), gram(X, X), mean(X), mean(X) + when return_gram is True and return_mean is True + """ + + from ._kernels._pca_sparse_kernel import ( + _copy_kernel, + _cov_kernel, + _gramm_kernel_csr, + ) + + gram_matrix = cp.zeros((x.shape[1], x.shape[1]), dtype=x.data.dtype) + + block = (128,) + grid = (x.shape[0],) + compute_mean_cov = _gramm_kernel_csr(x.data.dtype) + compute_mean_cov( + grid, + block, + ( + x.indptr, + x.indices, + x.data, + x.shape[0], + x.shape[1], + gram_matrix, + ), + ) + + copy_gram = _copy_kernel(x.data.dtype) + block = (32, 32) + grid = (math.ceil(x.shape[1] / block[0]), math.ceil(x.shape[1] / block[1])) + copy_gram( + grid, + block, + (gram_matrix, x.shape[1]), + ) + + mean_x = x.sum(axis=0) * (1 / x.shape[0]) + gram_matrix *= 1 / x.shape[0] + + if return_gram: + cov_result = cp.zeros( + (gram_matrix.shape[0], gram_matrix.shape[0]), + dtype=gram_matrix.dtype, + ) + else: + cov_result = gram_matrix + + compute_cov = _cov_kernel(x.dtype) + + block_size = (32, 32) + grid_size = (math.ceil(gram_matrix.shape[0] / 8),) * 2 + compute_cov( + grid_size, + block_size, + (cov_result, gram_matrix, mean_x, mean_x, gram_matrix.shape[0]), + ) + + if not return_gram and not return_mean: + return cov_result + elif return_gram and not return_mean: + return cov_result, gram_matrix + elif not return_gram and return_mean: + return cov_result, mean_x, mean_x + elif return_gram and return_mean: + return cov_result, gram_matrix, mean_x, mean_x + + +def _check_matrix_for_zero_genes(X): + gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) + + from ._kernels._pca_sparse_kernel import _zero_genes_kernel + + block = (32,) + grid = (int(math.ceil(X.nnz / block[0])),) + _zero_genes_kernel( + grid, + block, + ( + X.indices, + gene_ex, + X.nnz, + ), + ) + if cp.any(gene_ex == 0): + raise ValueError( + "There are genes with zero expression. " + "Please remove them before running PCA." + ) From b216890cfc2166694829ba22985728a7a316df5b Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 2 May 2024 18:20:27 +0200 Subject: [PATCH 07/69] pca update --- src/rapids_singlecell/preprocessing/_pca.py | 14 ++++++--- .../_sparse_pca/_dask_sparse_pca.py | 30 ++++++++++++++++--- .../_kernels/_pca_sparse_kernel.py | 22 ++++++++++++++ .../preprocessing/_sparse_pca/_sparse_pca.py | 16 ++++++++++ 4 files changed, 74 insertions(+), 8 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 61928831..47a5ae7c 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -5,7 +5,6 @@ import cupy as cp import numpy as np -from cuml.decomposition import TruncatedSVD from cuml.internals.input_utils import sparse_scipy_to_cp from cupyx.scipy.sparse import csr_matrix from cupyx.scipy.sparse import issparse as cpissparse @@ -16,6 +15,8 @@ from rapids_singlecell._compat import DaskArray, DaskClient from rapids_singlecell.get import _get_obs_rep +from ._utils import _check_gpu_X + if TYPE_CHECKING: from anndata import AnnData from numpy.typing import NDArray @@ -136,17 +137,18 @@ def pca( raise ValueError( "Dask arrays are not supported for chunked PCA computation." ) + _check_gpu_X(X, allow_dask=True) if isinstance(X._meta, cp.ndarray): from cuml.dask.decomposition import PCA if svd_solver == "auto": svd_solver = "jacobi" - pca_func = PCA(n_components=n_comps, svd_solver=svd_solver) + pca_func = PCA(n_components=n_comps, svd_solver=svd_solver, client=client) X_pca = pca_func.fit_transform(X) elif isinstance(X._meta, csr_matrix): - from cuml.dask.decomposition import PCA_sparse + from ._sparse_pca._dask_sparse_pca import PCA_sparse_dask - pca_func = PCA_sparse(n_components=n_comps) + pca_func = PCA_sparse_dask(n_components=n_comps, client=client) X_pca = pca_func.fit_transform(X) else: @@ -179,6 +181,8 @@ def pca( pca_func = PCA_sparse(n_components=n_comps) X_pca = pca_func.fit_transform(X) else: + from cuml.decomposition import PCA + pca_func = PCA( n_components=n_comps, svd_solver=svd_solver, @@ -188,6 +192,8 @@ def pca( X_pca = pca_func.fit_transform(X) elif not zero_center: + from cuml.decomposition import TruncatedSVD + pca_func = TruncatedSVD( n_components=n_comps, random_state=random_state, diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index ab2483e6..7ae60d59 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -6,7 +6,6 @@ import dask from cuml.dask.common.part_utils import _extract_partitions from cuml.internals.memory_utils import with_cupy_rmm -from cupyx import cusparse from rapids_singlecell._compat import ( _get_dask_client, @@ -60,8 +59,30 @@ def fit(self, x): return self def transform(self, X): + from ._kernels._pca_sparse_kernel import denser_kernel + + kernel = denser_kernel(X.dtype) + kernel.compile() + def _transform(X_part, mean_, components_): - dense = cusparse.csr2dense(X_part) + dense = cp.zeros(X_part.shape, dtype=X.dtype) + max_nnz = cp.diff(X_part.indptr).max() + tpb = (32, 32) + bpg_x = math.ceil(X_part.shape[0] / tpb[0]) + bpg_y = math.ceil(max_nnz / tpb[1]) + bpg = (bpg_x, bpg_y) + kernel( + bpg, + tpb, + ( + X_part.indptr, + X_part.indices, + X_part.data, + dense, + X_part.shape[0], + X_part.shape[1], + ), + ) dense = dense - mean_ X_pca = dense.dot(components_.T) return X_pca @@ -71,13 +92,14 @@ def _transform(X_part, mean_, components_): mean_=self.mean_, components_=self.components_, dtype=X.dtype, + chunks=(X.chunks[0], self.n_components_), meta=_meta_dense(X.dtype), ) self.components_ = self.components_.get() self.explained_variance_ = self.explained_variance_.get() self.explained_variance_ratio_ = self.explained_variance_ratio_.get() - return X_pca.persist() + return X_pca def fit_transform(self, X, y=None): return self.fit(X).transform(X) @@ -122,7 +144,7 @@ def _cov_sparse_dask(client, x, return_gram=False, return_mean=False): when return_gram is True and return_mean is True """ - from rapids_singlecell.preprocessing._kernels._pca_sparse_kernel import ( + from ._kernels._pca_sparse_kernel import ( _copy_kernel, _cov_kernel, _gramm_kernel_csr, diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py index 11aa197d..889fe350 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py @@ -61,6 +61,24 @@ } """ + +denser = r""" +(const int* indptr,const int *index,const {0} *data, + {0}* out, int major, int minor) { + int row = blockIdx.x*blockDim.x+threadIdx.x ; + int col = blockIdx.y*blockDim.y+threadIdx.y ; + if(row >= major){ + return; + } + int start = indptr[row]; + int stop = indptr[row+1]; + if (col>= (stop - start)){ + return; + } + int idx = index[start + col]; + out[row*minor+idx] = data[start + col];} +""" + _zero_genes_kernel = cp.RawKernel(check_zero_genes, "check_zero_genes") @@ -74,3 +92,7 @@ def _gramm_kernel_csr(dtype): def _copy_kernel(dtype): return cuda_kernel_factory(copy_kernel, (dtype,), "copy_kernel") + + +def denser_kernel(dtype): + return cuda_kernel_factory(denser, (dtype,), "denser") diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index 27ef71b6..eb0de780 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -54,6 +54,22 @@ def fit(self, x): return self def transform(self, X): + """ " + from ._kernels._pca_sparse_kernel import denser_kernel + dense_kernel = denser_kernel(X.dtype) + + dense = cp.zeros(X.shape,dtype=X.dtype) + max_nnz = cp.diff(X.indptr).max() + tpb = (32, 32) + bpg_x = math.ceil(X.shape[0] / tpb[0]) + bpg_y = math.ceil(max_nnz / tpb[1]) + bpg = (bpg_x, bpg_y) + dense_kernel(bpg, + tpb, + (X.indptr,X.indices, X.data, dense,X.shape[0],X.shape[1]),) + dense = dense - self.mean_ + X_transformed = dense.dot(self.components_.T) + """ X = X - self.mean_ X_transformed = X.dot(self.components_.T) self.components_ = self.components_.get() From cdffd336e358488a86c21249e6dea763514f0f62 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 3 May 2024 09:13:14 +0200 Subject: [PATCH 08/69] fix bug with csc matrix --- src/rapids_singlecell/preprocessing/_pca.py | 5 +++-- .../preprocessing/_sparse_pca/_sparse_pca.py | 20 +++++++++---------- 2 files changed, 12 insertions(+), 13 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 47a5ae7c..0323910a 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -6,7 +6,7 @@ import cupy as cp import numpy as np from cuml.internals.input_utils import sparse_scipy_to_cp -from cupyx.scipy.sparse import csr_matrix +from cupyx.scipy.sparse import csr_matrix, isspmatrix_csr from cupyx.scipy.sparse import issparse as cpissparse from scanpy._utils import Empty, _empty from scanpy.preprocessing._pca import _handle_mask_var @@ -175,9 +175,10 @@ def pca( if cpissparse(X) or issparse(X): if issparse(X): X = sparse_scipy_to_cp(X, dtype=X.dtype) - X = csr_matrix(X) from ._sparse_pca._sparse_pca import PCA_sparse + if not isspmatrix_csr(X): + X = X.tocsr() pca_func = PCA_sparse(n_components=n_comps) X_pca = pca_func.fit_transform(X) else: diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index eb0de780..22535384 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -3,7 +3,6 @@ import math import cupy as cp -from cupyx.scipy.sparse import isspmatrix_csr class PCA_sparse: @@ -18,8 +17,6 @@ def fit(self, x): else: self.n_components_ = self.n_components - if not isspmatrix_csr(x): - x = x.tocsr() _check_matrix_for_zero_genes(x) self.n_samples_ = x.shape[0] self.n_features_in_ = x.shape[1] if x.ndim == 2 else 1 @@ -54,24 +51,25 @@ def fit(self, x): return self def transform(self, X): - """ " from ._kernels._pca_sparse_kernel import denser_kernel + dense_kernel = denser_kernel(X.dtype) - dense = cp.zeros(X.shape,dtype=X.dtype) + dense = cp.zeros(X.shape, dtype=X.dtype) max_nnz = cp.diff(X.indptr).max() tpb = (32, 32) bpg_x = math.ceil(X.shape[0] / tpb[0]) bpg_y = math.ceil(max_nnz / tpb[1]) bpg = (bpg_x, bpg_y) - dense_kernel(bpg, - tpb, - (X.indptr,X.indices, X.data, dense,X.shape[0],X.shape[1]),) + dense_kernel( + bpg, + tpb, + (X.indptr, X.indices, X.data, dense, X.shape[0], X.shape[1]), + ) dense = dense - self.mean_ X_transformed = dense.dot(self.components_.T) - """ - X = X - self.mean_ - X_transformed = X.dot(self.components_.T) + # X = X - self.mean_ + # X_transformed = X.dot(self.components_.T) self.components_ = self.components_.get() self.explained_variance_ = self.explained_variance_.get() self.explained_variance_ratio_ = self.explained_variance_ratio_.get() From 177afa1c912d6c22153dd92d50a1a7fcc552743e Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 3 May 2024 09:19:35 +0200 Subject: [PATCH 09/69] add dask to docs --- docs/conf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index fb660361..9cf55043 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -62,7 +62,7 @@ autosummary_generate = True autodoc_member_order = "bysource" -autodoc_mock_imports = ["cudf", "cuml", "cugraph", "cupy", "cupyx", "pylibraft"] +autodoc_mock_imports = ["cudf", "cuml", "cugraph", "cupy", "cupyx", "pylibraft", "dask"] default_role = "literal" napoleon_google_docstring = False napoleon_numpy_docstring = True @@ -108,6 +108,7 @@ "rmm": ("https://docs.rapids.ai/api/rmm/stable/", None), "statsmodels": ("https://www.statsmodels.org/stable/", None), "omnipath": ("https://omnipath.readthedocs.io/en/latest/", None), + "dask": ("https://docs.dask.org/en/stable/", None), } # List of patterns, relative to source directory, that match files and From dd1377c950a69e8840b2fe163688eb0293ea71ec Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 3 May 2024 13:23:23 +0200 Subject: [PATCH 10/69] add tests --- src/rapids_singlecell/preprocessing/_hvg.py | 9 +- .../preprocessing/_normalize.py | 4 +- src/rapids_singlecell/preprocessing/_pca.py | 23 +++- src/rapids_singlecell/preprocessing/_qc.py | 12 +- src/rapids_singlecell/preprocessing/_utils.py | 64 ++++++++- tests/dask/conftest.py | 48 +++++++ tests/dask/test_dask_pca.py | 62 +++++++++ tests/dask/test_hvg.py | 125 ++++++++++++++++++ tests/dask/test_normalize.py | 56 ++++++++ tests/dask/test_qc.py | 54 ++++++++ 10 files changed, 432 insertions(+), 25 deletions(-) create mode 100644 tests/dask/conftest.py create mode 100644 tests/dask/test_dask_pca.py create mode 100644 tests/dask/test_hvg.py create mode 100644 tests/dask/test_normalize.py create mode 100644 tests/dask/test_qc.py diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 2e97332c..90541d9d 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -291,13 +291,12 @@ def _highly_variable_genes_single_batch( X = _get_obs_rep(adata, layer=layer) _check_gpu_X(X, allow_dask=True) if hasattr(X, "_view_args"): # AnnData array view - # For compatibility with anndata<0.9 - X = X.copy() # Doesn't actually copy memory, just removes View class wrapper + X = X.copy() if flavor == "seurat": if isinstance(X, DaskArray): if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(cp.expm1, meta=_meta_dense(X.dtype)) + X = X.map_blocks(lambda X: cp.expm1(X), meta=_meta_dense(X.dtype)) elif isinstance(X._meta, csr_matrix): X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) else: @@ -433,11 +432,11 @@ def _highly_variable_genes_batched( dfs = [] gene_list = adata.var_names for batch in batches: - adata_subset = adata[adata.obs[batch_key] == batch] + adata_subset = adata[adata.obs[batch_key] == batch].copy() calculate_qc_metrics(adata_subset, layer=layer, client=client) filt = adata_subset.var["n_cells_by_counts"].to_numpy() > 0 - adata_subset = adata_subset[:, filt] + adata_subset = adata_subset[:, filt].copy() hvg = _highly_variable_genes_single_batch( adata_subset, diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 8e92e9db..7f978403 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -135,7 +135,7 @@ def __mul(X_part): return X_part X = X.map_blocks(lambda X: __mul(X), meta=_meta_sparse(X.dtype)) - elif isinstance(X.meta, cp.ndarray): + elif isinstance(X._meta, cp.ndarray): from ._kernels._norm_kernel import _mul_dense mul_kernel = _mul_dense(X.dtype) @@ -269,7 +269,7 @@ def log1p( X = X.log1p() elif isinstance(X, DaskArray): if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(cp.log1p, meta=_meta_dense(X.dtype)) + X = X.map_blocks(lambda X: cp.log1p(X), meta=_meta_dense(X.dtype)) elif isinstance(X._meta, sparse.csr_matrix): X = X.map_blocks(lambda X: X.log1p(), meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 0323910a..96188e2f 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -143,13 +143,17 @@ def pca( if svd_solver == "auto": svd_solver = "jacobi" - pca_func = PCA(n_components=n_comps, svd_solver=svd_solver, client=client) + pca_func = PCA( + n_components=n_comps, svd_solver=svd_solver, whiten=False, client=client + ) X_pca = pca_func.fit_transform(X) + X_pca = X_pca.compute_chunk_sizes() elif isinstance(X._meta, csr_matrix): from ._sparse_pca._dask_sparse_pca import PCA_sparse_dask pca_func = PCA_sparse_dask(n_components=n_comps, client=client) - X_pca = pca_func.fit_transform(X) + pca_func = pca_func.fit(X) + X_pca = pca_func.transform(X) else: if chunked: @@ -213,16 +217,23 @@ def pca( "use_highly_variable": mask_var_param == "highly_variable", "mask_var": mask_var_param, }, - "variance": pca_func.explained_variance_, - "variance_ratio": pca_func.explained_variance_ratio_, + "variance": _as_numpy(pca_func.explained_variance_), + "variance_ratio": _as_numpy(pca_func.explained_variance_ratio_), } adata.uns["pca"] = uns_entry if layer is not None: adata.uns["pca"]["params"]["layer"] = layer if mask_var is not None: adata.varm["PCs"] = np.zeros(shape=(adata.n_vars, n_comps)) - adata.varm["PCs"][mask_var] = pca_func.components_.T + adata.varm["PCs"][mask_var] = _as_numpy(pca_func.components_.T) else: - adata.varm["PCs"] = pca_func.components_.T + adata.varm["PCs"] = _as_numpy(pca_func.components_.T) if copy: return adata + + +def _as_numpy(X): + if isinstance(X, cp.ndarray): + return X.get() + else: + return X diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index d87d8614..ba5ab9b9 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -223,8 +223,8 @@ def __qc_calc(X_part): elif isinstance(X._meta, cp.ndarray): from ._kernels._qc_kernels import _sparse_qc_dense - _sparse_qc_dense = _sparse_qc_dense(X.dtype) - _sparse_qc_dense.compile() + sparse_qc_dense = _sparse_qc_dense(X.dtype) + sparse_qc_dense.compile() def __qc_calc(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) @@ -238,7 +238,6 @@ def __qc_calc(X_part): int(math.ceil(X_part.shape[0] / block[0])), int(math.ceil(X_part.shape[1] / block[1])), ) - sparse_qc_dense = _sparse_qc_dense(X.dtype) sparse_qc_dense( grid, block, @@ -364,10 +363,10 @@ def __qc_calc(X_part): return sums_cells_sub elif isinstance(X._meta, cp.ndarray): - from ._kernels._qc_kernels import _sparse_qc_dense + from ._kernels._qc_kernels import _sparse_qc_dense_sub - _sparse_qc_dense = _sparse_qc_dense(X.dtype) - _sparse_qc_dense.compile() + sparse_qc_dense = _sparse_qc_dense_sub(X.dtype) + sparse_qc_dense.compile() def __qc_calc(X_part): sums_cells_sub = cp.zeros((X_part.shape[0]), dtype=X_part.dtype) @@ -378,7 +377,6 @@ def __qc_calc(X_part): int(math.ceil(X_part.shape[0] / block[0])), int(math.ceil(X_part.shape[1] / block[1])), ) - sparse_qc_dense = _sparse_qc_dense(X.dtype) sparse_qc_dense( grid, block, diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index a854a4f8..875bbec3 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -156,6 +156,53 @@ def __mean_var(X_part, minor, major): return mean, var +@with_cupy_rmm +def _mean_var_dense_dask(X, axis, client=None): + """ + Implements sum operation for dask array when the backend is cupy sparse csr matrix + """ + import dask.array as da + + client = _get_dask_client(client) + + def __mean_var(X_part, axis): + mean = X_part.sum(axis=axis) + var = (X_part**2).sum(axis=axis) + if axis == 0: + mean = mean.reshape(-1, 1) + var = var.reshape(-1, 1) + return mean, var + + parts = client.sync(_extract_partitions, X) + futures = [client.submit(__mean_var, part, axis, workers=[w]) for w, part in parts] + # Gather results from futures + results = client.gather(futures) + + # Initialize lists to hold the Dask arrays + means_objs = [] + var_objs = [] + + # Process each result + for means, vars_ in results: + # Append the arrays to their respective lists as Dask arrays + means_objs.append(da.from_array(means, chunks=means.shape)) + var_objs.append(da.from_array(vars_, chunks=vars_.shape)) + if axis == 0: + mean = da.concatenate(means_objs, axis=1).sum(axis=1) + var = da.concatenate(var_objs, axis=1).sum(axis=1) + else: + mean = da.concatenate(means_objs) + var = da.concatenate(var_objs) + + mean, var = da.compute(mean, var) + mean, var = mean.ravel(), var.ravel() + mean = mean / X.shape[axis] + var = var / X.shape[axis] + var -= cp.power(mean, 2) + var *= X.shape[axis] / (X.shape[axis] - 1) + return mean, var + + def _get_mean_var(X, axis=0, client=None): if issparse(X): if axis == 0: @@ -191,7 +238,8 @@ def _get_mean_var(X, axis=0, client=None): major = X.shape[0] minor = X.shape[1] mean, var = _mean_var_major_dask(X, major, minor, client) - + elif isinstance(X._meta, cp.ndarray): + mean, var = _mean_var_dense_dask(X, axis, client) else: mean = X.mean(axis=axis) var = X.var(axis=axis) @@ -215,7 +263,16 @@ def _check_nonnegative_integers(X): def _check_gpu_X(X, require_cf=False, allow_dask=False): - if isinstance(X, cp.ndarray): + if isinstance(X, DaskArray): + if allow_dask: + return _check_gpu_X(X._meta) + else: + raise TypeError( + "The input is a DaskArray. " + "Rapids-singlecell doesn't support DaskArray in this function, " + "so your input must be a CuPy ndarray or a CuPy sparse matrix. " + ) + elif isinstance(X, cp.ndarray): return True elif issparse(X): if not require_cf: @@ -225,9 +282,6 @@ def _check_gpu_X(X, require_cf=False, allow_dask=False): else: X.sort_indices() X.sum_duplicates() - elif allow_dask: - if isinstance(X, DaskArray): - return _check_gpu_X(X._meta) else: raise TypeError( "The input is not a CuPy ndarray or CuPy sparse matrix. " diff --git a/tests/dask/conftest.py b/tests/dask/conftest.py new file mode 100644 index 00000000..e3f7b610 --- /dev/null +++ b/tests/dask/conftest.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import cupy as cp +from cupyx.scipy import sparse as cusparse +from anndata.tests.helpers import as_dense_dask_array, as_sparse_dask_array +import pytest + +from dask_cuda import LocalCUDACluster +from dask_cuda.utils_test import IncreasedCloseTimeoutNanny +from dask.distributed import Client + + +def as_sparse_cupy_dask_array(X): + da = as_sparse_dask_array(X) + da = da.rechunk((da.shape[0]//2, da.shape[1])) + da = da.map_blocks(cusparse.csr_matrix, dtype = X.dtype) + return da + +def as_dense_cupy_dask_array(X): + X = as_dense_dask_array(X) + X = X.map_blocks(cp.array) + X = X.rechunk((X.shape[0]//2, X.shape[1])) + return X + +from dask_cuda import initialize +from dask_cuda import LocalCUDACluster +from dask_cuda.utils_test import IncreasedCloseTimeoutNanny +from dask.distributed import Client + +@pytest.fixture(scope="module") +def cluster(): + + cluster = LocalCUDACluster( + CUDA_VISIBLE_DEVICES ="0", + protocol="tcp", + scheduler_port=0, + worker_class=IncreasedCloseTimeoutNanny, + ) + yield cluster + cluster.close() + + +@pytest.fixture(scope="function") +def client(cluster): + + client = Client(cluster) + yield client + client.close() diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py new file mode 100644 index 00000000..f0fa9b6f --- /dev/null +++ b/tests/dask/test_dask_pca.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +from cupyx.scipy import sparse as cusparse +from scipy import sparse +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +import rapids_singlecell as rsc + +from scanpy.datasets import pbmc3k_processed + +def test_pca_sparse_dask(client): + sparse_ad = pbmc3k_processed() + default = pbmc3k_processed() + sparse_ad.X = sparse.csr_matrix(sparse_ad.X.astype(np.float64)) + default.X = as_sparse_cupy_dask_array(default.X.astype(np.float64)) + rsc.pp.pca(sparse_ad) + rsc.pp.pca(default) + + cp.testing.assert_allclose( + np.abs(sparse_ad.obsm["X_pca"]), + cp.abs(default.obsm["X_pca"].compute()), + rtol=1e-7, + atol=1e-6, + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.uns["pca"]["variance_ratio"]), + np.abs(default.uns["pca"]["variance_ratio"]), + rtol=1e-7, + atol=1e-6, + ) + +def test_pca_dense_dask(client): + sparse_ad = pbmc3k_processed() + default = pbmc3k_processed() + sparse_ad.X = cp.array(sparse_ad.X.astype(np.float64)) + default.X = as_dense_cupy_dask_array(default.X.astype(np.float64)) + rsc.pp.pca(sparse_ad, svd_solver="full") + rsc.pp.pca(default, svd_solver="full") + + cp.testing.assert_allclose( + np.abs(sparse_ad.obsm["X_pca"]), + cp.abs(default.obsm["X_pca"].compute()), + rtol=1e-7, + atol=1e-6, + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.uns["pca"]["variance_ratio"]), + np.abs(default.uns["pca"]["variance_ratio"]), + rtol=1e-7, + atol=1e-6, + ) diff --git a/tests/dask/test_hvg.py b/tests/dask/test_hvg.py new file mode 100644 index 00000000..8d817a49 --- /dev/null +++ b/tests/dask/test_hvg.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +from cupyx.scipy import sparse as cusparse +import scanpy as sc +import pandas as pd +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +import rapids_singlecell as rsc + +from scanpy.datasets import pbmc3k + +def _get_anndata(): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + sc.pp.normalize_total(adata) + sc.pp.log1p(adata) + return adata + +def test_seurat_sparse(client): + adata = _get_anndata() + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.highly_variable_genes(adata) + rsc.pp.highly_variable_genes(dask_data) + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + + +def test_seurat_sparse_batch(client): + adata = _get_anndata() + adata.obs["batch"] = ( + "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") + )[: adata.n_obs] + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.highly_variable_genes(adata, batch_key="batch") + rsc.pp.highly_variable_genes(dask_data,batch_key="batch") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + +def test_cr_sparse(client): + adata = _get_anndata() + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.highly_variable_genes(adata, flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + +def test_cr_sparse_batch(client): + adata = _get_anndata() + adata.obs["batch"] = ( + "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") + )[: adata.n_obs] + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data,batch_key="batch", flavor="cell_ranger") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + +def test_cr_dense(client): + adata = _get_anndata() + adata.X = adata.X.astype("float64") + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + rsc.pp.highly_variable_genes(adata, flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + +def test_seurat_dense(client): + adata = _get_anndata() + adata.X = adata.X.astype("float64") + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + rsc.pp.highly_variable_genes(adata) + rsc.pp.highly_variable_genes(dask_data) + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + + +def test_cr_dense_batch(client): + adata = _get_anndata() + adata.obs["batch"] = ( + "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") + )[: adata.n_obs] + adata.X = adata.X.astype("float64") + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data,batch_key="batch", flavor="cell_ranger") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + +def test_seurat_dense_batch(client): + adata = _get_anndata() + adata.obs["batch"] = ( + "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") + )[: adata.n_obs] + adata.X = adata.X.astype("float64") + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + rsc.pp.highly_variable_genes(adata, batch_key="batch") + rsc.pp.highly_variable_genes(dask_data,batch_key="batch") + cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) + cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) + cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) diff --git a/tests/dask/test_normalize.py b/tests/dask/test_normalize.py new file mode 100644 index 00000000..3c67d18d --- /dev/null +++ b/tests/dask/test_normalize.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +from cupyx.scipy import sparse as cusparse +import scanpy as sc +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +import rapids_singlecell as rsc + +from scanpy.datasets import pbmc3k + +def test_normalize_sparse(client): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.normalize_total(adata) + rsc.pp.normalize_total(dask_data) + cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + +def test_normalize_dense(client): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + rsc.pp.normalize_total(adata) + rsc.pp.normalize_total(dask_data) + cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + +def test_log1p_sparse(client): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + sc.pp.normalize_total(adata) + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.log1p(adata) + rsc.pp.log1p(dask_data) + cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + +def test_log1p_dense(client): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + sc.pp.normalize_total(adata) + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + rsc.pp.log1p(adata) + rsc.pp.log1p(dask_data) + cp.testing.assert_allclose(adata.X, dask_data.X.compute()) diff --git a/tests/dask/test_qc.py b/tests/dask/test_qc.py new file mode 100644 index 00000000..b79746e0 --- /dev/null +++ b/tests/dask/test_qc.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +from cupyx.scipy import sparse as cusparse +from scipy import sparse +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +import rapids_singlecell as rsc + +from scanpy.datasets import pbmc3k + +def test_qc_metrics_sparse(client): + adata = pbmc3k() + adata.var["mt"] = adata.var_names.str.startswith("MT-") + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p= True) + rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p= True) + np.testing.assert_allclose(adata.obs["n_genes_by_counts"] , dask_data.obs["n_genes_by_counts"]) + np.testing.assert_allclose(adata.obs["total_counts"] , dask_data.obs["total_counts"]) + np.testing.assert_allclose(adata.obs["log1p_n_genes_by_counts"] , dask_data.obs["log1p_n_genes_by_counts"]) + np.testing.assert_allclose(adata.obs["log1p_total_counts"] , dask_data.obs["log1p_total_counts"]) + np.testing.assert_allclose(adata.obs["pct_counts_mt"] , dask_data.obs["pct_counts_mt"]) + np.testing.assert_allclose(adata.obs["total_counts_mt"] , dask_data.obs["total_counts_mt"]) + np.testing.assert_allclose(adata.obs["log1p_total_counts_mt"] , dask_data.obs["log1p_total_counts_mt"]) + np.testing.assert_allclose(adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"]) + np.testing.assert_allclose(adata.var["total_counts"], dask_data.var["total_counts"]) + np.testing.assert_allclose(adata.var["mean_counts"], dask_data.var["mean_counts"]) + np.testing.assert_allclose(adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"]) + np.testing.assert_allclose(adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"]) + np.testing.assert_allclose(adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"]) + +def test_qc_metrics_dense(client): + adata = pbmc3k() + adata.var["mt"] = adata.var_names.str.startswith("MT-") + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p= True) + rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p= True) + np.testing.assert_allclose(adata.obs["n_genes_by_counts"] , dask_data.obs["n_genes_by_counts"]) + np.testing.assert_allclose(adata.obs["total_counts"] , dask_data.obs["total_counts"]) + np.testing.assert_allclose(adata.obs["log1p_n_genes_by_counts"] , dask_data.obs["log1p_n_genes_by_counts"]) + np.testing.assert_allclose(adata.obs["log1p_total_counts"] , dask_data.obs["log1p_total_counts"]) + np.testing.assert_allclose(adata.obs["pct_counts_mt"] , dask_data.obs["pct_counts_mt"]) + np.testing.assert_allclose(adata.obs["total_counts_mt"] , dask_data.obs["total_counts_mt"]) + np.testing.assert_allclose(adata.obs["log1p_total_counts_mt"] , dask_data.obs["log1p_total_counts_mt"]) + np.testing.assert_allclose(adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"]) + np.testing.assert_allclose(adata.var["total_counts"], dask_data.var["total_counts"]) + np.testing.assert_allclose(adata.var["mean_counts"], dask_data.var["mean_counts"]) + np.testing.assert_allclose(adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"]) + np.testing.assert_allclose(adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"]) + np.testing.assert_allclose(adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"]) From e25480050ca89ed4262fdd1473bdd9a00d0bf923 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 3 May 2024 13:50:50 +0200 Subject: [PATCH 11/69] update names --- tests/dask/{test_hvg.py => test_hvg_dask.py} | 0 tests/dask/{test_normalize.py => test_normalize_dask.py} | 0 tests/dask/{test_qc.py => test_qc_dask.py} | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename tests/dask/{test_hvg.py => test_hvg_dask.py} (100%) rename tests/dask/{test_normalize.py => test_normalize_dask.py} (100%) rename tests/dask/{test_qc.py => test_qc_dask.py} (100%) diff --git a/tests/dask/test_hvg.py b/tests/dask/test_hvg_dask.py similarity index 100% rename from tests/dask/test_hvg.py rename to tests/dask/test_hvg_dask.py diff --git a/tests/dask/test_normalize.py b/tests/dask/test_normalize_dask.py similarity index 100% rename from tests/dask/test_normalize.py rename to tests/dask/test_normalize_dask.py diff --git a/tests/dask/test_qc.py b/tests/dask/test_qc_dask.py similarity index 100% rename from tests/dask/test_qc.py rename to tests/dask/test_qc_dask.py From 77b3c345aa06de7b5d84fcba5cf2be53c4859e22 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Sat, 4 May 2024 12:36:43 +0200 Subject: [PATCH 12/69] get docs to work --- .../preprocessing/_normalize.py | 53 ++++++++++--------- src/rapids_singlecell/preprocessing/_pca.py | 2 +- 2 files changed, 30 insertions(+), 25 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 7f978403..d59415c7 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -2,7 +2,8 @@ import math import warnings -from functools import singledispatch + +# from functools import singledispatch from typing import TYPE_CHECKING import cupy as cp @@ -89,23 +90,26 @@ def normalize_total( return X -@singledispatch -def _normalize_total(X: cp.ndarray, target_sum: int, client=None) -> cp.ndarray: - from ._kernels._norm_kernel import _mul_dense +def _normalize_total(X: cp.ndarray, target_sum: int, client=None): + if isinstance(X, sparse.csr_matrix): + return _normalize_total_csr(X, target_sum) + elif isinstance(X, DaskArray): + return _normalize_total_dask(X, target_sum, client=client) + else: + from ._kernels._norm_kernel import _mul_dense - if not X.flags.c_contiguous: - X = cp.asarray(X, order="C") - mul_kernel = _mul_dense(X.dtype) - mul_kernel( - (math.ceil(X.shape[0] / 128),), - (128,), - (X, X.shape[0], X.shape[1], int(target_sum)), - ) - return X + if not X.flags.c_contiguous: + X = cp.asarray(X, order="C") + mul_kernel = _mul_dense(X.dtype) + mul_kernel( + (math.ceil(X.shape[0] / 128),), + (128,), + (X, X.shape[0], X.shape[1], int(target_sum)), + ) + return X -@_normalize_total.register(sparse.csr_matrix) -def _(X: sparse.csr_matrix, target_sum: int, client=None) -> sparse.csr_matrix: +def _normalize_total_csr(X: sparse.csr_matrix, target_sum: int) -> sparse.csr_matrix: from ._kernels._norm_kernel import _mul_csr mul_kernel = _mul_csr(X.dtype) @@ -117,8 +121,7 @@ def _(X: sparse.csr_matrix, target_sum: int, client=None) -> sparse.csr_matrix: return X -@_normalize_total.register(DaskArray) -def _(X: DaskArray, target_sum: int, client=None) -> DaskArray: +def _normalize_total_dask(X: DaskArray, target_sum: int, client=None) -> DaskArray: client = _get_dask_client(client) if isinstance(X._meta, sparse.csr_matrix): from ._kernels._norm_kernel import _mul_csr @@ -155,13 +158,16 @@ def __mul(X_part): return X -@singledispatch -def _get_target_sum(X: cp.ndarray, client=None) -> int: - return cp.median(X.sum(axis=1)) +def _get_target_sum(X, client=None) -> int: + if isinstance(X, sparse.csr_matrix): + return _get_target_sum_csr(X, client=client) + elif isinstance(X, DaskArray): + return _get_target_sum_dask(X, client=client) + else: + return cp.median(X.sum(axis=1)) -@_get_target_sum.register(sparse.csr_matrix) -def _(X: sparse.csr_matrix, client=None) -> int: +def _get_target_sum_csr(X: sparse.csr_matrix) -> int: from ._kernels._norm_kernel import _get_sparse_sum_major counts_per_cell = cp.zeros(X.shape[0], dtype=X.dtype) @@ -176,8 +182,7 @@ def _(X: sparse.csr_matrix, client=None) -> int: return target_sum -@_get_target_sum.register(DaskArray) -def _(X: DaskArray, client=None) -> int: +def _get_target_sum_dask(X: DaskArray, client=None) -> int: import dask.array as da client = _get_dask_client(client) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 96188e2f..0fdf9724 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -30,7 +30,7 @@ def pca( zero_center: bool = True, svd_solver: str = None, random_state: int | None = 0, - mask_var: NDArray[np.bool_] | str | None | Empty = _empty, + mask_var: NDArray[bool] | str | None | Empty = _empty, use_highly_variable: bool | None = None, dtype: str = "float32", copy: bool = False, From 36bebf9777323d887f5c9582a203b903a85bd2ca Mon Sep 17 00:00:00 2001 From: Intron7 Date: Sat, 4 May 2024 12:43:30 +0200 Subject: [PATCH 13/69] remove client from sparse calc --- src/rapids_singlecell/preprocessing/_normalize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index d59415c7..cb78f16e 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -160,7 +160,7 @@ def __mul(X_part): def _get_target_sum(X, client=None) -> int: if isinstance(X, sparse.csr_matrix): - return _get_target_sum_csr(X, client=client) + return _get_target_sum_csr(X) elif isinstance(X, DaskArray): return _get_target_sum_dask(X, client=client) else: From 82cc22c0dbc113c9b601140daa927e1c7b78f5f8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Sat, 4 May 2024 13:00:50 +0200 Subject: [PATCH 14/69] need dask for docs --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 4123d1da..c852c113 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ doc = [ "scanpydoc[typehints,theme]>=0.9.4", "readthedocs-sphinx-ext", "sphinx_copybutton", + "dask", ] test = [ "pytest", From e33821faf7c972410b29107c2536929a363e0a3e Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 7 May 2024 17:34:13 +0200 Subject: [PATCH 15/69] add scale --- .../preprocessing/_kernels/_scale_kernel.py | 2 +- src/rapids_singlecell/preprocessing/_pca.py | 2 +- src/rapids_singlecell/preprocessing/_scale.py | 217 +++++++++++++++++- src/rapids_singlecell/preprocessing/_utils.py | 1 + tests/dask/test_scale_dask.py | 61 +++++ 5 files changed, 278 insertions(+), 5 deletions(-) create mode 100644 tests/dask/test_scale_dask.py diff --git a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py index b0a59ed3..686bfde6 100644 --- a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py +++ b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py @@ -20,7 +20,7 @@ """ _csr_scale_diff_kernel = r""" -(const int *indptr, const int *indices, {0} *data, const double * std, const int *mask, {0} clipper,int nrows) { +(const int *indptr, const int *indices, {0} *data, const {0} * std, const int *mask, {0} clipper,int nrows) { int row = blockIdx.x; if(row >= nrows){ diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 0fdf9724..96188e2f 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -30,7 +30,7 @@ def pca( zero_center: bool = True, svd_solver: str = None, random_state: int | None = 0, - mask_var: NDArray[bool] | str | None | Empty = _empty, + mask_var: NDArray[np.bool_] | str | None | Empty = _empty, use_highly_variable: bool | None = None, dtype: str = "float32", copy: bool = False, diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index fac96d0f..a6c68271 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -8,9 +8,21 @@ from cupyx.scipy import sparse from scanpy._utils import view_to_actual +from rapids_singlecell._compat import ( + DaskArray, + DaskClient, + _get_dask_client, + _meta_dense, + _meta_sparse, +) from rapids_singlecell.get import _check_mask, _get_obs_rep, _set_obs_rep from rapids_singlecell.preprocessing._utils import _check_gpu_X, _get_mean_var +try: + import dask.array as da +except ImportError: + pass + def scale( adata: AnnData, @@ -22,6 +34,7 @@ def scale( obsm: str | None = None, mask_obs: np.ndarray | str | None = None, inplace: bool = True, + client: DaskClient | None = None, ) -> None | cp.ndarray: """ Scales matrix to unit variance and clips values @@ -56,6 +69,9 @@ def scale( inplace If True, update AnnData with results. Otherwise, return results. See below for details of what is returned. + client + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. + Returns ------- Returns a scaled copy or updates `adata` with a scaled version of the original `adata.X` and `adata.layers['layer']`, \ @@ -71,7 +87,7 @@ def scale( view_to_actual(adata) X = _get_obs_rep(adata, layer=layer, obsm=obsm) - _check_gpu_X(X) + _check_gpu_X(X, allow_dask=True) str_mean_std = ("mean", "std") if mask_obs is not None: @@ -81,7 +97,19 @@ def scale( str_mean_std = ("mean with mask", "std with mask") mask_obs = _check_mask(adata, mask_obs, "obs") - if isinstance(X, cp.ndarray): + if isinstance(X, DaskArray): + if client is None: + client = _get_dask_client(client=client) + X, means, std = _scale_dask( + X, + mask_obs=mask_obs, + zero_center=zero_center, + inplace=inplace, + max_value=max_value, + client=client, + ) + + elif isinstance(X, cp.ndarray): X, means, std = _scale_array( X, mask_obs=mask_obs, @@ -240,12 +268,195 @@ def _scale_sparse_csr( scale_csr( (X.shape[0],), (64,), - (X.indptr, X.indices, X.data, std, mask_array, max_value, X.shape[0]), + ( + X.indptr, + X.indices, + X.data, + std.astype(X.dtype), + mask_array, + max_value, + X.shape[0], + ), ) return X, mean, std +def _scale_dask( + X, *, mask_obs=None, zero_center=True, inplace=True, max_value=None, client=None +): + if not inplace: + X = X.copy() + if mask_obs is None: + mean, var = _get_mean_var(X, client=client) + mask_array = cp.ones(X.shape[0], dtype=cp.int32) + + else: + mean, var = _get_mean_var(X[mask_obs, :], client=client) + mask_array = cp.array(mask_obs).astype(cp.int32) + std = cp.sqrt(var) + std[std == 0] = 1 + max_value = _get_max_value(max_value, X.dtype) + + mask_array = da.from_array( + mask_array, chunks=(X.chunks[0],), meta=_meta_dense(mask_array.dtype) + ) + + if isinstance(X._meta, sparse.csr_matrix) and zero_center: + from ._sparse_pca._kernels._pca_sparse_kernel import denser_kernel + + kernel = denser_kernel(X.dtype) + kernel.compile() + + def __dense(X_part): + dense = cp.zeros(X_part.shape, dtype=X.dtype) + max_nnz = cp.diff(X_part.indptr).max() + tpb = (32, 32) + bpg_x = math.ceil(X_part.shape[0] / tpb[0]) + bpg_y = math.ceil(max_nnz / tpb[1]) + bpg = (bpg_x, bpg_y) + kernel( + bpg, + tpb, + ( + X_part.indptr, + X_part.indices, + X_part.data, + dense, + X_part.shape[0], + X_part.shape[1], + ), + ) + return dense + + X = X.map_blocks( + lambda x: __dense(x), + dtype=X.dtype, + meta=_meta_dense(X.dtype), + ) + return _scale_dask_array_zc( + X, mask_array=mask_array, mean=mean, std=std, max_value=max_value + ) + + elif isinstance(X._meta, sparse.csr_matrix) and not zero_center: + return _scale_sparse_csr_dask( + X, mask_array=mask_array, mean=mean, std=std, max_value=max_value + ) + + elif isinstance(X._meta, cp.ndarray) and zero_center: + return _scale_dask_array_zc( + X, mask_array=mask_array, mean=mean, std=std, max_value=max_value + ) + + elif isinstance(X._meta, cp.ndarray) and not zero_center: + return _scale_dask_array_nzc( + X, mask_array=mask_array, mean=mean, std=std, max_value=max_value + ) + + +def _scale_dask_array_zc(X, *, mask_array, mean, std, max_value): + from ._kernels._scale_kernel import _dense_center_scale_kernel + + scale_kernel_center = _dense_center_scale_kernel(X.dtype) + scale_kernel_center.compile() + + mean_ = mean.astype(X.dtype) + std_ = std.astype(X.dtype) + + def __scale_kernel_center(X_part, mask_part): + scale_kernel_center( + (math.ceil(X_part.shape[0] / 32), math.ceil(X_part.shape[1] / 32)), + (32, 32), + ( + X_part, + mean_, + std_, + mask_part, + max_value, + X_part.shape[0], + X_part.shape[1], + ), + ) + return X_part + + X = da.blockwise( + __scale_kernel_center, + "ij", + X, + "ij", + mask_array, + "i", + meta=_meta_dense(X.dtype), + dtype=X.dtype, + ) + return X, mean, std + + +def _scale_dask_array_nzc(X, *, mask_array, mean, std, max_value): + from ._kernels._scale_kernel import _dense_scale_kernel + + scale_kernel = _dense_scale_kernel(X.dtype) + scale_kernel.compile() + std_ = std.astype(X.dtype) + + def __scale_kernel(X_part, mask_part): + scale_kernel( + (math.ceil(X_part.shape[0] / 32), math.ceil(X_part.shape[1] / 32)), + (32, 32), + (X_part, std_, mask_part, max_value, X_part.shape[0], X_part.shape[1]), + ) + + return X_part + + X = da.blockwise( + __scale_kernel, + "ij", + X, + "ij", + mask_array, + "i", + meta=_meta_dense(X.dtype), + dtype=X.dtype, + ) + return X, mean, std + + +def _scale_sparse_csr_dask(X, *, mask_array, mean, std, max_value): + from ._kernels._scale_kernel import _csr_scale_kernel + + scale_kernel_csr = _csr_scale_kernel(X.dtype) + scale_kernel_csr.compile() + std_ = std.astype(X.dtype) + + def __scale_kernel_csr(X_part, mask_part): + scale_kernel_csr( + (X_part.shape[0],), + (64,), + ( + X_part.indptr, + X_part.indices, + X_part.data, + std_, + mask_part, + max_value, + X_part.shape[0], + ), + ) + return X_part + + X = da.blockwise( + __scale_kernel_csr, + "ij", + X, + "ij", + mask_array, + "i", + meta=_meta_sparse(X.dtype), + dtype=X.dtype, + ) + return X, mean, std + + def _get_max_value(val, dtype): if val is None: val = np.inf diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 875bbec3..b958b0f3 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -165,6 +165,7 @@ def _mean_var_dense_dask(X, axis, client=None): client = _get_dask_client(client) + # ToDo: get a 64bit version working without copying the data def __mean_var(X_part, axis): mean = X_part.sum(axis=axis) var = (X_part**2).sum(axis=axis) diff --git a/tests/dask/test_scale_dask.py b/tests/dask/test_scale_dask.py new file mode 100644 index 00000000..961c3876 --- /dev/null +++ b/tests/dask/test_scale_dask.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +import cupy as cp +import numpy as np +from cupyx.scipy import sparse as cusparse +from scipy import sparse +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +import rapids_singlecell as rsc +import scanpy as sc + +from scanpy.datasets import pbmc3k + + +def _get_anndata(): + adata = pbmc3k() + sc.pp.filter_cells(adata, min_genes=100) + sc.pp.filter_genes(adata, min_cells=3) + sc.pp.normalize_total(adata) + sc.pp.log1p(adata) + sc.pp.highly_variable_genes(adata, n_top_genes=1000, subset=True) + return adata.copy() + +def test_zc_sparse(client): + adata = _get_anndata() + mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X.astype(np.float64)) + adata.X = cusparse.csr_matrix(adata.X.astype(np.float64)) + rsc.pp.scale(adata, mask_obs = mask, max_value = 10) + rsc.pp.scale(dask_data, mask_obs = mask, max_value = 10) + cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + +def test_nzc_sparse(client): + adata = _get_anndata() + mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + dask_data = adata.copy() + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + rsc.pp.scale(adata, zero_center = False, mask_obs = mask, max_value = 10) + rsc.pp.scale(dask_data,zero_center = False, mask_obs = mask, max_value = 10) + cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + +def test_zc_dense(client): + adata = _get_anndata() + mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) + adata.X = cp.array(adata.X.toarray().astype(np.float64)) + rsc.pp.scale(adata, mask_obs = mask, max_value = 10) + rsc.pp.scale(dask_data, mask_obs = mask, max_value = 10) + cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + +def test_nzc_dense(client): + adata = _get_anndata() + mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + dask_data = adata.copy() + dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) + adata.X = cp.array(adata.X.toarray().astype(np.float64)) + rsc.pp.scale(adata, zero_center = False, mask_obs = mask, max_value = 10) + rsc.pp.scale(dask_data, zero_center = False, mask_obs = mask, max_value = 10) + cp.testing.assert_allclose(adata.X, dask_data.X.compute()) From e1e6c199f7f3a7a87ec8321b33d2a5cb7c983729 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 8 May 2024 14:59:26 +0200 Subject: [PATCH 16/69] int64 updates --- .../preprocessing/_kernels/_scale_kernel.py | 6 +++--- .../_sparse_pca/_kernels/_pca_sparse_kernel.py | 3 ++- .../preprocessing/_sparse_pca/_sparse_pca.py | 2 +- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py index 686bfde6..07f8e512 100644 --- a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py +++ b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py @@ -39,10 +39,10 @@ """ _dense_scale_center_diff_kernel = r""" -({0} *data, const {0} *mean, const {0} *std, const int *mask, {0} clipper,int nrows,int ncols) +({0} *data, const {0} *mean, const {0} *std, const int *mask, {0} clipper,long long int nrows,long long int ncols) { - int row = blockIdx.x * blockDim.x + threadIdx.x; - int col = blockIdx.y * blockDim.y + threadIdx.y; + long long int row = blockIdx.x * blockDim.x + threadIdx.x; + long long int col = blockIdx.y * blockDim.y + threadIdx.y; if (row < nrows && col < ncols) { if (mask[row]){ diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py index 889fe350..d4115c4c 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py @@ -76,7 +76,8 @@ return; } int idx = index[start + col]; - out[row*minor+idx] = data[start + col];} + long long int res_index = static_cast(row)*minor+idx; + out[res_index] = data[start + col];} """ _zero_genes_kernel = cp.RawKernel(check_zero_genes, "check_zero_genes") diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index 22535384..95e7bad5 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -66,7 +66,7 @@ def transform(self, X): tpb, (X.indptr, X.indices, X.data, dense, X.shape[0], X.shape[1]), ) - dense = dense - self.mean_ + dense -= self.mean_ X_transformed = dense.dot(self.components_.T) # X = X - self.mean_ # X_transformed = X.dot(self.components_.T) From 7da41e0d360165116320552bae47db30e96b58ff Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 8 May 2024 15:12:27 +0200 Subject: [PATCH 17/69] For main branch --- .../preprocessing/_kernels/_scale_kernel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py index 07f8e512..686bfde6 100644 --- a/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py +++ b/src/rapids_singlecell/preprocessing/_kernels/_scale_kernel.py @@ -39,10 +39,10 @@ """ _dense_scale_center_diff_kernel = r""" -({0} *data, const {0} *mean, const {0} *std, const int *mask, {0} clipper,long long int nrows,long long int ncols) +({0} *data, const {0} *mean, const {0} *std, const int *mask, {0} clipper,int nrows,int ncols) { - long long int row = blockIdx.x * blockDim.x + threadIdx.x; - long long int col = blockIdx.y * blockDim.y + threadIdx.y; + int row = blockIdx.x * blockDim.x + threadIdx.x; + int col = blockIdx.y * blockDim.y + threadIdx.y; if (row < nrows && col < ncols) { if (mask[row]){ From b6f436fffc1528d39c46cffa7d528ef3098b50ee Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 8 May 2024 16:00:04 +0200 Subject: [PATCH 18/69] test docs --- src/rapids_singlecell/preprocessing/_pca.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 96188e2f..f6a8da26 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -30,7 +30,7 @@ def pca( zero_center: bool = True, svd_solver: str = None, random_state: int | None = 0, - mask_var: NDArray[np.bool_] | str | None | Empty = _empty, + mask_var: NDArray[np.bool] | str | None | Empty = _empty, use_highly_variable: bool | None = None, dtype: str = "float32", copy: bool = False, From 4b22562a7f69e17b6aedcc95ec74ffde576a756a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 13 May 2024 13:57:57 +0000 Subject: [PATCH 19/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 0e042727..8a8f5618 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -42,7 +42,6 @@ def _mean_var_minor(X, major, minor): return mean, var - @with_cupy_rmm def _mean_var_minor_dask(X, major, minor, client=None): """ @@ -199,6 +198,7 @@ def __mean_var(X_part, axis): mean, var = da.compute(mean, var) mean, var = mean.ravel(), var.ravel() + def _mean_var_dense(X, axis): from ._kernels._mean_var_kernel import mean_sum, sq_sum @@ -211,7 +211,6 @@ def _mean_var_dense(X, axis): return mean, var - def _get_mean_var(X, axis=0, client=None): if issparse(X): if axis == 0: From 5ed8e68bbabb5dd588237122bc2c7db8d16ccc1e Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 13 May 2024 16:08:54 +0200 Subject: [PATCH 20/69] fix import --- src/rapids_singlecell/preprocessing/_utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 8a8f5618..4a0a1e81 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -197,6 +197,7 @@ def __mean_var(X_part, axis): mean, var = da.compute(mean, var) mean, var = mean.ravel(), var.ravel() + return mean, var def _mean_var_dense(X, axis): From b879ea407100e2e88be9d7c92f10af2ef9a9c169 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 13 May 2024 16:23:04 +0200 Subject: [PATCH 21/69] fix rebase --- src/rapids_singlecell/preprocessing/_utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 4a0a1e81..018d8e29 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -197,6 +197,10 @@ def __mean_var(X_part, axis): mean, var = da.compute(mean, var) mean, var = mean.ravel(), var.ravel() + mean = mean / X.shape[axis] + var = var / X.shape[axis] + var -= cp.power(mean, 2) + var *= X.shape[axis] / (X.shape[axis] - 1) return mean, var From f002b9c782d39b164014fb5eab96b11838f90ac5 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Tue, 16 Jul 2024 13:42:40 +0200 Subject: [PATCH 22/69] (fix): use `to_delayed` and `from_delayed` to submit gram matrix jobs (#210) * (fix): use `to_delayed` and `from_delayed` to submit gram matrix jobs * (refactor): use `map_blocks` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): correct first dimension size * (fix): add `x` as arg * (fix): `ncols` usages * (fix): try mapping block from x * (fix): matrix creation + cleaner `map_blocks` * (fix): `len(blocks)` -> `num_blocks` * (fix): don't need `dask.delayed` decorator * (fix): try some debugging * (fix): use `cp.sum` * (fix): revert to `to_delayed` * (refactor): use `n_cols` * (fix): remove `client` * (fix): `client` doc --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Severin Dicks <37635888+Intron7@users.noreply.github.com> --- src/rapids_singlecell/preprocessing/_pca.py | 4 +- .../_sparse_pca/_dask_sparse_pca.py | 40 +++++++++---------- 2 files changed, 21 insertions(+), 23 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index f6a8da26..c0b5ae1b 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -92,7 +92,7 @@ def pca( Required if `chunked=True` was passed. client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. + Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a dense Dask array. Returns ------- @@ -151,7 +151,7 @@ def pca( elif isinstance(X._meta, csr_matrix): from ._sparse_pca._dask_sparse_pca import PCA_sparse_dask - pca_func = PCA_sparse_dask(n_components=n_comps, client=client) + pca_func = PCA_sparse_dask(n_components=n_comps) pca_func = pca_func.fit(X) X_pca = pca_func.transform(X) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index 7ae60d59..b451ec57 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -4,20 +4,17 @@ import cupy as cp import dask -from cuml.dask.common.part_utils import _extract_partitions from cuml.internals.memory_utils import with_cupy_rmm from rapids_singlecell._compat import ( - _get_dask_client, _meta_dense, ) from rapids_singlecell.preprocessing._utils import _get_mean_var class PCA_sparse_dask: - def __init__(self, n_components, client) -> None: + def __init__(self, n_components) -> None: self.n_components = n_components - self.client = _get_dask_client(client) def fit(self, x): if self.n_components is None: @@ -30,7 +27,7 @@ def fit(self, x): self.n_samples_ = x.shape[0] self.n_features_in_ = x.shape[1] if x.ndim == 2 else 1 self.dtype = x.dtype - covariance, self.mean_, _ = _cov_sparse_dask(self.client, x=x, return_mean=True) + covariance, self.mean_, _ = _cov_sparse_dask(x=x, return_mean=True) self.explained_variance_, self.components_ = cp.linalg.eigh( covariance, UPLO="U" ) @@ -106,7 +103,7 @@ def fit_transform(self, X, y=None): @with_cupy_rmm -def _cov_sparse_dask(client, x, return_gram=False, return_mean=False): +def _cov_sparse_dask(x, return_gram=False, return_mean=False): """ Computes the mean and the covariance of matrix X of the form Cov(X, X) = E(XX) - E(X)E(X) @@ -152,8 +149,10 @@ def _cov_sparse_dask(client, x, return_gram=False, return_mean=False): compute_mean_cov = _gramm_kernel_csr(x.dtype) compute_mean_cov.compile() + n_cols = x.shape[1] - def __gram_block(x_part, n_cols): + @dask.delayed + def __gram_block(x_part): gram_matrix = cp.zeros((n_cols, n_cols), dtype=x.dtype) block = (128,) @@ -172,27 +171,26 @@ def __gram_block(x_part, n_cols): ) return gram_matrix - parts = client.sync(_extract_partitions, x) - futures = [ - client.submit(__gram_block, part, x.shape[1], workers=[w]) for w, part in parts - ] - # Gather results from futures - objs = [] - for i in range(len(futures)): - obj = dask.array.from_delayed( - futures[i], shape=(x.shape[1], x.shape[1]), dtype=x.dtype + blocks = x.to_delayed().ravel() + gram_chunk_matrices = [ + dask.array.from_delayed( + __gram_block(block), + shape=(n_cols, n_cols), + dtype=x.dtype, + meta=cp.array([]), ) - objs.append(obj) - gram_matrix = dask.array.stack(objs).sum(axis=0).compute() - mean_x, _ = _get_mean_var(x, client=client) + for block in blocks + ] + gram_matrix = dask.array.stack(gram_chunk_matrices).sum(axis=0).compute() + mean_x, _ = _get_mean_var(x) mean_x = mean_x.astype(x.dtype) copy_gram = _copy_kernel(x.dtype) block = (32, 32) - grid = (math.ceil(x.shape[1] / block[0]), math.ceil(x.shape[1] / block[1])) + grid = (math.ceil(n_cols / block[0]), math.ceil(n_cols / block[1])) copy_gram( grid, block, - (gram_matrix, x.shape[1]), + (gram_matrix, n_cols), ) gram_matrix *= 1 / x.shape[0] From a8b30d3c5e96b4639f6b94fb6ee99a981b3dc0c0 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Tue, 16 Jul 2024 13:43:24 +0200 Subject: [PATCH 23/69] (fix): use `map_blocks` for job submission in `_get_target_sum_dask` + `_second_pass_qc` (#211) * (fix): use `to_delayed` and `from_delayed` to submit gram matrix jobs * (refactor): use `map_blocks` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): correct first dimension size * (fix): change normalization sum * (fix): add `x` as arg * (fix): `ncols` usages * (fix): try mapping block from x * (fix): use `X` directly to `map_blocks` * (fix): matrix creation + cleaner `map_blocks` * (fix): `len(blocks)` -> `num_blocks` * (fix): don't need `dask.delayed` decorator * (fix): try some debugging * (fix): use `cp.sum` * (fix): revert to `to_delayed` * (refactor): use `n_cols` * (fix): remove `num_blocks` * (fix): need to specify `drop_axis` for reduction * (fix): `map_blocks` for `_second_pass_qc_dask` * (fix): remove `client` * (chore): remove `client` * (chore): remove in tests * (fix): `client` doc * (fix): return client to test context --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Severin Dicks <37635888+Intron7@users.noreply.github.com> --- .../preprocessing/_normalize.py | 47 +++++++------------ src/rapids_singlecell/preprocessing/_qc.py | 19 ++------ 2 files changed, 21 insertions(+), 45 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index cb78f16e..2e87ae77 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -7,14 +7,11 @@ from typing import TYPE_CHECKING import cupy as cp -from cuml.dask.common.part_utils import _extract_partitions from cupyx.scipy import sparse from scanpy.get import _get_obs_rep, _set_obs_rep from rapids_singlecell._compat import ( DaskArray, - DaskClient, - _get_dask_client, _meta_dense, _meta_sparse, ) @@ -32,7 +29,6 @@ def normalize_total( layer: int | str = None, inplace: bool = True, copy: bool = False, - client: DaskClient | None = None, ) -> AnnData | sparse.csr_matrix | cp.ndarray | None: """ Normalizes rows in matrix so they sum to `target_sum` @@ -53,10 +49,6 @@ def normalize_total( copy Whether to return a copy or update `adata`. Not compatible with inplace=False. - - client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. - Returns ------- Returns a normalized copy or updates `adata` with a normalized version of \ @@ -77,9 +69,9 @@ def normalize_total( if sparse.isspmatrix_csc(X): X = X.tocsr() if not target_sum: - target_sum = _get_target_sum(X, client=client) + target_sum = _get_target_sum(X) - X = _normalize_total(X, target_sum, client=client) + X = _normalize_total(X, target_sum) if inplace: _set_obs_rep(adata, X, layer=layer) @@ -90,11 +82,11 @@ def normalize_total( return X -def _normalize_total(X: cp.ndarray, target_sum: int, client=None): +def _normalize_total(X: cp.ndarray, target_sum: int): if isinstance(X, sparse.csr_matrix): return _normalize_total_csr(X, target_sum) elif isinstance(X, DaskArray): - return _normalize_total_dask(X, target_sum, client=client) + return _normalize_total_dask(X, target_sum) else: from ._kernels._norm_kernel import _mul_dense @@ -121,8 +113,7 @@ def _normalize_total_csr(X: sparse.csr_matrix, target_sum: int) -> sparse.csr_ma return X -def _normalize_total_dask(X: DaskArray, target_sum: int, client=None) -> DaskArray: - client = _get_dask_client(client) +def _normalize_total_dask(X: DaskArray, target_sum: int) -> DaskArray: if isinstance(X._meta, sparse.csr_matrix): from ._kernels._norm_kernel import _mul_csr @@ -158,11 +149,11 @@ def __mul(X_part): return X -def _get_target_sum(X, client=None) -> int: +def _get_target_sum(X) -> int: if isinstance(X, sparse.csr_matrix): return _get_target_sum_csr(X) elif isinstance(X, DaskArray): - return _get_target_sum_dask(X, client=client) + return _get_target_sum_dask(X) else: return cp.median(X.sum(axis=1)) @@ -182,11 +173,7 @@ def _get_target_sum_csr(X: sparse.csr_matrix) -> int: return target_sum -def _get_target_sum_dask(X: DaskArray, client=None) -> int: - import dask.array as da - - client = _get_dask_client(client) - +def _get_target_sum_dask(X: DaskArray) -> int: if isinstance(X._meta, sparse.csr_matrix): from ._kernels._norm_kernel import _get_sparse_sum_major @@ -208,16 +195,14 @@ def __sum(X_part): return X_part.sum(axis=1) else: raise ValueError(f"Cannot compute target sum for {type(X)}") - - parts = client.sync(_extract_partitions, X) - futures = [client.submit(__sum, part, workers=[w]) for w, part in parts] - # Gather results from futures - futures = client.gather(futures) - objs = [] - for i in futures: - objs.append(da.from_array(i, chunks=i.shape)) - - counts_per_cell = da.concatenate(objs).compute() + target_sum_chunk_matrices = X.map_blocks( + __sum, + meta=cp.array((1.0,), dtype=X.dtype), + dtype=X.dtype, + chunks=(X.chunksize[0],), + drop_axis=1, + ) + counts_per_cell = target_sum_chunk_matrices.compute() target_sum = cp.median(counts_per_cell) return target_sum diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index ba5ab9b9..b19fae9c 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -187,7 +187,6 @@ def _first_pass_qc(X, client=None): @with_cupy_rmm def _first_pass_qc_dask(X, client=None): import dask - from cuml.dask.common.part_utils import _extract_partitions client = _get_dask_client(client) @@ -334,8 +333,6 @@ def _second_pass_qc(X, mask, client=None): @with_cupy_rmm def _second_pass_qc_dask(X, mask, client=None): - import dask - client = _get_dask_client(client) if isinstance(X._meta, sparse.csr_matrix): @@ -345,7 +342,7 @@ def _second_pass_qc_dask(X, mask, client=None): sparse_qc_csr.compile() def __qc_calc(X_part): - sums_cells_sub = cp.zeros((X_part.shape[0]), dtype=X_part.dtype) + sums_cells_sub = cp.zeros(X_part.shape[0], dtype=X_part.dtype) block = (32,) grid = (int(math.ceil(X_part.shape[0] / block[0])),) sparse_qc_csr( @@ -369,7 +366,7 @@ def __qc_calc(X_part): sparse_qc_dense.compile() def __qc_calc(X_part): - sums_cells_sub = cp.zeros((X_part.shape[0]), dtype=X_part.dtype) + sums_cells_sub = cp.zeros(X_part.shape[0], dtype=X_part.dtype) if not X_part.flags.c_contiguous: X_part = cp.asarray(X_part, order="C") block = (16, 16) @@ -384,13 +381,7 @@ def __qc_calc(X_part): ) return sums_cells_sub - parts = client.sync(_extract_partitions, X) - futures = [client.submit(__qc_calc, part, workers=[w]) for w, part in parts] - # Gather results from futures - futures = client.gather(futures) - objs = [] - for i in futures: - objs.append(dask.array.from_array(i, chunks=i.shape)) - - sums_cells_sub = dask.array.concatenate(objs).compute() + sums_cells_sub = X.map_blocks( + __qc_calc, dtype=X.dtype, meta=cp.array([]), drop_axis=1 + ).compute() return sums_cells_sub.ravel() From b22834522bcb12bbb99b030f128d99934a662789 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Tue, 16 Jul 2024 13:43:50 +0200 Subject: [PATCH 24/69] (fix): remove `extract_partitions` from `mean`/`var` calculation (#221) * (fix): use `to_delayed` and `from_delayed` to submit gram matrix jobs * (refactor): use `map_blocks` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): correct first dimension size * (fix): change normalization sum * (fix): add `x` as arg * (fix): `ncols` usages * (fix): try mapping block from x * (fix): use `X` directly to `map_blocks` * (fix): matrix creation + cleaner `map_blocks` * (fix): `len(blocks)` -> `num_blocks` * (fix): don't need `dask.delayed` decorator * (fix): try some debugging * (fix): use `cp.sum` * (fix): revert to `to_delayed` * (refactor): use `n_cols` * (fix): remove `num_blocks` * (fix): need to specify `drop_axis` for reduction * (fix): try splitting mean/var directly * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): `x` to `X` * (fix) correct import * (fix): `delayed` decorator * (fix): get axis right * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): ravel mean/var * (fix): too many mistakes to count * (feat): same for other axis * (fix): resolve all small tolerance differences * (fix): try splitting first * (fix): `compute` once * (fix): stack `mean`/`var` * (fix): use `float64` for mean-var * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): remove unnecessary cast * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): add `dask` dep for major axis * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (refactor): use cleaner `zeros` * (fix): revert other `rtol` * (fix): remove last `extract_partitions` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): `map_blocks` for `_second_pass_qc_dask` * (fix): remove `client` * (fix): remove other instances * (fix): remove `client` * (chore): remove `client` * (chore): remove in tests * (fix): `client` doc * (fix): remove more `client` * (fix): return client to test context * (chore): re-add client * (fix): oops * (fiX): oops x 2 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Severin Dicks <37635888+Intron7@users.noreply.github.com> --- src/rapids_singlecell/preprocessing/_hvg.py | 2 +- src/rapids_singlecell/preprocessing/_scale.py | 12 +- src/rapids_singlecell/preprocessing/_utils.py | 138 ++++++++---------- 3 files changed, 61 insertions(+), 91 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index a6d40c02..27b859f2 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -306,7 +306,7 @@ def _highly_variable_genes_single_batch( else: X = cp.expm1(X) - mean, var = _get_mean_var(X, axis=0, client=client) + mean, var = _get_mean_var(X, axis=0) mean[mean == 0] = 1e-12 disp = var / mean if flavor == "seurat": # logarithmized mean as in Seurat diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index a6c68271..aca9b92d 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -11,7 +11,6 @@ from rapids_singlecell._compat import ( DaskArray, DaskClient, - _get_dask_client, _meta_dense, _meta_sparse, ) @@ -98,15 +97,12 @@ def scale( mask_obs = _check_mask(adata, mask_obs, "obs") if isinstance(X, DaskArray): - if client is None: - client = _get_dask_client(client=client) X, means, std = _scale_dask( X, mask_obs=mask_obs, zero_center=zero_center, inplace=inplace, max_value=max_value, - client=client, ) elif isinstance(X, cp.ndarray): @@ -282,17 +278,15 @@ def _scale_sparse_csr( return X, mean, std -def _scale_dask( - X, *, mask_obs=None, zero_center=True, inplace=True, max_value=None, client=None -): +def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=None): if not inplace: X = X.copy() if mask_obs is None: - mean, var = _get_mean_var(X, client=client) + mean, var = _get_mean_var(X) mask_array = cp.ones(X.shape[0], dtype=cp.int32) else: - mean, var = _get_mean_var(X[mask_obs, :], client=client) + mean, var = _get_mean_var(X[mask_obs, :]) mask_array = cp.array(mask_obs).astype(cp.int32) std = cp.sqrt(var) std[std == 0] = 1 diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 018d8e29..4320babf 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -3,11 +3,10 @@ import math import cupy as cp -from cuml.dask.common.part_utils import _extract_partitions from cuml.internals.memory_utils import with_cupy_rmm from cupyx.scipy.sparse import issparse, isspmatrix_csc, isspmatrix_csr -from rapids_singlecell._compat import DaskArray, _get_dask_client +from rapids_singlecell._compat import DaskArray def _mean_var_major(X, major, minor): @@ -43,72 +42,64 @@ def _mean_var_minor(X, major, minor): @with_cupy_rmm -def _mean_var_minor_dask(X, major, minor, client=None): +def _mean_var_minor_dask(X, major, minor): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ + import dask import dask.array as da from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( _get_mean_var_minor, ) - client = _get_dask_client(client) - get_mean_var_minor = _get_mean_var_minor(X.dtype) get_mean_var_minor.compile() + @dask.delayed def __mean_var(X_part, minor, major): - mean = cp.zeros((minor, 1), dtype=cp.float64) - var = cp.zeros((minor, 1), dtype=cp.float64) + mean = cp.zeros(minor, dtype=cp.float64) + var = cp.zeros(minor, dtype=cp.float64) block = (32,) grid = (int(math.ceil(X_part.nnz / block[0])),) get_mean_var_minor( grid, block, (X_part.indices, X_part.data, mean, var, major, X_part.nnz) ) - return mean, var - - parts = client.sync(_extract_partitions, X) - futures = [ - client.submit(__mean_var, part, minor, major, workers=[w]) for w, part in parts + return cp.vstack([mean, var]) + + blocks = X.to_delayed().ravel() + mean_var_blocks = [ + da.from_delayed( + __mean_var(block, minor, major), + shape=(2, minor), + dtype=cp.float64, + meta=cp.array([]), + ) + for block in blocks ] - # Gather results from futures - results = client.gather(futures) - - # Initialize lists to hold the Dask arrays - means_objs = [] - var_objs = [] - - # Process each result - for means, vars in results: - # Append the arrays to their respective lists as Dask arrays - means_objs.append(da.from_array(means, chunks=means.shape)) - var_objs.append(da.from_array(vars, chunks=vars.shape)) - mean = da.concatenate(means_objs, axis=1).sum(axis=1) - var = da.concatenate(var_objs, axis=1).sum(axis=1) - mean, var = da.compute(mean, var) - mean, var = mean.ravel(), var.ravel() + + mean, var = da.stack(mean_var_blocks, axis=1).sum(axis=1).compute() var = (var - mean**2) * (major / (major - 1)) return mean, var # todo: Implement this dynamically for csc matrix as well @with_cupy_rmm -def _mean_var_major_dask(X, major, minor, client=None): +def _mean_var_major_dask(X, major, minor): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ + import dask import dask.array as da from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( _get_mean_var_major, ) - client = _get_dask_client(client) - get_mean_var_major = _get_mean_var_major(X.dtype) get_mean_var_major.compile() + @dask.delayed def __mean_var(X_part, minor, major): mean = cp.zeros(X_part.shape[0], dtype=cp.float64) var = cp.zeros(X_part.shape[0], dtype=cp.float64) @@ -127,28 +118,21 @@ def __mean_var(X_part, minor, major): minor, ), ) - return mean, var - - parts = client.sync(_extract_partitions, X) - futures = [ - client.submit(__mean_var, part, minor, major, workers=[w]) for w, part in parts + return cp.vstack([mean, var]) + + blocks = X.to_delayed().ravel() + mean_var_blocks = [ + da.from_delayed( + __mean_var(block, minor, major), + shape=(2, X.chunks[0][ind]), + dtype=cp.float64, + meta=cp.array([]), + ) + for ind, block in enumerate(blocks) ] - # Gather results from futures - results = client.gather(futures) - - # Initialize lists to hold the Dask arrays - means_objs = [] - var_objs = [] - - # Process each result - for means, vars_ in results: - # Append the arrays to their respective lists as Dask arrays - means_objs.append(da.from_array(means, chunks=means.shape)) - var_objs.append(da.from_array(vars_, chunks=vars_.shape)) - mean = da.concatenate(means_objs) - var = da.concatenate(var_objs) - mean, var = da.compute(mean, var) - mean, var = mean.ravel(), var.ravel() + + mean, var = da.hstack(mean_var_blocks).compute() + mean = mean / minor var = var / minor var -= cp.power(mean, 2) @@ -157,46 +141,38 @@ def __mean_var(X_part, minor, major): @with_cupy_rmm -def _mean_var_dense_dask(X, axis, client=None): +def _mean_var_dense_dask(X, axis): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ + import dask import dask.array as da - client = _get_dask_client(client) - # ToDo: get a 64bit version working without copying the data + @dask.delayed def __mean_var(X_part, axis): mean = X_part.sum(axis=axis) var = (X_part**2).sum(axis=axis) if axis == 0: mean = mean.reshape(-1, 1) var = var.reshape(-1, 1) - return mean, var - - parts = client.sync(_extract_partitions, X) - futures = [client.submit(__mean_var, part, axis, workers=[w]) for w, part in parts] - # Gather results from futures - results = client.gather(futures) - - # Initialize lists to hold the Dask arrays - means_objs = [] - var_objs = [] - - # Process each result - for means, vars_ in results: - # Append the arrays to their respective lists as Dask arrays - means_objs.append(da.from_array(means, chunks=means.shape)) - var_objs.append(da.from_array(vars_, chunks=vars_.shape)) + return cp.vstack([mean.ravel(), var.ravel()]) + + blocks = X.to_delayed().ravel() + mean_var_blocks = [ + da.from_delayed( + __mean_var(block, axis=axis), + shape=(2, X.chunks[0][ind]) if axis else (2, X.shape[1]), + dtype=cp.float64, + meta=cp.array([]), + ) + for ind, block in enumerate(blocks) + ] if axis == 0: - mean = da.concatenate(means_objs, axis=1).sum(axis=1) - var = da.concatenate(var_objs, axis=1).sum(axis=1) + mean, var = da.stack(mean_var_blocks, axis=1).sum(axis=1).compute() else: - mean = da.concatenate(means_objs) - var = da.concatenate(var_objs) + mean, var = da.hstack(mean_var_blocks).compute() - mean, var = da.compute(mean, var) - mean, var = mean.ravel(), var.ravel() mean = mean / X.shape[axis] var = var / X.shape[axis] var -= cp.power(mean, 2) @@ -216,7 +192,7 @@ def _mean_var_dense(X, axis): return mean, var -def _get_mean_var(X, axis=0, client=None): +def _get_mean_var(X, axis=0): if issparse(X): if axis == 0: if isspmatrix_csr(X): @@ -246,13 +222,13 @@ def _get_mean_var(X, axis=0, client=None): if axis == 0: major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_minor_dask(X, major, minor, client) + mean, var = _mean_var_minor_dask(X, major, minor) if axis == 1: major = X.shape[0] minor = X.shape[1] - mean, var = _mean_var_major_dask(X, major, minor, client) + mean, var = _mean_var_major_dask(X, major, minor) elif isinstance(X._meta, cp.ndarray): - mean, var = _mean_var_dense_dask(X, axis, client) + mean, var = _mean_var_dense_dask(X, axis) else: mean, var = _mean_var_dense(X, axis) return mean, var From d7c34cec4463b88fa35b136a81e42d332735e333 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 16 Jul 2024 12:22:15 +0000 Subject: [PATCH 25/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_pca.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index cde51978..c0b5ae1b 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -232,7 +232,6 @@ def pca( return adata - def _as_numpy(X): if isinstance(X, cp.ndarray): return X.get() From 7390dd146adaacb8ddf3f6bf395f1a73fa8665e8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 16 Jul 2024 14:30:58 +0200 Subject: [PATCH 26/69] remove client from hvg --- src/rapids_singlecell/preprocessing/_hvg.py | 12 ++---------- src/rapids_singlecell/preprocessing/_utils.py | 5 ++--- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 27b859f2..cba1d07a 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -12,7 +12,7 @@ from cupyx.scipy.sparse import csr_matrix, issparse, isspmatrix_csc from scanpy.get import _get_obs_rep -from rapids_singlecell._compat import DaskArray, DaskClient, _meta_dense, _meta_sparse +from rapids_singlecell._compat import DaskArray, _meta_dense, _meta_sparse from ._qc import calculate_qc_metrics from ._utils import _check_gpu_X, _check_nonnegative_integers, _get_mean_var @@ -49,7 +49,6 @@ def highly_variable_genes( chunksize: int = 1000, n_samples: int = 10000, batch_key: str = None, - client: DaskClient | None = None, ) -> None: """\ Annotate highly variable genes. @@ -119,8 +118,6 @@ def highly_variable_genes( of enrichment of zeros for each gene (only for `flavor='poisson_gene_selection'`). batch_key If specified, highly-variable genes are selected within each batch separately and merged. - client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. Returns ------- @@ -198,7 +195,6 @@ def highly_variable_genes( cutoff=cutoff, n_bins=n_bins, flavor=flavor, - client=client, ) else: df = _highly_variable_genes_batched( @@ -208,7 +204,6 @@ def highly_variable_genes( cutoff=cutoff, n_bins=n_bins, flavor=flavor, - client=client, ) adata.uns["hvg"] = {"flavor": flavor} @@ -278,7 +273,6 @@ def _highly_variable_genes_single_batch( cutoff: _Cutoffs | int, n_bins: int = 20, flavor: Literal["seurat", "cell_ranger"] = "seurat", - client: DaskClient | None = None, ) -> pd.DataFrame: """\ See `highly_variable_genes`. @@ -425,7 +419,6 @@ def _highly_variable_genes_batched( n_bins: int, flavor: Literal["seurat", "cell_ranger"], cutoff: _Cutoffs | int, - client: DaskClient | None = None, ) -> pd.DataFrame: adata._sanitize() batches = adata.obs[batch_key].cat.categories @@ -434,7 +427,7 @@ def _highly_variable_genes_batched( for batch in batches: adata_subset = adata[adata.obs[batch_key] == batch].copy() - calculate_qc_metrics(adata_subset, layer=layer, client=client) + calculate_qc_metrics(adata_subset, layer=layer) filt = adata_subset.var["n_cells_by_counts"].to_numpy() > 0 adata_subset = adata_subset[:, filt].copy() @@ -444,7 +437,6 @@ def _highly_variable_genes_batched( cutoff=cutoff, n_bins=n_bins, flavor=flavor, - client=client, ) hvg.reset_index(drop=False, inplace=True, names=["gene"]) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 90327573..9c8b1b4a 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -7,6 +7,8 @@ from cuml.internals.memory_utils import with_cupy_rmm from cupyx.scipy.sparse import issparse, isspmatrix_csc, isspmatrix_csr, spmatrix +from rapids_singlecell._compat import DaskArray + def _sparse_to_dense(X: spmatrix, order: Literal["C", "F"] | None = None) -> cp.ndarray: if order is None: @@ -38,9 +40,6 @@ def _sparse_to_dense(X: spmatrix, order: Literal["C", "F"] | None = None) -> cp. return dense -from rapids_singlecell._compat import DaskArray - - def _mean_var_major(X, major, minor): from ._kernels._mean_var_kernel import _get_mean_var_major From 7fedc07cd6b6d262d3549a80b8721b1219fabc49 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 16 Jul 2024 15:49:55 +0200 Subject: [PATCH 27/69] remove client --- .../_kernels/_qc_kernels_dask.py | 102 +++++++++++ src/rapids_singlecell/preprocessing/_qc.py | 169 +++++++++++------- 2 files changed, 208 insertions(+), 63 deletions(-) create mode 100644 src/rapids_singlecell/preprocessing/_kernels/_qc_kernels_dask.py diff --git a/src/rapids_singlecell/preprocessing/_kernels/_qc_kernels_dask.py b/src/rapids_singlecell/preprocessing/_kernels/_qc_kernels_dask.py new file mode 100644 index 00000000..01172407 --- /dev/null +++ b/src/rapids_singlecell/preprocessing/_kernels/_qc_kernels_dask.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from cuml.common.kernel_utils import cuda_kernel_factory + +_sparse_qc_kernel_csr_dask_cells = r""" + (const int *indptr,const int *index,const {0} *data, + {0}* sums_cells, int* cell_ex, + int n_cells) { + int cell = blockDim.x * blockIdx.x + threadIdx.x; + if(cell >= n_cells){ + return; + } + int start_idx = indptr[cell]; + int stop_idx = indptr[cell+1]; + + {0} sums_cells_i = 0; + int cell_ex_i = 0; + for(int gene = start_idx; gene < stop_idx; gene++){ + {0} value = data[gene]; + int gene_number = index[gene]; + sums_cells_i += value; + cell_ex_i += 1; + } + sums_cells[cell] = sums_cells_i; + cell_ex[cell] = cell_ex_i; + } +""" + + +_sparse_qc_kernel_csr_dask_genes = r""" + (const int *index,const {0} *data, + {0}* sums_genes, int* gene_ex, + int nnz) { + int idx = blockDim.x * blockIdx.x + threadIdx.x; + if(idx >= nnz){ + return; + } + int minor_pos = index[idx]; + atomicAdd(&sums_genes[minor_pos], data[idx]); + atomicAdd(&gene_ex[minor_pos], 1); + } + """ + +_sparse_qc_kernel_dense_cells = r""" + (const {0} *data, + {0}* sums_cells, int* cell_ex, + int n_cells,int n_genes) { + int cell = blockDim.x * blockIdx.x + threadIdx.x; + int gene = blockDim.y * blockIdx.y + threadIdx.y; + if(cell >= n_cells || gene >=n_genes){ + return; + } + long long int index = static_cast(cell) * n_genes + gene; + {0} value = data[index]; + if (value>0.0){ + atomicAdd(&sums_cells[cell], value); + atomicAdd(&cell_ex[cell], 1); + } + } +""" + +_sparse_qc_kernel_dense_genes = r""" + (const {0} *data, + {0}* sums_genes,int* gene_ex, + int n_cells,int n_genes) { + int cell = blockDim.x * blockIdx.x + threadIdx.x; + int gene = blockDim.y * blockIdx.y + threadIdx.y; + if(cell >= n_cells || gene >=n_genes){ + return; + } + long long int index = static_cast(cell) * n_genes + gene; + {0} value = data[index]; + if (value>0.0){ + atomicAdd(&sums_genes[gene], value); + atomicAdd(&gene_ex[gene], 1); + } + } +""" + + +def _sparse_qc_csr_dask_cells(dtype): + return cuda_kernel_factory( + _sparse_qc_kernel_csr_dask_cells, (dtype,), "_sparse_qc_kernel_csr_dask_cells" + ) + + +def _sparse_qc_csr_dask_genes(dtype): + return cuda_kernel_factory( + _sparse_qc_kernel_csr_dask_genes, (dtype,), "_sparse_qc_kernel_csr_dask_genes" + ) + + +def _sparse_qc_dense_cells(dtype): + return cuda_kernel_factory( + _sparse_qc_kernel_dense_cells, (dtype,), "_sparse_qc_kernel_dense_cells" + ) + + +def _sparse_qc_dense_genes(dtype): + return cuda_kernel_factory( + _sparse_qc_kernel_dense_genes, (dtype,), "_sparse_qc_kernel_dense_genes" + ) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index b19fae9c..0bf4f443 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -4,12 +4,11 @@ from typing import TYPE_CHECKING import cupy as cp -from cuml.dask.common.part_utils import _extract_partitions from cuml.internals.memory_utils import with_cupy_rmm from cupyx.scipy import sparse from scanpy.get import _get_obs_rep -from rapids_singlecell._compat import DaskArray, DaskClient, _get_dask_client +from rapids_singlecell._compat import DaskArray from ._utils import _check_gpu_X @@ -25,7 +24,6 @@ def calculate_qc_metrics( qc_vars: str | list = None, log1p: bool = True, layer: str = None, - client: DaskClient | None = None, ) -> None: """\ Calculates basic qc Parameters. Calculates number of genes per cell (n_genes) and number of counts per cell (n_counts). @@ -46,8 +44,6 @@ def calculate_qc_metrics( Set to `False` to skip computing `log1p` transformed annotations. layer If provided, use :attr:`~anndata.AnnData.layers` for expression values instead of :attr:`~anndata.AnnData.X`. - client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. Returns ------- @@ -78,7 +74,7 @@ def calculate_qc_metrics( _check_gpu_X(X, allow_dask=True) - sums_cells, sums_genes, cell_ex, gene_ex = _first_pass_qc(X, client=client) + sums_cells, sums_genes, cell_ex, gene_ex = _first_pass_qc(X) # .var adata.var[f"n_cells_by_{expr_type}"] = cp.asnumpy(gene_ex) adata.var[f"total_{expr_type}"] = cp.asnumpy(sums_genes) @@ -102,7 +98,7 @@ def calculate_qc_metrics( qc_vars = [qc_vars] for qc_var in qc_vars: mask = cp.array(adata.var[qc_var], dtype=cp.bool_) - sums_cells_sub = _second_pass_qc(X, mask, client=client) + sums_cells_sub = _second_pass_qc(X, mask) adata.obs[f"total_{expr_type}_{qc_var}"] = cp.asnumpy(sums_cells_sub) adata.obs[f"pct_{expr_type}_{qc_var}"] = cp.asnumpy( @@ -114,9 +110,9 @@ def calculate_qc_metrics( ) -def _first_pass_qc(X, client=None): +def _first_pass_qc(X): if isinstance(X, DaskArray): - return _first_pass_qc_dask(X, client=client) + return _first_pass_qc_dask(X) sums_cells = cp.zeros(X.shape[0], dtype=X.dtype) sums_genes = cp.zeros(X.shape[1], dtype=X.dtype) @@ -185,25 +181,27 @@ def _first_pass_qc(X, client=None): @with_cupy_rmm -def _first_pass_qc_dask(X, client=None): +def _first_pass_qc_dask(X): import dask - - client = _get_dask_client(client) + import dask.array as da if isinstance(X._meta, sparse.csr_matrix): - from ._kernels._qc_kernels import _sparse_qc_csr + from ._kernels._qc_kernels_dask import ( + _sparse_qc_csr_dask_cells, + _sparse_qc_csr_dask_genes, + ) - sparse_qc_csr = _sparse_qc_csr(X.dtype) - sparse_qc_csr.compile() + sparse_qc_csr_cells = _sparse_qc_csr_dask_cells(X.dtype) + sparse_qc_csr_cells.compile() - def __qc_calc(X_part): + @dask.delayed + def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) - sums_genes = cp.zeros((X_part.shape[1], 1), dtype=X_part.dtype) cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) - gene_ex = cp.zeros((X_part.shape[1], 1), dtype=cp.int32) block = (32,) grid = (int(math.ceil(X_part.shape[0] / block[0])),) - sparse_qc_csr( + + sparse_qc_csr_cells( grid, block, ( @@ -211,25 +209,47 @@ def __qc_calc(X_part): X_part.indices, X_part.data, sums_cells, - sums_genes, cell_ex, - gene_ex, X_part.shape[0], ), ) - return sums_cells, sums_genes, cell_ex, gene_ex + return cp.vstack([sums_cells, cell_ex.astype(X_part.dtype)]) + + sparse_qc_csr_genes = _sparse_qc_csr_dask_genes(X.dtype) + sparse_qc_csr_genes.compile() + + @dask.delayed + def __qc_calc_2(X_part): + sums_genes = cp.zeros(X_part.shape[1], dtype=X_part.dtype) + gene_ex = cp.zeros(X_part.shape[1], dtype=cp.int32) + block = (32,) + grid = (int(math.ceil(X_part.nnz / block[0])),) + sparse_qc_csr_genes( + grid, + block, + ( + X_part.indices, + X_part.data, + sums_genes, + gene_ex, + X_part.nnz, + ), + ) + return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)]) elif isinstance(X._meta, cp.ndarray): - from ._kernels._qc_kernels import _sparse_qc_dense + from ._kernels._qc_kernels_dask import ( + _sparse_qc_dense_cells, + _sparse_qc_dense_genes, + ) - sparse_qc_dense = _sparse_qc_dense(X.dtype) - sparse_qc_dense.compile() + sparse_qc_dense_cells = _sparse_qc_dense_cells(X.dtype) + sparse_qc_dense_cells.compile() - def __qc_calc(X_part): + @dask.delayed + def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) - sums_genes = cp.zeros((X_part.shape[1], 1), dtype=X_part.dtype) cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) - gene_ex = cp.zeros((X_part.shape[1], 1), dtype=cp.int32) if not X_part.flags.c_contiguous: X_part = cp.asarray(X_part, order="C") block = (16, 16) @@ -237,60 +257,85 @@ def __qc_calc(X_part): int(math.ceil(X_part.shape[0] / block[0])), int(math.ceil(X_part.shape[1] / block[1])), ) - sparse_qc_dense( + sparse_qc_dense_cells( grid, block, ( X_part, sums_cells, - sums_genes, cell_ex, + X_part.shape[0], + X_part.shape[1], + ), + ) + return cp.vstack([sums_cells, cell_ex.astype(X_part.dtype)]) + + sparse_qc_dense_genes = _sparse_qc_dense_genes(X.dtype) + sparse_qc_dense_genes.compile() + + @dask.delayed + def __qc_calc_2(X_part): + sums_genes = cp.zeros((X_part.shape[1]), dtype=X_part.dtype) + gene_ex = cp.zeros((X_part.shape[1]), dtype=cp.int32) + if not X_part.flags.c_contiguous: + X_part = cp.asarray(X_part, order="C") + block = (16, 16) + grid = ( + int(math.ceil(X_part.shape[0] / block[0])), + int(math.ceil(X_part.shape[1] / block[1])), + ) + sparse_qc_dense_genes( + grid, + block, + ( + X_part, + sums_genes, gene_ex, X_part.shape[0], X_part.shape[1], ), ) - return sums_cells, sums_genes, cell_ex, gene_ex + return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)]) else: raise ValueError( "Please use a cupy csr_matrix or cp.ndarray. csc_matrix are not supported with dask." ) - parts = client.sync(_extract_partitions, X) - futures = [client.submit(__qc_calc, part, workers=[w]) for w, part in parts] - # Gather results from futures - results = client.gather(futures) - - # Initialize lists to hold the Dask arrays - sums_cells_objs = [] - sums_genes_objs = [] - cell_ex_objs = [] - gene_ex_objs = [] - - # Process each result - for sums_cells, sums_genes, cell_ex, gene_ex in results: - # Append the arrays to their respective lists as Dask arrays - sums_cells_objs.append( - dask.array.from_array(sums_cells, chunks=sums_cells.shape) + blocks = X.to_delayed().ravel() + cell_blocks = [ + da.from_delayed( + __qc_calc_1(block), + shape=(2, X.chunks[0][ind]), + dtype=X.dtype, + meta=cp.array([]), ) - sums_genes_objs.append( - dask.array.from_array(sums_genes, chunks=sums_genes.shape) + for ind, block in enumerate(blocks) + ] + + blocks = X.to_delayed().ravel() + gene_blocks = [ + da.from_delayed( + __qc_calc_2(block), + shape=(2, X.shape[0]), + dtype=X.dtype, + meta=cp.array([]), ) - cell_ex_objs.append(dask.array.from_array(cell_ex, chunks=cell_ex.shape)) - gene_ex_objs.append(dask.array.from_array(gene_ex, chunks=gene_ex.shape)) - sums_cells = dask.array.concatenate(sums_cells_objs) - sums_genes = dask.array.concatenate(sums_genes_objs, axis=1).sum(axis=1) - cell_ex = dask.array.concatenate(cell_ex_objs) - gene_ex = dask.array.concatenate(gene_ex_objs, axis=1).sum(axis=1) - sums_cells, sums_genes, cell_ex, gene_ex = dask.compute( - sums_cells, sums_genes, cell_ex, gene_ex + for ind, block in enumerate(blocks) + ] + sums_cells, cell_ex = da.hstack(cell_blocks).compute() + sums_genes, gene_ex = da.stack(gene_blocks, axis=1).sum(axis=1).compute() + + return ( + sums_cells.ravel(), + sums_genes.ravel(), + cell_ex.ravel().astype(cp.int32), + gene_ex.ravel().astype(cp.int32), ) - return sums_cells.ravel(), sums_genes.ravel(), cell_ex.ravel(), gene_ex.ravel() -def _second_pass_qc(X, mask, client=None): +def _second_pass_qc(X, mask): if isinstance(X, DaskArray): - return _second_pass_qc_dask(X, mask, client=client) + return _second_pass_qc_dask(X, mask) sums_cells_sub = cp.zeros(X.shape[0], dtype=X.dtype) if sparse.issparse(X): if sparse.isspmatrix_csr(X): @@ -332,9 +377,7 @@ def _second_pass_qc(X, mask, client=None): @with_cupy_rmm -def _second_pass_qc_dask(X, mask, client=None): - client = _get_dask_client(client) - +def _second_pass_qc_dask(X, mask): if isinstance(X._meta, sparse.csr_matrix): from ._kernels._qc_kernels import _sparse_qc_csr_sub From 36a081e9eabe71b76b194d5b55eb8373cc61a35d Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 16 Jul 2024 15:51:57 +0200 Subject: [PATCH 28/69] remove client from scale --- src/rapids_singlecell/preprocessing/_scale.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 5a374275..1c0c9c9a 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -10,7 +10,6 @@ from rapids_singlecell._compat import ( DaskArray, - DaskClient, _meta_dense, _meta_sparse, ) @@ -37,7 +36,6 @@ def scale( obsm: str | None = None, mask_obs: np.ndarray | str | None = None, inplace: bool = True, - client: DaskClient | None = None, ) -> None | cp.ndarray: """ Scales matrix to unit variance and clips values @@ -72,8 +70,6 @@ def scale( inplace If True, update AnnData with results. Otherwise, return results. See below for details of what is returned. - client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a Dask array. Returns ------- From 6be8d9b7145cd1d5137a95059aaaf80647720641 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 16 Jul 2024 16:57:22 +0200 Subject: [PATCH 29/69] update to fast transform --- .../_sparse_pca/_dask_sparse_pca.py | 25 +++---------------- .../preprocessing/_sparse_pca/_sparse_pca.py | 20 +++------------ 2 files changed, 7 insertions(+), 38 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index b451ec57..80cec8e3 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -62,27 +62,10 @@ def transform(self, X): kernel.compile() def _transform(X_part, mean_, components_): - dense = cp.zeros(X_part.shape, dtype=X.dtype) - max_nnz = cp.diff(X_part.indptr).max() - tpb = (32, 32) - bpg_x = math.ceil(X_part.shape[0] / tpb[0]) - bpg_y = math.ceil(max_nnz / tpb[1]) - bpg = (bpg_x, bpg_y) - kernel( - bpg, - tpb, - ( - X_part.indptr, - X_part.indices, - X_part.data, - dense, - X_part.shape[0], - X_part.shape[1], - ), - ) - dense = dense - mean_ - X_pca = dense.dot(components_.T) - return X_pca + pre_mean = mean_ @ components_.T + mean_impact = cp.ones((X_part.shape[0], 1)) @ pre_mean.reshape(1, -1) + X_transformed = X_part.dot(components_.T) - mean_impact + return X_transformed X_pca = X.map_blocks( _transform, diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index 95e7bad5..ac2b0030 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -51,23 +51,9 @@ def fit(self, x): return self def transform(self, X): - from ._kernels._pca_sparse_kernel import denser_kernel - - dense_kernel = denser_kernel(X.dtype) - - dense = cp.zeros(X.shape, dtype=X.dtype) - max_nnz = cp.diff(X.indptr).max() - tpb = (32, 32) - bpg_x = math.ceil(X.shape[0] / tpb[0]) - bpg_y = math.ceil(max_nnz / tpb[1]) - bpg = (bpg_x, bpg_y) - dense_kernel( - bpg, - tpb, - (X.indptr, X.indices, X.data, dense, X.shape[0], X.shape[1]), - ) - dense -= self.mean_ - X_transformed = dense.dot(self.components_.T) + precomputed_mean_impact = self.mean_ @ self.components_.T + mean_impact = cp.ones((X.shape[0], 1)) @ precomputed_mean_impact.reshape(1, -1) + X_transformed = X.dot(self.components_.T) - mean_impact # X = X - self.mean_ # X_transformed = X.dot(self.components_.T) self.components_ = self.components_.get() From 35da55287e3036ebb5c487aad130d7f8da6fa5f2 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Tue, 16 Jul 2024 17:07:07 +0200 Subject: [PATCH 30/69] (fix): `normalize_total` -> `log1p` -> `pca` with sparse (#217) * (fix): use `to_delayed` and `from_delayed` to submit gram matrix jobs * (refactor): use `map_blocks` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): correct first dimension size * (fix): change normalization sum * (fix): add `x` as arg * (fix): `ncols` usages * (fix): try mapping block from x * (fix): use `X` directly to `map_blocks` * (fix): matrix creation + cleaner `map_blocks` * (fix): `len(blocks)` -> `num_blocks` * (fix): don't need `dask.delayed` decorator * (fix): try some debugging * (fix): use `cp.sum` * (fix): revert to `to_delayed` * (refactor): use `n_cols` * (fix): remove `num_blocks` * (fix): need to specify `drop_axis` for reduction * (chore): add full pipline * (chore): use `cusparse` * (fix): use `scipy_sparse` * (fix): initialization * (fix): try splitting mean/var directly * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): `x` to `X` * (fix) correct import * (fix): `delayed` decorator * (fix): get axis right * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): ravel mean/var * (fix): too many mistakes to count * (feat): same for other axis * (fix): resolve all small tolerance differences * (fix): try splitting first * (fix): `compute` once * (fix): stack `mean`/`var` * (fix): use `float64` for mean-var * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): remove unnecessary cast * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): add `dask` dep for major axis * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (refactor): use cleaner `zeros` * (fix): revert other `rtol` * (fix): remove last `extract_partitions` * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * (fix): `map_blocks` for `_second_pass_qc_dask` * (fix): remove `client` * (fix): remove other instances * (fix): remove `client` * (chore): remove `client` * (chore): remove in tests * (fix): `client` doc * (fix): remove more `client` * (fix): return client to test context * (chore): re-add client * (fix): oops * (fiX): oops x 2 * (chore): add dense test, which also doesn't work? * (fix): corect filtering * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Severin Dicks <37635888+Intron7@users.noreply.github.com> --- tests/dask/test_dask_pca.py | 80 +++++++++++++++++++++++++++++++++++-- 1 file changed, 77 insertions(+), 3 deletions(-) diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py index f0fa9b6f..0648d8b4 100644 --- a/tests/dask/test_dask_pca.py +++ b/tests/dask/test_dask_pca.py @@ -2,12 +2,12 @@ import cupy as cp import numpy as np -from cupyx.scipy import sparse as cusparse from scipy import sparse +from cupyx.scipy import sparse as cusparse from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array import rapids_singlecell as rsc - -from scanpy.datasets import pbmc3k_processed +import numpy as np +from scanpy.datasets import pbmc3k_processed, pbmc3k def test_pca_sparse_dask(client): sparse_ad = pbmc3k_processed() @@ -35,6 +35,80 @@ def test_pca_sparse_dask(client): atol=1e-6, ) +def test_pca_dense_dask_full_pipeline(client): + dense = pbmc3k() + default = pbmc3k() + dense.X = cp.array(dense.X.astype(np.float64).toarray()) + default.X = as_dense_cupy_dask_array(default.X.astype(np.float64).toarray()) + + rsc.pp.filter_genes(dense,min_count=500) + rsc.pp.filter_genes(default,min_count=500) + + default.X = as_dense_cupy_dask_array(default.X.compute()) + + rsc.pp.normalize_total(dense, target_sum=1e4) + rsc.pp.normalize_total(default,target_sum=1e4) + + rsc.pp.log1p(dense) + rsc.pp.log1p(default) + + rsc.pp.pca(dense, svd_solver="full") + rsc.pp.pca(default, svd_solver="full") + + cp.testing.assert_allclose( + np.abs(dense.obsm["X_pca"]), + cp.abs(default.obsm["X_pca"].compute()), + rtol=1e-7, + atol=1e-6, + ) + + cp.testing.assert_allclose( + np.abs(dense.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 + ) + + cp.testing.assert_allclose( + np.abs(dense.uns["pca"]["variance_ratio"]), + np.abs(default.uns["pca"]["variance_ratio"]), + rtol=1e-7, + atol=1e-6, + ) + +def test_pca_sparse_dask_full_pipeline(client): + sparse_ad = pbmc3k() + default = pbmc3k() + sparse_ad.X = cusparse.csr_matrix(sparse.csr_matrix(sparse_ad.X.astype(np.float64))) + default.X = as_sparse_cupy_dask_array(default.X.astype(np.float64)) + + rsc.pp.filter_genes(sparse_ad,min_count=100) + rsc.pp.filter_genes(default,min_count=100) + + rsc.pp.normalize_total(sparse_ad, target_sum=1e4) + rsc.pp.normalize_total(default,target_sum=1e4) + + rsc.pp.log1p(sparse_ad) + rsc.pp.log1p(default) + + rsc.pp.pca(sparse_ad) + rsc.pp.pca(default) + + cp.testing.assert_allclose( + np.abs(sparse_ad.obsm["X_pca"]), + cp.abs(default.obsm["X_pca"].compute()), + rtol=1e-7, + atol=1e-6, + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 + ) + + cp.testing.assert_allclose( + np.abs(sparse_ad.uns["pca"]["variance_ratio"]), + np.abs(default.uns["pca"]["variance_ratio"]), + rtol=1e-7, + atol=1e-6, + ) + def test_pca_dense_dask(client): sparse_ad = pbmc3k_processed() default = pbmc3k_processed() From 30ccfd395676fc6903f52b6eb092880089a3e755 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 16 Jul 2024 17:43:58 +0200 Subject: [PATCH 31/69] fix taskgraph --- src/rapids_singlecell/preprocessing/_normalize.py | 14 +++++++++----- .../preprocessing/_sparse_pca/_dask_sparse_pca.py | 2 +- tests/dask/test_dask_pca.py | 2 -- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 2e87ae77..c81b34c9 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -122,13 +122,13 @@ def _normalize_total_dask(X: DaskArray, target_sum: int) -> DaskArray: def __mul(X_part): mul_kernel( - (math.ceil(X.shape[0] / 128),), - (128,), + (math.ceil(X_part.shape[0] / 32),), + (32,), (X_part.indptr, X_part.data, X_part.shape[0], int(target_sum)), ) return X_part - X = X.map_blocks(lambda X: __mul(X), meta=_meta_sparse(X.dtype)) + X = X.map_blocks(__mul, meta=_meta_sparse(X.dtype)) elif isinstance(X._meta, cp.ndarray): from ._kernels._norm_kernel import _mul_dense @@ -143,7 +143,7 @@ def __mul(X_part): ) return X_part - X = X.map_blocks(lambda X: __mul(X), meta=_meta_dense(X.dtype)) + X = X.map_blocks(__mul, meta=_meta_dense(X.dtype)) else: raise ValueError(f"Cannot normalize {type(X)}") return X @@ -261,7 +261,11 @@ def log1p( if isinstance(X._meta, cp.ndarray): X = X.map_blocks(lambda X: cp.log1p(X), meta=_meta_dense(X.dtype)) elif isinstance(X._meta, sparse.csr_matrix): - X = X.map_blocks(lambda X: X.log1p(), meta=_meta_sparse(X.dtype)) + + def __log1p(X_part): + return X_part.log1p() + + X = X.map_blocks(__log1p, meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} if inplace: _set_obs_rep(adata, X, layer=layer, obsm=obsm) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index 80cec8e3..e1b42e68 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -160,7 +160,7 @@ def __gram_block(x_part): __gram_block(block), shape=(n_cols, n_cols), dtype=x.dtype, - meta=cp.array([]), + meta=_meta_dense(x.dtype), ) for block in blocks ] diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py index 0648d8b4..3410d11f 100644 --- a/tests/dask/test_dask_pca.py +++ b/tests/dask/test_dask_pca.py @@ -44,8 +44,6 @@ def test_pca_dense_dask_full_pipeline(client): rsc.pp.filter_genes(dense,min_count=500) rsc.pp.filter_genes(default,min_count=500) - default.X = as_dense_cupy_dask_array(default.X.compute()) - rsc.pp.normalize_total(dense, target_sum=1e4) rsc.pp.normalize_total(default,target_sum=1e4) From 8cd60b492bc42bf02aa802c52c45156324ae7d90 Mon Sep 17 00:00:00 2001 From: Ilan Gold Date: Thu, 18 Jul 2024 10:38:11 +0200 Subject: [PATCH 32/69] (feat): use `map_blocks` in gram matrix calculation and and `mean_var` (#230) * (feat): use `map_blocks` instead of `to_delayed` for `gram_matrix` * (feat): same for `mean/var` * (fix): correct shape * (fix): add new axis to `_mean_var_dense_dask` + remove `dask.delayed` * (fix): chunks shape --- src/rapids_singlecell/preprocessing/_qc.py | 2 +- .../_sparse_pca/_dask_sparse_pca.py | 25 +++--- src/rapids_singlecell/preprocessing/_utils.py | 82 ++++++++----------- 3 files changed, 47 insertions(+), 62 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index 0bf4f443..159009f0 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -316,7 +316,7 @@ def __qc_calc_2(X_part): gene_blocks = [ da.from_delayed( __qc_calc_2(block), - shape=(2, X.shape[0]), + shape=(2, X.shape[1]), dtype=X.dtype, meta=cp.array([]), ) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index e1b42e68..8c8df9f0 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -3,7 +3,6 @@ import math import cupy as cp -import dask from cuml.internals.memory_utils import with_cupy_rmm from rapids_singlecell._compat import ( @@ -134,7 +133,6 @@ def _cov_sparse_dask(x, return_gram=False, return_mean=False): compute_mean_cov.compile() n_cols = x.shape[1] - @dask.delayed def __gram_block(x_part): gram_matrix = cp.zeros((n_cols, n_cols), dtype=x.dtype) @@ -152,19 +150,20 @@ def __gram_block(x_part): gram_matrix, ), ) - return gram_matrix - - blocks = x.to_delayed().ravel() - gram_chunk_matrices = [ - dask.array.from_delayed( - __gram_block(block), - shape=(n_cols, n_cols), + return gram_matrix[None, ...] # need new axis for summing + + n_blocks = len(x.to_delayed().ravel()) + gram_matrix = ( + x.map_blocks( + __gram_block, + new_axis=(1,), + chunks=((1,) * n_blocks, (x.shape[1],), (x.shape[1],)), + meta=cp.array([]), dtype=x.dtype, - meta=_meta_dense(x.dtype), ) - for block in blocks - ] - gram_matrix = dask.array.stack(gram_chunk_matrices).sum(axis=0).compute() + .sum(axis=0) + .compute() + ) mean_x, _ = _get_mean_var(x) mean_x = mean_x.astype(x.dtype) copy_gram = _copy_kernel(x.dtype) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 9c8b1b4a..78c80792 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -77,8 +77,6 @@ def _mean_var_minor_dask(X, major, minor): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ - import dask - import dask.array as da from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( _get_mean_var_minor, @@ -87,8 +85,7 @@ def _mean_var_minor_dask(X, major, minor): get_mean_var_minor = _get_mean_var_minor(X.dtype) get_mean_var_minor.compile() - @dask.delayed - def __mean_var(X_part, minor, major): + def __mean_var(X_part): mean = cp.zeros(minor, dtype=cp.float64) var = cp.zeros(minor, dtype=cp.float64) block = (32,) @@ -96,20 +93,20 @@ def __mean_var(X_part, minor, major): get_mean_var_minor( grid, block, (X_part.indices, X_part.data, mean, var, major, X_part.nnz) ) - return cp.vstack([mean, var]) - - blocks = X.to_delayed().ravel() - mean_var_blocks = [ - da.from_delayed( - __mean_var(block, minor, major), - shape=(2, minor), + return cp.vstack([mean, var])[None, ...] # new axis for summing + + n_blocks = len(X.to_delayed().ravel()) + mean, var = ( + X.map_blocks( + __mean_var, + new_axis=(1,), + chunks=((1,) * n_blocks, (2,), (minor,)), dtype=cp.float64, meta=cp.array([]), ) - for block in blocks - ] - - mean, var = da.stack(mean_var_blocks, axis=1).sum(axis=1).compute() + .sum(axis=0) + .compute() + ) var = (var - mean**2) * (major / (major - 1)) return mean, var @@ -120,8 +117,6 @@ def _mean_var_major_dask(X, major, minor): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ - import dask - import dask.array as da from rapids_singlecell.preprocessing._kernels._mean_var_kernel import ( _get_mean_var_major, @@ -130,8 +125,7 @@ def _mean_var_major_dask(X, major, minor): get_mean_var_major = _get_mean_var_major(X.dtype) get_mean_var_major.compile() - @dask.delayed - def __mean_var(X_part, minor, major): + def __mean_var(X_part): mean = cp.zeros(X_part.shape[0], dtype=cp.float64) var = cp.zeros(X_part.shape[0], dtype=cp.float64) block = (64,) @@ -151,18 +145,12 @@ def __mean_var(X_part, minor, major): ) return cp.vstack([mean, var]) - blocks = X.to_delayed().ravel() - mean_var_blocks = [ - da.from_delayed( - __mean_var(block, minor, major), - shape=(2, X.chunks[0][ind]), - dtype=cp.float64, - meta=cp.array([]), - ) - for ind, block in enumerate(blocks) - ] - - mean, var = da.hstack(mean_var_blocks).compute() + mean, var = X.map_blocks( + __mean_var, + chunks=((2,), X.chunks[0]), + dtype=cp.float64, + meta=cp.array([]), + ).compute() mean = mean / minor var = var / minor @@ -176,33 +164,31 @@ def _mean_var_dense_dask(X, axis): """ Implements sum operation for dask array when the backend is cupy sparse csr matrix """ - import dask - import dask.array as da # ToDo: get a 64bit version working without copying the data - @dask.delayed - def __mean_var(X_part, axis): + def __mean_var(X_part): mean = X_part.sum(axis=axis) var = (X_part**2).sum(axis=axis) if axis == 0: mean = mean.reshape(-1, 1) var = var.reshape(-1, 1) - return cp.vstack([mean.ravel(), var.ravel()]) + return cp.vstack([mean.ravel(), var.ravel()])[ + None if 1 - axis else slice(None, None), ... + ] + + n_blocks = len(X.to_delayed().ravel()) + mean_var = X.map_blocks( + __mean_var, + new_axis=(1,) if axis - 1 else None, + chunks=((2,), X.chunks[0]) if axis else ((1,) * n_blocks, (2,), (X.shape[1],)), + dtype=cp.float64, + meta=cp.array([]), + ) - blocks = X.to_delayed().ravel() - mean_var_blocks = [ - da.from_delayed( - __mean_var(block, axis=axis), - shape=(2, X.chunks[0][ind]) if axis else (2, X.shape[1]), - dtype=cp.float64, - meta=cp.array([]), - ) - for ind, block in enumerate(blocks) - ] if axis == 0: - mean, var = da.stack(mean_var_blocks, axis=1).sum(axis=1).compute() + mean, var = mean_var.sum(axis=0).compute() else: - mean, var = da.hstack(mean_var_blocks).compute() + mean, var = mean_var.compute() mean = mean / X.shape[axis] var = var / X.shape[axis] From 4bfd5a1a49c7e3a236f08a5a58bbd7d407be0774 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 18 Jul 2024 16:42:24 +0200 Subject: [PATCH 33/69] update pca --- src/rapids_singlecell/preprocessing/_pca.py | 10 ++-------- .../preprocessing/_sparse_pca/_dask_sparse_pca.py | 5 ----- 2 files changed, 2 insertions(+), 13 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index c0b5ae1b..26bd5a8e 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -12,7 +12,7 @@ from scanpy.preprocessing._pca import _handle_mask_var from scipy.sparse import issparse -from rapids_singlecell._compat import DaskArray, DaskClient +from rapids_singlecell._compat import DaskArray from rapids_singlecell.get import _get_obs_rep from ._utils import _check_gpu_X @@ -36,7 +36,6 @@ def pca( copy: bool = False, chunked: bool = False, chunk_size: int = None, - client: DaskClient | None = None, ) -> None | AnnData: """ Performs PCA using the cuml decomposition function. @@ -91,9 +90,6 @@ def pca( Number of observations to include in each chunk. \ Required if `chunked=True` was passed. - client - Dask client to use for computation. If `None`, the default client is used. Only used if `X` is a dense Dask array. - Returns ------- adds fields to `adata`: @@ -143,9 +139,7 @@ def pca( if svd_solver == "auto": svd_solver = "jacobi" - pca_func = PCA( - n_components=n_comps, svd_solver=svd_solver, whiten=False, client=client - ) + pca_func = PCA(n_components=n_comps, svd_solver=svd_solver, whiten=False) X_pca = pca_func.fit_transform(X) X_pca = X_pca.compute_chunk_sizes() elif isinstance(X._meta, csr_matrix): diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index 8c8df9f0..cef18de6 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -55,11 +55,6 @@ def fit(self, x): return self def transform(self, X): - from ._kernels._pca_sparse_kernel import denser_kernel - - kernel = denser_kernel(X.dtype) - kernel.compile() - def _transform(X_part, mean_, components_): pre_mean = mean_ @ components_.T mean_impact = cp.ones((X_part.shape[0], 1)) @ pre_mean.reshape(1, -1) From 8ddba491641382cdc23e1c87ba7be89d013cb967 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 19 Jul 2024 11:23:16 +0200 Subject: [PATCH 34/69] use lambda --- src/rapids_singlecell/preprocessing/_normalize.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index c81b34c9..772c5cb8 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -261,11 +261,7 @@ def log1p( if isinstance(X._meta, cp.ndarray): X = X.map_blocks(lambda X: cp.log1p(X), meta=_meta_dense(X.dtype)) elif isinstance(X._meta, sparse.csr_matrix): - - def __log1p(X_part): - return X_part.log1p() - - X = X.map_blocks(__log1p, meta=_meta_sparse(X.dtype)) + X = X.map_blocks(lambda x: x.log1p(), meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} if inplace: _set_obs_rep(adata, X, layer=layer, obsm=obsm) From dea8d3d47105231ec8cf780660251f18be946923 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 19 Jul 2024 12:24:45 +0200 Subject: [PATCH 35/69] remove unused kernel --- .../_kernels/_pca_sparse_kernel.py | 22 ------------------- 1 file changed, 22 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py index d4115c4c..878d79de 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_kernels/_pca_sparse_kernel.py @@ -62,24 +62,6 @@ } """ -denser = r""" -(const int* indptr,const int *index,const {0} *data, - {0}* out, int major, int minor) { - int row = blockIdx.x*blockDim.x+threadIdx.x ; - int col = blockIdx.y*blockDim.y+threadIdx.y ; - if(row >= major){ - return; - } - int start = indptr[row]; - int stop = indptr[row+1]; - if (col>= (stop - start)){ - return; - } - int idx = index[start + col]; - long long int res_index = static_cast(row)*minor+idx; - out[res_index] = data[start + col];} -""" - _zero_genes_kernel = cp.RawKernel(check_zero_genes, "check_zero_genes") @@ -93,7 +75,3 @@ def _gramm_kernel_csr(dtype): def _copy_kernel(dtype): return cuda_kernel_factory(copy_kernel, (dtype,), "copy_kernel") - - -def denser_kernel(dtype): - return cuda_kernel_factory(denser, (dtype,), "denser") From a71116d711f0a628ebda169f70b48cd522a99ef8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 19 Jul 2024 13:28:25 +0200 Subject: [PATCH 36/69] test removed kernel --- src/rapids_singlecell/preprocessing/_scale.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 1c0c9c9a..109eae20 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -297,29 +297,24 @@ def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=N ) if isinstance(X._meta, sparse.csr_matrix) and zero_center: - from ._sparse_pca._kernels._pca_sparse_kernel import denser_kernel + from ._kernels._sparse2dense import _sparse2densekernel - kernel = denser_kernel(X.dtype) + kernel = _sparse2densekernel(X.dtype) kernel.compile() def __dense(X_part): - dense = cp.zeros(X_part.shape, dtype=X.dtype) + major, minor = X_part.shape + dense = cp.zeros(X_part.shape, order="C", dtype=X_part.dtype) max_nnz = cp.diff(X_part.indptr).max() tpb = (32, 32) - bpg_x = math.ceil(X_part.shape[0] / tpb[0]) + bpg_x = math.ceil(major / tpb[0]) bpg_y = math.ceil(max_nnz / tpb[1]) bpg = (bpg_x, bpg_y) + kernel( bpg, tpb, - ( - X_part.indptr, - X_part.indices, - X_part.data, - dense, - X_part.shape[0], - X_part.shape[1], - ), + (X_part.indptr, X_part.indices, X_part.data, dense, major, minor, 1), ) return dense From 3680a68ee619663e08a1c58d5b6a8c9154bc5f00 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 5 Aug 2024 20:45:26 +0000 Subject: [PATCH 37/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index eebf4196..3f151136 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -13,7 +13,6 @@ from anndata import AnnData - def _sparse_to_dense(X: spmatrix, order: Literal["C", "F"] | None = None) -> cp.ndarray: if order is None: order = "C" From 9c0e09a9dc5167d019e3441aeaf881a7ef31b63c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 13 Aug 2024 08:31:33 +0000 Subject: [PATCH 38/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_normalize.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 16545acb..6fc409b1 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -2,10 +2,8 @@ import math import warnings - from typing import TYPE_CHECKING, Union - import cupy as cp from cupyx.scipy import sparse from scanpy.get import _get_obs_rep, _set_obs_rep From 383bafb778be73fd18a0137899a9ceba16b4a9d0 Mon Sep 17 00:00:00 2001 From: Severin Dicks <37635888+Intron7@users.noreply.github.com> Date: Tue, 13 Aug 2024 16:50:52 +0200 Subject: [PATCH 39/69] add outside compute (#245) --- src/rapids_singlecell/preprocessing/_hvg.py | 4 ++++ src/rapids_singlecell/preprocessing/_qc.py | 8 ++++++-- src/rapids_singlecell/preprocessing/_scale.py | 3 +++ .../_sparse_pca/_dask_sparse_pca.py | 20 +++++++++---------- src/rapids_singlecell/preprocessing/_utils.py | 2 +- 5 files changed, 23 insertions(+), 14 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index cba1d07a..30d9831c 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -301,6 +301,10 @@ def _highly_variable_genes_single_batch( X = cp.expm1(X) mean, var = _get_mean_var(X, axis=0) + if isinstance(X, DaskArray): + import dask + + mean, var = dask.compute(mean, var) mean[mean == 0] = 1e-12 disp = var / mean if flavor == "seurat": # logarithmized mean as in Seurat diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index 159009f0..c095e314 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -322,8 +322,12 @@ def __qc_calc_2(X_part): ) for ind, block in enumerate(blocks) ] - sums_cells, cell_ex = da.hstack(cell_blocks).compute() - sums_genes, gene_ex = da.stack(gene_blocks, axis=1).sum(axis=1).compute() + sums_cells, cell_ex = da.hstack(cell_blocks) + sums_genes, gene_ex = da.stack(gene_blocks, axis=1).sum(axis=1) + + sums_cells, cell_ex, sums_genes, gene_ex = dask.compute( + sums_cells, cell_ex, sums_genes, gene_ex + ) return ( sums_cells.ravel(), diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 91b12331..4e2946f3 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -280,6 +280,8 @@ def _scale_sparse_csr( def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=None): + import dask + if not inplace: X = X.copy() if mask_obs is None: @@ -289,6 +291,7 @@ def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=N else: mean, var = _get_mean_var(X[mask_obs, :]) mask_array = cp.array(mask_obs).astype(cp.int32) + mean, var = dask.compute(mean, var) std = cp.sqrt(var) std[std == 0] = 1 max_value = _get_max_value(max_value, X.dtype) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index cef18de6..003f8393 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -117,6 +117,7 @@ def _cov_sparse_dask(x, return_gram=False, return_mean=False): cov(X, X), gram(X, X), mean(X), mean(X) when return_gram is True and return_mean is True """ + import dask from ._kernels._pca_sparse_kernel import ( _copy_kernel, @@ -148,18 +149,15 @@ def __gram_block(x_part): return gram_matrix[None, ...] # need new axis for summing n_blocks = len(x.to_delayed().ravel()) - gram_matrix = ( - x.map_blocks( - __gram_block, - new_axis=(1,), - chunks=((1,) * n_blocks, (x.shape[1],), (x.shape[1],)), - meta=cp.array([]), - dtype=x.dtype, - ) - .sum(axis=0) - .compute() - ) + gram_matrix = x.map_blocks( + __gram_block, + new_axis=(1,), + chunks=((1,) * n_blocks, (x.shape[1],), (x.shape[1],)), + meta=cp.array([]), + dtype=x.dtype, + ).sum(axis=0) mean_x, _ = _get_mean_var(x) + gram_matrix, mean_x = dask.compute(gram_matrix, mean_x) mean_x = mean_x.astype(x.dtype) copy_gram = _copy_kernel(x.dtype) block = (32, 32) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 3f151136..208d940b 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -153,7 +153,7 @@ def __mean_var(X_part): chunks=((2,), X.chunks[0]), dtype=cp.float64, meta=cp.array([]), - ).compute() + ) mean = mean / minor var = var / minor From 4f88abce2407fd35bfcc2549181f1c1e2472518c Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 26 Sep 2024 13:55:31 +0200 Subject: [PATCH 40/69] update utils for lazy compute --- src/rapids_singlecell/preprocessing/_utils.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 208d940b..ca556140 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -99,17 +99,13 @@ def __mean_var(X_part): return cp.vstack([mean, var])[None, ...] # new axis for summing n_blocks = len(X.to_delayed().ravel()) - mean, var = ( - X.map_blocks( - __mean_var, - new_axis=(1,), - chunks=((1,) * n_blocks, (2,), (minor,)), - dtype=cp.float64, - meta=cp.array([]), - ) - .sum(axis=0) - .compute() - ) + mean, var = X.map_blocks( + __mean_var, + new_axis=(1,), + chunks=((1,) * n_blocks, (2,), (minor,)), + dtype=cp.float64, + meta=cp.array([]), + ).sum(axis=0) var = (var - mean**2) * (major / (major - 1)) return mean, var @@ -189,9 +185,9 @@ def __mean_var(X_part): ) if axis == 0: - mean, var = mean_var.sum(axis=0).compute() + mean, var = mean_var.sum(axis=0) else: - mean, var = mean_var.compute() + mean, var = mean_var mean = mean / X.shape[axis] var = var / X.shape[axis] From 602f845abec1e4c91b324a03568b0d2ce3af2d4f Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 26 Sep 2024 14:39:35 +0200 Subject: [PATCH 41/69] update utils --- src/rapids_singlecell/preprocessing/_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index ca556140..2f667e3f 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -161,13 +161,13 @@ def __mean_var(X_part): @with_cupy_rmm def _mean_var_dense_dask(X, axis): """ - Implements sum operation for dask array when the backend is cupy sparse csr matrix + Implements sum operation for dask array when the backend is cupy dense matrix """ + from ._kernels._mean_var_kernel import mean_sum, sq_sum - # ToDo: get a 64bit version working without copying the data def __mean_var(X_part): - mean = X_part.sum(axis=axis) - var = (X_part**2).sum(axis=axis) + var = sq_sum(X_part, axis=axis) + mean = mean_sum(X_part, axis=axis) if axis == 0: mean = mean.reshape(-1, 1) var = var.reshape(-1, 1) @@ -191,7 +191,7 @@ def __mean_var(X_part): mean = mean / X.shape[axis] var = var / X.shape[axis] - var -= cp.power(mean, 2) + var -= mean**2 var *= X.shape[axis] / (X.shape[axis] - 1) return mean, var From 98441ea1101c46d636e9e027dbb9e7658216bb2a Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 26 Sep 2024 15:01:02 +0000 Subject: [PATCH 42/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/dask/conftest.py | 22 +++---- tests/dask/test_dask_pca.py | 28 +++++---- tests/dask/test_hvg_dask.py | 56 +++++++++++------ tests/dask/test_normalize_dask.py | 9 ++- tests/dask/test_qc_dask.py | 99 ++++++++++++++++++++++--------- tests/dask/test_scale_dask.py | 37 ++++++------ 6 files changed, 159 insertions(+), 92 deletions(-) diff --git a/tests/dask/conftest.py b/tests/dask/conftest.py index e3f7b610..84dda0b7 100644 --- a/tests/dask/conftest.py +++ b/tests/dask/conftest.py @@ -1,37 +1,32 @@ from __future__ import annotations import cupy as cp -from cupyx.scipy import sparse as cusparse -from anndata.tests.helpers import as_dense_dask_array, as_sparse_dask_array import pytest - +from anndata.tests.helpers import as_dense_dask_array, as_sparse_dask_array +from cupyx.scipy import sparse as cusparse +from dask.distributed import Client from dask_cuda import LocalCUDACluster from dask_cuda.utils_test import IncreasedCloseTimeoutNanny -from dask.distributed import Client def as_sparse_cupy_dask_array(X): da = as_sparse_dask_array(X) - da = da.rechunk((da.shape[0]//2, da.shape[1])) - da = da.map_blocks(cusparse.csr_matrix, dtype = X.dtype) + da = da.rechunk((da.shape[0] // 2, da.shape[1])) + da = da.map_blocks(cusparse.csr_matrix, dtype=X.dtype) return da + def as_dense_cupy_dask_array(X): X = as_dense_dask_array(X) X = X.map_blocks(cp.array) - X = X.rechunk((X.shape[0]//2, X.shape[1])) + X = X.rechunk((X.shape[0] // 2, X.shape[1])) return X -from dask_cuda import initialize -from dask_cuda import LocalCUDACluster -from dask_cuda.utils_test import IncreasedCloseTimeoutNanny -from dask.distributed import Client @pytest.fixture(scope="module") def cluster(): - cluster = LocalCUDACluster( - CUDA_VISIBLE_DEVICES ="0", + CUDA_VISIBLE_DEVICES="0", protocol="tcp", scheduler_port=0, worker_class=IncreasedCloseTimeoutNanny, @@ -42,7 +37,6 @@ def cluster(): @pytest.fixture(scope="function") def client(cluster): - client = Client(cluster) yield client client.close() diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py index 3410d11f..9e360f3f 100644 --- a/tests/dask/test_dask_pca.py +++ b/tests/dask/test_dask_pca.py @@ -2,12 +2,13 @@ import cupy as cp import numpy as np -from scipy import sparse -from cupyx.scipy import sparse as cusparse from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +from cupyx.scipy import sparse as cusparse +from scanpy.datasets import pbmc3k, pbmc3k_processed +from scipy import sparse + import rapids_singlecell as rsc -import numpy as np -from scanpy.datasets import pbmc3k_processed, pbmc3k + def test_pca_sparse_dask(client): sparse_ad = pbmc3k_processed() @@ -35,17 +36,18 @@ def test_pca_sparse_dask(client): atol=1e-6, ) + def test_pca_dense_dask_full_pipeline(client): dense = pbmc3k() default = pbmc3k() dense.X = cp.array(dense.X.astype(np.float64).toarray()) default.X = as_dense_cupy_dask_array(default.X.astype(np.float64).toarray()) - rsc.pp.filter_genes(dense,min_count=500) - rsc.pp.filter_genes(default,min_count=500) + rsc.pp.filter_genes(dense, min_count=500) + rsc.pp.filter_genes(default, min_count=500) - rsc.pp.normalize_total(dense, target_sum=1e4) - rsc.pp.normalize_total(default,target_sum=1e4) + rsc.pp.normalize_total(dense, target_sum=1e4) + rsc.pp.normalize_total(default, target_sum=1e4) rsc.pp.log1p(dense) rsc.pp.log1p(default) @@ -71,17 +73,18 @@ def test_pca_dense_dask_full_pipeline(client): atol=1e-6, ) + def test_pca_sparse_dask_full_pipeline(client): sparse_ad = pbmc3k() default = pbmc3k() sparse_ad.X = cusparse.csr_matrix(sparse.csr_matrix(sparse_ad.X.astype(np.float64))) default.X = as_sparse_cupy_dask_array(default.X.astype(np.float64)) - rsc.pp.filter_genes(sparse_ad,min_count=100) - rsc.pp.filter_genes(default,min_count=100) + rsc.pp.filter_genes(sparse_ad, min_count=100) + rsc.pp.filter_genes(default, min_count=100) - rsc.pp.normalize_total(sparse_ad, target_sum=1e4) - rsc.pp.normalize_total(default,target_sum=1e4) + rsc.pp.normalize_total(sparse_ad, target_sum=1e4) + rsc.pp.normalize_total(default, target_sum=1e4) rsc.pp.log1p(sparse_ad) rsc.pp.log1p(default) @@ -107,6 +110,7 @@ def test_pca_sparse_dask_full_pipeline(client): atol=1e-6, ) + def test_pca_dense_dask(client): sparse_ad = pbmc3k_processed() default = pbmc3k_processed() diff --git a/tests/dask/test_hvg_dask.py b/tests/dask/test_hvg_dask.py index 8d817a49..e0fe98ba 100644 --- a/tests/dask/test_hvg_dask.py +++ b/tests/dask/test_hvg_dask.py @@ -1,14 +1,14 @@ from __future__ import annotations import cupy as cp -import numpy as np -from cupyx.scipy import sparse as cusparse -import scanpy as sc import pandas as pd +import scanpy as sc from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +from cupyx.scipy import sparse as cusparse +from scanpy.datasets import pbmc3k + import rapids_singlecell as rsc -from scanpy.datasets import pbmc3k def _get_anndata(): adata = pbmc3k() @@ -18,6 +18,7 @@ def _get_anndata(): sc.pp.log1p(adata) return adata + def test_seurat_sparse(client): adata = _get_anndata() dask_data = adata.copy() @@ -27,7 +28,9 @@ def test_seurat_sparse(client): rsc.pp.highly_variable_genes(dask_data) cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) def test_seurat_sparse_batch(client): @@ -39,10 +42,13 @@ def test_seurat_sparse_batch(client): dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() adata.X = cusparse.csr_matrix(adata.X) rsc.pp.highly_variable_genes(adata, batch_key="batch") - rsc.pp.highly_variable_genes(dask_data,batch_key="batch") + rsc.pp.highly_variable_genes(dask_data, batch_key="batch") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) + def test_cr_sparse(client): adata = _get_anndata() @@ -53,7 +59,10 @@ def test_cr_sparse(client): rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) + def test_cr_sparse_batch(client): adata = _get_anndata() @@ -64,10 +73,13 @@ def test_cr_sparse_batch(client): dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() adata.X = cusparse.csr_matrix(adata.X) rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data,batch_key="batch", flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data, batch_key="batch", flavor="cell_ranger") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) + def test_cr_dense(client): adata = _get_anndata() @@ -79,7 +91,10 @@ def test_cr_dense(client): rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) + def test_seurat_dense(client): adata = _get_anndata() @@ -91,7 +106,9 @@ def test_seurat_dense(client): rsc.pp.highly_variable_genes(dask_data) cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) def test_cr_dense_batch(client): @@ -104,10 +121,13 @@ def test_cr_dense_batch(client): dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() adata.X = cp.array(adata.X.toarray()) rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data,batch_key="batch", flavor="cell_ranger") + rsc.pp.highly_variable_genes(dask_data, batch_key="batch", flavor="cell_ranger") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) + def test_seurat_dense_batch(client): adata = _get_anndata() @@ -117,9 +137,11 @@ def test_seurat_dense_batch(client): adata.X = adata.X.astype("float64") dask_data = adata.copy() dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() - adata.X = cp.array(adata.X.toarray()) + adata.X = cp.array(adata.X.toarray()) rsc.pp.highly_variable_genes(adata, batch_key="batch") - rsc.pp.highly_variable_genes(dask_data,batch_key="batch") + rsc.pp.highly_variable_genes(dask_data, batch_key="batch") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose(adata.var["dispersions_norm"], dask_data.var["dispersions_norm"]) + cp.testing.assert_allclose( + adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] + ) diff --git a/tests/dask/test_normalize_dask.py b/tests/dask/test_normalize_dask.py index 3c67d18d..a9771c5f 100644 --- a/tests/dask/test_normalize_dask.py +++ b/tests/dask/test_normalize_dask.py @@ -1,13 +1,13 @@ from __future__ import annotations import cupy as cp -import numpy as np -from cupyx.scipy import sparse as cusparse import scanpy as sc from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +from cupyx.scipy import sparse as cusparse +from scanpy.datasets import pbmc3k + import rapids_singlecell as rsc -from scanpy.datasets import pbmc3k def test_normalize_sparse(client): adata = pbmc3k() @@ -20,6 +20,7 @@ def test_normalize_sparse(client): rsc.pp.normalize_total(dask_data) cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + def test_normalize_dense(client): adata = pbmc3k() sc.pp.filter_cells(adata, min_genes=100) @@ -31,6 +32,7 @@ def test_normalize_dense(client): rsc.pp.normalize_total(dask_data) cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + def test_log1p_sparse(client): adata = pbmc3k() sc.pp.filter_cells(adata, min_genes=100) @@ -43,6 +45,7 @@ def test_log1p_sparse(client): rsc.pp.log1p(dask_data) cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + def test_log1p_dense(client): adata = pbmc3k() sc.pp.filter_cells(adata, min_genes=100) diff --git a/tests/dask/test_qc_dask.py b/tests/dask/test_qc_dask.py index b79746e0..1801596a 100644 --- a/tests/dask/test_qc_dask.py +++ b/tests/dask/test_qc_dask.py @@ -2,12 +2,12 @@ import cupy as cp import numpy as np -from cupyx.scipy import sparse as cusparse -from scipy import sparse from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +from cupyx.scipy import sparse as cusparse +from scanpy.datasets import pbmc3k + import rapids_singlecell as rsc -from scanpy.datasets import pbmc3k def test_qc_metrics_sparse(client): adata = pbmc3k() @@ -15,21 +15,42 @@ def test_qc_metrics_sparse(client): dask_data = adata.copy() dask_data.X = as_sparse_cupy_dask_array(dask_data.X) adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p= True) - rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p= True) - np.testing.assert_allclose(adata.obs["n_genes_by_counts"] , dask_data.obs["n_genes_by_counts"]) - np.testing.assert_allclose(adata.obs["total_counts"] , dask_data.obs["total_counts"]) - np.testing.assert_allclose(adata.obs["log1p_n_genes_by_counts"] , dask_data.obs["log1p_n_genes_by_counts"]) - np.testing.assert_allclose(adata.obs["log1p_total_counts"] , dask_data.obs["log1p_total_counts"]) - np.testing.assert_allclose(adata.obs["pct_counts_mt"] , dask_data.obs["pct_counts_mt"]) - np.testing.assert_allclose(adata.obs["total_counts_mt"] , dask_data.obs["total_counts_mt"]) - np.testing.assert_allclose(adata.obs["log1p_total_counts_mt"] , dask_data.obs["log1p_total_counts_mt"]) - np.testing.assert_allclose(adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"]) + rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p=True) + rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p=True) + np.testing.assert_allclose( + adata.obs["n_genes_by_counts"], dask_data.obs["n_genes_by_counts"] + ) + np.testing.assert_allclose(adata.obs["total_counts"], dask_data.obs["total_counts"]) + np.testing.assert_allclose( + adata.obs["log1p_n_genes_by_counts"], dask_data.obs["log1p_n_genes_by_counts"] + ) + np.testing.assert_allclose( + adata.obs["log1p_total_counts"], dask_data.obs["log1p_total_counts"] + ) + np.testing.assert_allclose( + adata.obs["pct_counts_mt"], dask_data.obs["pct_counts_mt"] + ) + np.testing.assert_allclose( + adata.obs["total_counts_mt"], dask_data.obs["total_counts_mt"] + ) + np.testing.assert_allclose( + adata.obs["log1p_total_counts_mt"], dask_data.obs["log1p_total_counts_mt"] + ) + np.testing.assert_allclose( + adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"] + ) np.testing.assert_allclose(adata.var["total_counts"], dask_data.var["total_counts"]) np.testing.assert_allclose(adata.var["mean_counts"], dask_data.var["mean_counts"]) - np.testing.assert_allclose(adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"]) - np.testing.assert_allclose(adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"]) - np.testing.assert_allclose(adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"]) + np.testing.assert_allclose( + adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"] + ) + np.testing.assert_allclose( + adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"] + ) + np.testing.assert_allclose( + adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"] + ) + def test_qc_metrics_dense(client): adata = pbmc3k() @@ -37,18 +58,38 @@ def test_qc_metrics_dense(client): dask_data = adata.copy() dask_data.X = as_dense_cupy_dask_array(dask_data.X) adata.X = cp.array(adata.X.toarray()) - rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p= True) - rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p= True) - np.testing.assert_allclose(adata.obs["n_genes_by_counts"] , dask_data.obs["n_genes_by_counts"]) - np.testing.assert_allclose(adata.obs["total_counts"] , dask_data.obs["total_counts"]) - np.testing.assert_allclose(adata.obs["log1p_n_genes_by_counts"] , dask_data.obs["log1p_n_genes_by_counts"]) - np.testing.assert_allclose(adata.obs["log1p_total_counts"] , dask_data.obs["log1p_total_counts"]) - np.testing.assert_allclose(adata.obs["pct_counts_mt"] , dask_data.obs["pct_counts_mt"]) - np.testing.assert_allclose(adata.obs["total_counts_mt"] , dask_data.obs["total_counts_mt"]) - np.testing.assert_allclose(adata.obs["log1p_total_counts_mt"] , dask_data.obs["log1p_total_counts_mt"]) - np.testing.assert_allclose(adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"]) + rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p=True) + rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p=True) + np.testing.assert_allclose( + adata.obs["n_genes_by_counts"], dask_data.obs["n_genes_by_counts"] + ) + np.testing.assert_allclose(adata.obs["total_counts"], dask_data.obs["total_counts"]) + np.testing.assert_allclose( + adata.obs["log1p_n_genes_by_counts"], dask_data.obs["log1p_n_genes_by_counts"] + ) + np.testing.assert_allclose( + adata.obs["log1p_total_counts"], dask_data.obs["log1p_total_counts"] + ) + np.testing.assert_allclose( + adata.obs["pct_counts_mt"], dask_data.obs["pct_counts_mt"] + ) + np.testing.assert_allclose( + adata.obs["total_counts_mt"], dask_data.obs["total_counts_mt"] + ) + np.testing.assert_allclose( + adata.obs["log1p_total_counts_mt"], dask_data.obs["log1p_total_counts_mt"] + ) + np.testing.assert_allclose( + adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"] + ) np.testing.assert_allclose(adata.var["total_counts"], dask_data.var["total_counts"]) np.testing.assert_allclose(adata.var["mean_counts"], dask_data.var["mean_counts"]) - np.testing.assert_allclose(adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"]) - np.testing.assert_allclose(adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"]) - np.testing.assert_allclose(adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"]) + np.testing.assert_allclose( + adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"] + ) + np.testing.assert_allclose( + adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"] + ) + np.testing.assert_allclose( + adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"] + ) diff --git a/tests/dask/test_scale_dask.py b/tests/dask/test_scale_dask.py index 961c3876..ecd656f2 100644 --- a/tests/dask/test_scale_dask.py +++ b/tests/dask/test_scale_dask.py @@ -2,14 +2,13 @@ import cupy as cp import numpy as np -from cupyx.scipy import sparse as cusparse -from scipy import sparse -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array -import rapids_singlecell as rsc import scanpy as sc - +from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array +from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k +import rapids_singlecell as rsc + def _get_anndata(): adata = pbmc3k() @@ -20,42 +19,46 @@ def _get_anndata(): sc.pp.highly_variable_genes(adata, n_top_genes=1000, subset=True) return adata.copy() + def test_zc_sparse(client): adata = _get_anndata() - mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) dask_data = adata.copy() dask_data.X = as_sparse_cupy_dask_array(dask_data.X.astype(np.float64)) adata.X = cusparse.csr_matrix(adata.X.astype(np.float64)) - rsc.pp.scale(adata, mask_obs = mask, max_value = 10) - rsc.pp.scale(dask_data, mask_obs = mask, max_value = 10) + rsc.pp.scale(adata, mask_obs=mask, max_value=10) + rsc.pp.scale(dask_data, mask_obs=mask, max_value=10) cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + def test_nzc_sparse(client): adata = _get_anndata() - mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) dask_data = adata.copy() dask_data.X = as_sparse_cupy_dask_array(dask_data.X) adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.scale(adata, zero_center = False, mask_obs = mask, max_value = 10) - rsc.pp.scale(dask_data,zero_center = False, mask_obs = mask, max_value = 10) + rsc.pp.scale(adata, zero_center=False, mask_obs=mask, max_value=10) + rsc.pp.scale(dask_data, zero_center=False, mask_obs=mask, max_value=10) cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + def test_zc_dense(client): adata = _get_anndata() - mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) dask_data = adata.copy() dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) adata.X = cp.array(adata.X.toarray().astype(np.float64)) - rsc.pp.scale(adata, mask_obs = mask, max_value = 10) - rsc.pp.scale(dask_data, mask_obs = mask, max_value = 10) + rsc.pp.scale(adata, mask_obs=mask, max_value=10) + rsc.pp.scale(dask_data, mask_obs=mask, max_value=10) cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + def test_nzc_dense(client): adata = _get_anndata() - mask = np.random.randint(0,2,adata.shape[0],dtype = np.bool_) + mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) dask_data = adata.copy() dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) adata.X = cp.array(adata.X.toarray().astype(np.float64)) - rsc.pp.scale(adata, zero_center = False, mask_obs = mask, max_value = 10) - rsc.pp.scale(dask_data, zero_center = False, mask_obs = mask, max_value = 10) + rsc.pp.scale(adata, zero_center=False, mask_obs=mask, max_value=10) + rsc.pp.scale(dask_data, zero_center=False, mask_obs=mask, max_value=10) cp.testing.assert_allclose(adata.X, dask_data.X.compute()) From 118a37abb48670d10eda5e1aae729b9fdd20399f Mon Sep 17 00:00:00 2001 From: Phil Schaf Date: Thu, 26 Sep 2024 18:02:58 +0200 Subject: [PATCH 43/69] move test helpers --- src/rapids_singlecell/_testing.py | 19 ++++++++++++++++++- tests/dask/conftest.py | 17 ----------------- tests/dask/test_dask_pca.py | 5 ++++- tests/dask/test_hvg_dask.py | 5 ++++- tests/dask/test_normalize_dask.py | 5 ++++- tests/dask/test_qc_dask.py | 5 ++++- tests/dask/test_scale_dask.py | 5 ++++- 7 files changed, 38 insertions(+), 23 deletions(-) diff --git a/src/rapids_singlecell/_testing.py b/src/rapids_singlecell/_testing.py index eee65f78..b18590eb 100644 --- a/src/rapids_singlecell/_testing.py +++ b/src/rapids_singlecell/_testing.py @@ -4,8 +4,10 @@ from typing import TYPE_CHECKING +import cupy as cp import pytest -from anndata.tests.helpers import asarray +from anndata.tests.helpers import as_dense_dask_array, as_sparse_dask_array, asarray +from cupyx.scipy import sparse as cusparse from scipy import sparse if TYPE_CHECKING: @@ -38,3 +40,18 @@ def param_with( ARRAY_TYPES_MEM = tuple( at for (strg, _), ats in MAP_ARRAY_TYPES.items() if strg == "mem" for at in ats ) + + +def as_sparse_cupy_dask_array(X): + da = as_sparse_dask_array(X) + da = da.rechunk((da.shape[0] // 2, da.shape[1])) + da = da.map_blocks(cusparse.csr_matrix, dtype=X.dtype) + return da + + +def as_dense_cupy_dask_array(X): + X = as_dense_dask_array(X) + X = X.map_blocks(cp.array) + X = X.rechunk((X.shape[0] // 2, X.shape[1])) + return X + diff --git a/tests/dask/conftest.py b/tests/dask/conftest.py index 84dda0b7..c2ded911 100644 --- a/tests/dask/conftest.py +++ b/tests/dask/conftest.py @@ -1,28 +1,11 @@ from __future__ import annotations -import cupy as cp import pytest -from anndata.tests.helpers import as_dense_dask_array, as_sparse_dask_array -from cupyx.scipy import sparse as cusparse from dask.distributed import Client from dask_cuda import LocalCUDACluster from dask_cuda.utils_test import IncreasedCloseTimeoutNanny -def as_sparse_cupy_dask_array(X): - da = as_sparse_dask_array(X) - da = da.rechunk((da.shape[0] // 2, da.shape[1])) - da = da.map_blocks(cusparse.csr_matrix, dtype=X.dtype) - return da - - -def as_dense_cupy_dask_array(X): - X = as_dense_dask_array(X) - X = X.map_blocks(cp.array) - X = X.rechunk((X.shape[0] // 2, X.shape[1])) - return X - - @pytest.fixture(scope="module") def cluster(): cluster = LocalCUDACluster( diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py index 9e360f3f..95801d96 100644 --- a/tests/dask/test_dask_pca.py +++ b/tests/dask/test_dask_pca.py @@ -2,12 +2,15 @@ import cupy as cp import numpy as np -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k, pbmc3k_processed from scipy import sparse import rapids_singlecell as rsc +from rapids_singlecell._testing import ( + as_dense_cupy_dask_array, + as_sparse_cupy_dask_array, +) def test_pca_sparse_dask(client): diff --git a/tests/dask/test_hvg_dask.py b/tests/dask/test_hvg_dask.py index e0fe98ba..1b47c392 100644 --- a/tests/dask/test_hvg_dask.py +++ b/tests/dask/test_hvg_dask.py @@ -3,11 +3,14 @@ import cupy as cp import pandas as pd import scanpy as sc -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k import rapids_singlecell as rsc +from rapids_singlecell._testing import ( + as_dense_cupy_dask_array, + as_sparse_cupy_dask_array, +) def _get_anndata(): diff --git a/tests/dask/test_normalize_dask.py b/tests/dask/test_normalize_dask.py index a9771c5f..dcdca21c 100644 --- a/tests/dask/test_normalize_dask.py +++ b/tests/dask/test_normalize_dask.py @@ -2,11 +2,14 @@ import cupy as cp import scanpy as sc -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k import rapids_singlecell as rsc +from rapids_singlecell._testing import ( + as_dense_cupy_dask_array, + as_sparse_cupy_dask_array, +) def test_normalize_sparse(client): diff --git a/tests/dask/test_qc_dask.py b/tests/dask/test_qc_dask.py index 1801596a..f1fd5f9c 100644 --- a/tests/dask/test_qc_dask.py +++ b/tests/dask/test_qc_dask.py @@ -2,11 +2,14 @@ import cupy as cp import numpy as np -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k import rapids_singlecell as rsc +from rapids_singlecell._testing import ( + as_dense_cupy_dask_array, + as_sparse_cupy_dask_array, +) def test_qc_metrics_sparse(client): diff --git a/tests/dask/test_scale_dask.py b/tests/dask/test_scale_dask.py index ecd656f2..6a9a7351 100644 --- a/tests/dask/test_scale_dask.py +++ b/tests/dask/test_scale_dask.py @@ -3,11 +3,14 @@ import cupy as cp import numpy as np import scanpy as sc -from conftest import as_dense_cupy_dask_array, as_sparse_cupy_dask_array from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k import rapids_singlecell as rsc +from rapids_singlecell._testing import ( + as_dense_cupy_dask_array, + as_sparse_cupy_dask_array, +) def _get_anndata(): From b6c268902b9a9b3dcfd77f00168cf4a38c4f99a9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 26 Sep 2024 16:04:57 +0000 Subject: [PATCH 44/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/_testing.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/rapids_singlecell/_testing.py b/src/rapids_singlecell/_testing.py index b18590eb..4e73261d 100644 --- a/src/rapids_singlecell/_testing.py +++ b/src/rapids_singlecell/_testing.py @@ -54,4 +54,3 @@ def as_dense_cupy_dask_array(X): X = X.map_blocks(cp.array) X = X.rechunk((X.shape[0] // 2, X.shape[1])) return X - From ea57084a4052a156120c21a90f9ceda3714de162 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 14 Oct 2024 15:09:00 +0200 Subject: [PATCH 45/69] update typing --- src/rapids_singlecell/_compat.py | 24 +--------------------- src/rapids_singlecell/_utils/__init__.py | 7 +++++++ src/rapids_singlecell/preprocessing/_qc.py | 14 +++++++++---- 3 files changed, 18 insertions(+), 27 deletions(-) diff --git a/src/rapids_singlecell/_compat.py b/src/rapids_singlecell/_compat.py index ecb20551..f569504b 100644 --- a/src/rapids_singlecell/_compat.py +++ b/src/rapids_singlecell/_compat.py @@ -2,29 +2,7 @@ import cupy as cp from cupyx.scipy.sparse import csr_matrix - -try: - from dask.array import Array as DaskArray -except ImportError: - - class DaskArray: - pass - - -try: - from dask.distributed import Client as DaskClient -except ImportError: - - class DaskClient: - pass - - -def _get_dask_client(client=None): - from dask.distributed import default_client - - if client is None or not isinstance(client, DaskClient): - return default_client() - return client +from dask.array import Array as DaskArray # noqa: F401 def _meta_dense(dtype): diff --git a/src/rapids_singlecell/_utils/__init__.py b/src/rapids_singlecell/_utils/__init__.py index 201809b6..8076b7f1 100644 --- a/src/rapids_singlecell/_utils/__init__.py +++ b/src/rapids_singlecell/_utils/__init__.py @@ -2,6 +2,13 @@ from typing import TYPE_CHECKING, Union +import cupy as cp import numpy as np +from cupyx.scipy.sparse import csc_matrix, csr_matrix +from dask.array import Array as DaskArray AnyRandom = Union[int, np.random.RandomState, None] # noqa: UP007 + + +ArrayTypes = Union[cp.ndarray, csc_matrix, csr_matrix] # noqa: UP007 +ArrayTypesDask = Union[cp.ndarray, csc_matrix, csr_matrix, DaskArray] # noqa: UP007 diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index c095e314..bd6bdd2c 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -15,6 +15,8 @@ if TYPE_CHECKING: from anndata import AnnData + from rapids_singlecell._utils import ArrayTypesDask + def calculate_qc_metrics( adata: AnnData, @@ -110,7 +112,9 @@ def calculate_qc_metrics( ) -def _first_pass_qc(X): +def _first_pass_qc( + X: ArrayTypesDask, +) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]: if isinstance(X, DaskArray): return _first_pass_qc_dask(X) @@ -181,7 +185,9 @@ def _first_pass_qc(X): @with_cupy_rmm -def _first_pass_qc_dask(X): +def _first_pass_qc_dask( + X: DaskArray, +) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]: import dask import dask.array as da @@ -337,7 +343,7 @@ def __qc_calc_2(X_part): ) -def _second_pass_qc(X, mask): +def _second_pass_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: if isinstance(X, DaskArray): return _second_pass_qc_dask(X, mask) sums_cells_sub = cp.zeros(X.shape[0], dtype=X.dtype) @@ -381,7 +387,7 @@ def _second_pass_qc(X, mask): @with_cupy_rmm -def _second_pass_qc_dask(X, mask): +def _second_pass_qc_dask(X: DaskArray, mask: cp.ndarray) -> cp.ndarray: if isinstance(X._meta, sparse.csr_matrix): from ._kernels._qc_kernels import _sparse_qc_csr_sub From 13760b7f413b4d1549e4b257fd094866d0d83f18 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 15 Oct 2024 12:26:58 +0200 Subject: [PATCH 46/69] update normalize --- src/rapids_singlecell/preprocessing/_normalize.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 6fc409b1..9e33af60 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -20,6 +20,8 @@ from anndata import AnnData from cupyx.scipy.sparse import csr_matrix, spmatrix + from rapids_singlecell._utils import ArrayTypesDask + def normalize_total( adata: AnnData, @@ -81,7 +83,7 @@ def normalize_total( return X -def _normalize_total(X: cp.ndarray, target_sum: int): +def _normalize_total(X: ArrayTypesDask, target_sum: int): if isinstance(X, sparse.csr_matrix): return _normalize_total_csr(X, target_sum) elif isinstance(X, DaskArray): @@ -148,7 +150,7 @@ def __mul(X_part): return X -def _get_target_sum(X) -> int: +def _get_target_sum(X: ArrayTypesDask) -> int: if isinstance(X, sparse.csr_matrix): return _get_target_sum_csr(X) elif isinstance(X, DaskArray): @@ -167,7 +169,7 @@ def _get_target_sum_csr(X: sparse.csr_matrix) -> int: (64,), (X.indptr, X.data, counts_per_cell, X.shape[0]), ) - + counts_per_cell = counts_per_cell[counts_per_cell > 0] target_sum = cp.median(counts_per_cell) return target_sum @@ -202,6 +204,7 @@ def __sum(X_part): drop_axis=1, ) counts_per_cell = target_sum_chunk_matrices.compute() + counts_per_cell = counts_per_cell[counts_per_cell > 0] target_sum = cp.median(counts_per_cell) return target_sum @@ -255,7 +258,7 @@ def log1p( X = X.log1p() elif isinstance(X, DaskArray): if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(lambda X: cp.log1p(X), meta=_meta_dense(X.dtype)) + X = X.map_blocks(cp.log1p, meta=_meta_dense(X.dtype)) elif isinstance(X._meta, sparse.csr_matrix): X = X.map_blocks(lambda x: x.log1p(), meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} From b9e49312311c491db953418487b2b6417ac93adc Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 15 Oct 2024 15:19:42 +0200 Subject: [PATCH 47/69] go back to lambda --- src/rapids_singlecell/preprocessing/_normalize.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_normalize.py b/src/rapids_singlecell/preprocessing/_normalize.py index 9e33af60..1acf9dd3 100644 --- a/src/rapids_singlecell/preprocessing/_normalize.py +++ b/src/rapids_singlecell/preprocessing/_normalize.py @@ -258,7 +258,7 @@ def log1p( X = X.log1p() elif isinstance(X, DaskArray): if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(cp.log1p, meta=_meta_dense(X.dtype)) + X = X.map_blocks(lambda x: cp.log1p(x), meta=_meta_dense(X.dtype)) elif isinstance(X._meta, sparse.csr_matrix): X = X.map_blocks(lambda x: x.log1p(), meta=_meta_sparse(X.dtype)) adata.uns["log1p"] = {"base": None} From 9308e211ec52429693a24df457c8cdf326cdacce Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 15 Oct 2024 16:37:28 +0200 Subject: [PATCH 48/69] slim down tests --- tests/dask/test_dask_pca.py | 135 ++++++++++-------------------- tests/dask/test_normalize_dask.py | 68 ++++++++------- tests/dask/test_qc_dask.py | 56 +++---------- tests/dask/test_scale_dask.py | 62 +++++--------- 4 files changed, 119 insertions(+), 202 deletions(-) diff --git a/tests/dask/test_dask_pca.py b/tests/dask/test_dask_pca.py index 95801d96..08d6d5d9 100644 --- a/tests/dask/test_dask_pca.py +++ b/tests/dask/test_dask_pca.py @@ -2,6 +2,7 @@ import cupy as cp import numpy as np +import pytest from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k, pbmc3k_processed from scipy import sparse @@ -13,129 +14,85 @@ ) -def test_pca_sparse_dask(client): - sparse_ad = pbmc3k_processed() - default = pbmc3k_processed() - sparse_ad.X = sparse.csr_matrix(sparse_ad.X.astype(np.float64)) - default.X = as_sparse_cupy_dask_array(default.X.astype(np.float64)) - rsc.pp.pca(sparse_ad) - rsc.pp.pca(default) +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +def test_pca_dask(client, data_kind): + adata_1 = pbmc3k_processed() + adata_2 = pbmc3k_processed() - cp.testing.assert_allclose( - np.abs(sparse_ad.obsm["X_pca"]), - cp.abs(default.obsm["X_pca"].compute()), - rtol=1e-7, - atol=1e-6, - ) + if data_kind == "sparse": + adata_1.X = sparse.csr_matrix(adata_1.X.astype(np.float64)) + adata_2.X = as_sparse_cupy_dask_array(adata_2.X.astype(np.float64)) + elif data_kind == "dense": + adata_1.X = cp.array(adata_1.X.astype(np.float64)) + adata_2.X = as_dense_cupy_dask_array(adata_2.X.astype(np.float64)) + else: + raise ValueError(f"Unknown data_kind {data_kind}") - cp.testing.assert_allclose( - np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 - ) + rsc.pp.pca(adata_1, svd_solver="full") + rsc.pp.pca(adata_2, svd_solver="full") cp.testing.assert_allclose( - np.abs(sparse_ad.uns["pca"]["variance_ratio"]), - np.abs(default.uns["pca"]["variance_ratio"]), + np.abs(adata_1.obsm["X_pca"]), + cp.abs(adata_2.obsm["X_pca"].compute()), rtol=1e-7, atol=1e-6, ) - -def test_pca_dense_dask_full_pipeline(client): - dense = pbmc3k() - default = pbmc3k() - dense.X = cp.array(dense.X.astype(np.float64).toarray()) - default.X = as_dense_cupy_dask_array(default.X.astype(np.float64).toarray()) - - rsc.pp.filter_genes(dense, min_count=500) - rsc.pp.filter_genes(default, min_count=500) - - rsc.pp.normalize_total(dense, target_sum=1e4) - rsc.pp.normalize_total(default, target_sum=1e4) - - rsc.pp.log1p(dense) - rsc.pp.log1p(default) - - rsc.pp.pca(dense, svd_solver="full") - rsc.pp.pca(default, svd_solver="full") - cp.testing.assert_allclose( - np.abs(dense.obsm["X_pca"]), - cp.abs(default.obsm["X_pca"].compute()), + np.abs(adata_1.varm["PCs"]), + np.abs(adata_2.varm["PCs"]), rtol=1e-7, atol=1e-6, ) cp.testing.assert_allclose( - np.abs(dense.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 - ) - - cp.testing.assert_allclose( - np.abs(dense.uns["pca"]["variance_ratio"]), - np.abs(default.uns["pca"]["variance_ratio"]), + np.abs(adata_1.uns["pca"]["variance_ratio"]), + np.abs(adata_2.uns["pca"]["variance_ratio"]), rtol=1e-7, atol=1e-6, ) -def test_pca_sparse_dask_full_pipeline(client): - sparse_ad = pbmc3k() - default = pbmc3k() - sparse_ad.X = cusparse.csr_matrix(sparse.csr_matrix(sparse_ad.X.astype(np.float64))) - default.X = as_sparse_cupy_dask_array(default.X.astype(np.float64)) +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +def test_pca_dask_full_pipeline(client, data_kind): + adata_1 = pbmc3k() + adata_2 = pbmc3k() - rsc.pp.filter_genes(sparse_ad, min_count=100) - rsc.pp.filter_genes(default, min_count=100) + if data_kind == "sparse": + adata_1.X = cusparse.csr_matrix(sparse.csr_matrix(adata_1.X.astype(np.float64))) + adata_2.X = as_sparse_cupy_dask_array(adata_2.X.astype(np.float64)) + elif data_kind == "dense": + adata_1.X = cp.array(adata_1.X.astype(np.float64).toarray()) + adata_2.X = as_dense_cupy_dask_array(adata_2.X.astype(np.float64).toarray()) + else: + raise ValueError(f"Unknown data_kind {data_kind}") - rsc.pp.normalize_total(sparse_ad, target_sum=1e4) - rsc.pp.normalize_total(default, target_sum=1e4) + rsc.pp.filter_genes(adata_1, min_count=500) + rsc.pp.filter_genes(adata_2, min_count=500) - rsc.pp.log1p(sparse_ad) - rsc.pp.log1p(default) - - rsc.pp.pca(sparse_ad) - rsc.pp.pca(default) - - cp.testing.assert_allclose( - np.abs(sparse_ad.obsm["X_pca"]), - cp.abs(default.obsm["X_pca"].compute()), - rtol=1e-7, - atol=1e-6, - ) - - cp.testing.assert_allclose( - np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 - ) - - cp.testing.assert_allclose( - np.abs(sparse_ad.uns["pca"]["variance_ratio"]), - np.abs(default.uns["pca"]["variance_ratio"]), - rtol=1e-7, - atol=1e-6, - ) + rsc.pp.normalize_total(adata_1, target_sum=1e4) + rsc.pp.normalize_total(adata_2, target_sum=1e4) + rsc.pp.log1p(adata_1) + rsc.pp.log1p(adata_2) -def test_pca_dense_dask(client): - sparse_ad = pbmc3k_processed() - default = pbmc3k_processed() - sparse_ad.X = cp.array(sparse_ad.X.astype(np.float64)) - default.X = as_dense_cupy_dask_array(default.X.astype(np.float64)) - rsc.pp.pca(sparse_ad, svd_solver="full") - rsc.pp.pca(default, svd_solver="full") + rsc.pp.pca(adata_1, svd_solver="full") + rsc.pp.pca(adata_2, svd_solver="full") cp.testing.assert_allclose( - np.abs(sparse_ad.obsm["X_pca"]), - cp.abs(default.obsm["X_pca"].compute()), + np.abs(adata_1.obsm["X_pca"]), + cp.abs(adata_2.obsm["X_pca"].compute()), rtol=1e-7, atol=1e-6, ) cp.testing.assert_allclose( - np.abs(sparse_ad.varm["PCs"]), np.abs(default.varm["PCs"]), rtol=1e-7, atol=1e-6 + np.abs(adata_1.varm["PCs"]), np.abs(adata_2.varm["PCs"]), rtol=1e-7, atol=1e-6 ) cp.testing.assert_allclose( - np.abs(sparse_ad.uns["pca"]["variance_ratio"]), - np.abs(default.uns["pca"]["variance_ratio"]), + np.abs(adata_1.uns["pca"]["variance_ratio"]), + np.abs(adata_2.uns["pca"]["variance_ratio"]), rtol=1e-7, atol=1e-6, ) diff --git a/tests/dask/test_normalize_dask.py b/tests/dask/test_normalize_dask.py index dcdca21c..5a76111a 100644 --- a/tests/dask/test_normalize_dask.py +++ b/tests/dask/test_normalize_dask.py @@ -1,6 +1,7 @@ from __future__ import annotations import cupy as cp +import pytest import scanpy as sc from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k @@ -12,51 +13,60 @@ ) -def test_normalize_sparse(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +def test_normalize_total(client, data_kind): adata = pbmc3k() sc.pp.filter_cells(adata, min_genes=100) sc.pp.filter_genes(adata, min_cells=3) dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X) - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.normalize_total(adata) - rsc.pp.normalize_total(dask_data) - cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + if data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + elif data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + else: + raise ValueError(f"Unknown data_kind {data_kind}") -def test_normalize_dense(client): - adata = pbmc3k() - sc.pp.filter_cells(adata, min_genes=100) - sc.pp.filter_genes(adata, min_cells=3) - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X) - adata.X = cp.array(adata.X.toarray()) rsc.pp.normalize_total(adata) rsc.pp.normalize_total(dask_data) - cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + if data_kind == "sparse": + adata_X = adata.X.toarray() + dask_X = dask_data.X.compute().toarray() + else: + adata_X = adata.X + dask_X = dask_data.X.compute() -def test_log1p_sparse(client): - adata = pbmc3k() - sc.pp.filter_cells(adata, min_genes=100) - sc.pp.filter_genes(adata, min_cells=3) - sc.pp.normalize_total(adata) - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X) - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.log1p(adata) - rsc.pp.log1p(dask_data) - cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) + cp.testing.assert_allclose(adata_X, dask_X) -def test_log1p_dense(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +def test_log1p(client, data_kind): adata = pbmc3k() sc.pp.filter_cells(adata, min_genes=100) sc.pp.filter_genes(adata, min_cells=3) sc.pp.normalize_total(adata) dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X) - adata.X = cp.array(adata.X.toarray()) + + if data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + elif data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + else: + raise ValueError(f"Unknown data_kind {data_kind}") + rsc.pp.log1p(adata) rsc.pp.log1p(dask_data) - cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + + if data_kind == "sparse": + adata_X = adata.X.toarray() + dask_X = dask_data.X.compute().toarray() + else: + adata_X = adata.X + dask_X = dask_data.X.compute() + + cp.testing.assert_allclose(adata_X, dask_X) diff --git a/tests/dask/test_qc_dask.py b/tests/dask/test_qc_dask.py index f1fd5f9c..2beafc85 100644 --- a/tests/dask/test_qc_dask.py +++ b/tests/dask/test_qc_dask.py @@ -2,6 +2,7 @@ import cupy as cp import numpy as np +import pytest from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k @@ -12,55 +13,20 @@ ) -def test_qc_metrics_sparse(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +def test_qc_metrics_sparse(client, data_kind): adata = pbmc3k() adata.var["mt"] = adata.var_names.str.startswith("MT-") dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X) - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p=True) - rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p=True) - np.testing.assert_allclose( - adata.obs["n_genes_by_counts"], dask_data.obs["n_genes_by_counts"] - ) - np.testing.assert_allclose(adata.obs["total_counts"], dask_data.obs["total_counts"]) - np.testing.assert_allclose( - adata.obs["log1p_n_genes_by_counts"], dask_data.obs["log1p_n_genes_by_counts"] - ) - np.testing.assert_allclose( - adata.obs["log1p_total_counts"], dask_data.obs["log1p_total_counts"] - ) - np.testing.assert_allclose( - adata.obs["pct_counts_mt"], dask_data.obs["pct_counts_mt"] - ) - np.testing.assert_allclose( - adata.obs["total_counts_mt"], dask_data.obs["total_counts_mt"] - ) - np.testing.assert_allclose( - adata.obs["log1p_total_counts_mt"], dask_data.obs["log1p_total_counts_mt"] - ) - np.testing.assert_allclose( - adata.var["n_cells_by_counts"], dask_data.var["n_cells_by_counts"] - ) - np.testing.assert_allclose(adata.var["total_counts"], dask_data.var["total_counts"]) - np.testing.assert_allclose(adata.var["mean_counts"], dask_data.var["mean_counts"]) - np.testing.assert_allclose( - adata.var["pct_dropout_by_counts"], dask_data.var["pct_dropout_by_counts"] - ) - np.testing.assert_allclose( - adata.var["log1p_total_counts"], dask_data.var["log1p_total_counts"] - ) - np.testing.assert_allclose( - adata.var["log1p_mean_counts"], dask_data.var["log1p_mean_counts"] - ) + if data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X) + adata.X = cusparse.csr_matrix(adata.X) + elif data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X) + adata.X = cp.array(adata.X.toarray()) + else: + raise ValueError(f"Unknown data_kind {data_kind}") - -def test_qc_metrics_dense(client): - adata = pbmc3k() - adata.var["mt"] = adata.var_names.str.startswith("MT-") - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X) - adata.X = cp.array(adata.X.toarray()) rsc.pp.calculate_qc_metrics(adata, qc_vars=["mt"], log1p=True) rsc.pp.calculate_qc_metrics(dask_data, qc_vars=["mt"], log1p=True) np.testing.assert_allclose( diff --git a/tests/dask/test_scale_dask.py b/tests/dask/test_scale_dask.py index 6a9a7351..d8c3eb1f 100644 --- a/tests/dask/test_scale_dask.py +++ b/tests/dask/test_scale_dask.py @@ -2,6 +2,7 @@ import cupy as cp import numpy as np +import pytest import scanpy as sc from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k @@ -23,45 +24,28 @@ def _get_anndata(): return adata.copy() -def test_zc_sparse(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +@pytest.mark.parametrize("zero_center", [True, False]) +def test_scale(client, data_kind, zero_center): adata = _get_anndata() - mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) + mask = np.random.randint(0, 2, adata.shape[0], dtype=bool) dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X.astype(np.float64)) - adata.X = cusparse.csr_matrix(adata.X.astype(np.float64)) - rsc.pp.scale(adata, mask_obs=mask, max_value=10) - rsc.pp.scale(dask_data, mask_obs=mask, max_value=10) - cp.testing.assert_allclose(adata.X, dask_data.X.compute()) - -def test_nzc_sparse(client): - adata = _get_anndata() - mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X) - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.scale(adata, zero_center=False, mask_obs=mask, max_value=10) - rsc.pp.scale(dask_data, zero_center=False, mask_obs=mask, max_value=10) - cp.testing.assert_allclose(adata.X.toarray(), dask_data.X.compute().toarray()) - - -def test_zc_dense(client): - adata = _get_anndata() - mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) - adata.X = cp.array(adata.X.toarray().astype(np.float64)) - rsc.pp.scale(adata, mask_obs=mask, max_value=10) - rsc.pp.scale(dask_data, mask_obs=mask, max_value=10) - cp.testing.assert_allclose(adata.X, dask_data.X.compute()) - - -def test_nzc_dense(client): - adata = _get_anndata() - mask = np.random.randint(0, 2, adata.shape[0], dtype=np.bool_) - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) - adata.X = cp.array(adata.X.toarray().astype(np.float64)) - rsc.pp.scale(adata, zero_center=False, mask_obs=mask, max_value=10) - rsc.pp.scale(dask_data, zero_center=False, mask_obs=mask, max_value=10) - cp.testing.assert_allclose(adata.X, dask_data.X.compute()) + if data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X.astype(np.float64)) + adata.X = cusparse.csr_matrix(adata.X.astype(np.float64)) + elif data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X.astype(np.float64)) + adata.X = cp.array(adata.X.toarray().astype(np.float64)) + else: + raise ValueError(f"Unknown data_kind {data_kind}") + + rsc.pp.scale(adata, zero_center=zero_center, mask_obs=mask, max_value=10) + rsc.pp.scale(dask_data, zero_center=zero_center, mask_obs=mask, max_value=10) + if data_kind == "sparse" and not zero_center: + adata_X = adata.X.toarray() + dask_X = dask_data.X.compute().toarray() + else: + adata_X = adata.X + dask_X = dask_data.X.compute() + cp.testing.assert_allclose(adata_X, dask_X) From f8d6269100c0b27c5035fd9b10c7f257d2045eca Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 15 Oct 2024 16:43:26 +0200 Subject: [PATCH 49/69] run tests on rapids-24.08 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 4505e519..071d475a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ [project.optional-dependencies] rapids11 = ["cupy-cuda11x","cudf-cu11==24.10.*", "cuml-cu11==24.10.*", "cugraph-cu11==24.10.*"] -rapids12 = ["cupy-cuda12x","cudf-cu12==24.10.*", "cuml-cu12==24.10.*", "cugraph-cu12==24.10.*"] +rapids12 = ["cupy-cuda12x","cudf-cu12==24.8.*", "cuml-cu12==24.8.*", "cugraph-cu12==24.8.*"] doc = [ "sphinx>=4.5.0", "sphinx-copybutton", From 65f941ae88eac35c180a9e4bcfb2902534d65d64 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 16 Oct 2024 15:17:44 +0200 Subject: [PATCH 50/69] compress hvg tests --- tests/dask/test_hvg_dask.py | 126 ++++++++---------------------------- 1 file changed, 26 insertions(+), 100 deletions(-) diff --git a/tests/dask/test_hvg_dask.py b/tests/dask/test_hvg_dask.py index 1b47c392..2cc5913d 100644 --- a/tests/dask/test_hvg_dask.py +++ b/tests/dask/test_hvg_dask.py @@ -2,6 +2,7 @@ import cupy as cp import pandas as pd +import pytest import scanpy as sc from cupyx.scipy import sparse as cusparse from scanpy.datasets import pbmc3k @@ -22,91 +23,25 @@ def _get_anndata(): return adata -def test_seurat_sparse(client): - adata = _get_anndata() - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.highly_variable_genes(adata) - rsc.pp.highly_variable_genes(dask_data) - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) - - -def test_seurat_sparse_batch(client): - adata = _get_anndata() - adata.obs["batch"] = ( - "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") - )[: adata.n_obs] - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.highly_variable_genes(adata, batch_key="batch") - rsc.pp.highly_variable_genes(dask_data, batch_key="batch") - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) - - -def test_cr_sparse(client): - adata = _get_anndata() - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.highly_variable_genes(adata, flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) - - -def test_cr_sparse_batch(client): - adata = _get_anndata() - adata.obs["batch"] = ( - "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") - )[: adata.n_obs] - dask_data = adata.copy() - dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() - adata.X = cusparse.csr_matrix(adata.X) - rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data, batch_key="batch", flavor="cell_ranger") - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) - - -def test_cr_dense(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +@pytest.mark.parametrize("flavor", ["seurat", "cell_ranger"]) +def test_highly_variable_genes(client, data_kind, flavor): adata = _get_anndata() adata.X = adata.X.astype("float64") dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() - adata.X = cp.array(adata.X.toarray()) - rsc.pp.highly_variable_genes(adata, flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data, flavor="cell_ranger") - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) + if data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + elif data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + else: + raise ValueError(f"Unknown data_kind: {data_kind}") + + rsc.pp.highly_variable_genes(adata, flavor=flavor) + rsc.pp.highly_variable_genes(dask_data, flavor=flavor) -def test_seurat_dense(client): - adata = _get_anndata() - adata.X = adata.X.astype("float64") - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() - adata.X = cp.array(adata.X.toarray()) - rsc.pp.highly_variable_genes(adata) - rsc.pp.highly_variable_genes(dask_data) cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) cp.testing.assert_allclose( @@ -114,33 +49,24 @@ def test_seurat_dense(client): ) -def test_cr_dense_batch(client): +@pytest.mark.parametrize("data_kind", ["sparse", "dense"]) +@pytest.mark.parametrize("flavor", ["seurat", "cell_ranger"]) +def test_highly_variable_genes_batched(client, data_kind, flavor): adata = _get_anndata() adata.obs["batch"] = ( "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") )[: adata.n_obs] - adata.X = adata.X.astype("float64") dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() - adata.X = cp.array(adata.X.toarray()) - rsc.pp.highly_variable_genes(adata, batch_key="batch", flavor="cell_ranger") - rsc.pp.highly_variable_genes(dask_data, batch_key="batch", flavor="cell_ranger") - cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) - cp.testing.assert_allclose(adata.var["dispersions"], dask_data.var["dispersions"]) - cp.testing.assert_allclose( - adata.var["dispersions_norm"], dask_data.var["dispersions_norm"] - ) + if data_kind == "dense": + dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() + adata.X = cp.array(adata.X.toarray()) + elif data_kind == "sparse": + dask_data.X = as_sparse_cupy_dask_array(dask_data.X).persist() + adata.X = cusparse.csr_matrix(adata.X) + else: + raise ValueError(f"Unknown data_kind: {data_kind}") -def test_seurat_dense_batch(client): - adata = _get_anndata() - adata.obs["batch"] = ( - "source_" + pd.array([*range(1, 6), 5]).repeat(500).astype("string") - )[: adata.n_obs] - adata.X = adata.X.astype("float64") - dask_data = adata.copy() - dask_data.X = as_dense_cupy_dask_array(dask_data.X).persist() - adata.X = cp.array(adata.X.toarray()) rsc.pp.highly_variable_genes(adata, batch_key="batch") rsc.pp.highly_variable_genes(dask_data, batch_key="batch") cp.testing.assert_allclose(adata.var["means"], dask_data.var["means"]) From 17ca2efbb966f45f471c15fad69fbc64cfe873a7 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 16 Oct 2024 16:59:53 +0200 Subject: [PATCH 51/69] remove .todelayed --- .../preprocessing/_sparse_pca/_dask_sparse_pca.py | 2 +- src/rapids_singlecell/preprocessing/_utils.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index 003f8393..2a59cbfe 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -148,7 +148,7 @@ def __gram_block(x_part): ) return gram_matrix[None, ...] # need new axis for summing - n_blocks = len(x.to_delayed().ravel()) + n_blocks = x.blocks.size gram_matrix = x.map_blocks( __gram_block, new_axis=(1,), diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 2f667e3f..cc0a4b59 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -98,7 +98,7 @@ def __mean_var(X_part): ) return cp.vstack([mean, var])[None, ...] # new axis for summing - n_blocks = len(X.to_delayed().ravel()) + n_blocks = X.blocks.size mean, var = X.map_blocks( __mean_var, new_axis=(1,), @@ -175,7 +175,7 @@ def __mean_var(X_part): None if 1 - axis else slice(None, None), ... ] - n_blocks = len(X.to_delayed().ravel()) + n_blocks = X.blocks.size mean_var = X.map_blocks( __mean_var, new_axis=(1,) if axis - 1 else None, From 06ce8e584b4ba0d078daa657ee260bf13423466b Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 16 Oct 2024 17:04:08 +0200 Subject: [PATCH 52/69] remove dask.delayed --- src/rapids_singlecell/preprocessing/_qc.py | 103 ++++++++------------- tests/dask/test_qc_dask.py | 2 +- 2 files changed, 42 insertions(+), 63 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index bd6bdd2c..2199e962 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -128,43 +128,33 @@ def _first_pass_qc( block = (32,) grid = (int(math.ceil(X.shape[0] / block[0])),) - sparse_qc_csr = _sparse_qc_csr(X.data.dtype) - sparse_qc_csr( - grid, - block, - ( - X.indptr, - X.indices, - X.data, - sums_cells, - sums_genes, - cell_ex, - gene_ex, - X.shape[0], - ), - ) + call_shape = X.shape[0] + sparse_qc_kernel = _sparse_qc_csr(X.data.dtype) + elif sparse.isspmatrix_csc(X): from ._kernels._qc_kernels import _sparse_qc_csc block = (32,) grid = (int(math.ceil(X.shape[1] / block[0])),) - sparse_qc_csc = _sparse_qc_csc(X.data.dtype) - sparse_qc_csc( - grid, - block, - ( - X.indptr, - X.indices, - X.data, - sums_cells, - sums_genes, - cell_ex, - gene_ex, - X.shape[1], - ), - ) + call_shape = X.shape[1] + sparse_qc_kernel = _sparse_qc_csc(X.data.dtype) + else: raise ValueError("Please use a csr or csc matrix") + sparse_qc_kernel( + grid, + block, + ( + X.indptr, + X.indices, + X.data, + sums_cells, + sums_genes, + cell_ex, + gene_ex, + call_shape, + ), + ) else: from ._kernels._qc_kernels import _sparse_qc_dense @@ -189,7 +179,6 @@ def _first_pass_qc_dask( X: DaskArray, ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]: import dask - import dask.array as da if isinstance(X._meta, sparse.csr_matrix): from ._kernels._qc_kernels_dask import ( @@ -200,7 +189,6 @@ def _first_pass_qc_dask( sparse_qc_csr_cells = _sparse_qc_csr_dask_cells(X.dtype) sparse_qc_csr_cells.compile() - @dask.delayed def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) @@ -219,12 +207,11 @@ def __qc_calc_1(X_part): X_part.shape[0], ), ) - return cp.vstack([sums_cells, cell_ex.astype(X_part.dtype)]) + return cp.stack([sums_cells, cell_ex.astype(X_part.dtype)], axis=1) sparse_qc_csr_genes = _sparse_qc_csr_dask_genes(X.dtype) sparse_qc_csr_genes.compile() - @dask.delayed def __qc_calc_2(X_part): sums_genes = cp.zeros(X_part.shape[1], dtype=X_part.dtype) gene_ex = cp.zeros(X_part.shape[1], dtype=cp.int32) @@ -241,7 +228,7 @@ def __qc_calc_2(X_part): X_part.nnz, ), ) - return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)]) + return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)])[None, ...] elif isinstance(X._meta, cp.ndarray): from ._kernels._qc_kernels_dask import ( @@ -252,7 +239,6 @@ def __qc_calc_2(X_part): sparse_qc_dense_cells = _sparse_qc_dense_cells(X.dtype) sparse_qc_dense_cells.compile() - @dask.delayed def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) @@ -274,12 +260,11 @@ def __qc_calc_1(X_part): X_part.shape[1], ), ) - return cp.vstack([sums_cells, cell_ex.astype(X_part.dtype)]) + return cp.stack([sums_cells, cell_ex.astype(X_part.dtype)], axis=1) sparse_qc_dense_genes = _sparse_qc_dense_genes(X.dtype) sparse_qc_dense_genes.compile() - @dask.delayed def __qc_calc_2(X_part): sums_genes = cp.zeros((X_part.shape[1]), dtype=X_part.dtype) gene_ex = cp.zeros((X_part.shape[1]), dtype=cp.int32) @@ -301,35 +286,29 @@ def __qc_calc_2(X_part): X_part.shape[1], ), ) - return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)]) + return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)])[None, ...] else: raise ValueError( "Please use a cupy csr_matrix or cp.ndarray. csc_matrix are not supported with dask." ) - blocks = X.to_delayed().ravel() - cell_blocks = [ - da.from_delayed( - __qc_calc_1(block), - shape=(2, X.chunks[0][ind]), - dtype=X.dtype, - meta=cp.array([]), - ) - for ind, block in enumerate(blocks) - ] - - blocks = X.to_delayed().ravel() - gene_blocks = [ - da.from_delayed( - __qc_calc_2(block), - shape=(2, X.shape[1]), - dtype=X.dtype, - meta=cp.array([]), - ) - for ind, block in enumerate(blocks) - ] - sums_cells, cell_ex = da.hstack(cell_blocks) - sums_genes, gene_ex = da.stack(gene_blocks, axis=1).sum(axis=1) + cell_results = X.map_blocks( + __qc_calc_1, + chunks=(X.chunks[0], (2,)), + dtype=X.dtype, + meta=cp.empty((0, 2), dtype=X.dtype), + ) + sums_cells = cell_results[:, 0] + cell_ex = cell_results[:, 1] + + n_blocks = X.blocks.size + sums_genes, gene_ex = X.map_blocks( + __qc_calc_2, + new_axis=(1,), + chunks=((1,) * n_blocks, (2,), (X.shape[1],)), + dtype=X.dtype, + meta=cp.array([]), + ).sum(axis=0) sums_cells, cell_ex, sums_genes, gene_ex = dask.compute( sums_cells, cell_ex, sums_genes, gene_ex diff --git a/tests/dask/test_qc_dask.py b/tests/dask/test_qc_dask.py index 2beafc85..3b350dbe 100644 --- a/tests/dask/test_qc_dask.py +++ b/tests/dask/test_qc_dask.py @@ -14,7 +14,7 @@ @pytest.mark.parametrize("data_kind", ["sparse", "dense"]) -def test_qc_metrics_sparse(client, data_kind): +def test_qc_metrics(client, data_kind): adata = pbmc3k() adata.var["mt"] = adata.var_names.str.startswith("MT-") dask_data = adata.copy() From 6d948355a4b387379f64ad1148329960e3caa84d Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 17 Oct 2024 10:46:48 +0200 Subject: [PATCH 53/69] update qc --- src/rapids_singlecell/preprocessing/_qc.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index 2199e962..843e2866 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -332,24 +332,22 @@ def _second_pass_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: block = (32,) grid = (int(math.ceil(X.shape[0] / block[0])),) - sparse_qc_csr_sub = _sparse_qc_csr_sub(X.data.dtype) - sparse_qc_csr_sub( - grid, - block, - (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[0]), - ) + call_shape = X.shape[0] + sparse_qc_sub = _sparse_qc_csr_sub(X.data.dtype) + elif sparse.isspmatrix_csc(X): from ._kernels._qc_kernels import _sparse_qc_csc_sub block = (32,) grid = (int(math.ceil(X.shape[1] / block[0])),) - sparse_qc_csc_sub = _sparse_qc_csc_sub(X.data.dtype) - sparse_qc_csc_sub( - grid, - block, - (X.indptr, X.indices, X.data, sums_cells_sub, mask, X.shape[1]), - ) + call_shape = X.shape[1] + sparse_qc_sub = _sparse_qc_csc_sub(X.data.dtype) + sparse_qc_sub( + grid, + block, + (X.indptr, X.indices, X.data, sums_cells_sub, mask, call_shape), + ) else: from ._kernels._qc_kernels import _sparse_qc_dense_sub From b733ab3e92cc14892a65da1b5db655afb44d115d Mon Sep 17 00:00:00 2001 From: Severin Dicks <37635888+Intron7@users.noreply.github.com> Date: Tue, 22 Oct 2024 13:35:33 +0200 Subject: [PATCH 54/69] Update src/rapids_singlecell/preprocessing/_pca.py Co-authored-by: Philipp A. --- src/rapids_singlecell/preprocessing/_pca.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 7d970a37..f3c11a32 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -201,7 +201,7 @@ def pca( ) X_pca = pca_func.fit_transform(X) - elif not zero_center: + else: # not zero_center from cuml.decomposition import TruncatedSVD pca_func = TruncatedSVD( From 75bbbb80a3cca0ca703d8814a3c4ad9dd1104946 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 22 Oct 2024 11:35:39 +0000 Subject: [PATCH 55/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_pca.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index f3c11a32..26cc3021 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -201,7 +201,7 @@ def pca( ) X_pca = pca_func.fit_transform(X) - else: # not zero_center + else: # not zero_center from cuml.decomposition import TruncatedSVD pca_func = TruncatedSVD( From d2d2e45264ba7da87dd202143883c181cb60c312 Mon Sep 17 00:00:00 2001 From: Severin Dicks <37635888+Intron7@users.noreply.github.com> Date: Tue, 12 Nov 2024 14:25:56 +0100 Subject: [PATCH 56/69] Update src/rapids_singlecell/preprocessing/_scale.py Co-authored-by: Philipp A. --- src/rapids_singlecell/preprocessing/_scale.py | 22 ++++++------------- 1 file changed, 7 insertions(+), 15 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 4e2946f3..db6278c2 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -327,24 +327,16 @@ def __dense(X_part): dtype=X.dtype, meta=_meta_dense(X.dtype), ) - return _scale_dask_array_zc( - X, mask_array=mask_array, mean=mean, std=std, max_value=max_value - ) - + scale = _scale_dask_array_zc elif isinstance(X._meta, sparse.csr_matrix) and not zero_center: - return _scale_sparse_csr_dask( - X, mask_array=mask_array, mean=mean, std=std, max_value=max_value - ) - + scale = _scale_sparse_csr_dask elif isinstance(X._meta, cp.ndarray) and zero_center: - return _scale_dask_array_zc( - X, mask_array=mask_array, mean=mean, std=std, max_value=max_value - ) - + scale = _scale_dask_array_zc elif isinstance(X._meta, cp.ndarray) and not zero_center: - return _scale_dask_array_nzc( - X, mask_array=mask_array, mean=mean, std=std, max_value=max_value - ) + scale = _scale_dask_array_nzc + else: + TODO # raise error + return scale(X, mask_array=mask_array, mean=mean, std=std, max_value=max_value) def _scale_dask_array_zc(X, *, mask_array, mean, std, max_value): From d3421caa78905596612e49a072c2eac2cffc32c8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 12 Nov 2024 14:43:42 +0100 Subject: [PATCH 57/69] add error --- src/rapids_singlecell/preprocessing/_scale.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index db6278c2..093f6898 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -335,7 +335,9 @@ def __dense(X_part): elif isinstance(X._meta, cp.ndarray) and not zero_center: scale = _scale_dask_array_nzc else: - TODO # raise error + raise ValueError( + "Invalid `._meta` type only supports `cupyx.scipy.sparse.csr_matrix` and `cp.ndarray`" + ) return scale(X, mask_array=mask_array, mean=mean, std=std, max_value=max_value) From 8b32156eff668d8e336ef09e07dc2f26be5ca2d6 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Tue, 12 Nov 2024 15:02:39 +0100 Subject: [PATCH 58/69] update tree pca --- src/rapids_singlecell/preprocessing/_pca.py | 64 ++++++++++----------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 26cc3021..89b29ed7 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -145,6 +145,8 @@ def pca( "Dask arrays are not supported for chunked PCA computation." ) _check_gpu_X(X, allow_dask=True) + if not zero_center: + raise ValueError("Dask arrays do not support non-zero centered PCA.") if isinstance(X._meta, cp.ndarray): from cuml.dask.decomposition import PCA @@ -160,7 +162,7 @@ def pca( pca_func = pca_func.fit(X) X_pca = pca_func.transform(X) - else: + elif zero_center: if chunked: from cuml.decomposition import IncrementalPCA @@ -179,38 +181,36 @@ def pca( if issparse(chunk) or cpissparse(chunk): chunk = chunk.toarray() X_pca[start_idx:stop_idx] = pca_func.transform(chunk) + elif cpissparse(X) or issparse(X): + if issparse(X): + X = sparse_scipy_to_cp(X, dtype=X.dtype) + from ._sparse_pca._sparse_pca import PCA_sparse + + if not isspmatrix_csr(X): + X = X.tocsr() + pca_func = PCA_sparse(n_components=n_comps) + X_pca = pca_func.fit_transform(X) else: - if zero_center: - if cpissparse(X) or issparse(X): - if issparse(X): - X = sparse_scipy_to_cp(X, dtype=X.dtype) - from ._sparse_pca._sparse_pca import PCA_sparse - - if not isspmatrix_csr(X): - X = X.tocsr() - pca_func = PCA_sparse(n_components=n_comps) - X_pca = pca_func.fit_transform(X) - else: - from cuml.decomposition import PCA - - pca_func = PCA( - n_components=n_comps, - svd_solver=svd_solver, - random_state=random_state, - output_type="numpy", - ) - X_pca = pca_func.fit_transform(X) - - else: # not zero_center - from cuml.decomposition import TruncatedSVD - - pca_func = TruncatedSVD( - n_components=n_comps, - random_state=random_state, - algorithm=svd_solver, - output_type="numpy", - ) - X_pca = pca_func.fit_transform(X) + from cuml.decomposition import PCA + + pca_func = PCA( + n_components=n_comps, + svd_solver=svd_solver, + random_state=random_state, + output_type="numpy", + ) + X_pca = pca_func.fit_transform(X) + + else: # not zero_center + from cuml.decomposition import TruncatedSVD + + pca_func = TruncatedSVD( + n_components=n_comps, + random_state=random_state, + algorithm=svd_solver, + output_type="numpy", + ) + X_pca = pca_func.fit_transform(X) if X_pca.dtype.descr != np.dtype(dtype).descr: X_pca = X_pca.astype(dtype) From bb10cda5eea817be7790a8fdc18e7980d1187d9d Mon Sep 17 00:00:00 2001 From: Severin Dicks <37635888+Intron7@users.noreply.github.com> Date: Wed, 13 Nov 2024 12:02:35 +0100 Subject: [PATCH 59/69] Update src/rapids_singlecell/preprocessing/_scale.py Co-authored-by: Ilan Gold --- src/rapids_singlecell/preprocessing/_scale.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 093f6898..479182fa 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -284,14 +284,11 @@ def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=N if not inplace: X = X.copy() + mean, var = dask.compute(*_get_mean_var(X[mask_obs if mask_obs is not None else slice(None), :])) if mask_obs is None: - mean, var = _get_mean_var(X) mask_array = cp.ones(X.shape[0], dtype=cp.int32) - else: - mean, var = _get_mean_var(X[mask_obs, :]) mask_array = cp.array(mask_obs).astype(cp.int32) - mean, var = dask.compute(mean, var) std = cp.sqrt(var) std[std == 0] = 1 max_value = _get_max_value(max_value, X.dtype) From 1ae74d7e2e123108ae7078e989643b4fcad101d7 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 13 Nov 2024 12:45:01 +0100 Subject: [PATCH 60/69] add note --- src/rapids_singlecell/preprocessing/_pca.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/rapids_singlecell/preprocessing/_pca.py b/src/rapids_singlecell/preprocessing/_pca.py index 89b29ed7..c970dc7b 100644 --- a/src/rapids_singlecell/preprocessing/_pca.py +++ b/src/rapids_singlecell/preprocessing/_pca.py @@ -154,6 +154,7 @@ def pca( svd_solver = "jacobi" pca_func = PCA(n_components=n_comps, svd_solver=svd_solver, whiten=False) X_pca = pca_func.fit_transform(X) + # cuml-issue #5883 X_pca = X_pca.compute_chunk_sizes() elif isinstance(X._meta, csr_matrix): from ._sparse_pca._dask_sparse_pca import PCA_sparse_dask From 38e4ad0fe10e99dfdc1def4d6b4f0c124de658e8 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 13 Nov 2024 12:45:21 +0100 Subject: [PATCH 61/69] dask import --- src/rapids_singlecell/preprocessing/_scale.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_scale.py b/src/rapids_singlecell/preprocessing/_scale.py index 479182fa..13678258 100644 --- a/src/rapids_singlecell/preprocessing/_scale.py +++ b/src/rapids_singlecell/preprocessing/_scale.py @@ -4,6 +4,8 @@ from typing import Union import cupy as cp +import dask +import dask.array as da import numpy as np from anndata import AnnData from cupyx.scipy import sparse @@ -21,11 +23,6 @@ _sparse_to_dense, ) -try: - import dask.array as da -except ImportError: - pass - def scale( adata: AnnData, @@ -280,11 +277,11 @@ def _scale_sparse_csr( def _scale_dask(X, *, mask_obs=None, zero_center=True, inplace=True, max_value=None): - import dask - if not inplace: X = X.copy() - mean, var = dask.compute(*_get_mean_var(X[mask_obs if mask_obs is not None else slice(None), :])) + mean, var = dask.compute( + *_get_mean_var(X[mask_obs if mask_obs is not None else slice(None), :]) + ) if mask_obs is None: mask_array = cp.ones(X.shape[0], dtype=cp.int32) else: From 97556415c5e853f8846167bf2d8159b7f06b515b Mon Sep 17 00:00:00 2001 From: Intron7 Date: Wed, 13 Nov 2024 12:45:34 +0100 Subject: [PATCH 62/69] update qc names --- src/rapids_singlecell/preprocessing/_qc.py | 86 +++++++++++++--------- 1 file changed, 50 insertions(+), 36 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_qc.py b/src/rapids_singlecell/preprocessing/_qc.py index 843e2866..e14618c3 100644 --- a/src/rapids_singlecell/preprocessing/_qc.py +++ b/src/rapids_singlecell/preprocessing/_qc.py @@ -76,23 +76,25 @@ def calculate_qc_metrics( _check_gpu_X(X, allow_dask=True) - sums_cells, sums_genes, cell_ex, gene_ex = _first_pass_qc(X) + sums_cells, sums_genes, genes_per_cell, cells_per_gene = _basic_qc(X) # .var - adata.var[f"n_cells_by_{expr_type}"] = cp.asnumpy(gene_ex) + adata.var[f"n_cells_by_{expr_type}"] = cp.asnumpy(cells_per_gene) adata.var[f"total_{expr_type}"] = cp.asnumpy(sums_genes) mean_array = sums_genes / adata.n_obs adata.var[f"mean_{expr_type}"] = cp.asnumpy(mean_array) adata.var[f"pct_dropout_by_{expr_type}"] = cp.asnumpy( - (1 - gene_ex / adata.n_obs) * 100 + (1 - cells_per_gene / adata.n_obs) * 100 ) if log1p: adata.var[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_genes)) adata.var[f"log1p_mean_{expr_type}"] = cp.asnumpy(cp.log1p(mean_array)) # .obs - adata.obs[f"n_{var_type}_by_{expr_type}"] = cp.asnumpy(cell_ex) + adata.obs[f"n_{var_type}_by_{expr_type}"] = cp.asnumpy(genes_per_cell) adata.obs[f"total_{expr_type}"] = cp.asnumpy(sums_cells) if log1p: - adata.obs[f"log1p_n_{var_type}_by_{expr_type}"] = cp.asnumpy(cp.log1p(cell_ex)) + adata.obs[f"log1p_n_{var_type}_by_{expr_type}"] = cp.asnumpy( + cp.log1p(genes_per_cell) + ) adata.obs[f"log1p_total_{expr_type}"] = cp.asnumpy(cp.log1p(sums_cells)) if qc_vars: @@ -100,7 +102,7 @@ def calculate_qc_metrics( qc_vars = [qc_vars] for qc_var in qc_vars: mask = cp.array(adata.var[qc_var], dtype=cp.bool_) - sums_cells_sub = _second_pass_qc(X, mask) + sums_cells_sub = _geneset_qc(X, mask) adata.obs[f"total_{expr_type}_{qc_var}"] = cp.asnumpy(sums_cells_sub) adata.obs[f"pct_{expr_type}_{qc_var}"] = cp.asnumpy( @@ -112,16 +114,16 @@ def calculate_qc_metrics( ) -def _first_pass_qc( +def _basic_qc( X: ArrayTypesDask, ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]: if isinstance(X, DaskArray): - return _first_pass_qc_dask(X) + return _basic_qc_dask(X) sums_cells = cp.zeros(X.shape[0], dtype=X.dtype) sums_genes = cp.zeros(X.shape[1], dtype=X.dtype) - cell_ex = cp.zeros(X.shape[0], dtype=cp.int32) - gene_ex = cp.zeros(X.shape[1], dtype=cp.int32) + genes_per_cell = cp.zeros(X.shape[0], dtype=cp.int32) + cells_per_gene = cp.zeros(X.shape[1], dtype=cp.int32) if sparse.issparse(X): if sparse.isspmatrix_csr(X): from ._kernels._qc_kernels import _sparse_qc_csr @@ -150,8 +152,8 @@ def _first_pass_qc( X.data, sums_cells, sums_genes, - cell_ex, - gene_ex, + genes_per_cell, + cells_per_gene, call_shape, ), ) @@ -169,13 +171,21 @@ def _first_pass_qc( sparse_qc_dense( grid, block, - (X, sums_cells, sums_genes, cell_ex, gene_ex, X.shape[0], X.shape[1]), + ( + X, + sums_cells, + sums_genes, + genes_per_cell, + cells_per_gene, + X.shape[0], + X.shape[1], + ), ) - return sums_cells, sums_genes, cell_ex, gene_ex + return sums_cells, sums_genes, genes_per_cell, cells_per_gene @with_cupy_rmm -def _first_pass_qc_dask( +def _basic_qc_dask( X: DaskArray, ) -> tuple[cp.ndarray, cp.ndarray, cp.ndarray, cp.ndarray]: import dask @@ -191,7 +201,7 @@ def _first_pass_qc_dask( def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) - cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) + genes_per_cell = cp.zeros(X_part.shape[0], dtype=cp.int32) block = (32,) grid = (int(math.ceil(X_part.shape[0] / block[0])),) @@ -203,18 +213,18 @@ def __qc_calc_1(X_part): X_part.indices, X_part.data, sums_cells, - cell_ex, + genes_per_cell, X_part.shape[0], ), ) - return cp.stack([sums_cells, cell_ex.astype(X_part.dtype)], axis=1) + return cp.stack([sums_cells, genes_per_cell.astype(X_part.dtype)], axis=1) sparse_qc_csr_genes = _sparse_qc_csr_dask_genes(X.dtype) sparse_qc_csr_genes.compile() def __qc_calc_2(X_part): sums_genes = cp.zeros(X_part.shape[1], dtype=X_part.dtype) - gene_ex = cp.zeros(X_part.shape[1], dtype=cp.int32) + cells_per_gene = cp.zeros(X_part.shape[1], dtype=cp.int32) block = (32,) grid = (int(math.ceil(X_part.nnz / block[0])),) sparse_qc_csr_genes( @@ -224,11 +234,13 @@ def __qc_calc_2(X_part): X_part.indices, X_part.data, sums_genes, - gene_ex, + cells_per_gene, X_part.nnz, ), ) - return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)])[None, ...] + return cp.vstack([sums_genes, cells_per_gene.astype(X_part.dtype)])[ + None, ... + ] elif isinstance(X._meta, cp.ndarray): from ._kernels._qc_kernels_dask import ( @@ -241,7 +253,7 @@ def __qc_calc_2(X_part): def __qc_calc_1(X_part): sums_cells = cp.zeros(X_part.shape[0], dtype=X_part.dtype) - cell_ex = cp.zeros(X_part.shape[0], dtype=cp.int32) + genes_per_cell = cp.zeros(X_part.shape[0], dtype=cp.int32) if not X_part.flags.c_contiguous: X_part = cp.asarray(X_part, order="C") block = (16, 16) @@ -255,19 +267,19 @@ def __qc_calc_1(X_part): ( X_part, sums_cells, - cell_ex, + genes_per_cell, X_part.shape[0], X_part.shape[1], ), ) - return cp.stack([sums_cells, cell_ex.astype(X_part.dtype)], axis=1) + return cp.stack([sums_cells, genes_per_cell.astype(X_part.dtype)], axis=1) sparse_qc_dense_genes = _sparse_qc_dense_genes(X.dtype) sparse_qc_dense_genes.compile() def __qc_calc_2(X_part): sums_genes = cp.zeros((X_part.shape[1]), dtype=X_part.dtype) - gene_ex = cp.zeros((X_part.shape[1]), dtype=cp.int32) + cells_per_gene = cp.zeros((X_part.shape[1]), dtype=cp.int32) if not X_part.flags.c_contiguous: X_part = cp.asarray(X_part, order="C") block = (16, 16) @@ -281,12 +293,14 @@ def __qc_calc_2(X_part): ( X_part, sums_genes, - gene_ex, + cells_per_gene, X_part.shape[0], X_part.shape[1], ), ) - return cp.vstack([sums_genes, gene_ex.astype(X_part.dtype)])[None, ...] + return cp.vstack([sums_genes, cells_per_gene.astype(X_part.dtype)])[ + None, ... + ] else: raise ValueError( "Please use a cupy csr_matrix or cp.ndarray. csc_matrix are not supported with dask." @@ -299,10 +313,10 @@ def __qc_calc_2(X_part): meta=cp.empty((0, 2), dtype=X.dtype), ) sums_cells = cell_results[:, 0] - cell_ex = cell_results[:, 1] + genes_per_cell = cell_results[:, 1] n_blocks = X.blocks.size - sums_genes, gene_ex = X.map_blocks( + sums_genes, cells_per_gene = X.map_blocks( __qc_calc_2, new_axis=(1,), chunks=((1,) * n_blocks, (2,), (X.shape[1],)), @@ -310,21 +324,21 @@ def __qc_calc_2(X_part): meta=cp.array([]), ).sum(axis=0) - sums_cells, cell_ex, sums_genes, gene_ex = dask.compute( - sums_cells, cell_ex, sums_genes, gene_ex + sums_cells, genes_per_cell, sums_genes, cells_per_gene = dask.compute( + sums_cells, genes_per_cell, sums_genes, cells_per_gene ) return ( sums_cells.ravel(), sums_genes.ravel(), - cell_ex.ravel().astype(cp.int32), - gene_ex.ravel().astype(cp.int32), + genes_per_cell.ravel().astype(cp.int32), + cells_per_gene.ravel().astype(cp.int32), ) -def _second_pass_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: +def _geneset_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: if isinstance(X, DaskArray): - return _second_pass_qc_dask(X, mask) + return _geneset_qc_dask(X, mask) sums_cells_sub = cp.zeros(X.shape[0], dtype=X.dtype) if sparse.issparse(X): if sparse.isspmatrix_csr(X): @@ -364,7 +378,7 @@ def _second_pass_qc(X: ArrayTypesDask, mask: cp.ndarray) -> cp.ndarray: @with_cupy_rmm -def _second_pass_qc_dask(X: DaskArray, mask: cp.ndarray) -> cp.ndarray: +def _geneset_qc_dask(X: DaskArray, mask: cp.ndarray) -> cp.ndarray: if isinstance(X._meta, sparse.csr_matrix): from ._kernels._qc_kernels import _sparse_qc_csr_sub From 5fd5a97192f5682179a0bcc890cd2f4825fe909c Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 14 Nov 2024 09:03:29 +0100 Subject: [PATCH 63/69] update --- docs/release-notes/0.11.0.md | 14 ++++++++++ src/rapids_singlecell/preprocessing/_hvg.py | 31 ++++++++++++--------- 2 files changed, 32 insertions(+), 13 deletions(-) create mode 100644 docs/release-notes/0.11.0.md diff --git a/docs/release-notes/0.11.0.md b/docs/release-notes/0.11.0.md new file mode 100644 index 00000000..9d1bf581 --- /dev/null +++ b/docs/release-notes/0.11.0.md @@ -0,0 +1,14 @@ +### 0.10.11 {small}`the-future` + +```{rubric} Features +``` +* Adds support for Multi-GPU out-of-core support through Dask {pr}`179` {smaller}`S Dicks, I Gold & P Angerer` +```{rubric} Performance +``` + + +```{rubric} Bug fixes +``` + +```{rubric} Misc +``` diff --git a/src/rapids_singlecell/preprocessing/_hvg.py b/src/rapids_singlecell/preprocessing/_hvg.py index 30d9831c..6db78c74 100644 --- a/src/rapids_singlecell/preprocessing/_hvg.py +++ b/src/rapids_singlecell/preprocessing/_hvg.py @@ -266,6 +266,21 @@ def in_bounds( ) +def _hvg_expm1(X): + if isinstance(X, DaskArray): + if isinstance(X._meta, cp.ndarray): + X = X.map_blocks(lambda X: cp.expm1(X), meta=_meta_dense(X.dtype)) + elif isinstance(X._meta, csr_matrix): + X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) + else: + X = X.copy() + if issparse(X): + X = X.expm1() + else: + X = cp.expm1(X) + return X + + def _highly_variable_genes_single_batch( adata: AnnData, *, @@ -288,17 +303,7 @@ def _highly_variable_genes_single_batch( X = X.copy() if flavor == "seurat": - if isinstance(X, DaskArray): - if isinstance(X._meta, cp.ndarray): - X = X.map_blocks(lambda X: cp.expm1(X), meta=_meta_dense(X.dtype)) - elif isinstance(X._meta, csr_matrix): - X = X.map_blocks(lambda X: X.expm1(), meta=_meta_sparse(X.dtype)) - else: - X = X.copy() - if issparse(X): - X = X.expm1() - else: - X = cp.expm1(X) + X = _hvg_expm1(X) mean, var = _get_mean_var(X, axis=0) if isinstance(X, DaskArray): @@ -429,11 +434,11 @@ def _highly_variable_genes_batched( dfs = [] gene_list = adata.var_names for batch in batches: - adata_subset = adata[adata.obs[batch_key] == batch].copy() + adata_subset = adata[adata.obs[batch_key] == batch] calculate_qc_metrics(adata_subset, layer=layer) filt = adata_subset.var["n_cells_by_counts"].to_numpy() > 0 - adata_subset = adata_subset[:, filt].copy() + adata_subset = adata_subset[:, filt] hvg = _highly_variable_genes_single_batch( adata_subset, From af1faf5983537156c63f9214dbac8119b4acb4cc Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 14 Nov 2024 12:11:50 +0100 Subject: [PATCH 64/69] update _check_gpu_X --- src/rapids_singlecell/preprocessing/_utils.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index cc0a4b59..65b02fd3 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -260,10 +260,10 @@ def _check_nonnegative_integers(X): return True -def _check_gpu_X(X, require_cf=False, allow_dask=False): +def _check_gpu_X(X, require_cf=False, allow_dask=False, allow_csc=True): if isinstance(X, DaskArray): if allow_dask: - return _check_gpu_X(X._meta) + return _check_gpu_X(X._meta, allow_csc=False) else: raise TypeError( "The input is a DaskArray. " @@ -272,8 +272,13 @@ def _check_gpu_X(X, require_cf=False, allow_dask=False): ) elif isinstance(X, cp.ndarray): return True - elif issparse(X): - if not require_cf: + elif isspmatrix_csc(X) or isspmatrix_csr(X): + if not allow_csc and isspmatrix_csc(X): + raise TypeError( + "Dask only supports CSR matrices and cp.ndarray. " + "Please convert your data to CSR format before passing it to this function." + ) + elif not require_cf: return True elif X.has_canonical_format: return True From 73662007d8d734bfb152df25cfc436e8dce472eb Mon Sep 17 00:00:00 2001 From: Intron7 Date: Thu, 14 Nov 2024 13:26:44 +0100 Subject: [PATCH 65/69] update docs --- docs/release-notes/index.md | 4 ++++ src/rapids_singlecell/preprocessing/_utils.py | 4 ++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/docs/release-notes/index.md b/docs/release-notes/index.md index faa775ff..65f29b12 100644 --- a/docs/release-notes/index.md +++ b/docs/release-notes/index.md @@ -2,6 +2,10 @@ # Release notes +## Version 0.11.0 +```{include} /release-notes/0.11.0.md +``` + ## Version 0.10.0 ```{include} /release-notes/0.10.11.md ``` diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 65b02fd3..f375d268 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -275,8 +275,8 @@ def _check_gpu_X(X, require_cf=False, allow_dask=False, allow_csc=True): elif isspmatrix_csc(X) or isspmatrix_csr(X): if not allow_csc and isspmatrix_csc(X): raise TypeError( - "Dask only supports CSR matrices and cp.ndarray. " - "Please convert your data to CSR format before passing it to this function." + "When using Dask, only CuPy ndarrays and CSR matrices are supported as " + "meta arrays. Please convert your data to CSR format if it is in CSC." ) elif not require_cf: return True From d1a6344de8bf308624ab3fd36229a728ff811f72 Mon Sep 17 00:00:00 2001 From: Intron7 Date: Fri, 15 Nov 2024 12:14:31 +0100 Subject: [PATCH 66/69] docs update --- docs/release-notes/index.md | 7 ------- 1 file changed, 7 deletions(-) diff --git a/docs/release-notes/index.md b/docs/release-notes/index.md index 65f29b12..9d05efa7 100644 --- a/docs/release-notes/index.md +++ b/docs/release-notes/index.md @@ -35,27 +35,20 @@ ## Version 0.9.0 ```{include} /release-notes/0.9.6.md ``` - ```{include} /release-notes/0.9.5.md ``` - ```{include} /release-notes/0.9.4.md ``` - ```{include} /release-notes/0.9.3.md ``` - ```{include} /release-notes/0.9.2.md ``` - ```{include} /release-notes/0.9.1.md ``` - ```{include} /release-notes/0.9.0.md ``` ## Version 0.8.0 - ```{include} /release-notes/0.8.1.md ``` ```{include} /release-notes/0.8.0.md From c65585dc176e54373b9cf533ffe4a94b0935a4fa Mon Sep 17 00:00:00 2001 From: Intron7 Date: Mon, 25 Nov 2024 11:49:12 +0100 Subject: [PATCH 67/69] make sure dtype is correct PCA --- .../preprocessing/_sparse_pca/_dask_sparse_pca.py | 4 +++- .../preprocessing/_sparse_pca/_sparse_pca.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py index 2a59cbfe..41001eb7 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_dask_sparse_pca.py @@ -57,7 +57,9 @@ def fit(self, x): def transform(self, X): def _transform(X_part, mean_, components_): pre_mean = mean_ @ components_.T - mean_impact = cp.ones((X_part.shape[0], 1)) @ pre_mean.reshape(1, -1) + mean_impact = cp.ones( + (X_part.shape[0], 1), dtype=X_part.dtype + ) @ pre_mean.reshape(1, -1) X_transformed = X_part.dot(components_.T) - mean_impact return X_transformed diff --git a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py index ac2b0030..2f9d5117 100644 --- a/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py +++ b/src/rapids_singlecell/preprocessing/_sparse_pca/_sparse_pca.py @@ -52,7 +52,9 @@ def fit(self, x): def transform(self, X): precomputed_mean_impact = self.mean_ @ self.components_.T - mean_impact = cp.ones((X.shape[0], 1)) @ precomputed_mean_impact.reshape(1, -1) + mean_impact = cp.ones( + (X.shape[0], 1), dtype=cp.float32 + ) @ precomputed_mean_impact.reshape(1, -1) X_transformed = X.dot(self.components_.T) - mean_impact # X = X - self.mean_ # X_transformed = X.dot(self.components_.T) From e7a11184be95c7eddaeb80f23fca79f1d6fc51e5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 26 Nov 2024 10:37:23 +0000 Subject: [PATCH 68/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- src/rapids_singlecell/preprocessing/_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/rapids_singlecell/preprocessing/_utils.py b/src/rapids_singlecell/preprocessing/_utils.py index 66b1e0cf..09e8eb43 100644 --- a/src/rapids_singlecell/preprocessing/_utils.py +++ b/src/rapids_singlecell/preprocessing/_utils.py @@ -4,9 +4,8 @@ from typing import TYPE_CHECKING, Literal import cupy as cp - -from cuml.internals.memory_utils import with_cupy_rmm import numpy as np +from cuml.internals.memory_utils import with_cupy_rmm from cupyx.scipy.sparse import issparse, isspmatrix_csc, isspmatrix_csr, spmatrix from rapids_singlecell._compat import DaskArray From 03e601a552046e5d1fd81857f8e59018556a21f5 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 6 Dec 2024 13:17:43 +0000 Subject: [PATCH 69/69] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/conf.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/docs/conf.py b/docs/conf.py index 732aa93e..a86fbc13 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -62,7 +62,16 @@ autosummary_generate = True autodoc_member_order = "bysource" -autodoc_mock_imports = ["cudf", "cuml", "cugraph", "cupy", "cupyx", "pylibraft", "dask","cuvs"] +autodoc_mock_imports = [ + "cudf", + "cuml", + "cugraph", + "cupy", + "cupyx", + "pylibraft", + "dask", + "cuvs", +] default_role = "literal" napoleon_google_docstring = False napoleon_numpy_docstring = True