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

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

sclab/preprocess/_subset.py

@@ -0,0 +1,208 @@
+from typing import Sequence
+
+import numpy as np
+import pandas as pd
+from anndata import AnnData
+
+
+def subset_obs(
+    adata: AnnData,
+    subset: pd.Index | Sequence[str | int | bool] | str,
+) -> None:
+    """Subset observations (rows) in an AnnData object.
+
+    This function modifies the AnnData object in-place by selecting a subset of observations
+    based on the provided subset parameter. The subsetting can be done using observation
+    names, integer indices, a boolean mask, a query string, or a pandas Index.
+
+    Parameters
+    ----------
+    adata : AnnData
+        The annotated data matrix to subset. Will be modified in-place.
+    subset : pd.Index | Sequence[str | int | bool] | str
+        The subset specification. Can be one of:
+        * A pandas Index containing observation names
+        * A sequence of observation names (strings)
+        * A sequence of integer indices
+        * A boolean mask of length `adata.n_obs`
+        * A query string to match observations by their metadata columns
+
+    Examples
+    --------
+    >>> # Create an example AnnData object
+    >>> import anndata
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>>
+    >>> obs = pd.DataFrame(
+    ...     index=['A', 'B', 'C'],
+    ...     data={'cell_type': ['type1', 'type2', 'type2']})
+    >>> adata_ = anndata.AnnData(obs=obs)
+    >>>
+    >>> # Subset using pandas Index
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, pd.Index(['B', 'C']))
+    >>> adata.obs_names.tolist()
+    ['B', 'C']
+    >>>
+    >>> # Subset using observation names
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, ['A', 'B'])
+    >>> adata.obs_names.tolist()
+    ['A', 'B']
+    >>>
+    >>> # Subset using integer indices
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, [0, 1])
+    >>> adata.obs_names.tolist()
+    ['A', 'B']
+    >>>
+    >>> # Subset using boolean mask
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, [True, False, True])
+    >>> adata.obs_names.tolist()
+    ['A', 'C']
+    >>>
+    >>> # Subset using query string
+    >>> adata = adata_.copy()
+    >>> subset_obs(adata, 'cell_type == "type2"')
+    >>> adata.obs_names.tolist()
+    ['B', 'C']
+
+    Notes
+    -----
+    - The function modifies the AnnData object in-place
+    - When using a boolean mask, its length must match the number of observations
+    - When using integer indices, they must be valid indices for the observations
+    - Invalid observation names or indices will raise KeyError or IndexError respectively
+    - The order of observations in the output will match the order in the subset parameter
+    """
+    if isinstance(subset, str):
+        subset = adata.obs.query(subset).index
+
+    if not isinstance(subset, pd.Index):
+        subset = np.asarray(subset)
+
+    # Handle boolean mask
+    if subset.dtype.kind == "b":
+        if len(subset) != adata.n_obs:
+            raise IndexError(
+                f"Boolean mask length ({len(subset)}) does not match number of "
+                f"observations ({adata.n_obs})"
+            )
+        subset = adata.obs_names[subset]
+
+    # Handle integer indices
+    elif subset.dtype.kind in "iu":
+        if np.any(subset < 0) or np.any(subset >= adata.n_obs):
+            raise IndexError(f"Integer indices must be between 0 and {adata.n_obs - 1}")
+        subset = adata.obs_names[subset]
+
+    if adata.n_obs == subset.size and (subset == adata.obs_names).all():
+        # No need to subset, avoid making a copy. Useful for large AnnData objects
+        return
+
+    adata._inplace_subset_obs(subset)
+
+
+def subset_var(
+    adata: AnnData,
+    subset: pd.Index | Sequence[str | int | bool] | str,
+) -> None:
+    """Subset variables (columns) in an AnnData object.
+
+    This function modifies the AnnData object in-place by selecting a subset of variables
+    based on the provided subset parameter. The subsetting can be done using variable
+    names, integer indices, a boolean mask, a query string, or a pandas Index.
+
+    Parameters
+    ----------
+    adata : AnnData
+        The annotated data matrix to subset. Will be modified in-place.
+    subset : pd.Index | Sequence[str | int | bool] | str
+        The subset specification. Can be one of:
+        * A pandas Index containing variable names
+        * A sequence of variable names (strings)
+        * A sequence of integer indices
+        * A boolean mask of length `adata.n_vars`
+        * A query string to match variables by their metadata columns
+
+    Examples
+    --------
+    >>> # Create an example AnnData object
+    >>> import anndata
+    >>> import pandas as pd
+    >>> import numpy as np
+    >>>
+    >>> var = pd.DataFrame(
+    ...     index=['gene1', 'gene2', 'gene3'],
+    ...     data={'gene_type': ['type1', 'type2', 'type1']})
+    >>> adata_ = anndata.AnnData(var=var)
+    >>>
+    >>> # Subset using pandas Index
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, pd.Index(['gene2', 'gene3']))
+    >>> adata.var_names.tolist()
+    ['gene2', 'gene3']
+    >>>
+    >>> # Subset using variable names
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, ['gene1', 'gene2'])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene2']
+    >>>
+    >>> # Subset using integer indices
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, [0, 1])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene2']
+    >>>
+    >>> # Subset using boolean mask
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, [True, False, True])
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene3']
+    >>>
+    >>> # Subset using query string
+    >>> adata = adata_.copy()
+    >>> subset_var(adata, 'gene_type == "type1"')
+    >>> adata.var_names.tolist()
+    ['gene1', 'gene3']
+
+    Notes
+    -----
+    - The function modifies the AnnData object in-place
+    - When using a boolean mask, its length must match the number of variables
+    - When using integer indices, they must be valid indices for the variables
+    - Invalid variable names or indices will raise KeyError or IndexError respectively
+    - The order of variables in the output will match the order in the subset parameter
+    """
+
+    if isinstance(subset, str):
+        subset = adata.var.query(subset).index
+
+    if not isinstance(subset, pd.Index):
+        subset = np.asarray(subset)
+
+    # Handle boolean mask
+    if subset.dtype.kind == "b":
+        if len(subset) != adata.n_vars:
+            raise IndexError(
+                f"Boolean mask length ({len(subset)}) does not match number of "
+                f"variables ({adata.n_vars})"
+            )
+        subset = adata.var_names[subset]
+
+    # Handle integer indices
+    elif subset.dtype.kind in "iu":
+        if np.any(subset < 0) or np.any(subset >= adata.n_vars):
+            raise IndexError(
+                f"Integer indices must be between 0 and {adata.n_vars - 1}"
+            )
+        subset = adata.var_names[subset]
+
+    if adata.n_vars == subset.size and (subset == adata.var_names).all():
+        # No need to subset, avoid making a copy. Useful for large AnnData objects
+        return
+
+    adata._inplace_subset_var(subset)