pertpy 0.6.0__py3-none-any.whl → 0.8.0__py3-none-any.whl

pertpy/tools/_cinemaot.py

@@ -2,14 +2,17 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import scanpy as sc
 import scipy.stats as ss
+import seaborn as sns
 import sklearn.metrics
 from ott.geometry import pointcloud
 from ott.problems.linear import linear_problem
 from ott.solvers.linear import sinkhorn, sinkhorn_lr
+from scanpy.plotting import _utils
 from scipy.sparse import issparse
 from sklearn.decomposition import FastICA
 from sklearn.linear_model import LinearRegression
@@ -17,6 +20,7 @@ from sklearn.neighbors import NearestNeighbors
 
 if TYPE_CHECKING:
     from anndata import AnnData
+    from matplotlib.axes import Axes
     from statsmodels.tools.typing import ArrayLike
 
 
@@ -67,6 +71,14 @@ class Cinemaot:
             Returns an AnnData object that contains the single-cell level treatment effect as de.X and the
             corresponding low dimensional embedding in de.obsm['X_embedding'], and optional matching matrix
             stored in the de.obsm['ot']. Also puts the confounding variation in adata.obsm[cf_rep].
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> out_adata = model.causaleffect(
+            >>>         adata, pert_key="perturbation", control="No stimulation", return_matching=True,
+            >>>         thres=0.5, smoothness=1e-5, eps=1e-3, solver="Sinkhorn", preweight_label="cell_type0528")
         """
         available_solvers = ["Sinkhorn", "LRSinkhorn"]
         if solver not in available_solvers:
@@ -225,6 +237,14 @@ class Cinemaot:
             Returns an anndata object that contains the single-cell level treatment effect as de.X and the
             corresponding low dimensional embedding in de.obsm['X_embedding'], and optional matching matrix
             stored in the de.obsm['ot']. Also puts the confounding variation in adata.obsm[cf_rep].
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> ad, de = model.causaleffect_weighted(
+            >>>              adata, pert_key="perturbation", control="No stimulation", return_matching=True,
+            >>>              thres=0.5, smoothness=1e-5, eps=1e-3, solver="Sinkhorn")
         """
         available_solvers = ["Sinkhorn", "LRSinkhorn"]
         assert solver in available_solvers, (
@@ -288,6 +308,16 @@ class Cinemaot:
 
         Returns:
             Returns an anndata object that contains aggregated pseudobulk profiles and associated metadata.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> de = model.causaleffect(
+            >>>         adata, pert_key="perturbation", control="No stimulation", return_matching=True, thres=0.5,
+            >>>         smoothness=1e-5, eps=1e-3, solver="Sinkhorn", preweight_label="cell_type0528")
+            >>> adata_pb = model.generate_pseudobulk(
+            >>>         adata, de, pert_key="perturbation", control="No stimulation", label_list=None)
         """
         sc.pp.neighbors(de, use_rep=de_rep)
         sc.tl.leiden(de, resolution=de_resolution)
@@ -308,7 +338,7 @@ class Cinemaot:
             sc.tl.leiden(adata, resolution=cf_resolution)
             df["ct"] = adata.obs["leiden"].astype(str)
         df["ptb"] = "control"
-        df["ptb"][adata.obs[pert_key] != control] = de.obs["leiden"].astype(str)
+        df.loc[adata.obs[pert_key] != control, "ptb"] = de.obs["leiden"].astype(str)
         label_list.append("ptb")
         df = df.groupby(label_list).sum()
         new_index = df.index.map(lambda x: "_".join(map(str, x)))
@@ -336,6 +366,12 @@ class Cinemaot:
 
         Returns:
             Returns the estimated dimension number.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> dim = model.get_dim(adata)
         """
         sk = SinkhornKnopp()
         if issparse(adata.raw.X):
@@ -369,6 +405,12 @@ class Cinemaot:
 
         Returns:
             Returns the indices.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> idx = model.get_weightidx(adata, pert_key="perturbation", control="No stimulation")
         """
         adata_ = adata.copy()
         X_pca1 = adata_.obsm[use_rep][adata_.obs[pert_key] == control, :]
@@ -390,7 +432,7 @@ class Cinemaot:
         expr_label = "control"
 
         adata_.obs["ct"] = ref_label
-        adata_.obs["ct"][adata_.obs[pert_key] == control] = expr_label
+        adata_.obs.loc[adata_.obs[pert_key] == control, "ct"] = expr_label
         pert_key = "ct"
         z = np.zeros(adata_.shape[0]) + 1
 
@@ -485,6 +527,15 @@ class Cinemaot:
 
         Returns:
             Returns an AnnData object that contains the single-cell level synergy matrix de.X and the embedding.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.dong_2023()
+            >>> sc.pp.pca(adata)
+            >>> model = pt.tl.Cinemaot()
+            >>> combo = model.synergy(adata, pert_key='perturbation', base='No stimulation', A='IFNb', B='IFNg',
+            >>>                   AB='IFNb+ IFNg', thres=0.5, smoothness=1e-5, eps=1e-3, solver='Sinkhorn')
+
         """
         adata1 = adata[adata.obs[pert_key].isin([base, A]), :].copy()
         adata2 = adata[adata.obs[pert_key].isin([B, AB]), :].copy()
@@ -552,6 +603,12 @@ class Cinemaot:
 
         Returns:
             Returns the confounder effect (c_effect) and the residual effect (s_effect).
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> model = pt.tl.Cinemaot()
+            >>> c_effect, s_effect = model.attribution_scatter(adata, pert_key="perturbation", control="No stimulation")
         """
         cf = adata.obsm[cf_rep]
         if use_raw:
@@ -582,6 +639,84 @@ class Cinemaot:
         s_effect = (np.linalg.norm(e1, axis=0) + 1e-6) / (np.linalg.norm(e0, axis=0) + 1e-6)
         return c_effect, s_effect
 
+    def plot_vis_matching(
+        self,
+        adata: AnnData,
+        de: AnnData,
+        pert_key: str,
+        control: str,
+        de_label: str,
+        source_label: str,
+        matching_rep: str = "ot",
+        resolution: float = 0.5,
+        normalize: str = "col",
+        title: str = "CINEMA-OT matching matrix",
+        min_val: float = 0.01,
+        show: bool = True,
+        save: str | None = None,
+        ax: Axes | None = None,
+        **kwargs,
+    ) -> None:
+        """Visualize the CINEMA-OT matching matrix.
+
+        Args:
+            adata: the original anndata after running cinemaot.causaleffect or cinemaot.causaleffect_weighted.
+            de: The anndata output from Cinemaot.causaleffect() or Cinemaot.causaleffect_weighted().
+            pert_key: The column  of `.obs` with perturbation categories, should also contain `control`.
+            control: Control category from the `pert_key` column.
+            de_label: the label for differential response. If none, use leiden cluster labels at resolution 1.0.
+            source_label: the confounder / cell type label.
+            matching_rep: the place that stores the matching matrix. default de.obsm['ot'].
+            normalize: normalize the coarse-grained matching matrix by row / column.
+            title: the title for the figure.
+            min_val: The min value to truncate the matching matrix.
+            show: Show the plot, do not return axis.
+            save: If `True` or a `str`, save the figure. A string is appended to the default filename.
+                Infer the filetype if ending on {`'.pdf'`, `'.png'`, `'.svg'`}.
+            **kwargs: Other parameters to input for seaborn.heatmap.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.cinemaot_example()
+            >>> cot = pt.tl.Cinemaot()
+            >>> de = cot.causaleffect(
+            >>>         adata, pert_key="perturbation", control="No stimulation", return_matching=True,
+            >>>         thres=0.5, smoothness=1e-5, eps=1e-3, solver="Sinkhorn", preweight_label="cell_type0528")
+            >>> cot.plot_vis_matching(
+            >>>         adata, de, pert_key="perturbation",control="No stimulation", de_label=None, source_label="cell_type0528")
+        """
+        adata_ = adata[adata.obs[pert_key] == control]
+
+        df = pd.DataFrame(de.obsm[matching_rep])
+        if de_label is None:
+            de_label = "leiden"
+            sc.pp.neighbors(de, use_rep="X_embedding")
+            sc.tl.leiden(de, resolution=resolution)
+        df["de_label"] = de.obs[de_label].astype(str).values
+        df["de_label"] = "Response " + df["de_label"]
+        df = df.groupby("de_label").sum().T
+        df["source_label"] = adata_.obs[source_label].astype(str).values
+        df = df.groupby("source_label").sum()
+
+        if normalize == "col":
+            df = df / df.sum(axis=0)
+        else:
+            df = (df.T / df.sum(axis=1)).T
+        df = df.clip(lower=min_val) - min_val
+        if normalize == "col":
+            df = df / df.sum(axis=0)
+        else:
+            df = (df.T / df.sum(axis=1)).T
+
+        g = sns.heatmap(df, annot=True, ax=ax, **kwargs)
+        plt.title(title)
+        _utils.savefig_or_show("matching_heatmap", show=show, save=save)
+        if not show:
+            if ax is not None:
+                return ax
+            else:
+                return g
+
 
 class Xi:
     """
@@ -610,7 +745,7 @@ class Xi:
         # same as pandas rank method 'first'
         rankdata = ss.rankdata(randomized, method="ordinal")
         # Reindexing based on pairs of indices before and after
-        unrandomized = [rankdata[j] for i, j in sorted(zip(randomized_indices, range(len_x)))]
+        unrandomized = [rankdata[j] for i, j in sorted(zip(randomized_indices, range(len_x), strict=False))]
         return unrandomized
 
     @property
@@ -648,6 +783,7 @@ class Xi:
                         for x, y in zip(
                             x1,
                             x2,
+                            strict=False,
                         )
                     ]
                 )
@@ -694,13 +830,16 @@ class Xi:
         ind = [i + 1 for i in range(self.sample_size)]
         ind2 = [2 * self.sample_size - 2 * ind[i - 1] + 1 for i in ind]
 
-        a = np.mean([i * j * j for i, j in zip(ind2, sorted_ordered_x_rank)]) / self.sample_size
+        a = np.mean([i * j * j for i, j in zip(ind2, sorted_ordered_x_rank, strict=False)]) / self.sample_size
 
-        c = np.mean([i * j for i, j in zip(ind2, sorted_ordered_x_rank)]) / self.sample_size
+        c = np.mean([i * j for i, j in zip(ind2, sorted_ordered_x_rank, strict=False)]) / self.sample_size
 
         cq = np.cumsum(sorted_ordered_x_rank)
 
-        m = [(i + (self.sample_size - j) * k) / self.sample_size for i, j, k in zip(cq, ind, sorted_ordered_x_rank)]
+        m = [
+            (i + (self.sample_size - j) * k) / self.sample_size
+            for i, j, k in zip(cq, ind, sorted_ordered_x_rank, strict=False)
+        ]
 
         b = np.mean([np.square(i) for i in m])
         v = (a - 2 * b + np.square(c)) / np.square(self.inverse_g_mean)