pertpy 0.9.4__py3-none-any.whl → 0.9.5__py3-none-any.whl

pertpy/tools/_mixscape.py

@@ -18,6 +18,7 @@ from scipy.sparse import csr_matrix, issparse, spmatrix
 from sklearn.mixture import GaussianMixture
 
 import pertpy as pt
+from pertpy._doc import _doc_params, doc_common_plot_args
 
 if TYPE_CHECKING:
     from collections.abc import Sequence
@@ -25,6 +26,7 @@ if TYPE_CHECKING:
     from anndata import AnnData
     from matplotlib.axes import Axes
     from matplotlib.colors import Colormap
+    from matplotlib.pyplot import Figure
     from scipy import sparse
 
 
@@ -102,7 +104,7 @@ class Mixscape:
             control_mask_split = control_mask & split_mask
 
             R_split = representation[split_mask]
-            R_control = representation[control_mask_split]
+            R_control = representation[np.asarray(control_mask_split)]
 
             from pynndescent import NNDescent
 
@@ -110,7 +112,7 @@ class Mixscape:
             nn_index = NNDescent(R_control, **kwargs)
             indices, _ = nn_index.query(R_split, k=n_neighbors, epsilon=eps)
 
-            X_control = np.expm1(adata.X[control_mask_split])
+            X_control = np.expm1(adata.X[np.asarray(control_mask_split)])
 
             n_split = split_mask.sum()
             n_control = X_control.shape[0]
@@ -254,7 +256,7 @@ class Mixscape:
                 else:
                     de_genes = perturbation_markers[(category, gene)]
                     de_genes_indices = self._get_column_indices(adata, list(de_genes))
-                    dat = X[all_cells][:, de_genes_indices]
+                    dat = X[np.asarray(all_cells)][:, de_genes_indices]
                     converged = False
                     n_iter = 0
                     old_classes = adata.obs[labels][all_cells]
@@ -264,8 +266,8 @@ class Mixscape:
                         # get average value for each gene over all selected cells
                         # all cells in current split&Gene minus all NT cells in current split
                         # Each row is for each cell, each column is for each gene, get mean for each column
-                        vec = np.mean(X[guide_cells][:, de_genes_indices], axis=0) - np.mean(
-                            X[nt_cells][:, de_genes_indices], axis=0
+                        vec = np.mean(X[np.asarray(guide_cells)][:, de_genes_indices], axis=0) - np.mean(
+                            X[np.asarray(nt_cells)][:, de_genes_indices], axis=0
                         )
                         # project cells onto the perturbation vector
                         if isinstance(dat, spmatrix):
@@ -506,21 +508,23 @@ class Mixscape:
 
         return [mu, sd]
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_barplot(  # pragma: no cover
         self,
         adata: AnnData,
         guide_rna_column: str,
+        *,
         mixscape_class_global: str = "mixscape_class_global",
         axis_text_x_size: int = 8,
         axis_text_y_size: int = 6,
         axis_title_size: int = 8,
         legend_title_size: int = 8,
         legend_text_size: int = 8,
-        return_fig: bool | None = None,
-        ax: Axes | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
-    ):
+        legend_bbox_to_anchor: tuple[float, float] = None,
+        figsize: tuple[float, float] = (25, 25),
+        show: bool = True,
+        return_fig: bool = False,
+    ) -> Figure | None:
         """Barplot to visualize perturbation scores calculated by the `mixscape` function.
 
         Args:
@@ -528,12 +532,17 @@ class Mixscape:
             guide_rna_column: The column of `.obs` with guide RNA labels. The target gene labels.
                               The format must be <gene_target>g<#>. Examples are 'STAT2g1' and 'ATF2g1'.
             mixscape_class_global: The column of `.obs` with mixscape global classification result (perturbed, NP or NT).
-            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'}.
+            axis_text_x_size: Size of the x-axis text.
+            axis_text_y_size: Size of the y-axis text.
+            axis_title_size: Size of the axis title.
+            legend_title_size: Size of the legend title.
+            legend_text_size: Size of the legend text.
+            legend_bbox_to_anchor: The bbox that the legend will be anchored.
+            figsize: The size of the figure.
+            {common_plot_args}
 
         Returns:
-            If `show==False`, return a :class:`~matplotlib.axes.Axes.
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
 
         Examples:
             >>> import pertpy as pt
@@ -565,63 +574,66 @@ class Mixscape:
         all_cells_percentage["guide_number"] = "g" + all_cells_percentage["guide_number"]
         NP_KO_cells = all_cells_percentage[all_cells_percentage["gene"] != "NT"]
 
-        if show:
-            color_mapping = {"KO": "salmon", "NP": "lightgray", "NT": "grey"}
-            unique_genes = NP_KO_cells["gene"].unique()
-            fig, axs = plt.subplots(int(len(unique_genes) / 5), 5, figsize=(25, 25), sharey=True)
-            for i, gene in enumerate(unique_genes):
-                ax = axs[int(i / 5), i % 5]
-                grouped_df = (
-                    NP_KO_cells[NP_KO_cells["gene"] == gene]
-                    .groupby(["guide_number", "mixscape_class_global"], observed=False)["value"]
-                    .sum()
-                    .unstack()
-                )
-                grouped_df.plot(
-                    kind="bar",
-                    stacked=True,
-                    color=[color_mapping[col] for col in grouped_df.columns],
-                    ax=ax,
-                    width=0.8,
-                    legend=False,
-                )
-                ax.set_title(
-                    gene, bbox={"facecolor": "white", "edgecolor": "black", "pad": 1}, fontsize=axis_title_size
-                )
-                ax.set(xlabel="sgRNA", ylabel="% of cells")
-                sns.despine(ax=ax, top=True, right=True, left=False, bottom=False)
-                ax.set_xticklabels(ax.get_xticklabels(), rotation=0, ha="right", fontsize=axis_text_x_size)
-                ax.set_yticklabels(ax.get_yticklabels(), rotation=0, fontsize=axis_text_y_size)
-            fig.subplots_adjust(right=0.8)
-            fig.subplots_adjust(hspace=0.5, wspace=0.5)
+        color_mapping = {"KO": "salmon", "NP": "lightgray", "NT": "grey"}
+        unique_genes = NP_KO_cells["gene"].unique()
+        fig, axs = plt.subplots(int(len(unique_genes) / 5), 5, figsize=figsize, sharey=True)
+        for i, gene in enumerate(unique_genes):
+            ax = axs[int(i / 5), i % 5]
+            grouped_df = (
+                NP_KO_cells[NP_KO_cells["gene"] == gene]
+                .groupby(["guide_number", "mixscape_class_global"], observed=False)["value"]
+                .sum()
+                .unstack()
+            )
+            grouped_df.plot(
+                kind="bar",
+                stacked=True,
+                color=[color_mapping[col] for col in grouped_df.columns],
+                ax=ax,
+                width=0.8,
+                legend=False,
+            )
+            ax.set_title(gene, bbox={"facecolor": "white", "edgecolor": "black", "pad": 1}, fontsize=axis_title_size)
+            ax.set(xlabel="sgRNA", ylabel="% of cells")
+            sns.despine(ax=ax, top=True, right=True, left=False, bottom=False)
+            ax.set_xticks(ax.get_xticks(), ax.get_xticklabels(), rotation=0, ha="right", fontsize=axis_text_x_size)
+            ax.set_yticks(ax.get_yticks(), ax.get_yticklabels(), rotation=0, fontsize=axis_text_y_size)
             ax.legend(
-                title="mixscape_class_global",
+                title="Mixscape Class",
                 loc="center right",
-                bbox_to_anchor=(2.2, 3.5),
+                bbox_to_anchor=legend_bbox_to_anchor,
                 frameon=True,
                 fontsize=legend_text_size,
                 title_fontsize=legend_title_size,
             )
 
+        fig.subplots_adjust(right=0.8)
+        fig.subplots_adjust(hspace=0.5, wspace=0.5)
         plt.tight_layout()
-        _utils.savefig_or_show("mixscape_barplot", show=show, save=save)
 
+        if show:
+            plt.show()
+        if return_fig:
+            return fig
+        return None
+
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_heatmap(  # pragma: no cover
         self,
         adata: AnnData,
         labels: str,
         target_gene: str,
         control: str,
+        *,
         layer: str | None = None,
         method: str | None = "wilcoxon",
         subsample_number: int | None = 900,
         vmin: float | None = -2,
         vmax: float | None = 2,
-        return_fig: bool | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
+        show: bool = True,
+        return_fig: bool = False,
         **kwds,
-    ) -> Axes | None:
+    ) -> Figure | None:
         """Heatmap plot using mixscape results. Requires `pt.tl.mixscape()` to be run first.
 
         Args:
@@ -634,14 +646,11 @@ class Mixscape:
             subsample_number: Subsample to this number of observations.
             vmin: The value representing the lower limit of the color scale. Values smaller than vmin are plotted with the same color as vmin.
             vmax: The value representing the upper limit of the color scale. Values larger than vmax are plotted with the same color as vmax.
-            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'`}.
-            ax: A matplotlib axes object. Only works if plotting a single component.
+            {common_plot_args}
             **kwds: Additional arguments to `scanpy.pl.rank_genes_groups_heatmap`.
 
         Returns:
-            If `show==False`, return a :class:`~matplotlib.axes.Axes`.
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
 
         Examples:
             >>> import pertpy as pt
@@ -663,35 +672,39 @@ class Mixscape:
         sc.pp.scale(adata_subset, max_value=vmax)
         sc.pp.subsample(adata_subset, n_obs=subsample_number)
 
-        return sc.pl.rank_genes_groups_heatmap(
+        fig = sc.pl.rank_genes_groups_heatmap(
             adata_subset,
             groupby="mixscape_class",
             vmin=vmin,
             vmax=vmax,
             n_genes=20,
             groups=["NT"],
-            return_fig=return_fig,
-            show=show,
-            save=save,
+            show=False,
             **kwds,
         )
 
+        if show:
+            plt.show()
+        if return_fig:
+            return fig
+        return None
+
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_perturbscore(  # pragma: no cover
         self,
         adata: AnnData,
         labels: str,
         target_gene: str,
+        *,
         mixscape_class: str = "mixscape_class",
         color: str = "orange",
         palette: dict[str, str] = None,
         split_by: str = None,
         before_mixscape: bool = False,
         perturbation_type: str = "KO",
-        return_fig: bool | None = None,
-        ax: Axes | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
-    ) -> None:
+        show: bool = True,
+        return_fig: bool = False,
+    ) -> Figure | None:
         """Density plots to visualize perturbation scores calculated by the `pt.tl.mixscape` function.
 
         Requires `pt.tl.mixscape` to be run first.
@@ -710,6 +723,10 @@ class Mixscape:
             before_mixscape: Option to split densities based on mixscape classification (default) or original target gene classification.
                              Default is set to NULL and plots cells by original class ID.
             perturbation_type: Specify type of CRISPR perturbation expected for labeling mixscape classifications.
+            {common_plot_args}
+
+        Returns:
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
 
         Examples:
             Visualizing the perturbation scores for the cells in a dataset:
@@ -778,15 +795,6 @@ class Mixscape:
                 plt.legend(title="gene_target", title_fontsize=14, fontsize=12)
                 sns.despine()
 
-            if save:
-                plt.savefig(save, bbox_inches="tight")
-            if show:
-                plt.show()
-            if return_fig:
-                return plt.gcf()
-            if not (show or save):
-                return plt.gca()
-
         # If before_mixscape is False, split densities based on mixscape classifications
         else:
             if palette is None:
@@ -843,19 +851,18 @@ class Mixscape:
                 plt.legend(title="mixscape class", title_fontsize=14, fontsize=12)
                 sns.despine()
 
-            if save:
-                plt.savefig(save, bbox_inches="tight")
-            if show:
-                plt.show()
-            if return_fig:
-                return plt.gcf()
-            if not (show or save):
-                return plt.gca()
+        if show:
+            plt.show()
+        if return_fig:
+            return plt.gcf()
+        return None
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_violin(  # pragma: no cover
         self,
         adata: AnnData,
         target_gene_idents: str | list[str],
+        *,
         keys: str | Sequence[str] = "mixscape_class_p_ko",
         groupby: str | None = "mixscape_class",
         log: bool = False,
@@ -872,10 +879,10 @@ class Mixscape:
         ylabel: str | Sequence[str] | None = None,
         rotation: float | None = None,
         ax: Axes | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
+        show: bool = True,
+        return_fig: bool = False,
         **kwargs,
-    ):
+    ) -> Axes | Figure | None:
         """Violin plot using mixscape results.
 
         Requires `pt.tl.mixscape` to be run first.
@@ -892,14 +899,12 @@ class Mixscape:
             xlabel: Label of the x-axis. Defaults to `groupby` if `rotation` is `None`, otherwise, no label is shown.
             ylabel: Label of the y-axis. If `None` and `groupby` is `None`, defaults to `'value'`.
                     If `None` and `groubpy` is not `None`, defaults to `keys`.
-            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'`}.
             ax: A matplotlib axes object. Only works if plotting a single component.
+            {common_plot_args}
             **kwargs: Additional arguments to `seaborn.violinplot`.
 
         Returns:
-            A :class:`~matplotlib.axes.Axes` object if `ax` is `None` else `None`.
+            If `return_fig` is `True`, returns the figure (as Axes list if it's a multi-panel plot), otherwise `None`.
 
         Examples:
             >>> import pertpy as pt
@@ -1045,20 +1050,24 @@ class Mixscape:
         show = settings.autoshow if show is None else show
         if hue is not None and stripplot is True:
             plt.legend(handles, labels)
-        _utils.savefig_or_show("mixscape_violin", show=show, save=save)
 
-        if not show:
+        if show:
+            plt.show()
+        if return_fig:
             if multi_panel and groupby is None and len(ys) == 1:
                 return g
             elif len(axs) == 1:
                 return axs[0]
             else:
                 return axs
+        return None
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_lda(  # pragma: no cover
         self,
         adata: AnnData,
         control: str,
+        *,
         mixscape_class: str = "mixscape_class",
         mixscape_class_global: str = "mixscape_class_global",
         perturbation_type: str | None = "KO",
@@ -1066,12 +1075,11 @@ class Mixscape:
         n_components: int | None = None,
         color_map: Colormap | str | None = None,
         palette: str | Sequence[str] | None = None,
-        return_fig: bool | None = None,
         ax: Axes | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
+        show: bool = True,
+        return_fig: bool = False,
         **kwds,
-    ) -> None:
+    ) -> Figure | None:
         """Visualizing perturbation responses with Linear Discriminant Analysis. Requires `pt.tl.mixscape()` to be run first.
 
         Args:
@@ -1082,9 +1090,7 @@ class Mixscape:
             perturbation_type: Specify type of CRISPR perturbation expected for labeling mixscape classifications.
             lda_key: If not specified, lda looks .uns["mixscape_lda"] for the LDA results.
             n_components: The number of dimensions of the embedding.
-            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'`}.
+            {common_plot_args}
             **kwds: Additional arguments to `scanpy.pl.umap`.
 
         Examples:
@@ -1112,14 +1118,19 @@ class Mixscape:
             n_components = adata_subset.uns[lda_key].shape[1]
         sc.pp.neighbors(adata_subset, use_rep=lda_key)
         sc.tl.umap(adata_subset, n_components=n_components)
-        sc.pl.umap(
+        fig = sc.pl.umap(
             adata_subset,
             color=mixscape_class,
             palette=palette,
             color_map=color_map,
             return_fig=return_fig,
-            show=show,
-            save=save,
+            show=False,
             ax=ax,
             **kwds,
         )
+
+        if show:
+            plt.show()
+        if return_fig:
+            return fig
+        return None

pertpy/tools/_perturbation_space/_perturbation_space.py

@@ -7,6 +7,7 @@ import pandas as pd
 from anndata import AnnData
 from lamin_utils import logger
 from rich import print
+from scipy.stats import entropy
 
 if TYPE_CHECKING:
     from collections.abc import Iterable
@@ -41,7 +42,7 @@ class PerturbationSpace:
         Args:
             adata: Anndata object of size cells x genes.
             target_col: .obs column name that stores the label of the perturbation applied to each cell.
-            group_col: .obs column name that stores the label of the group of eah cell. If None, ignore groups.
+            group_col: .obs column name that stores the label of the group of each cell. If None, ignore groups.
             reference_key: The key of the control values.
             layer_key: Key of the AnnData layer to use for computation.
             new_layer_key: the results are stored in the given layer.
@@ -364,50 +365,58 @@ class PerturbationSpace:
         self,
         adata: AnnData,
         column: str = "perturbation",
+        column_uncertainty_score_key: str = "perturbation_transfer_uncertainty",
         target_val: str = "unknown",
-        n_neighbors: int = 5,
-        use_rep: str = "X_umap",
+        neighbors_key: str = "neighbors",
     ) -> None:
         """Impute missing values in the specified column using KNN imputation in the space defined by `use_rep`.
 
+        Uncertainty is calculated as the entropy of the label distribution in the neighborhood of the target cell.
+        In other words, a cell where all neighbors have the same set of labels will have an uncertainty of 0, whereas a cell
+        where all neighbors have many different labels will have high uncertainty.
+
         Args:
             adata: The AnnData object containing single-cell data.
-            column: The column name in AnnData object to perform imputation on.
+            column: The column name in adata.obs to perform imputation on.
+            column_uncertainty_score_key: The column name in adata.obs to store the uncertainty score of the label transfer.
             target_val: The target value to impute.
-            n_neighbors: Number of neighbors to use for imputation.
-            use_rep: The key in `adata.obsm` where the embedding (UMAP, PCA, etc.) is stored.
+            neighbors_key: The key in adata.uns where the neighbors are stored.
 
         Examples:
             >>> import pertpy as pt
             >>> import scanpy as sc
             >>> import numpy as np
             >>> adata = sc.datasets.pbmc68k_reduced()
-            >>> rng = np.random.default_rng()
-            >>> adata.obs["perturbation"] = rng.choice(
-            ...     ["A", "B", "C", "unknown"], size=adata.n_obs, p=[0.33, 0.33, 0.33, 0.01]
-            ... )
+            >>> # randomly dropout 10% of the data annotations
+            >>> adata.obs["perturbation"] = adata.obs["louvain"].astype(str).copy()
+            >>> random_cells = np.random.choice(adata.obs.index, int(adata.obs.shape[0] * 0.1), replace=False)
+            >>> adata.obs.loc[random_cells, "perturbation"] = "unknown"
             >>> sc.pp.neighbors(adata)
             >>> sc.tl.umap(adata)
             >>> ps = pt.tl.PseudobulkSpace()
-            >>> ps.label_transfer(adata, n_neighbors=5, use_rep="X_umap")
+            >>> ps.label_transfer(adata)
         """
-        if use_rep not in adata.obsm:
-            raise ValueError(f"Representation {use_rep} not found in the AnnData object.")
-
-        embedding = adata.obsm[use_rep]
-
-        from pynndescent import NNDescent
-
-        nnd = NNDescent(embedding, n_neighbors=n_neighbors)
-        indices, _ = nnd.query(embedding, k=n_neighbors)
-
-        perturbations = np.array(adata.obs[column])
-        missing_mask = perturbations == target_val
-
-        for idx in np.where(missing_mask)[0]:
-            neighbor_indices = indices[idx]
-            neighbor_categories = perturbations[neighbor_indices]
-            most_common = pd.Series(neighbor_categories).mode()[0]
-            perturbations[idx] = most_common
-
-        adata.obs[column] = perturbations
+        if neighbors_key not in adata.uns:
+            raise ValueError(f"Key {neighbors_key} not found in adata.uns. Please run `sc.pp.neighbors` first.")
+
+        labels = adata.obs[column].astype(str)
+        target_cells = labels == target_val
+
+        connectivities = adata.obsp[adata.uns[neighbors_key]["connectivities_key"]]
+        # convert labels to an incidence matrix
+        one_hot_encoded_labels = adata.obs[column].astype(str).str.get_dummies()
+        # convert to distance-weighted neighborhood incidence matrix
+        weighted_label_occurence = pd.DataFrame(
+            (one_hot_encoded_labels.values.T * connectivities).T,
+            index=adata.obs_names,
+            columns=one_hot_encoded_labels.columns,
+        )
+        # choose best label for each target cell
+        best_labels = weighted_label_occurence.drop(target_val, axis=1)[target_cells].idxmax(axis=1)
+        adata.obs[column] = labels
+        adata.obs.loc[target_cells, column] = best_labels
+
+        # calculate uncertainty
+        uncertainty = np.zeros(adata.n_obs)
+        uncertainty[target_cells] = entropy(weighted_label_occurence.drop(target_val, axis=1)[target_cells], axis=1)
+        adata.obs[column_uncertainty_score_key] = uncertainty

pertpy/tools/_perturbation_space/_simple.py

@@ -1,13 +1,20 @@
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import decoupler as dc
+import matplotlib.pyplot as plt
 import numpy as np
 from anndata import AnnData
 from sklearn.cluster import DBSCAN, KMeans
 
+from pertpy._doc import _doc_params, doc_common_plot_args
 from pertpy.tools._perturbation_space._clustering import ClusteringSpace
 from pertpy.tools._perturbation_space._perturbation_space import PerturbationSpace
 
+if TYPE_CHECKING:
+    from matplotlib.pyplot import Figure
+
 
 class CentroidSpace(PerturbationSpace):
     """Computes the centroids per perturbation of a pre-computed embedding."""
@@ -168,6 +175,49 @@ class PseudobulkSpace(PerturbationSpace):
 
         return ps_adata
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
+    def plot_psbulk_samples(
+        self,
+        adata: AnnData,
+        groupby: str,
+        *,
+        show: bool = True,
+        return_fig: bool = False,
+        **kwargs,
+    ) -> Figure | None:
+        """Plot the pseudobulk samples of an AnnData object.
+
+        Plot the count number vs. the number of cells per pseudobulk sample.
+
+        Args:
+            adata: Anndata containing pseudobulk samples.
+            groupby: `.obs` column to color the samples by.
+            {common_plot_args}
+            **kwargs: Are passed to decoupler's plot_psbulk_samples.
+
+        Returns:
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> adata = pt.dt.zhang_2021()
+            >>> ps = pt.tl.PseudobulkSpace()
+            >>> pdata = ps.compute(
+            ...     adata, target_col="Patient", groups_col="Cluster", mode="sum", min_cells=10, min_counts=1000
+            ... )
+            >>> ps.plot_psbulk_samples(pdata, groupby=["Patient", "Major celltype"], figsize=(12, 4))
+
+        Preview:
+            .. image:: /_static/docstring_previews/pseudobulk_samples.png
+        """
+        fig = dc.plot_psbulk_samples(adata, groupby, return_fig=True, **kwargs)
+
+        if show:
+            plt.show()
+        if return_fig:
+            return fig
+        return None
+
 
 class KMeansSpace(ClusteringSpace):
     """Computes K-Means clustering of the expression values."""