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

pertpy/tools/_milo.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 import random
 import re
 from typing import TYPE_CHECKING, Literal
@@ -14,6 +13,8 @@ from anndata import AnnData
 from lamin_utils import logger
 from mudata import MuData
 
+from pertpy._doc import _doc_params, doc_common_plot_args
+
 if TYPE_CHECKING:
     from collections.abc import Sequence
 
@@ -125,7 +126,7 @@ class Milo:
             try:
                 use_rep = adata.uns["neighbors"]["params"]["use_rep"]
             except KeyError:
-                logging.warning("Using X_pca as default embedding")
+                logger.warning("Using X_pca as default embedding")
                 use_rep = "X_pca"
             try:
                 knn_graph = adata.obsp["connectivities"].copy()
@@ -136,7 +137,7 @@ class Milo:
             try:
                 use_rep = adata.uns[neighbors_key]["params"]["use_rep"]
             except KeyError:
-                logging.warning("Using X_pca as default embedding")
+                logger.warning("Using X_pca as default embedding")
                 use_rep = "X_pca"
             knn_graph = adata.obsp[neighbors_key + "_connectivities"].copy()
 
@@ -182,7 +183,7 @@ class Milo:
             knn_dists = adata.obsp[neighbors_key + "_distances"]
 
         nhood_ixs = adata.obs["nhood_ixs_refined"] == 1
-        dist_mat = knn_dists[nhood_ixs, :]
+        dist_mat = knn_dists[np.asarray(nhood_ixs), :]
         k_distances = dist_mat.max(1).toarray().ravel()
         adata.obs["nhood_kth_distance"] = 0
         adata.obs["nhood_kth_distance"] = adata.obs["nhood_kth_distance"].astype(float)
@@ -703,8 +704,8 @@ class Milo:
         pvalues = sample_adata.var["PValue"]
         keep_nhoods = ~pvalues.isna()  # Filtering in case of test on subset of nhoods
         o = pvalues[keep_nhoods].argsort()
-        pvalues = pvalues[keep_nhoods][o]
-        w = w[keep_nhoods][o]
+        pvalues = pvalues.loc[keep_nhoods].iloc[o]
+        w = w.loc[keep_nhoods].iloc[o]
 
         adjp = np.zeros(shape=len(o))
         adjp[o] = (sum(w) * pvalues / np.cumsum(w))[::-1].cummin()[::-1]
@@ -713,9 +714,11 @@ class Milo:
         sample_adata.var["SpatialFDR"] = np.nan
         sample_adata.var.loc[keep_nhoods, "SpatialFDR"] = adjp
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_nhood_graph(
         self,
         mdata: MuData,
+        *,
         alpha: float = 0.1,
         min_logFC: float = 0,
         min_size: int = 10,
@@ -724,10 +727,9 @@ class Milo:
         color_map: Colormap | str | None = None,
         palette: str | Sequence[str] | None = None,
         ax: Axes | None = None,
-        show: bool | None = None,
-        save: bool | str | None = None,
+        return_fig: bool = False,
         **kwargs,
-    ) -> None:
+    ) -> Figure | None:
         """Visualize DA results on abstracted graph (wrapper around sc.pl.embedding)
 
         Args:
@@ -737,9 +739,7 @@ class Milo:
             min_size: Minimum size of nodes in visualization. (default: 10)
             plot_edges: If edges for neighbourhood overlaps whould be plotted.
             title: Plot title.
-            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}
             **kwargs: Additional arguments to `scanpy.pl.embedding`.
 
         Examples:
@@ -782,7 +782,7 @@ class Milo:
         vmax = np.max([nhood_adata.obs["graph_color"].max(), abs(nhood_adata.obs["graph_color"].min())])
         vmin = -vmax
 
-        sc.pl.embedding(
+        fig = sc.pl.embedding(
             nhood_adata,
             "X_milo_graph",
             color="graph_color",
@@ -798,33 +798,40 @@ class Milo:
             color_map=color_map,
             palette=palette,
             ax=ax,
-            show=show,
-            save=save,
+            show=False,
             **kwargs,
         )
 
+        if return_fig:
+            return fig
+        plt.show()
+        return None
+
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_nhood(
         self,
         mdata: MuData,
         ix: int,
+        *,
         feature_key: str | None = "rna",
         basis: str = "X_umap",
         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,
+        return_fig: bool = False,
         **kwargs,
-    ) -> None:
+    ) -> Figure | None:
         """Visualize cells in a neighbourhood.
 
         Args:
             mdata: MuData object with feature_key slot, storing neighbourhood assignments in `mdata[feature_key].obsm['nhoods']`
             ix: index of neighbourhood to visualize
+            feature_key: Key in mdata to the cell-level AnnData object.
             basis: Embedding to use for visualization.
-            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'}.
+            color_map: Colormap to use for coloring.
+            palette: Color palette to use for coloring.
+            ax: Axes to plot on.
+            {common_plot_args}
             **kwargs: Additional arguments to `scanpy.pl.embedding`.
 
         Examples:
@@ -842,7 +849,7 @@ class Milo:
             .. image:: /_static/docstring_previews/milo_nhood.png
         """
         mdata[feature_key].obs["Nhood"] = mdata[feature_key].obsm["nhoods"][:, ix].toarray().ravel()
-        sc.pl.embedding(
+        fig = sc.pl.embedding(
             mdata[feature_key],
             basis,
             color="Nhood",
@@ -852,32 +859,41 @@ class Milo:
             palette=palette,
             return_fig=return_fig,
             ax=ax,
-            show=show,
-            save=save,
+            show=False,
             **kwargs,
         )
 
+        if return_fig:
+            return fig
+        plt.show()
+        return None
+
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_da_beeswarm(
         self,
         mdata: MuData,
+        *,
         feature_key: str | None = "rna",
         anno_col: str = "nhood_annotation",
         alpha: float = 0.1,
         subset_nhoods: list[str] = None,
         palette: str | Sequence[str] | dict[str, str] | None = None,
-        return_fig: bool | None = None,
-        save: bool | str | None = None,
-        show: bool | None = None,
-    ) -> Figure | Axes | None:
+        return_fig: bool = False,
+    ) -> Figure | None:
         """Plot beeswarm plot of logFC against nhood labels
 
         Args:
             mdata: MuData object
+            feature_key: Key in mdata to the cell-level AnnData object.
             anno_col: Column in adata.uns['nhood_adata'].obs to use as annotation. (default: 'nhood_annotation'.)
             alpha: Significance threshold. (default: 0.1)
             subset_nhoods: List of nhoods to plot. If None, plot all nhoods.
             palette: Name of Seaborn color palette for violinplots.
                      Defaults to pre-defined category colors for violinplots.
+            {common_plot_args}
+
+        Returns:
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
 
         Examples:
             >>> import pertpy as pt
@@ -973,29 +989,21 @@ class Milo:
         plt.legend(loc="upper left", title=f"< {int(alpha * 100)}% SpatialFDR", bbox_to_anchor=(1, 1), frameon=False)
         plt.axvline(x=0, ymin=0, ymax=1, color="black", linestyle="--")
 
-        if save:
-            plt.savefig(save, bbox_inches="tight")
-            return None
-        if show:
-            plt.show()
-            return None
         if return_fig:
             return plt.gcf()
-        if (not show and not save) or (show is None and save is None):
-            return plt.gca()
-
+        plt.show()
         return None
 
+    @_doc_params(common_plot_args=doc_common_plot_args)
     def plot_nhood_counts_by_cond(
         self,
         mdata: MuData,
         test_var: str,
+        *,
         subset_nhoods: list[str] = None,
         log_counts: bool = False,
-        return_fig: bool | None = None,
-        save: bool | str | None = None,
-        show: bool | None = None,
-    ) -> Figure | Axes | None:
+        return_fig: bool = False,
+    ) -> Figure | None:
         """Plot boxplot of cell numbers vs condition of interest.
 
         Args:
@@ -1003,6 +1011,10 @@ class Milo:
             test_var: Name of column in adata.obs storing condition of interest (y-axis for boxplot)
             subset_nhoods: List of obs_names for neighbourhoods to include in plot. If None, plot all nhoods.
             log_counts: Whether to plot log1p of cell counts.
+            {common_plot_args}
+
+        Returns:
+            If `return_fig` is `True`, returns the figure, otherwise `None`.
         """
         try:
             nhood_adata = mdata["milo"].T.copy()
@@ -1031,15 +1043,7 @@ class Milo:
         plt.xticks(rotation=90)
         plt.xlabel(test_var)
 
-        if save:
-            plt.savefig(save, bbox_inches="tight")
-            return None
-        if show:
-            plt.show()
-            return None
         if return_fig:
             return plt.gcf()
-        if not (show or save):
-            return plt.gca()
-
+        plt.show()
         return None