sclab 0.3.2__py3-none-any.whl → 0.3.4__py3-none-any.whl

sclab/__init__.py

@@ -6,4 +6,4 @@ __all__ = [
     "SCLabDashboard",
 ]
 
-__version__ = "0.3.2"
+__version__ = "0.3.4"

sclab/examples/processor_steps/_integration.py

@@ -63,7 +63,7 @@ class Integration(ProcessorStepBase):
 
     def function(
         self,
-        use_rep: str,
+        use_rep: str | None,
         group_by: str,
         flavor: str,
         reference_batch: str | None,
@@ -71,6 +71,9 @@ class Integration(ProcessorStepBase):
     ):
         adata = self.parent.dataset.adata
 
+        if use_rep is None:
+            use_rep = "X"
+
         key_added = f"{use_rep}_{flavor}"
         kvargs = {
             "adata": adata,

sclab/examples/processor_steps/_preprocess.py

@@ -106,25 +106,41 @@ class Preprocess(ProcessorStepBase):
         )
         pbar.update(10)
 
-        sc.pp.highly_variable_genes(
-            adata,
-            layer=f"{layer}_log1p",
-            flavor="seurat",
-            batch_key=group_by,
-        )
-        hvg_seurat = adata.var["highly_variable"]
-        sc.pp.highly_variable_genes(
-            adata,
-            layer=layer,
-            flavor="seurat_v3_paper",
-            batch_key=group_by,
-            n_top_genes=hvg_seurat.sum(),
-        )
-        hvg_seurat_v3 = adata.var["highly_variable"]
+        if group_by is not None:
+            adata.var["highly_variable"] = False
+            for name, idx in adata.obs.groupby(group_by, observed=True).groups.items():
+                hvg_seurat = sc.pp.highly_variable_genes(
+                    adata[idx],
+                    layer=f"{layer}_log1p",
+                    flavor="seurat",
+                    inplace=False,
+                )["highly_variable"]
+
+                hvg_seurat_v3 = sc.pp.highly_variable_genes(
+                    adata[idx],
+                    layer=layer,
+                    flavor="seurat_v3_paper",
+                    n_top_genes=hvg_seurat.sum(),
+                    inplace=False,
+                )["highly_variable"]
+
+                adata.var[f"highly_variable_{name}"] = hvg_seurat | hvg_seurat_v3
+                adata.var["highly_variable"] |= adata.var[f"highly_variable_{name}"]
+
+        else:
+            sc.pp.highly_variable_genes(adata, layer=f"{layer}_log1p", flavor="seurat")
+            hvg_seurat = adata.var["highly_variable"]
+
+            sc.pp.highly_variable_genes(
+                adata,
+                layer=layer,
+                flavor="seurat_v3_paper",
+                n_top_genes=hvg_seurat.sum(),
+            )
+            hvg_seurat_v3 = adata.var["highly_variable"]
+
+            adata.var["highly_variable"] = hvg_seurat | hvg_seurat_v3
 
-        adata.var["highly_variable"] = hvg_seurat | hvg_seurat_v3
-        adata.var["highly_variable_seurat"] = hvg_seurat
-        adata.var["highly_variable_seurat_v3"] = hvg_seurat_v3
         pbar.update(10)
         pbar.update(10)
 

sclab/preprocess/__init__.py

@@ -2,8 +2,11 @@ from ._cca_integrate import cca_integrate, cca_integrate_pair
 from ._filter_obs import filter_obs
 from ._harmony_integrate import harmony_integrate
 from ._normalize_weighted import normalize_weighted
+from ._pca import pca
+from ._preprocess import preprocess
+from ._qc import qc
 from ._subset import subset_obs, subset_var
-from ._transfer_metadata import transfer_metadata
+from ._transfer_metadata import propagate_metadata, transfer_metadata
 from ._transform import pool_neighbors
 
 __all__ = [
@@ -12,7 +15,11 @@ __all__ = [
     "filter_obs",
     "harmony_integrate",
     "normalize_weighted",
+    "pca",
     "pool_neighbors",
+    "preprocess",
+    "propagate_metadata",
+    "qc",
     "subset_obs",
     "subset_var",
     "transfer_metadata",

sclab/preprocess/_cca.py

@@ -1,24 +1,31 @@
 import logging
+import os
 from typing import Literal
 
 import numpy as np
+from joblib import Parallel, delayed
 from numpy import matrix
 from numpy.typing import NDArray
 from scipy.linalg import svd
 from scipy.sparse import csc_matrix, csr_matrix, issparse
+from scipy.sparse import vstack as sparse_vstack
 from scipy.sparse.linalg import svds
 from sklearn.utils.extmath import randomized_svd
 
 logger = logging.getLogger(__name__)
 
 
+N_CPUS = os.cpu_count()
+
+
 def cca(
     X: NDArray | csr_matrix | csc_matrix,
     Y: NDArray | csr_matrix | csc_matrix,
     n_components=None,
-    svd_solver: Literal["full", "partial", "randomized"] = "partial",
+    svd_solver: Literal["full", "partial", "randomized"] = "randomized",
     normalize: bool = False,
     random_state=42,
+    n_jobs: int = N_CPUS,
 ) -> tuple[NDArray, NDArray, NDArray]:
     """
     CCA-style integration for two single-cell matrices with unequal numbers of cells.
@@ -50,7 +57,7 @@ def cca(
     k = n_components or min(n1, n2)
 
     if issparse(X):
-        C = _cross_covariance_sparse(X, Y)
+        C = _cross_covariance_sparse(X, Y, n_jobs=n_jobs)
     else:
         C = _cross_covariance_dense(X, Y)
 
@@ -103,7 +110,7 @@ def _svd_decomposition(
     return Uc, s, Vct
 
 
-def _cross_covariance_sparse(X: csr_matrix, Y: csr_matrix) -> NDArray:
+def _cross_covariance_sparse(X: csr_matrix, Y: csr_matrix, n_jobs=N_CPUS) -> NDArray:
     _, p1 = X.shape
     _, p2 = Y.shape
     if p1 != p2:
@@ -118,7 +125,7 @@ def _cross_covariance_sparse(X: csr_matrix, Y: csr_matrix) -> NDArray:
     mux: matrix = X.mean(axis=0)
     muy: matrix = Y.mean(axis=0)
 
-    XYt: csr_matrix = X.dot(Y.T)
+    XYt: csr_matrix = _spmm_parallel(X, Y.T, n_jobs=n_jobs)
     Xmuyt: matrix = X.dot(muy.T)
     muxYt: matrix = Y.dot(mux.T).T
     muxmuyt: float = (mux @ muy.T)[0, 0]
@@ -152,3 +159,18 @@ def _dense_scale(A: NDArray) -> NDArray:
     A = np.asarray(A)
     eps = np.finfo(A.dtype).eps
     return A / (A.std(axis=0, ddof=1, keepdims=True) + eps)
+
+
+def _spmm_chunk(A_csr, X, start, stop):
+    return A_csr[start:stop, :] @ X
+
+
+def _spmm_parallel(A_csr: csr_matrix, X_csc: csc_matrix, n_jobs=N_CPUS):
+    n = A_csr.shape[0]
+
+    bounds = np.linspace(0, n, n_jobs + 1, dtype=int)
+    Ys = Parallel(n_jobs=n_jobs, prefer="processes")(
+        delayed(_spmm_chunk)(A_csr, X_csc, bounds[i], bounds[i + 1])
+        for i in range(n_jobs)
+    )
+    return sparse_vstack(Ys)  # result is sparse if X is sparse, dense otherwise

sclab/preprocess/_cca_integrate.py

@@ -13,8 +13,8 @@ def cca_integrate(
     reference_batch: str | list[str] | None = None,
     mask_var: str | None = None,
     n_components: int = 30,
-    svd_solver: str = "partial",
-    normalize: bool = False,
+    svd_solver: str = "randomized",
+    normalize: bool = True,
     random_state: int | None = None,
 ):
     n_groups = adata.obs[key].nunique()
@@ -46,8 +46,8 @@ def cca_integrate_pair(
     adjusted_basis: str | None = None,
     mask_var: str | None = None,
     n_components: int = 30,
-    svd_solver: str = "partial",
-    normalize: bool = False,
+    svd_solver: str = "randomized",
+    normalize: bool = True,
     random_state: int | None = None,
 ):
     if basis is None:

sclab/preprocess/_normalize_weighted.py

@@ -9,6 +9,7 @@ def normalize_weighted(
     adata: AnnData,
     target_scale: float | None = None,
     batch_key: str | None = None,
+    q: float = 0.99,
 ) -> None:
     if batch_key is not None:
         for _, idx in adata.obs.groupby(batch_key, observed=True).groups.items():
@@ -22,6 +23,8 @@ def normalize_weighted(
 
         return
 
+    target_scale = None
+
     X: csr_matrix
     Y: csr_matrix
     Z: csr_matrix
@@ -38,6 +41,7 @@ def normalize_weighted(
     Y.eliminate_zeros()
     Y.data = -Y.data * np.log(Y.data)
     entropy = Y.sum(axis=0)
+    entropy[:, entropy.A1 < np.quantile(entropy.A1, q)] *= 0.0
 
     Z = X.multiply(entropy)
     Z = Z.tocsr()
@@ -48,7 +52,7 @@ def normalize_weighted(
             "ignore", category=RuntimeWarning, message="divide by zero"
         )
         scale = Z.sum(axis=1)
-        Z = Z.multiply(1 / scale)
+        Z = X.multiply(1 / scale)
     Z = Z.tocsr()
 
     if target_scale is None:

sclab/preprocess/_pca.py

@@ -0,0 +1,51 @@
+from anndata import AnnData
+
+
+def pca(
+    adata: AnnData,
+    layer: str | None = None,
+    n_comps: int = 30,
+    mask_var: str | None = None,
+    batch_key: str | None = None,
+    reference_batch: str | None = None,
+    zero_center: bool = False,
+):
+    import scanpy as sc
+
+    pca_kwargs = dict(
+        n_comps=n_comps,
+        layer=layer,
+        mask_var=mask_var,
+        svd_solver="arpack",
+    )
+
+    if reference_batch:
+        obs_mask = adata.obs[batch_key] == reference_batch
+        adata_ref = adata[obs_mask].copy()
+        if mask_var == "highly_variable":
+            sc.pp.highly_variable_genes(
+                adata_ref, layer=f"{layer if layer else 'X'}_log1p", flavor="seurat"
+            )
+            hvg_seurat = adata_ref.var["highly_variable"]
+            sc.pp.highly_variable_genes(
+                adata_ref,
+                layer=layer,
+                flavor="seurat_v3_paper",
+                n_top_genes=hvg_seurat.sum(),
+            )
+            hvg_seurat_v3 = adata_ref.var["highly_variable"]
+            adata_ref.var["highly_variable"] = hvg_seurat | hvg_seurat_v3
+
+        sc.pp.pca(adata_ref, **pca_kwargs)
+        uns_pca = adata_ref.uns["pca"]
+        uns_pca["reference_batch"] = reference_batch
+        PCs = adata_ref.varm["PCs"]
+        adata.obsm["X_pca"] = adata.X.dot(PCs)
+        adata.uns["pca"] = uns_pca
+        adata.varm["PCs"] = PCs
+    else:
+        sc.pp.pca(adata, **pca_kwargs)
+        adata.obsm["X_pca"] = adata.X.dot(adata.varm["PCs"])
+
+    if zero_center:
+        adata.obsm["X_pca"] -= adata.obsm["X_pca"].mean(axis=0, keepdims=True)

sclab/preprocess/_preprocess.py

@@ -0,0 +1,155 @@
+import warnings
+from typing import Literal
+
+import numpy as np
+from anndata import AnnData, ImplicitModificationWarning
+from tqdm.auto import tqdm
+
+
+def preprocess(
+    adata: AnnData,
+    counts_layer: str = "counts",
+    group_by: str | None = None,
+    min_cells: int = 5,
+    min_genes: int = 5,
+    compute_hvg: bool = True,
+    regress_total_counts: bool = False,
+    regress_n_genes: bool = False,
+    normalization_method: Literal["library", "weighted", "none"] = "library",
+    target_scale: float = 1e4,
+    weighted_norm_quantile: float = 0.9,
+    log1p: bool = True,
+    scale: bool = True,
+):
+    import scanpy as sc
+
+    from ._normalize_weighted import normalize_weighted
+
+    with tqdm(total=100, bar_format="{percentage:3.0f}%|{bar}|") as pbar:
+        if counts_layer not in adata.layers:
+            adata.layers[counts_layer] = adata.X.copy()
+
+        if f"{counts_layer}_log1p" not in adata.layers:
+            adata.layers[f"{counts_layer}_log1p"] = sc.pp.log1p(
+                adata.layers[counts_layer].copy()
+            )
+        pbar.update(10)
+
+        adata.X = adata.layers[counts_layer].copy()
+        sc.pp.calculate_qc_metrics(
+            adata,
+            percent_top=None,
+            log1p=False,
+            inplace=True,
+        )
+        sc.pp.filter_cells(adata, min_genes=min_genes)
+        sc.pp.filter_genes(adata, min_cells=min_cells)
+        pbar.update(10)
+
+        sc.pp.calculate_qc_metrics(
+            adata,
+            percent_top=None,
+            log1p=False,
+            inplace=True,
+        )
+        pbar.update(10)
+
+        if compute_hvg:
+            if group_by is not None:
+                adata.var["highly_variable"] = False
+                for name, idx in adata.obs.groupby(
+                    group_by, observed=True
+                ).groups.items():
+                    hvg_seurat = sc.pp.highly_variable_genes(
+                        adata[idx],
+                        layer=f"{counts_layer}_log1p",
+                        flavor="seurat",
+                        inplace=False,
+                    )["highly_variable"]
+
+                    hvg_seurat_v3 = sc.pp.highly_variable_genes(
+                        adata[idx],
+                        layer=counts_layer,
+                        flavor="seurat_v3_paper",
+                        n_top_genes=hvg_seurat.sum(),
+                        inplace=False,
+                    )["highly_variable"]
+
+                    adata.var[f"highly_variable_{name}"] = hvg_seurat | hvg_seurat_v3
+                    adata.var["highly_variable"] |= adata.var[f"highly_variable_{name}"]
+
+            else:
+                sc.pp.highly_variable_genes(
+                    adata, layer=f"{counts_layer}_log1p", flavor="seurat"
+                )
+                hvg_seurat = adata.var["highly_variable"]
+
+                sc.pp.highly_variable_genes(
+                    adata,
+                    layer=counts_layer,
+                    flavor="seurat_v3_paper",
+                    n_top_genes=hvg_seurat.sum(),
+                )
+                hvg_seurat_v3 = adata.var["highly_variable"]
+
+                adata.var["highly_variable"] = hvg_seurat | hvg_seurat_v3
+
+        pbar.update(10)
+        pbar.update(10)
+
+        new_layer = counts_layer
+        if normalization_method == "library":
+            new_layer += "_normt"
+            sc.pp.normalize_total(adata, target_sum=target_scale)
+        elif normalization_method == "weighted":
+            new_layer += "_normw"
+            normalize_weighted(
+                adata,
+                target_scale=target_scale,
+                batch_key=group_by,
+                q=weighted_norm_quantile,
+            )
+
+        pbar.update(10)
+        pbar.update(10)
+
+        if log1p:
+            new_layer += "_log1p"
+            adata.uns.pop("log1p", None)
+            sc.pp.log1p(adata)
+        pbar.update(10)
+
+        vars_to_regress = []
+        if regress_n_genes:
+            vars_to_regress.append("n_genes_by_counts")
+
+        if regress_total_counts and log1p:
+            adata.obs["log1p_total_counts"] = np.log1p(adata.obs["total_counts"])
+            vars_to_regress.append("log1p_total_counts")
+        elif regress_total_counts:
+            vars_to_regress.append("total_counts")
+
+        if vars_to_regress:
+            new_layer += "_regr"
+            sc.pp.regress_out(adata, keys=vars_to_regress, n_jobs=1)
+        pbar.update(10)
+
+        if scale:
+            new_layer += "_scale"
+            if group_by is not None:
+                for _, idx in adata.obs.groupby(group_by, observed=True).groups.items():
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings(
+                            "ignore",
+                            category=ImplicitModificationWarning,
+                            message="Modifying `X` on a view results in data being overridden",
+                        )
+                        adata[idx].X = sc.pp.scale(adata[idx].X, zero_center=False)
+            else:
+                sc.pp.scale(adata, zero_center=False)
+
+        adata.layers[new_layer] = adata.X.copy()
+
+        pbar.update(10)
+
+        adata.X = adata.X.astype(np.float32)

sclab/preprocess/_qc.py

@@ -0,0 +1,38 @@
+import numpy as np
+from anndata import AnnData
+
+
+def qc(
+    adata: AnnData,
+    counts_layer: str = "counts",
+    min_counts: int = 50,
+    min_genes: int = 5,
+    min_cells: int = 5,
+    max_rank: int = 0,
+):
+    import scanpy as sc
+
+    if counts_layer not in adata.layers:
+        adata.layers[counts_layer] = adata.X.copy()
+
+    adata.layers["qc_tmp_current_X"] = adata.X
+    adata.X = adata.layers[counts_layer].copy()
+    rowsums = np.asarray(adata.X.sum(axis=1)).squeeze()
+
+    obs_idx = adata.obs_names[rowsums >= min_counts]
+    adata._inplace_subset_obs(obs_idx)
+
+    sc.pp.calculate_qc_metrics(adata, percent_top=None, log1p=False, inplace=True)
+
+    sc.pp.filter_cells(adata, min_genes=min_genes)
+    sc.pp.filter_genes(adata, min_cells=min_cells)
+    sc.pp.calculate_qc_metrics(adata, percent_top=None, log1p=False, inplace=True)
+    adata.obs["barcode_rank"] = adata.obs["total_counts"].rank(ascending=False)
+
+    # Restore original X
+    adata.X = adata.layers.pop("qc_tmp_current_X")
+
+    if max_rank > 0:
+        series = adata.obs["barcode_rank"]
+        index = series.loc[series < max_rank].index
+        adata._inplace_subset_obs(index)

sclab/preprocess/_rpca.py

@@ -0,0 +1,116 @@
+import numpy as np
+from anndata import AnnData
+from numpy.typing import NDArray
+
+
+def rpca(
+    adata: AnnData,
+    key: str,
+    *,
+    basis: str = "X",
+    adjusted_basis: str | None = None,
+    reference_batch: str | list[str] | None = None,
+    mask_var: str | None = None,
+    n_components: int = 30,
+    min_variance_ratio: float = 0.0005,
+    svd_solver: str = "arpack",
+    normalize: bool = True,
+):
+    if basis is None:
+        basis = "X"
+
+    if adjusted_basis is None:
+        adjusted_basis = basis + "_rpca"
+
+    if mask_var is not None:
+        mask = adata.var[mask_var].values
+    else:
+        mask = np.ones(adata.n_vars, dtype=bool)
+
+    X = _get_basis(adata[:, mask], basis)
+    uns = {}
+
+    groups = adata.obs.groupby(key, observed=True).groups
+    if reference_batch is None:
+        reference_batch = list(groups.keys())
+    elif isinstance(reference_batch, str):
+        reference_batch = [reference_batch]
+
+    for gr, idx in groups.items():
+        if gr not in reference_batch:
+            continue
+
+        ref_basis_key = f"{adjusted_basis}_{gr}"
+        ref_PCs_key = f"{adjusted_basis}_{gr}_PCs"
+
+        X_reference = _get_basis(adata[idx, mask], basis)
+        proj_result = pca_projection(
+            X,
+            X_reference,
+            n_components=n_components,
+            min_variance_ratio=min_variance_ratio,
+            svd_solver=svd_solver,
+            normalize=normalize,
+        )
+        res_ncomps = proj_result[0].shape[1]
+        components = np.zeros((res_ncomps, adata.n_vars))
+        components[:, mask] = proj_result[1]
+
+        adata.obsm[ref_basis_key] = proj_result[0]
+        adata.varm[ref_PCs_key] = components.T
+
+        uns[gr] = {
+            "n_components": res_ncomps,
+            "explained_variance_ratio": proj_result[2],
+            "explained_variance": proj_result[3],
+        }
+
+    adata.uns[adjusted_basis] = uns
+
+
+def pca_projection(
+    X: NDArray,
+    X_reference: NDArray,
+    n_components: int = 30,
+    min_variance_ratio: float = 0.0005,
+    svd_solver: str = "arpack",
+    normalize: bool = False,
+) -> tuple[NDArray, NDArray, NDArray, NDArray]:
+    import scanpy as sc
+
+    pca_kwargs = dict(
+        n_comps=n_components,
+        svd_solver=svd_solver,
+        return_info=True,
+    )
+
+    pca_result = sc.pp.pca(X_reference, **pca_kwargs)
+    _, components, explained_variance_ratio, explained_variance = pca_result
+
+    components_mask = explained_variance_ratio > min_variance_ratio
+    components = components[components_mask]
+    explained_variance_ratio = explained_variance_ratio[components_mask]
+    explained_variance = explained_variance[components_mask]
+
+    X_pca = X.dot(components.T)
+
+    if normalize:
+        X_pca = X_pca / np.linalg.norm(X_pca, axis=1, keepdims=True)
+
+    return X_pca, components, explained_variance_ratio, explained_variance
+
+
+def _get_basis(adata: AnnData, basis: str):
+    if basis == "X":
+        X = adata.X
+
+    elif basis in adata.layers:
+        X = adata.layers[basis]
+
+    elif basis in adata.obsm:
+        X = adata.obsm[basis]
+
+    else:
+        raise ValueError(f"Unknown basis {basis}")
+
+    return X