sclab 0.2.5__py3-none-any.whl → 0.3.1__py3-none-any.whl

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

sclab/preprocess/_harmony_integrate.py

@@ -0,0 +1,53 @@
+"""Use harmony to integrate cells from different experiments.
+
+Note: code adapted from scanpy to use a custom version of harmonypy
+
+Harmony:
+Korsunsky, I., Millard, N., Fan, J. et al. Fast, sensitive and accurate integration of single-cell data with Harmony.
+Nat Methods 16, 1289-1296 (2019). https://doi.org/10.1038/s41592-019-0619-0
+
+Scanpy:
+Wolf, F., Angerer, P. & Theis, F. SCANPY: large-scale single-cell gene expression data analysis.
+Genome Biol 19, 15 (2018). https://doi.org/10.1186/s13059-017-1382-0
+
+Scverse:
+Virshup, I., Bredikhin, D., Heumos, L. et al. The scverse project provides a computational ecosystem for single-cell omics data analysis.
+Nat Biotechnol 41, 604-606 (2023). https://doi.org/10.1038/s41587-023-01733-8
+"""
+
+from collections.abc import Sequence
+
+import numpy as np
+from anndata import AnnData
+
+from ._harmony import run_harmony
+
+
+def harmony_integrate(
+    adata: AnnData,
+    key: str | Sequence[str],
+    *,
+    basis: str = "X_pca",
+    adjusted_basis: str | None = None,
+    reference_batch: str | list[str] | None = None,
+    **kwargs,
+):
+    """Use harmonypy :cite:p:`Korsunsky2019` to integrate different experiments."""
+
+    if adjusted_basis is None:
+        adjusted_basis = f"{basis}_harmony"
+
+    if isinstance(reference_batch, str):
+        reference_batch = [reference_batch]
+
+    if reference_batch is not None:
+        reference_values = np.zeros(adata.n_obs, dtype=bool)
+        for batch in reference_batch:
+            reference_values |= adata.obs[key].values == batch
+        kwargs["reference_values"] = reference_values
+
+    X = adata.obsm[basis].astype(np.float64)
+
+    harmony_out = run_harmony(X, adata.obs, key, **kwargs)
+
+    adata.obsm[adjusted_basis] = harmony_out.Z_corr.T

sclab/preprocess/_normalize_weighted.py

@@ -0,0 +1,61 @@
+import warnings
+
+import numpy as np
+from anndata import AnnData, ImplicitModificationWarning
+from scipy.sparse import csr_matrix, issparse
+
+
+def normalize_weighted(
+    adata: AnnData,
+    target_scale: float | None = None,
+    batch_key: str | None = None,
+) -> None:
+    if batch_key is not None:
+        for _, idx in adata.obs.groupby(batch_key, observed=True).groups.items():
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    category=ImplicitModificationWarning,
+                    message="Modifying `X` on a view results in data being overridden",
+                )
+                normalize_weighted(adata[idx], target_scale, None)
+
+        return
+
+    X: csr_matrix
+    Y: csr_matrix
+    Z: csr_matrix
+
+    X = adata.X
+    assert issparse(X)
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=RuntimeWarning, message="divide by zero"
+        )
+        Y = X.multiply(1 / X.sum(axis=0))
+    Y = Y.tocsr()
+    Y.eliminate_zeros()
+    Y.data = -Y.data * np.log(Y.data)
+    entropy = Y.sum(axis=0)
+
+    Z = X.multiply(entropy)
+    Z = Z.tocsr()
+    Z.eliminate_zeros()
+
+    with warnings.catch_warnings():
+        warnings.filterwarnings(
+            "ignore", category=RuntimeWarning, message="divide by zero"
+        )
+        scale = Z.sum(axis=1)
+        Z = Z.multiply(1 / scale)
+    Z = Z.tocsr()
+
+    if target_scale is None:
+        target_scale = np.median(scale.A1[scale.A1 > 0])
+
+    Z = Z * target_scale
+
+    adata.X = Z
+
+    return