"""Plotting for graph functions."""
from __future__ import annotations
import warnings
from collections.abc import Mapping, Sequence
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Literal
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import seaborn as sns
from anndata import AnnData
from matplotlib.axes import Axes
from squidpy._compat import old_positionals
from squidpy._constants._constants import RipleyStat
from squidpy._constants._pkg_constants import Key
from squidpy._docs import d
from squidpy._validators import assert_non_empty_sequence, get_valid_values
from squidpy.gr._utils import _assert_categorical_obs
from squidpy.pl._color_utils import Palette_t, _get_palette, _maybe_set_colors
from squidpy.pl._utils import _heatmap, save_fig
__all__ = [
"centrality_scores",
"interaction_matrix",
"nhood_enrichment",
"nhood_enrichment_dotplot",
"ripley",
"co_occurrence",
]
def _get_data(adata: AnnData, *, cluster_key: str, func_name: str, attr: str = "uns", **kwargs: Any) -> Any:
key = getattr(Key.uns, func_name)(cluster_key, **kwargs)
try:
if attr == "uns":
return adata.uns[key]
elif attr == "obsm":
return adata.obsm[key]
else:
raise ValueError(f"attr must be either 'uns' or 'obsm', got {attr}")
except KeyError:
raise KeyError(
f"Unable to get the data from `adata.uns[{key!r}]`. "
f"Please run `squidpy.gr.{func_name}(..., cluster_key={cluster_key!r})` first."
) from None
[docs]
@d.dedent
@old_positionals("cluster_key", "score", "legend_kwargs", "palette", "figsize", "dpi", "save")
def centrality_scores(
adata: AnnData,
*,
cluster_key: str,
score: str | Sequence[str] | None = None,
legend_kwargs: Mapping[str, Any] = MappingProxyType({}),
palette: Palette_t = None,
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
save: str | Path | None = None,
**kwargs: Any,
) -> None:
"""
Plot centrality scores.
The centrality scores are computed by :func:`squidpy.gr.centrality_scores`.
Parameters
----------
%(adata)s
%(cluster_key)s
score
Whether to plot all scores or only selected ones.
legend_kwargs
Keyword arguments for :func:`matplotlib.pyplot.legend`.
%(cat_plotting)s
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
df = _get_data(adata, cluster_key=cluster_key, func_name="centrality_scores").copy()
legend_kwargs = dict(legend_kwargs)
if "loc" not in legend_kwargs:
legend_kwargs["loc"] = "center left"
legend_kwargs.setdefault("bbox_to_anchor", (1, 0.5))
scores = df.columns.values
df[cluster_key] = df.index.values
clusters = adata.obs[cluster_key].cat.categories
palette = _get_palette(adata, cluster_key=cluster_key, categories=clusters, palette=palette)
score = scores if score is None else score
score = assert_non_empty_sequence(score, name="centrality scores")
score = sorted(get_valid_values(score, scores))
fig, axs = plt.subplots(1, len(score), figsize=figsize, dpi=dpi, constrained_layout=True)
axs = np.ravel(axs) # make into iterable
for g, ax in zip(score, axs, strict=False):
sns.scatterplot(
x=g,
y=cluster_key,
data=df,
hue=cluster_key,
hue_order=clusters,
palette=palette,
ax=ax,
**kwargs,
)
ax.set_title(str(g).replace("_", " ").capitalize())
ax.set_xlabel("value")
ax.set_yticks([])
ax.legend(**legend_kwargs)
if save is not None:
save_fig(fig, path=save)
[docs]
@d.dedent
@old_positionals(
"cluster_key", "annotate", "method", "title", "cmap", "palette", "cbar_kwargs", "figsize", "dpi", "save", "ax"
)
def interaction_matrix(
adata: AnnData,
*,
cluster_key: str,
annotate: bool = False,
method: str | None = None,
title: str | None = None,
cmap: str = "viridis",
palette: Palette_t = None,
cbar_kwargs: Mapping[str, Any] = MappingProxyType({}),
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
save: str | Path | None = None,
ax: Axes | None = None,
**kwargs: Any,
) -> None:
"""
Plot cluster interaction matrix.
The interaction matrix is computed by :func:`squidpy.gr.interaction_matrix`.
Parameters
----------
%(adata)s
%(cluster_key)s
%(heatmap_plotting)s
kwargs
Keyword arguments for :func:`matplotlib.pyplot.text`.
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
array = _get_data(adata, cluster_key=cluster_key, func_name="interaction_matrix")
ad = AnnData(X=array, obs={cluster_key: pd.Categorical(adata.obs[cluster_key].cat.categories)})
_maybe_set_colors(source=adata, target=ad, key=cluster_key, palette=palette)
if title is None:
title = "Interaction matrix"
fig = _heatmap(
ad,
key=cluster_key,
title=title,
method=method,
cont_cmap=cmap,
annotate=annotate,
figsize=(2 * ad.n_obs // 3, 2 * ad.n_obs // 3) if figsize is None else figsize,
dpi=dpi,
cbar_kwargs=cbar_kwargs,
ax=ax,
**kwargs,
)
if save is not None:
save_fig(fig, path=save)
[docs]
@d.dedent
@old_positionals(
"cluster_key",
"mode",
"annotate",
"method",
"title",
"cmap",
"palette",
"cbar_kwargs",
"figsize",
"dpi",
"save",
"ax",
)
def nhood_enrichment(
adata: AnnData,
*,
cluster_key: str,
mode: Literal["zscore", "count"] = "zscore",
annotate: bool = False,
method: str | None = None,
title: str | None = None,
cmap: str = "viridis",
palette: Palette_t = None,
cbar_kwargs: Mapping[str, Any] = MappingProxyType({}),
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
save: str | Path | None = None,
ax: Axes | None = None,
**kwargs: Any,
) -> None:
"""
Plot neighborhood enrichment.
The enrichment is computed by :func:`squidpy.gr.nhood_enrichment`.
Parameters
----------
%(adata)s
%(cluster_key)s
mode
Which :func:`squidpy.gr.nhood_enrichment` result to plot. Valid options are:
- `'zscore'` - z-score values of enrichment statistic.
- `'count'` - enrichment count.
%(heatmap_plotting)s
kwargs
Keyword arguments for :func:`matplotlib.pyplot.text`.
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
array = _get_data(adata, cluster_key=cluster_key, func_name="nhood_enrichment")[mode]
ad = AnnData(X=array, obs={cluster_key: pd.Categorical(adata.obs[cluster_key].cat.categories)})
_maybe_set_colors(source=adata, target=ad, key=cluster_key, palette=palette)
if title is None:
title = "Neighborhood enrichment"
fig = _heatmap(
ad,
key=cluster_key,
title=title,
method=method,
cont_cmap=cmap,
annotate=annotate,
figsize=(2 * ad.n_obs // 3, 2 * ad.n_obs // 3) if figsize is None else figsize,
dpi=dpi,
cbar_kwargs=cbar_kwargs,
ax=ax,
**kwargs,
)
if save is not None:
save_fig(fig, path=save)
[docs]
@d.dedent
def nhood_enrichment_dotplot(
adata: AnnData,
*,
cluster_key: str,
annotate: bool = False,
title: str | None = None,
cmap: str = "RdBu_r",
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
size_range: tuple[float, float] = (10, 200),
save: str | Path | None = None,
ax: Axes | None = None,
**kwargs: Any,
) -> None:
"""
Dot plot of neighborhood enrichment.
The enrichment is computed by :func:`squidpy.gr.nhood_enrichment`. Color shows the z-score,
dot size the conditional cell ratio (CCR) from ``normalization='conditional'`` -- the fraction
of cells of the row type having at least one neighbor of the column type.
Parameters
----------
%(adata)s
%(cluster_key)s
annotate
Whether to annotate the dots with their CCR values.
title
The title of the plot.
cmap
Diverging colormap for the z-score, centered on zero whenever the data spans both signs.
size_range
Smallest and largest dot size, which the CCR is scaled onto.
%(plotting)s
ax
Axes, :class:`matplotlib.axes.Axes`.
kwargs
Keyword arguments for :meth:`scanpy.pl.DotPlot.style`.
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
enrichment = _get_data(adata, cluster_key=cluster_key, func_name="nhood_enrichment")
zscore = enrichment["zscore"]
ccr = enrichment.get("conditional_ratio")
if ccr is None:
warnings.warn(
"'conditional_ratio' is None in nhood_enrichment results. Please run nhood_enrichment "
"with normalization = 'conditional'. Dot size will not reflect conditional cell ratios.",
UserWarning,
stacklevel=2,
)
ccr = np.ones_like(zscore)
annotate = False # the values are a placeholder; printing them would look like real data
cats = list(adata.obs[cluster_key].cat.categories)
# `DotPlot` reads its values from ``dot_color_df``/``dot_size_df``, so this stand-in only has to
# carry the right groups and var names; ``X`` is never touched.
grid = AnnData(
np.zeros((len(cats), len(cats)), dtype=np.float32),
obs=pd.DataFrame({cluster_key: pd.Categorical(cats, categories=cats)}, index=cats),
var=pd.DataFrame(index=pd.Index(cats)),
)
# A z-score is signed, so anchor the colormap at zero. Only when the data actually straddles
# zero, since TwoSlopeNorm requires vmin < vcenter < vmax -- and only with explicit bounds,
# because undefined enrichments are NaN by default and autoscaling over them yields NaN.
finite = zscore[np.isfinite(zscore)]
if finite.size and finite.min() < 0 < finite.max():
vmin, vmax, vcenter = float(finite.min()), float(finite.max()), 0.0
else:
vmin = vmax = vcenter = None
dp = (
sc.pl.DotPlot(
grid,
var_names=cats,
groupby=cluster_key,
dot_color_df=pd.DataFrame(zscore, index=cats, columns=cats),
dot_size_df=pd.DataFrame(ccr, index=cats, columns=cats),
title=title,
figsize=figsize,
vmin=vmin,
vmax=vmax,
vcenter=vcenter,
ax=ax,
)
.style(cmap=cmap, smallest_dot=size_range[0], largest_dot=size_range[1], **kwargs)
.legend(colorbar_title="Z-score", size_title="CCR")
)
dp.make_figure()
main_ax = dp.get_axes()["mainplot_ax"]
main_ax.set_xlabel("Neighbor cell type")
main_ax.set_ylabel("Index cell type")
if annotate:
for i in range(len(cats)):
for j in range(len(cats)):
main_ax.text(j + 0.5, i + 0.5, f"{ccr[i, j]:.2f}", ha="center", va="center")
fig = main_ax.figure
if dpi is not None:
fig.set_dpi(dpi)
if save is not None:
# `set_dpi` is not honoured by `savefig`, so the saved file needs it passed through
save_fig(fig, path=save, **({} if dpi is None else {"dpi": dpi}))
[docs]
@d.dedent
@old_positionals("cluster_key", "mode", "plot_sims", "palette", "figsize", "dpi", "save", "ax", "legend_kwargs")
def ripley(
adata: AnnData,
*,
cluster_key: str,
mode: Literal["F", "G", "L"] = "F",
plot_sims: bool = True,
palette: Palette_t = None,
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
save: str | Path | None = None,
ax: Axes | None = None,
legend_kwargs: Mapping[str, Any] = MappingProxyType({}),
**kwargs: Any,
) -> None:
"""
Plot Ripley's statistics for each cluster.
The estimate is computed by :func:`squidpy.gr.ripley`.
Parameters
----------
%(adata)s
%(cluster_key)s
mode
Ripley's statistics to be plotted.
plot_sims
Whether to overlay simulations in the plot.
%(cat_plotting)s
ax
Axes, :class:`matplotlib.axes.Axes`.
legend_kwargs
Keyword arguments for :func:`matplotlib.pyplot.legend`.
kwargs
Keyword arguments for :func:`seaborn.lineplot`.
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
res = _get_data(adata, cluster_key=cluster_key, func_name="ripley", mode=mode)
mode = RipleyStat(mode) # type: ignore[assignment]
if TYPE_CHECKING:
assert isinstance(mode, RipleyStat)
legend_kwargs = dict(legend_kwargs)
if "loc" not in legend_kwargs:
legend_kwargs["loc"] = "center left"
legend_kwargs.setdefault("bbox_to_anchor", (1, 0.5))
categories = adata.obs[cluster_key].cat.categories
palette = _get_palette(adata, cluster_key=cluster_key, categories=categories, palette=palette)
if ax is None:
fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
else:
fig = ax.figure
sns.lineplot(
y="stats",
x="bins",
hue=cluster_key,
data=res[f"{mode.s}_stat"],
hue_order=categories,
palette=palette,
ax=ax,
**kwargs,
)
if plot_sims:
sns.lineplot(y="stats", x="bins", errorbar="sd", alpha=0.01, color="gray", data=res["sims_stat"], ax=ax)
ax.legend(**legend_kwargs)
ax.set_ylabel("value")
ax.set_title(f"Ripley's {mode.s}")
if save is not None:
save_fig(fig, path=save)
[docs]
@d.dedent
@old_positionals("cluster_key", "palette", "clusters", "figsize", "dpi", "save", "legend_kwargs")
def co_occurrence(
adata: AnnData,
*,
cluster_key: str,
palette: Palette_t = None,
clusters: str | Sequence[str] | None = None,
figsize: tuple[float, float] | None = None,
dpi: int | None = None,
save: str | Path | None = None,
legend_kwargs: Mapping[str, Any] = MappingProxyType({}),
**kwargs: Any,
) -> None:
"""
Plot co-occurrence probability ratio for each cluster.
The co-occurrence is computed by :func:`squidpy.gr.co_occurrence`.
Parameters
----------
%(adata)s
%(cluster_key)s
clusters
Cluster instances for which to plot conditional probability.
%(cat_plotting)s
legend_kwargs
Keyword arguments for :func:`matplotlib.pyplot.legend`.
kwargs
Keyword arguments for :func:`seaborn.lineplot`.
Returns
-------
%(plotting_returns)s
"""
_assert_categorical_obs(adata, key=cluster_key)
occurrence_data = _get_data(adata, cluster_key=cluster_key, func_name="co_occurrence")
legend_kwargs = dict(legend_kwargs)
if "loc" not in legend_kwargs:
legend_kwargs["loc"] = "center left"
legend_kwargs.setdefault("bbox_to_anchor", (1, 0.5))
out = occurrence_data["occ"]
interval = occurrence_data["interval"][1:]
categories = adata.obs[cluster_key].cat.categories
clusters = categories if clusters is None else clusters
clusters = assert_non_empty_sequence(clusters, name="clusters")
clusters = sorted(get_valid_values(clusters, categories))
palette = _get_palette(adata, cluster_key=cluster_key, categories=categories, palette=palette)
fig, axs = plt.subplots(
1,
len(clusters),
figsize=(5 * len(clusters), 5) if figsize is None else figsize,
dpi=dpi,
constrained_layout=True,
)
axs = np.ravel(axs) # make into iterable
for g, ax in zip(clusters, axs, strict=False):
idx = np.where(categories == g)[0][0]
df = pd.DataFrame(out[idx, :, :].T, columns=categories).melt(var_name=cluster_key, value_name="probability")
df["distance"] = np.tile(interval, len(categories))
sns.lineplot(
x="distance",
y="probability",
data=df,
dashes=False,
hue=cluster_key,
hue_order=categories,
palette=palette,
ax=ax,
**kwargs,
)
ax.legend(**legend_kwargs)
ax.set_title(rf"$\frac{{p(exp|{g})}}{{p(exp)}}$")
ax.set_ylabel("value")
if save is not None:
save_fig(fig, path=save)