sclab 0.2.4__py3-none-any.whl → 0.3.0__py3-none-any.whl

sclab/preprocess/_cca.py

@@ -0,0 +1,154 @@
+import logging
+from typing import Literal
+
+import numpy as np
+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.linalg import svds
+from sklearn.utils.extmath import randomized_svd
+
+logger = logging.getLogger(__name__)
+
+
+def cca(
+    X: NDArray | csr_matrix | csc_matrix,
+    Y: NDArray | csr_matrix | csc_matrix,
+    n_components=None,
+    svd_solver: Literal["full", "partial", "randomized"] = "partial",
+    normalize: bool = False,
+    random_state=42,
+):
+    """
+    CCA-style integration for two single-cell matrices with unequal numbers of cells.
+
+    Parameters
+    ----------
+    X, Y : array-like, shape (n_cells, n_features)
+        feature-by-cell matrices with same column space (variable genes/pcs) in the same order.
+    n_components : int or None
+        Dimensionality of the canonical space (default = all that the smaller
+        dataset allows).
+    svd_solver : {'full', 'partial', 'randomized'}
+        'randomized' uses Halko et al. algorithm (`sklearn.utils.extmath.randomized_svd`)
+        and is strongly recommended when only the leading few components are needed.
+    random_state : int or None
+        Passed through to the randomized SVD for reproducibility.
+
+    Returns
+    -------
+    U : (n_cells(X), k) ndarray
+    V : (n_cells(Y), k) ndarray
+        Cell-level canonical variables.
+    """
+    n1, p1 = X.shape
+    n2, p2 = Y.shape
+    if p1 != p2:
+        raise ValueError("The two matrices must have the same number of features.")
+
+    k = n_components or min(n1, n2)
+
+    if issparse(X):
+        C = _cross_covariance_sparse(X, Y)
+    else:
+        C = _cross_covariance_dense(X, Y)
+
+    logger.info(f"Cross-covariance computed. Shape: {C.shape}")
+
+    Uc, s, Vct = _svd_decomposition(C, k, svd_solver, random_state)
+
+    # canonical variables
+    # Left and right singular vectors are cell embeddings
+    U = Uc  # (n1 x k)
+    V = Vct.T  # (n2 x k)
+
+    if normalize:
+        logger.info("Normalizing canonical variables...")
+        U = U / np.linalg.norm(U, axis=1, keepdims=True)
+        V = V / np.linalg.norm(V, axis=1, keepdims=True)
+
+    logger.info("Done.")
+
+    return U, s, V
+
+
+def _svd_decomposition(
+    C: NDArray,
+    k: int,
+    svd_solver: Literal["full", "partial", "randomized"],
+    random_state: int | None,
+) -> tuple[NDArray, NDArray, NDArray]:
+    if svd_solver == "full":
+        logger.info("SVD decomposition with full SVD...")
+        Uc, s, Vct = svd(C, full_matrices=False)
+        Uc, s, Vct = Uc[:, :k], s[:k], Vct[:k, :]
+
+    elif svd_solver == "partial":
+        logger.info("SVD decomposition with partial SVD...")
+        Uc, s, Vct = svds(C, k=k)
+
+    elif svd_solver == "randomized":
+        logger.info("SVD decomposition with randomized SVD...")
+        Uc, s, Vct = randomized_svd(C, n_components=k, random_state=random_state)
+
+    else:
+        raise ValueError("svd_solver must be 'full' or 'partial'.")
+
+    order = np.argsort(-s)
+    s = s[order]
+    Uc = Uc[:, order]
+    Vct = Vct[order, :]
+
+    return Uc, s, Vct
+
+
+def _cross_covariance_sparse(X: csr_matrix, Y: csr_matrix) -> NDArray:
+    _, p1 = X.shape
+    _, p2 = Y.shape
+    if p1 != p2:
+        raise ValueError("The two matrices must have the same number of features.")
+
+    p = p1
+
+    # TODO: incorporate sparse scaling
+
+    logger.info("Computing cross-covariance on sparse matrices...")
+
+    mux: matrix = X.mean(axis=0)
+    muy: matrix = Y.mean(axis=0)
+
+    XYt: csr_matrix = X.dot(Y.T)
+    Xmuyt: matrix = X.dot(muy.T)
+    muxYt: matrix = Y.dot(mux.T).T
+    muxmuyt: float = (mux @ muy.T)[0, 0]
+
+    C = (XYt - Xmuyt - muxYt + muxmuyt) / (p - 1)
+
+    return np.asarray(C)
+
+
+def _cross_covariance_dense(X: NDArray, Y: NDArray) -> NDArray:
+    _, p1 = X.shape
+    _, p2 = Y.shape
+    if p1 != p2:
+        raise ValueError("The two matrices must have the same number of features.")
+
+    p = p1
+
+    logger.info("Computing cross-covariance on dense matrices...")
+    X = _dense_scale(X)
+    Y = _dense_scale(Y)
+
+    X = X - X.mean(axis=0, keepdims=True)
+    Y = Y - Y.mean(axis=0, keepdims=True)
+
+    C: NDArray = (X @ Y.T) / (p - 1)
+
+    return C
+
+
+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)

sclab/preprocess/_cca_integrate.py

@@ -0,0 +1,77 @@
+import numpy as np
+from anndata import AnnData
+
+from ._cca import cca
+
+
+def cca_integrate_pair(
+    adata: AnnData,
+    key: str,
+    group1: str,
+    group2: str,
+    *,
+    basis: str | None = None,
+    adjusted_basis: str | None = None,
+    mask_var: str | None = None,
+    n_components: int = 50,
+    svd_solver: str = "partial",
+    normalize: bool = False,
+    random_state: int | None = None,
+):
+    if basis is None:
+        basis = "X"
+
+    if adjusted_basis is None:
+        adjusted_basis = basis + "_cca"
+
+    if mask_var is not None:
+        mask = adata.var[mask_var].values
+    else:
+        mask = np.ones(adata.n_vars, dtype=bool)
+
+    Xs = {}
+    groups = adata.obs.groupby(key, observed=True).groups
+    for gr, idx in groups.items():
+        Xs[gr] = _get_basis(adata[idx, mask], basis)
+
+    Ys = {}
+    Ys[group1], sigma, Ys[group2] = cca(
+        Xs[group1],
+        Xs[group2],
+        n_components=n_components,
+        svd_solver=svd_solver,
+        normalize=normalize,
+        random_state=random_state,
+    )
+
+    if (
+        adjusted_basis not in adata.obsm
+        or adata.obsm[adjusted_basis].shape[1] != n_components
+    ):
+        adata.obsm[adjusted_basis] = np.full((adata.n_obs, n_components), np.nan)
+
+    if adjusted_basis not in adata.uns:
+        adata.uns[adjusted_basis] = {}
+
+    uns = adata.uns[adjusted_basis]
+    uns[f"{group1}-{group2}"] = {"sigma": sigma}
+    for gr, obs_names in groups.items():
+        idx = adata.obs_names.get_indexer(obs_names)
+        adata.obsm[adjusted_basis][idx] = Ys[gr]
+        uns[gr] = Ys[gr]
+
+
+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/_filter_obs.py

@@ -0,0 +1,42 @@
+import numpy as np
+from anndata import AnnData
+from scipy.stats import rankdata
+
+
+def filter_obs(
+    adata: AnnData,
+    *,
+    layer: str | None = None,
+    min_counts: int | None = None,
+    min_genes: int | None = None,
+    max_counts: int | None = None,
+    max_cells: int | None = None,
+) -> None:
+    if layer is not None:
+        X = adata.layers[layer]
+    else:
+        X = adata.X
+
+    remove_mask = np.zeros(X.shape[0], dtype=bool)
+
+    if min_genes is not None:
+        M = X > 0
+        rowsums = np.asarray(M.sum(axis=1)).squeeze()
+        remove_mask[rowsums < min_genes] = True
+
+    if min_counts is not None or max_counts is not None or max_cells is not None:
+        rowsums = np.asarray(X.sum(axis=1)).squeeze()
+
+        if min_counts is not None:
+            remove_mask[rowsums < min_counts] = True
+
+        if max_counts is not None:
+            remove_mask[rowsums > max_counts] = True
+
+        if max_cells is not None:
+            ranks = rankdata(-rowsums, method="min")
+            remove_mask[ranks > max_cells] = True
+
+    if remove_mask.any():
+        obs_idx = adata.obs_names[~remove_mask]
+        adata._inplace_subset_obs(obs_idx)

sclab/preprocess/_harmony.py

@@ -0,0 +1,421 @@
+# harmonypy - A data alignment algorithm.
+# Copyright (C) 2018  Ilya Korsunsky
+#               2019  Kamil Slowikowski <kslowikowski@gmail.com>
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from functools import partial
+import pandas as pd
+import numpy as np
+from sklearn.cluster import KMeans
+import logging
+
+# create logger
+logger = logging.getLogger("harmonypy")
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+# from IPython.core.debugger import set_trace
+
+
+def run_harmony(
+    data_mat: np.ndarray,
+    meta_data: pd.DataFrame,
+    vars_use,
+    theta=None,
+    lamb=None,
+    sigma=0.1,
+    nclust=None,
+    tau=0,
+    block_size=0.05,
+    max_iter_harmony=10,
+    max_iter_kmeans=20,
+    epsilon_cluster=1e-5,
+    epsilon_harmony=1e-4,
+    plot_convergence=False,
+    verbose=True,
+    reference_values=None,
+    cluster_prior=None,
+    random_state=0,
+    cluster_fn="kmeans",
+):
+    """Run Harmony."""
+
+    # theta = None
+    # lamb = None
+    # sigma = 0.1
+    # nclust = None
+    # tau = 0
+    # block_size = 0.05
+    # epsilon_cluster = 1e-5
+    # epsilon_harmony = 1e-4
+    # plot_convergence = False
+    # verbose = True
+    # reference_values = None
+    # cluster_prior = None
+    # random_state = 0
+    # cluster_fn = 'kmeans'. Also accepts a callable object with data, num_clusters parameters
+
+    N = meta_data.shape[0]
+    if data_mat.shape[1] != N:
+        data_mat = data_mat.T
+
+    assert data_mat.shape[1] == N, (
+        "data_mat and meta_data do not have the same number of cells"
+    )
+
+    if nclust is None:
+        nclust = np.min([np.round(N / 30.0), 100]).astype(int)
+
+    if type(sigma) is float and nclust > 1:
+        sigma = np.repeat(sigma, nclust)
+
+    if isinstance(vars_use, str):
+        vars_use = [vars_use]
+
+    phi = pd.get_dummies(meta_data[vars_use]).to_numpy().T
+    phi_n = meta_data[vars_use].describe().loc["unique"].to_numpy().astype(int)
+
+    if theta is None:
+        theta = np.repeat([1] * len(phi_n), phi_n)
+    elif isinstance(theta, float) or isinstance(theta, int):
+        theta = np.repeat([theta] * len(phi_n), phi_n)
+    elif len(theta) == len(phi_n):
+        theta = np.repeat([theta], phi_n)
+
+    assert len(theta) == np.sum(phi_n), "each batch variable must have a theta"
+
+    if lamb is None:
+        lamb = np.repeat([1] * len(phi_n), phi_n)
+    elif isinstance(lamb, float) or isinstance(lamb, int):
+        lamb = np.repeat([lamb] * len(phi_n), phi_n)
+    elif len(lamb) == len(phi_n):
+        lamb = np.repeat([lamb], phi_n)
+
+    assert len(lamb) == np.sum(phi_n), "each batch variable must have a lambda"
+
+    # Number of items in each category.
+    N_b = phi.sum(axis=1)
+    # Proportion of items in each category.
+    Pr_b = N_b / N
+
+    if tau > 0:
+        theta = theta * (1 - np.exp(-((N_b / (nclust * tau)) ** 2)))
+
+    lamb_mat = np.diag(np.insert(lamb, 0, 0))
+
+    phi_moe = np.vstack((np.repeat(1, N), phi))
+
+    np.random.seed(random_state)
+
+    ho = Harmony(
+        data_mat,
+        phi,
+        phi_moe,
+        Pr_b,
+        sigma,
+        theta,
+        max_iter_harmony,
+        max_iter_kmeans,
+        epsilon_cluster,
+        epsilon_harmony,
+        nclust,
+        block_size,
+        lamb_mat,
+        verbose,
+        random_state,
+        cluster_fn,
+        reference_values,
+    )
+
+    return ho
+
+
+class Harmony(object):
+    def __init__(
+        self,
+        Z,
+        Phi,
+        Phi_moe,
+        Pr_b,
+        sigma,
+        theta,
+        max_iter_harmony,
+        max_iter_kmeans,
+        epsilon_kmeans,
+        epsilon_harmony,
+        K,
+        block_size,
+        lamb,
+        verbose,
+        random_state=None,
+        cluster_fn="kmeans",
+        frozen_values=None,
+    ):
+        self.Z_corr = np.array(Z)
+        self.Z_orig = np.array(Z)
+
+        self.Z_cos = self.Z_orig / self.Z_orig.max(axis=0)
+        self.Z_cos = self.Z_cos / np.linalg.norm(self.Z_cos, ord=2, axis=0)
+
+        self.Phi = Phi
+        self.Phi_moe = Phi_moe
+        self.N = self.Z_corr.shape[1]
+        self.Pr_b = Pr_b
+        self.B = self.Phi.shape[0]  # number of batch variables
+        self.d = self.Z_corr.shape[0]
+        self.window_size = 3
+        self.epsilon_kmeans = epsilon_kmeans
+        self.epsilon_harmony = epsilon_harmony
+        self.reference_values = frozen_values
+
+        self.lamb = lamb
+        self.sigma = sigma
+        self.sigma_prior = sigma
+        self.block_size = block_size
+        self.K = K  # number of clusters
+        self.max_iter_harmony = max_iter_harmony
+        self.max_iter_kmeans = max_iter_kmeans
+        self.verbose = verbose
+        self.theta = theta
+
+        self.objective_harmony = []
+        self.objective_kmeans = []
+        self.objective_kmeans_dist = []
+        self.objective_kmeans_entropy = []
+        self.objective_kmeans_cross = []
+        self.kmeans_rounds = []
+
+        self.allocate_buffers()
+        if cluster_fn == "kmeans":
+            cluster_fn = partial(Harmony._cluster_kmeans, random_state=random_state)
+        self.init_cluster(cluster_fn)
+        self.harmonize(self.max_iter_harmony, self.verbose)
+
+    def result(self):
+        return self.Z_corr
+
+    def allocate_buffers(self):
+        self._scale_dist = np.zeros((self.K, self.N))
+        self.dist_mat = np.zeros((self.K, self.N))
+        self.O = np.zeros((self.K, self.B))
+        self.E = np.zeros((self.K, self.B))
+        self.W = np.zeros((self.B + 1, self.d))
+        self.Phi_Rk = np.zeros((self.B + 1, self.N))
+
+    @staticmethod
+    def _cluster_kmeans(data, K, random_state):
+        # Start with cluster centroids
+        logger.info("Computing initial centroids with sklearn.KMeans...")
+        model = KMeans(
+            n_clusters=K,
+            init="k-means++",
+            n_init=10,
+            max_iter=25,
+            random_state=random_state,
+        )
+        model.fit(data)
+        km_centroids, km_labels = model.cluster_centers_, model.labels_
+        logger.info("sklearn.KMeans initialization complete.")
+        return km_centroids
+
+    def init_cluster(self, cluster_fn):
+        self.Y = cluster_fn(self.Z_cos.T, self.K).T
+        # (1) Normalize
+        self.Y = self.Y / np.linalg.norm(self.Y, ord=2, axis=0)
+        # (2) Assign cluster probabilities
+        self.dist_mat = 2 * (1 - np.dot(self.Y.T, self.Z_cos))
+        self.R = -self.dist_mat
+        self.R = self.R / self.sigma[:, None]
+        self.R -= np.max(self.R, axis=0)
+        self.R = np.exp(self.R)
+        self.R = self.R / np.sum(self.R, axis=0)
+        # (3) Batch diversity statistics
+        self.E = np.outer(np.sum(self.R, axis=1), self.Pr_b)
+        self.O = np.inner(self.R, self.Phi)
+        self.compute_objective()
+        # Save results
+        self.objective_harmony.append(self.objective_kmeans[-1])
+
+    def compute_objective(self):
+        kmeans_error = np.sum(self.R * self.dist_mat)
+        # Entropy
+        _entropy = np.sum(safe_entropy(self.R) * self.sigma[:, np.newaxis])
+        # Cross Entropy
+        x = self.R * self.sigma[:, np.newaxis]
+        y = np.tile(self.theta[:, np.newaxis], self.K).T
+        z = np.log((self.O + 1) / (self.E + 1))
+        w = np.dot(y * z, self.Phi)
+        _cross_entropy = np.sum(x * w)
+        # Save results
+        # print(f"{kmeans_error=}, {_entropy=}, {_cross_entropy=}")
+        self.objective_kmeans.append(kmeans_error + _entropy + _cross_entropy)
+        self.objective_kmeans_dist.append(kmeans_error)
+        self.objective_kmeans_entropy.append(_entropy)
+        self.objective_kmeans_cross.append(_cross_entropy)
+
+    def harmonize(self, iter_harmony=10, verbose=True):
+        converged = False
+        for i in range(1, iter_harmony + 1):
+            if verbose:
+                # logger.info("Iteration {} of {}".format(i, iter_harmony))
+                pass
+            # STEP 1: Clustering
+            self.cluster()
+            # STEP 2: Regress out covariates
+            # self.moe_correct_ridge()
+            self.Z_cos, self.Z_corr, self.W, self.Phi_Rk = moe_correct_ridge(
+                self.Z_orig,
+                self.Z_cos,
+                self.Z_corr,
+                self.R,
+                self.W,
+                self.K,
+                self.Phi_Rk,
+                self.Phi_moe,
+                self.lamb,
+                self.reference_values,
+            )
+            # STEP 3: Check for convergence
+            converged = self.check_convergence(1)
+            if converged:
+                if verbose:
+                    logger.info(
+                        "Converged after {} iteration{}".format(i, "s" if i > 1 else "")
+                    )
+                break
+        if verbose and not converged:
+            logger.info("Stopped before convergence")
+        return 0
+
+    def cluster(self):
+        # Z_cos has changed
+        # R is assumed to not have changed
+        # Update Y to match new integrated data
+        self.dist_mat = 2 * (1 - np.dot(self.Y.T, self.Z_cos))
+        for i in range(self.max_iter_kmeans):
+            # print("kmeans {}".format(i))
+            # STEP 1: Update Y
+            self.Y = np.dot(self.Z_cos, self.R.T)
+            self.Y = self.Y / np.linalg.norm(self.Y, ord=2, axis=0)
+            # STEP 2: Update dist_mat
+            self.dist_mat = 2 * (1 - np.dot(self.Y.T, self.Z_cos))
+            # STEP 3: Update R
+            self.update_R()
+            # STEP 4: Check for convergence
+            self.compute_objective()
+            if i > self.window_size:
+                converged = self.check_convergence(0)
+                if converged:
+                    break
+        self.kmeans_rounds.append(i)
+        self.objective_harmony.append(self.objective_kmeans[-1])
+        return 0
+
+    def update_R(self):
+        self._scale_dist = -self.dist_mat
+        self._scale_dist = self._scale_dist / self.sigma[:, None]
+        self._scale_dist -= np.max(self._scale_dist, axis=0)
+        self._scale_dist = np.exp(self._scale_dist)
+        # Update cells in blocks
+        update_order = np.arange(self.N)
+        np.random.shuffle(update_order)
+        n_blocks = np.ceil(1 / self.block_size).astype(int)
+        blocks = np.array_split(update_order, n_blocks)
+        for b in blocks:
+            # STEP 1: Remove cells
+            self.E -= np.outer(np.sum(self.R[:, b], axis=1), self.Pr_b)
+            self.O -= np.dot(self.R[:, b], self.Phi[:, b].T)
+            # STEP 2: Recompute R for removed cells
+            self.R[:, b] = self._scale_dist[:, b]
+            self.R[:, b] = np.multiply(
+                self.R[:, b],
+                np.dot(
+                    np.power((self.E + 1) / (self.O + 1), self.theta), self.Phi[:, b]
+                ),
+            )
+            self.R[:, b] = self.R[:, b] / np.linalg.norm(self.R[:, b], ord=1, axis=0)
+            # STEP 3: Put cells back
+            self.E += np.outer(np.sum(self.R[:, b], axis=1), self.Pr_b)
+            self.O += np.dot(self.R[:, b], self.Phi[:, b].T)
+        return 0
+
+    def check_convergence(self, i_type):
+        obj_old = 0.0
+        obj_new = 0.0
+        # Clustering, compute new window mean
+        if i_type == 0:
+            okl = len(self.objective_kmeans)
+            for i in range(self.window_size):
+                obj_old += self.objective_kmeans[okl - 2 - i]
+                obj_new += self.objective_kmeans[okl - 1 - i]
+            if (score := (abs(obj_old - obj_new) / abs(obj_old))) < self.epsilon_kmeans:
+                return True
+            # logger.info("Score: {} >= {}".format(score, self.epsilon_kmeans))
+            return False
+        # Harmony
+        if i_type == 1:
+            obj_old = self.objective_harmony[-2]
+            obj_new = self.objective_harmony[-1]
+            if (
+                score := (abs(obj_old - obj_new) / abs(obj_old))
+            ) < self.epsilon_harmony:
+                # logger.info("Score: {} >= {}".format(score, self.epsilon_harmony))
+                return True
+            # logger.info("Score: {} >= {}".format(score, self.epsilon_harmony))
+            return False
+        return True
+
+
+def safe_entropy(x: np.array):
+    y = np.multiply(x, np.log(x))
+    y[~np.isfinite(y)] = 0.0
+    return y
+
+
+def moe_correct_ridge(
+    Z_orig, Z_cos, Z_corr, R, W, K, Phi_Rk, Phi_moe, lamb, frozen_values=None
+):
+    """
+    Z_orig, Z_cos, Z_corr: DxN
+    R: KxN
+    W: (B+1)xD
+    Phi_moe: (B+1)xN
+    lamb: (B+1)x(B+1) diag matrix
+    """
+    Z_corr = Z_orig.copy()
+
+    if frozen_values is not None:
+        update_mask = ~frozen_values
+    else:
+        update_mask = np.ones(Z_orig.shape[1], dtype=bool)
+
+    for i in range(K):
+        # standard design
+        Phi_Rk = Phi_moe * R[i, :]
+
+        # ridge regression to get W
+        x = Phi_Rk @ Phi_moe.T + lamb
+        W = np.linalg.inv(x) @ Phi_Rk @ Z_orig.T
+        W[0, :] = 0  # don’t remove intercept
+
+        # apply correction
+        Z_corr[:, update_mask] -= (W.T @ Phi_Rk)[:, update_mask]
+
+    Z_cos = Z_corr / np.linalg.norm(Z_corr, ord=2, axis=0)
+    return Z_cos, Z_corr, W, Phi_Rk