"""Cell-type composition plots: stacked bars and clustered heatmaps.
All functions here work from a sample-by-cell-type proportion table
(``pd.crosstab`` under the hood), then offer several sample orderings:
as-is, hierarchically clustered, or clustered within groups of a metadata
column.
"""
from __future__ import annotations
import math
from collections.abc import Sequence
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from matplotlib import font_manager
from matplotlib import patches as mpatches
from scipy.cluster.hierarchy import leaves_list, linkage
from scipy.spatial.distance import pdist
from ._config import PlotConfig
from ._palettes import resolve_palette
from ._style import save_figure
from ._utils import require_columns, resolve_sample_column
__all__ = [
"cluster_samples",
"cluster_grouped_samples",
"stacked_barplots",
"stacked_barplots_multi_meta",
"composition_heatmap",
"abundance_bubble_grid",
]
[docs]
def cluster_samples(ct_props: pd.DataFrame) -> pd.Index:
"""Order samples by hierarchical (Ward) clustering of their proportions."""
dist_matrix = pdist(ct_props.values, metric="euclidean")
linkage_matrix = linkage(dist_matrix, method="ward")
return ct_props.index[leaves_list(linkage_matrix)]
[docs]
def cluster_grouped_samples(
ct_props: pd.DataFrame, adata, sample_column: str, order_by_column: str, ascending: bool = True
) -> list:
"""Order samples by clustering within each group of ``order_by_column``.
Groups themselves are ordered by their first observed value, ascending
or descending as requested; samples within a group are hierarchically
clustered (or left as-is if the group has a single sample).
"""
group_labels = (
adata.obs.groupby(sample_column, observed=True)[order_by_column].first().sort_values(ascending=ascending)
)
ordered = []
for group_value in group_labels.unique():
group_sample_ids = group_labels[group_labels == group_value].index
group_props = ct_props.loc[group_sample_ids]
if len(group_props) > 1:
d = pdist(group_props.values, metric="euclidean")
link = linkage(d, method="ward")
ordered.extend(group_props.index[leaves_list(link)])
else:
ordered.extend(group_props.index)
return ordered
def _clean_label(label: str, strip_prefix: str | None) -> str:
return label.replace(strip_prefix, "") if strip_prefix else label
def _fit_legends_in_figure(fig, ax, legends):
"""Grow the figure canvas, if needed, so ``legends`` fit on it.
``bbox_inches="tight"`` (used by :func:`~scplotkit._style.save_figure`)
only rescues legends registered figure-level via ``fig.legend()``; legends
kept alive via ``ax.legend()`` + ``ax.add_artist`` are invisible to it,
and Jupyter's inline display uses the same tight-bbox machinery. Given
figure-level legends, growing the canvas itself so it physically contains
them fixes both the saved file and the notebook-inline rendering --
legends anchored right of the axes push the canvas wider, legends
anchored below it push the canvas taller.
"""
fig.canvas.draw()
renderer = fig.canvas.get_renderer()
dpi = fig.dpi
fig_width_in, fig_height_in = fig.get_figwidth(), fig.get_figheight()
fig_width_px, fig_height_px = fig_width_in * dpi, fig_height_in * dpi
extents = [legend.get_window_extent(renderer) for legend in legends]
max_right_px = max(e.x1 for e in extents)
min_bottom_px = min(e.y0 for e in extents)
ax_pos = ax.get_position()
new_width_in, new_height_in = fig_width_in, fig_height_in
adjust_kwargs = {}
if max_right_px > fig_width_px:
extra_in = (max_right_px - fig_width_px) / dpi + 0.15
new_width_in = fig_width_in + extra_in
adjust_kwargs["right"] = (ax_pos.x1 * fig_width_in) / new_width_in
if min_bottom_px < 0:
extra_in = -min_bottom_px / dpi + 0.15
new_height_in = fig_height_in + extra_in
adjust_kwargs["bottom"] = (ax_pos.y0 * fig_height_in + extra_in) / new_height_in
adjust_kwargs["top"] = (ax_pos.y1 * fig_height_in + extra_in) / new_height_in
if not adjust_kwargs:
return
fig.set_size_inches(new_width_in, new_height_in, forward=True)
fig.subplots_adjust(**adjust_kwargs)
def _order_color_strip(ax, order_values, order_color_dict, bar_width, y_offset, height):
for i, val in enumerate(order_values):
ax.add_patch(
plt.Rectangle(
(i - bar_width / 2, y_offset),
bar_width,
height,
linewidth=0,
edgecolor=order_color_dict[val],
facecolor=order_color_dict[val],
transform=ax.transData,
clip_on=False,
)
)
def _plot_stacked_bar(
ct_props, color_dict, title, ylabel, figsize, sample_column,
order_values=None, order_color_dict=None, order_by_column=None,
strip_prefix=None, config: PlotConfig = None,
):
general = config.general
fig, ax = plt.subplots(figsize=figsize)
bottom = np.zeros(len(ct_props))
bar_width = 1
x_positions = np.arange(len(ct_props))
for cell_type in ct_props.columns:
values = ct_props[cell_type].values
ax.bar(
x_positions, values, bottom=bottom, width=bar_width,
color=color_dict.get(cell_type, general["fallback_color"]),
label=_clean_label(cell_type, strip_prefix),
edgecolor="white", linewidth=1.0,
)
bottom += values
ax.set_xlabel("")
ax.set_ylabel(ylabel, fontsize=12)
ax.set_xticks([])
ax.set_xticklabels([])
ax.set_title(title, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"], pad=20)
ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
ax.set_yticklabels(["0%", "25%", "50%", "75%", "100%"])
ax.set_xlim(-0.5, len(ct_props) - 0.5)
ax.set_ylim(0, 1.0)
ax.set_axisbelow(True)
if order_values is not None and order_color_dict is not None:
_order_color_strip(
ax, order_values, order_color_dict, bar_width, y_offset=-0.035, height=0.025
)
handles = [
plt.Rectangle((0, 0), 1, 1, color=color_dict.get(k, general["fallback_color"])) for k in ct_props.columns
]
labels = [_clean_label(k, strip_prefix) for k in ct_props.columns]
# fig.legend() (rather than ax.legend() + ax.add_artist()) so this legend is picked up by
# matplotlib's own tight-bbox accounting -- ax.add_artist'd legends are invisible to it, which
# is what let this legend get silently cropped by bbox_inches="tight" (the default for both
# save_figure()'s save and Jupyter's inline display) whenever it overflowed the axes.
legend1 = fig.legend(
handles, labels, bbox_to_anchor=(1.01, 1), bbox_transform=ax.transAxes, loc="upper left",
title="Cell Types", title_fontproperties=font_manager.FontProperties(weight="bold"),
fontsize=general["legend_fontsize"], frameon=False,
)
legends = [legend1]
if order_values is not None and order_color_dict is not None and order_by_column:
order_handles = [plt.Rectangle((0, 0), 1, 1, color=order_color_dict[v]) for v in order_color_dict]
order_labels = [str(v) for v in order_color_dict]
legend2 = fig.legend(
order_handles, order_labels, bbox_to_anchor=(1.01, 0), bbox_transform=ax.transAxes, loc="lower left",
title=order_by_column.replace("_", " ").title(), frameon=False,
)
legend2.get_title().set_fontweight("bold")
legends.append(legend2)
fig.subplots_adjust(right=0.78)
_fit_legends_in_figure(fig, ax, legends)
return fig
[docs]
def stacked_barplots(
adata,
level_column: str,
sample_column: str | None = "Sample_ID",
subset_level: str | None = None,
subset_value: str | Sequence | None = None,
order_by_column: str | None = None,
order_ascending: bool = True,
figsize: tuple = (16, 8),
palette: dict | None = None,
order_palette: dict | None = None,
strip_prefix: str | None = None,
save_name_prefix: str = "composition",
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
) -> dict:
"""Stacked bar plots of cell-type proportion per sample.
Returns a dict with up to three variants -- ``"basic"``, ``"clustered"``
(samples ordered by hierarchical clustering), and ``"clustered_grouped"``
(present only if ``order_by_column`` is given: samples clustered within
each group of that column).
"""
config = PlotConfig.load(config)
sample_column = resolve_sample_column(adata, sample_column)
require_columns(adata, level_column, sample_column, subset_level, order_by_column)
if subset_level and subset_value is not None:
subset_value = [subset_value] if isinstance(subset_value, str) else subset_value
adata = adata[adata.obs[subset_level].isin(subset_value), :]
ct_data = pd.crosstab(adata.obs[sample_column], adata.obs[level_column])
ct_props = ct_data.div(ct_data.sum(axis=1), axis=0)
color_dict = resolve_palette(config, level_column, ct_props.columns, palette)
ylabel = f"Cell Type Proportion ({subset_value[0]}s)" if subset_value else "Cell Type Proportion"
title_base = f"Cell Type Composition by {sample_column.replace('_', ' ').title()}"
order_color_dict = None
if order_by_column:
order_vals_all = adata.obs.groupby(sample_column, observed=True)[order_by_column].first()
order_color_dict = resolve_palette(config, order_by_column, order_vals_all.unique(), order_palette)
def order_values_for(index):
if not order_by_column:
return None
return [adata.obs[adata.obs[sample_column] == s][order_by_column].iloc[0] for s in index]
figs = {
"basic": _plot_stacked_bar(
ct_props, color_dict, title_base, ylabel, figsize, sample_column,
order_values_for(ct_props.index), order_color_dict, order_by_column, strip_prefix, config,
)
}
clustered_order = cluster_samples(ct_props)
figs["clustered"] = _plot_stacked_bar(
ct_props.loc[clustered_order], color_dict, f"{title_base} (Clustered)", ylabel, figsize, sample_column,
order_values_for(clustered_order), order_color_dict, order_by_column, strip_prefix, config,
)
if order_by_column:
grouped_order = cluster_grouped_samples(ct_props, adata, sample_column, order_by_column, order_ascending)
figs["clustered_grouped"] = _plot_stacked_bar(
ct_props.loc[grouped_order], color_dict, f"{title_base} (Clustered per Group)", ylabel, figsize,
sample_column, order_values_for(grouped_order), order_color_dict, order_by_column,
strip_prefix, config,
)
for key, fig in figs.items():
save_figure(
fig, output_dir, "compositional_plot", f"{save_name_prefix}_{key}.png", config, include_legends=True
)
return figs
def _plot_stacked_bar_multi_meta(
ct_props, color_dict, title, ylabel, figsize, sample_column,
meta_annotations_list, meta_color_dicts_list, meta_columns_list,
xlabel=None, strip_prefix=None, config: PlotConfig = None,
):
general = config.general
fig, ax = plt.subplots(figsize=figsize)
bottom = np.zeros(len(ct_props))
bar_width = 1
x_positions = np.arange(len(ct_props))
for cell_type in ct_props.columns:
values = ct_props[cell_type].values
ax.bar(
x_positions, values, bottom=bottom, width=bar_width,
color=color_dict.get(cell_type, general["fallback_color"]),
label=_clean_label(cell_type, strip_prefix), edgecolor="white", linewidth=1.0,
)
bottom += values
ax.set_xlabel("", fontsize=12)
ax.set_ylabel(ylabel, fontsize=12)
ax.set_xticks([])
ax.set_xticklabels([])
ax.set_title(title, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"], pad=20)
ax.set_yticks([0, 0.25, 0.5, 0.75, 1.0])
ax.set_yticklabels(["0%", "25%", "50%", "75%", "100%"])
ax.set_xlim(-0.5, len(ct_props) - 0.5)
ax.set_ylim(0, 1.0)
ax.set_axisbelow(True)
if meta_annotations_list and meta_color_dicts_list and meta_columns_list:
bar_height = 0.025
y_offset_base = -bar_height - 0.01
for idx, (annotations, color_dict_meta, _col) in enumerate(
zip(meta_annotations_list, meta_color_dicts_list, meta_columns_list)
):
y_offset = y_offset_base - idx * (bar_height + 0.005)
for i, val in enumerate(annotations):
color = color_dict_meta.get(val, general["fallback_color"])
ax.add_patch(
mpatches.Rectangle(
(i - bar_width / 2, y_offset), bar_width, bar_height,
linewidth=0, edgecolor=color, facecolor=color,
transform=ax.transData, clip_on=False,
)
)
handles = [
plt.Rectangle((0, 0), 1, 1, color=color_dict.get(k, general["fallback_color"])) for k in ct_props.columns
]
labels = [_clean_label(k, strip_prefix) for k in ct_props.columns]
legend1 = fig.legend(
handles, labels, bbox_to_anchor=(1.01, 1), bbox_transform=ax.transAxes, loc="upper left",
title="Cell Types", title_fontproperties=font_manager.FontProperties(weight="bold"),
fontsize=general["legend_fontsize"], frameon=False,
)
legends = [legend1]
if meta_annotations_list and meta_color_dicts_list and meta_columns_list:
current_y_anchor = -0.15
for idx, (color_dict_meta, column_name) in enumerate(zip(meta_color_dicts_list, meta_columns_list)):
present_values = sorted(set(meta_annotations_list[idx]), key=str)
order_handles = [
plt.Rectangle((0, 0), 1, 1, color=color_dict_meta[v]) for v in present_values if v in color_dict_meta
]
order_labels = [str(v) for v in present_values if v in color_dict_meta]
if not order_handles:
continue
# fig.legend() (rather than ax.legend() + ax.add_artist()) so this legend is picked up
# by matplotlib's own tight-bbox accounting -- see _fit_legends_in_figure.
legend = fig.legend(
order_handles, order_labels, bbox_to_anchor=(0.0, current_y_anchor), bbox_transform=ax.transAxes,
loc="upper left", title=column_name.replace("_", " ").title(), frameon=False,
fontsize=general["legend_fontsize"], ncol=1,
)
legend.get_title().set_fontweight("bold")
legends.append(legend)
fig.canvas.draw()
bbox = legend.get_window_extent(fig.canvas.get_renderer())
figure_height = fig.get_figheight() * fig.dpi
current_y_anchor -= (bbox.height / figure_height) + 0.05
fig.subplots_adjust(right=0.78)
_fit_legends_in_figure(fig, ax, legends)
return fig
[docs]
def composition_heatmap(
adata,
level_column: str,
sample_column: str | None = "Sample_ID",
subset_level: str | None = None,
subset_value: str | Sequence | None = None,
order_by_column: str | None = None,
order_palette: dict | None = None,
figsize: tuple = (16, 6),
title: str | None = None,
cmap: str | None = None,
save_name: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
**heatmap_kwargs,
):
"""Clustered heatmap of cell-type proportion per sample."""
import seaborn as sns
config = PlotConfig.load(config)
sample_column = resolve_sample_column(adata, sample_column)
require_columns(adata, level_column, sample_column, subset_level, order_by_column)
if subset_level is not None and subset_value is not None:
subset_value = [subset_value] if isinstance(subset_value, str) else subset_value
mask = adata.obs[subset_level].isin(subset_value)
adata = adata[mask, :]
print(f"Subset to {mask.sum()} cells with {subset_level} in {subset_value}")
ct_data = pd.crosstab(adata.obs[sample_column], adata.obs[level_column])
ct_props = ct_data.div(ct_data.sum(axis=1), axis=0)
dist_matrix = pdist(ct_props.values, metric="euclidean")
linkage_matrix = linkage(dist_matrix, method="ward")
ct_props = ct_props.iloc[leaves_list(linkage_matrix)]
order_color_dict = None
order_bar_colors = None
unique_vals = None
if order_by_column:
order_values = [adata.obs[adata.obs[sample_column] == s][order_by_column].iloc[0] for s in ct_props.index]
unique_vals = list(dict.fromkeys(order_values))
order_color_dict = resolve_palette(config, order_by_column, unique_vals, order_palette)
order_bar_colors = [order_color_dict[v] for v in order_values]
general = config.general
fig, ax = plt.subplots(figsize=figsize)
sns.heatmap(
ct_props.T, annot=False, cmap=cmap or config["continuous"]["cmap"], ax=ax, cbar=True,
linewidths=0, linecolor="white", xticklabels=False, yticklabels=True, square=False,
**heatmap_kwargs,
)
ax.set_aspect("auto")
ax.set_title(
title or "Cell Type Composition (Clustered)",
fontsize=general["title_fontsize"], fontweight=general["title_fontweight"],
)
ax.set_xlabel("")
ax.set_ylabel(f"Cell Types ({subset_value[0]}s)" if subset_value else "Cell Types", fontsize=12)
ax.set_xticks([])
ax.set_xticklabels([])
ax.set_yticklabels(ct_props.columns, rotation=0)
if order_by_column:
bar_height = 0.5
for idx, color in enumerate(order_bar_colors):
ax.add_patch(
plt.Rectangle(
(idx, ct_props.shape[1] + 0.05), 1, bar_height,
linewidth=0, color=color, transform=ax.transData, clip_on=False,
)
)
legend_handles = [plt.Rectangle((0, 0), 1, 1, color=order_color_dict[v]) for v in unique_vals]
legend_labels = [str(v) for v in unique_vals]
legend = fig.legend(
legend_handles, legend_labels, title=order_by_column.replace("_", " ").title(),
loc="lower center", bbox_to_anchor=(0.5, -0.12), ncol=len(unique_vals),
fontsize=general["legend_fontsize"], title_fontsize=general["legend_fontsize"], frameon=False,
)
legend.get_title().set_fontweight("bold")
plt.tight_layout(rect=[0, 0.01, 1, 1])
if save_name is not None:
save_figure(fig, output_dir, "compositional_plot", f"{save_name}.png", config)
return fig
def _size_legend_values(max_count: int, n: int = 4) -> list[int]:
"""Pick ``n`` round reference counts spanning roughly 0..``max_count`` for a size legend."""
if max_count <= 0:
return [0]
values = set()
for v in np.linspace(max_count / n, max_count, n):
magnitude = 10 ** math.floor(math.log10(v)) if v >= 1 else 1
values.add(int(round(v / magnitude) * magnitude))
return sorted(values)
[docs]
def abundance_bubble_grid(
adata,
level_column: str,
sample_column: str | None = "Sample_ID",
subset_level: str | None = None,
subset_value: str | Sequence | None = None,
dot_color: str = "#37474f",
min_size: float = 20.0,
max_size: float = 500.0,
figsize: tuple | None = None,
title: str | None = None,
xlabel: str | None = None,
save_name: str = "abundance_bubble_grid",
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Dot grid of cell counts per (cell type, sample/dataset) pair.
Rows are ``level_column`` categories, columns are ``sample_column``
values (any per-sample or per-dataset column works); dot area encodes
cell count. A single neutral ``dot_color`` is used throughout rather
than a per-category palette, so size alone carries the value. Pairs
with zero cells are drawn as a faint open ring rather than left blank,
so gaps in atlas coverage/balance across studies are visible rather
than ambiguous with "not shown". Rows and columns are both ordered by
descending total count, so the best-covered cell types and
samples/datasets end up in the top-left.
Purely descriptive -- raw cell counts only, no proportions and no
statistical framing.
Parameters
----------
adata:
Annotated data.
level_column:
``adata.obs`` column defining the rows (e.g. cell type).
sample_column:
``adata.obs`` column defining the columns (sample or dataset ID).
Defaults to the first object/category column if not found.
subset_level, subset_value:
Optionally restrict to cells where ``subset_level`` is in
``subset_value`` before counting (e.g. one lineage).
dot_color:
Fill color used for every filled dot.
min_size, max_size:
Marker area (in points^2) for the smallest nonzero and the largest
count, respectively.
save_name:
Filename (without extension) to save under. Pass ``None`` to skip saving.
config:
Active :class:`~scplotkit.PlotConfig`. Uses defaults if omitted.
output_dir:
Directory figures are written under.
"""
config = PlotConfig.load(config)
sample_column = resolve_sample_column(adata, sample_column)
require_columns(adata, level_column, sample_column, subset_level)
if subset_level and subset_value is not None:
subset_value = [subset_value] if isinstance(subset_value, str) else subset_value
adata = adata[adata.obs[subset_level].isin(subset_value), :]
counts = pd.crosstab(adata.obs[level_column], adata.obs[sample_column])
counts = counts.loc[
counts.sum(axis=1).sort_values(ascending=False).index,
counts.sum(axis=0).sort_values(ascending=False).index,
]
n_rows, n_cols = counts.shape
max_count = int(counts.values.max())
def size_for(count: int) -> float:
return min_size + (max_size - min_size) * (count / max_count) if max_count > 0 else min_size
xs, ys, sizes, zero_xs, zero_ys = [], [], [], [], []
for yi, cell_type in enumerate(counts.index):
for xi, sample in enumerate(counts.columns):
count = counts.loc[cell_type, sample]
if count > 0:
xs.append(xi)
ys.append(yi)
sizes.append(size_for(count))
else:
zero_xs.append(xi)
zero_ys.append(yi)
general = config.general
figsize = figsize or (max(6.0, n_cols * 0.45 + 3.0), max(4.0, n_rows * 0.4 + 2.0))
fig, ax = plt.subplots(figsize=figsize)
ax.set_axisbelow(True)
ax.grid(True, color="#e5e5e5", linewidth=0.6, zorder=0)
zero_color = general.get("fallback_color", "lightgrey")
if zero_xs:
ax.scatter(zero_xs, zero_ys, s=min_size, facecolors="none", edgecolors=zero_color, linewidths=0.9, zorder=2)
ax.scatter(xs, ys, s=sizes, color=dot_color, alpha=0.85, edgecolors="white", linewidths=0.6, zorder=3)
ax.set_xlim(-0.6, n_cols - 0.4)
ax.set_ylim(-0.6, n_rows - 0.4)
ax.set_xticks(range(n_cols))
ax.set_xticklabels(counts.columns, rotation=90)
ax.set_yticks(range(n_rows))
ax.set_yticklabels(counts.index)
ax.invert_yaxis() # most abundant cell type at the top
for spine in ax.spines.values():
spine.set_visible(False)
ax.set_xlabel(xlabel or sample_column.replace("_", " ").title(), fontsize=12)
ax.set_ylabel(level_column.replace("_", " ").title(), fontsize=12)
ax.set_title(
title or "Cell Type Abundance Across Samples",
fontsize=general["title_fontsize"], fontweight=general["title_fontweight"], pad=20,
)
legend_values = _size_legend_values(max_count)
legend_handles = [
plt.scatter([], [], s=size_for(v), color=dot_color, alpha=0.85, edgecolors="white", linewidths=0.6)
for v in legend_values
] + [plt.scatter([], [], s=min_size, facecolors="none", edgecolors=zero_color, linewidths=0.9)]
legend_labels = [str(v) for v in legend_values] + ["0"]
legend1 = fig.legend(
legend_handles, legend_labels, bbox_to_anchor=(1.01, 1), bbox_transform=ax.transAxes, loc="upper left",
title="Cell Count", title_fontproperties=font_manager.FontProperties(weight="bold"),
fontsize=general["legend_fontsize"], frameon=False, labelspacing=1.4,
)
fig.subplots_adjust(right=0.78)
_fit_legends_in_figure(fig, ax, [legend1])
if save_name is not None:
save_figure(fig, output_dir, "abundance_bubble_grid", f"{save_name}.png", config, include_legends=True)
return fig
