Add examples to plotting functions

pertpy/plot/_coda.py

@@ -112,6 +112,9 @@ def stacked_barplot(  # pragma: no cover
             >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
                 sample_identifier="batch", covariate_obs=["condition"])
             >>> sccoda.plot_stacked_barplot(mdata, feature_name="samples")
+
+        Preview:
+            .. image:: ../_static/docstring_previews/sccoda_stacked_barplot.png
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -185,6 +188,9 @@ def effects_barplot(  # pragma: no cover
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.plot_effects_barplot(mdata)
+
+        Preview:
+            .. image:: ../_static/docstring_previews/sccoda_effects_barplot.png
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -262,6 +268,9 @@ def boxplots(  # pragma: no cover
             >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
                 sample_identifier="batch", covariate_obs=["condition"])
             >>> sccoda.plot_boxplots(mdata, feature_name="condition", add_dots=True)
+
+        Preview:
+            .. image:: ../_static/docstring_previews/sccoda_boxplots.png
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -332,6 +341,9 @@ def rel_abundance_dispersion_plot(  # pragma: no cover
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.plot_rel_abundance_dispersion_plot(mdata)
+
+        Preview:
+            .. image:: ../_static/docstring_previews/sccoda_rel_abundance_dispersion_plot.png
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -412,6 +424,8 @@ def draw_tree(  # pragma: no cover
             >>> )
             >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
             >>> tasccoda.plot_draw_tree(mdata, tree="lineage")
+
+        Preview: #TODO: Add preview
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -498,6 +512,8 @@ def draw_effects(  # pragma: no cover
             >>> )
             >>> tasccoda.run_nuts(mdata, num_samples=1000, num_warmup=100, rng_key=42)
             >>> tasccoda.plot_draw_effects(mdata, covariate="Health[T.Inflamed]", tree="lineage")
+
+        Preview: #TODO: Add preview
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"
@@ -585,6 +601,8 @@ def effects_umap(  # pragma: no cover
             >>>                                      "effect_df_condition[T.Hpoly.Day10]"],
             >>>                                       cluster_key="nsbm_level_1",
             >>>                         )
+
+        Preview: #TODO: Add preview
         """
         warnings.warn(
             "This function is deprecated and will be removed in pertpy 0.8.0!"

pertpy/tools/_augur.py

@@ -969,7 +969,7 @@ def predict_differential_prioritization(
         return delta
 
     def plot_dp_scatter(
-        self, results: pd.DataFrame, top_n=None, ax: Axes = None, return_figure: bool = False
+        self, results: pd.DataFrame, top_n: int = None, ax: Axes = None, return_figure: bool = False
     ) -> Figure | Axes:
         """Plot scatterplot of differential prioritization.
 
@@ -1027,7 +1027,12 @@ def plot_dp_scatter(
         return fig if return_figure else ax
 
     def plot_important_features(
-        self, data: dict[str, Any], key: str = "augurpy_results", top_n=10, ax: Axes = None, return_figure: bool = False
+        self,
+        data: dict[str, Any],
+        key: str = "augurpy_results",
+        top_n: int = 10,
+        ax: Axes = None,
+        return_figure: bool = False,
     ) -> Figure | Axes:
         """Plot a lollipop plot of the n features with largest feature importances.
 
@@ -1125,7 +1130,7 @@ def plot_lollipop(
         return fig if return_figure else ax
 
     def plot_scatterplot(
-        self, results1: dict[str, Any], results2: dict[str, Any], top_n=None, return_figure: bool = False
+        self, results1: dict[str, Any], results2: dict[str, Any], top_n: int = None, return_figure: bool = False
     ) -> Figure | Axes:
         """Create scatterplot with two augur results.
 

pertpy/tools/_coda/_base_coda.py

@@ -1225,6 +1225,9 @@ def plot_stacked_barplot(  # pragma: no cover
             >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
                 sample_identifier="batch", covariate_obs=["condition"])
             >>> sccoda.plot_stacked_barplot(mdata, feature_name="samples")
+
+        Preview:
+            .. image:: /_static/docstring_previews/sccoda_stacked_barplot.png
         """
         if isinstance(data, MuData):
             data = data[modality_key]
@@ -1327,6 +1330,9 @@ def plot_effects_barplot(  # pragma: no cover
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.plot_effects_barplot(mdata)
+
+        Preview:
+            .. image:: /_static/docstring_previews/sccoda_effects_barplot.png
         """
         if args_barplot is None:
             args_barplot = {}
@@ -1507,6 +1513,9 @@ def plot_boxplots(  # pragma: no cover
             >>> mdata = sccoda.load(haber_cells, type="cell_level", generate_sample_level=True, cell_type_identifier="cell_label", \
                 sample_identifier="batch", covariate_obs=["condition"])
             >>> sccoda.plot_boxplots(mdata, feature_name="condition", add_dots=True)
+
+        Preview:
+            .. image:: /_static/docstring_previews/sccoda_boxplots.png
         """
         if args_boxplot is None:
             args_boxplot = {}
@@ -1698,7 +1707,7 @@ def plot_rel_abundance_dispersion_plot(  # pragma: no cover
         If the count of the cell type is larger than 0 in more than abundant_threshold percent of all samples, the cell type will be marked in a different color.
 
         Args:
-            data: AnnData object or MuData object.
+            data: AnnData or MuData object.
             modality_key: If data is a MuData object, specify which modality to use. Defaults to "coda".
                           Defaults to "coda".
             abundant_threshold: Presence threshold for abundant cell types. Defaults to 0.9.
@@ -1723,6 +1732,9 @@ def plot_rel_abundance_dispersion_plot(  # pragma: no cover
             >>> mdata = sccoda.prepare(mdata, formula="condition", reference_cell_type="Endocrine")
             >>> sccoda.run_nuts(mdata, num_warmup=100, num_samples=1000, rng_key=42)
             >>> sccoda.plot_rel_abundance_dispersion_plot(mdata)
+
+        Preview:
+            .. image:: /_static/docstring_previews/sccoda_rel_abundance_dispersion_plot.png
         """
         if isinstance(data, MuData):
             data = data[modality_key]
@@ -2083,7 +2095,7 @@ def plot_effects_umap(  # pragma: no cover
         show: bool = None,
         ax: Axes = None,
         **kwargs,
-    ):
+    ) -> plt.Axes | None:
         """Plot a UMAP visualization colored by effect strength.
 
         Effect results in .varm of aggregated sample-level AnnData (default is data['coda']) are assigned to cell-level AnnData

pertpy/tools/_dialogue.py

@@ -1064,7 +1064,7 @@ def plot_split_violins(
         self,
         adata: AnnData,
         split_key: str,
-        celltype_key=str,
+        celltype_key: str,
         split_which: tuple[str, str] = None,
         mcp: str = "mcp_0",
     ) -> Axes:
@@ -1091,6 +1091,9 @@ def plot_split_violins(
                 n_counts_key = "nCount_RNA", n_mpcs = 3)
             >>> adata, mcps, ws, ct_subs = dl.calculate_multifactor_PMD(adata, normalize=True)
             >>> dl.plot_split_violins(adata, split_key='gender', celltype_key='cell.subtypes')
+
+        Preview:
+            .. image:: /_static/docstring_previews/dialogue_violin.png
         """
         df = sc.get.obs_df(adata, [celltype_key, mcp, split_key])
         if split_which is None:
@@ -1131,6 +1134,9 @@ def plot_pairplot(
                 n_counts_key = "nCount_RNA", n_mpcs = 3)
             >>> adata, mcps, ws, ct_subs = dl.calculate_multifactor_PMD(adata, normalize=True)
             >>> dl.plot_pairplot(adata, celltype_key="cell.subtypes", color="gender", sample_id="clinical.status")
+
+        Preview:
+            .. image:: /_static/docstring_previews/dialogue_pairplot.png
         """
         mean_mcps = adata.obs.groupby([sample_id, celltype_key])[mcp].mean()
         mean_mcps = mean_mcps.reset_index()

pertpy/tools/_enrichment.py

@@ -318,8 +318,15 @@ def plot_dotplot(
 
         Examples:
             >>> import pertpy as pt
+            >>> import scanpy as sc
             >>> pt_enrichment = pt.tl.Enrichment()
-            >>> pt_enrichment.plot_dotplot(adata, categories=["B01","B02","B03"], groupby="leiden")
+            >>> adata = sc.datasets.pbmc3k_processed()
+            >>> pt_enrichment.score(adata)
+            >>> sc.tl.rank_genes_groups(adata, method="wilcoxon", groupby="louvain")
+            >>> pt_enrichment.plot_dotplot(adata, categories=["B01", "B02", "B03"], groupby="louvain")
+
+        Preview:
+            .. image:: /_static/docstring_previews/enrichment_dotplot.png
         """
         if categories is not None:
             if isinstance(categories, str):
@@ -384,6 +391,19 @@ def plot_gsea(
             n: How many top scores to show for each group. Defaults to 10.
             key: GSEA results key in `uns`. Defaults to "pertpy_enrichment_gsea".
             interactive_plot: Whether to plot interactively or not. Defaults to False.
+
+        Examples:
+            >>> import pertpy as pt
+            >>> import scanpy as sc
+            >>> pt_enrichment = pt.tl.Enrichment()
+            >>> adata = sc.datasets.pbmc3k_processed()
+            >>> pt_enrichment.score(adata)
+            >>> sc.tl.rank_genes_groups(adata, method="wilcoxon", groupby="louvain")
+            >>> enrichment = pt_enrichment.gsea(adata)
+            >>> pt_enrichment.plot_gsea(adata, enrichment, interactive_plot=True)
+
+        Preview:
+            .. image:: /_static/docstring_previews/enrichment_gsea.png
         """
         for cluster in enrichment:
             fig = blitzgsea.plot.top_table(

pertpy/tools/_milo.py

@@ -734,6 +734,9 @@ def plot_nhood_graph(
             >>>            model_contrasts='labelwithdraw_15d_Cocaine-labelwithdraw_48h_Cocaine')
             >>> milo.build_nhood_graph(mdata)
             >>> milo.plot_nhood_graph(mdata)
+
+        Preview:
+            .. image:: /_static/docstring_previews/milo_nhood_graph.png
         """
         nhood_adata = mdata["milo"].T.copy()
 
@@ -779,7 +782,7 @@ def plot_nhood(
         mdata: MuData,
         ix: int,
         feature_key: str | None = "rna",
-        basis="X_umap",
+        basis: str = "X_umap",
         show: bool | None = None,
         save: bool | str | None = None,
         **kwargs,
@@ -804,6 +807,9 @@ def plot_nhood(
             >>> sc.tl.umap(mdata["rna"])
             >>> milo.make_nhoods(mdata["rna"])
             >>> milo.plot_nhood(mdata, ix=0)
+
+        Preview:
+            .. image:: /_static/docstring_previews/milo_nhood.png
         """
         mdata[feature_key].obs["Nhood"] = mdata[feature_key].obsm["nhoods"][:, ix].toarray().ravel()
         sc.pl.embedding(
@@ -818,7 +824,7 @@ def plot_da_beeswarm(
         alpha: float = 0.1,
         subset_nhoods: list[str] = None,
         palette: str | Sequence[str] | dict[str, str] | None = None,
-    ):
+    ) -> None:
         """Plot beeswarm plot of logFC against nhood labels
 
         Args:
@@ -928,7 +934,7 @@ def plot_nhood_counts_by_cond(
         test_var: str,
         subset_nhoods: list[str] = None,
         log_counts: bool = False,
-    ):
+    ) -> None:
         """Plot boxplot of cell numbers vs condition of interest.
 
         Args:

pertpy/tools/_mixscape.py

@@ -443,7 +443,7 @@ def _get_perturbation_markers(
         min_de_genes: float,
         logfc_threshold: float,
     ) -> dict[tuple, np.ndarray]:
-        """determine gene sets across all splits/groups through differential gene expression
+        """Determine gene sets across all splits/groups through differential gene expression
 
         Args:
             adata: :class:`~anndata.AnnData` object
@@ -503,7 +503,7 @@ def plot_barplot(  # pragma: no cover
         self,
         adata: AnnData,
         guide_rna_column: str,
-        mixscape_class_global="mixscape_class_global",
+        mixscape_class_global: str = "mixscape_class_global",
         axis_text_x_size: int = 8,
         axis_text_y_size: int = 6,
         axis_title_size: int = 8,
@@ -524,7 +524,7 @@ def plot_barplot(  # pragma: no cover
                   Infer the filetype if ending on {'.pdf', '.png', '.svg'}.
 
         Returns:
-            If show is False, return ggplot object used to draw the plot.
+            If `show==False`, return a :class:`~matplotlib.axes.Axes.
 
         Examples:
             >>> import pertpy as pt
@@ -533,6 +533,9 @@ def plot_barplot(  # pragma: no cover
             >>> ms_pt.perturbation_signature(mdata['rna'], 'perturbation', 'NT', 'replicate')
             >>> ms_pt.mixscape(adata = mdata['rna'], control = 'NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.plot_barplot(mdata['rna'], guide_rna_column='NT')
+
+        Preview:
+            .. image:: /_static/docstring_previews/mixscape_barplot.png
         """
         if mixscape_class_global not in adata.obs:
             raise ValueError("Please run the `mixscape` function first.")
@@ -608,7 +611,7 @@ def plot_heatmap(  # pragma: no cover
         show: bool | None = None,
         save: bool | str | None = None,
         **kwds,
-    ):
+    ) -> Axes | None:
         """Heatmap plot using mixscape results. Requires `pt.tl.mixscape()` to be run first.
 
         Args:
@@ -627,13 +630,19 @@ def plot_heatmap(  # pragma: no cover
             ax: A matplotlib axes object. Only works if plotting a single component.
             **kwds: Additional arguments to `scanpy.pl.rank_genes_groups_heatmap`.
 
+        Returns:
+            If `show==False`, return a :class:`~matplotlib.axes.Axes`.
+
         Examples:
             >>> import pertpy as pt
             >>> mdata = pt.dt.papalexi_2021()
             >>> ms_pt = pt.tl.Mixscape()
             >>> ms_pt.perturbation_signature(mdata['rna'], 'perturbation', 'NT', 'replicate')
             >>> ms_pt.mixscape(adata = mdata['rna'], control = 'NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.plot_heatmap(adata = mdata['rna'], labels='gene_target', target_gene='IFNGR2', layer='X_pert', control='NT')
+
+        Preview:
+            .. image:: /_static/docstring_previews/mixscape_heatmap.png
         """
         if "mixscape_class" not in adata.obs:
             raise ValueError("Please run `pt.tl.mixscape` first.")
@@ -686,9 +695,6 @@ def plot_perturbscore(  # pragma: no cover
             perturbation_type: Specify type of CRISPR perturbation expected for labeling mixscape classifications.
                                Defaults to `KO`.
 
-        Returns:
-            The ggplot object used for drawn.
-
         Examples:
             Visualizing the perturbation scores for the cells in a dataset:
 
@@ -698,6 +704,9 @@ def plot_perturbscore(  # pragma: no cover
             >>> ms_pt.perturbation_signature(mdata['rna'], 'perturbation', 'NT', 'replicate')
             >>> ms_pt.mixscape(adata = mdata['rna'], control = 'NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.plot_perturbscore(adata = mdata['rna'], labels='gene_target', target_gene='IFNGR2', color = 'orange')
+
+        Preview:
+            .. image:: /_static/docstring_previews/mixscape_perturbscore.png
         """
         if "mixscape" not in adata.uns:
             raise ValueError("Please run the `mixscape` function first.")
@@ -864,6 +873,9 @@ def plot_violin(  # pragma: no cover
             >>> ms_pt.perturbation_signature(mdata['rna'], 'perturbation', 'NT', 'replicate')
             >>> ms_pt.mixscape(adata = mdata['rna'], control = 'NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.plot_violin(adata = mdata['rna'], target_gene_idents=['NT', 'IFNGR2 NP', 'IFNGR2 KO'], groupby='mixscape_class')
+
+        Preview:
+            .. image:: /_static/docstring_previews/mixscape_violin.png
         """
         if isinstance(target_gene_idents, str):
             mixscape_class_mask = adata.obs[groupby] == target_gene_idents
@@ -1010,8 +1022,8 @@ def plot_lda(  # pragma: no cover
         self,
         adata: AnnData,
         control: str,
-        mixscape_class="mixscape_class",
-        mixscape_class_global="mixscape_class_global",
+        mixscape_class: str = "mixscape_class",
+        mixscape_class_global: str = "mixscape_class_global",
         perturbation_type: str | None = "KO",
         lda_key: str | None = "mixscape_lda",
         n_components: int | None = None,
@@ -1043,6 +1055,9 @@ def plot_lda(  # pragma: no cover
             >>> ms_pt.mixscape(adata = mdata['rna'], control = 'NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.lda(adata=mdata['rna'], control='NT', labels='gene_target', layer='X_pert')
             >>> ms_pt.plot_lda(adata=mdata['rna'], control='NT')
+
+        Preview:
+            .. image:: /_static/docstring_previews/mixscape_lda.png
         """
         if mixscape_class not in adata.obs:
             raise ValueError(f'Did not find `.obs["{mixscape_class!r}"]`. Please run the `mixscape` function first.')