"""Dataset overview plots: sample/cell counts, abundance, and comparisons."""
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from matplotlib import gridspec
from matplotlib import patches as mpatches
from ._config import PlotConfig
from ._palettes import auto_palette, resolve_palette
from ._style import save_figure
from ._utils import require_columns
__all__ = [
"sample_and_cell_counts_barplot",
"sample_and_cell_counts_barplot_break_axis",
"cells_per_patient_boxplot",
"cell_abundance_barplot",
"compare_cell_abundance_barplot",
"compare_cell_abundance_boxplot",
]
_POINT_MARKERS = ["o", "X", "s", "^", "D", "P", "*"]
def _category_colors(config: PlotConfig, column: str, categories: Sequence, palette: dict | None) -> dict:
"""Explicit/configured palette if available, else the original pastel default."""
if palette or config.palette_for(column):
return resolve_palette(config, column, categories, palette)
return dict(zip(categories, sns.color_palette("pastel", len(categories))))
def _style_axes(ax, general):
ax.tick_params(axis="x", rotation=90)
ax.tick_params(axis="both", width=1.8)
for sp in ax.spines.values():
sp.set_linewidth(1.8)
ax.grid(False)
[docs]
def sample_and_cell_counts_barplot(
adata,
level_column: str,
sample_column: str = "Sample_ID",
figsize: tuple = (10, 5),
title: str | None = None,
save_name: str | None = None,
palette: dict | None = None,
xlabel: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Side-by-side bars of samples-per-category and cells-per-category.
Returns ``(fig, fig_grey)``: a colored version and a grayscale twin
(handy for figures that need to stay monochrome for print).
"""
config = PlotConfig.load(config)
require_columns(adata, level_column, sample_column)
general = config.general
sample_counts = adata.obs.groupby(level_column, observed=True)[sample_column].nunique()
cell_counts = adata.obs[level_column].value_counts()
sorted_categories = sample_counts.sort_values(ascending=False).index.tolist()
sample_counts = sample_counts.reindex(sorted_categories)
cell_counts = cell_counts.reindex(sorted_categories)
color_dict = _category_colors(config, level_column, sorted_categories, palette)
colors = [color_dict[c] for c in sorted_categories]
xlabel = xlabel or level_column.replace("_", " ").title()
def _build(bar_colors):
fig = plt.figure(figsize=figsize)
gs = gridspec.GridSpec(2, 2, height_ratios=[1, 2])
ax = fig.add_subplot(gs[:, 0])
ax.bar(sorted_categories, sample_counts.values, color=bar_colors)
ax.set_title("Samples per Category", fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
ax.set_ylabel("Number of Samples", fontsize=12)
ax.set_xlabel(xlabel, fontsize=12, labelpad=10)
_style_axes(ax, general)
ax = fig.add_subplot(gs[:, 1])
ax.bar(sorted_categories, cell_counts.values, color=bar_colors)
ax.set_title("Cells per Category", fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
ax.set_ylabel("Number of Cells", fontsize=12)
ax.set_xlabel(xlabel, fontsize=12, labelpad=10)
_style_axes(ax, general)
plt.tight_layout(rect=[0, 0, 1, 0.95])
return fig
fig = _build(colors)
fig_grey = _build("grey")
if save_name is not None:
save_figure(fig, output_dir, "sample_cell_counts", f"{save_name}.png", config)
save_figure(fig_grey, output_dir, "sample_cell_counts", f"{save_name}_all_grey.png", config)
return fig, fig_grey
[docs]
def sample_and_cell_counts_barplot_break_axis(
adata,
level_column: str,
sample_column: str = "Sample_ID",
figsize: tuple = (5, 5),
title: str | None = None,
save_name: str | None = None,
palette: dict | None = None,
xlabel: str | None = None,
break_point: float | None = None,
break_ratio: float = 0.3,
break_gap: float = 0.02,
plot_type: str = "samples",
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Like :func:`sample_and_cell_counts_barplot`, with a broken y-axis so a
few outlier categories don't compress the rest of the bars.
``plot_type`` selects between ``"samples"`` (unique sample count per
category) and ``"cells"`` (total cell count per category).
"""
config = PlotConfig.load(config)
require_columns(adata, level_column, sample_column)
if plot_type not in ("samples", "cells"):
raise ValueError("plot_type must be either 'samples' or 'cells'")
general = config.general
if plot_type == "samples":
counts = adata.obs.groupby(level_column, observed=True)[sample_column].nunique()
ylabel, plot_title = "Number of Samples", "Samples per Category"
else:
counts = adata.obs[level_column].value_counts()
ylabel, plot_title = "Number of Cells", "Cells per Category"
sorted_categories = counts.sort_values(ascending=False).index.tolist()
counts = counts.reindex(sorted_categories)
if break_point is None:
break_point = np.percentile(counts.values, 75)
color_dict = _category_colors(config, level_column, sorted_categories, palette)
colors = [color_dict[c] for c in sorted_categories]
xlabel = xlabel or level_column.replace("_", " ").title()
def _build(bar_colors):
fig = plt.figure(figsize=figsize)
gs = gridspec.GridSpec(2, 1, height_ratios=[break_ratio, 1 - break_ratio], hspace=break_gap)
ax_upper = fig.add_subplot(gs[0])
ax_lower = fig.add_subplot(gs[1], sharex=ax_upper)
ax_upper.bar(sorted_categories, counts.values, color=bar_colors)
ax_lower.bar(sorted_categories, counts.values, color=bar_colors)
max_val = np.max(counts.values)
ax_upper.set_ylim(break_point, max_val * 1.15)
ax_lower.set_ylim(0, break_point * 0.9)
ax_upper.spines["bottom"].set_visible(False)
ax_lower.spines["top"].set_visible(False)
ax_upper.tick_params(bottom=False, top=False, labelbottom=False, labeltop=False)
ax_lower.xaxis.tick_bottom()
d = 0.5
kwargs = dict(
marker=[(-1, -d), (1, d)], markersize=8, linestyle="none", color="k", mec="k", mew=1, clip_on=False
)
ax_upper.plot([0, 1], [0, 0], transform=ax_upper.transAxes, **kwargs)
ax_lower.plot([0, 1], [1, 1], transform=ax_lower.transAxes, **kwargs)
ax_upper.set_title(
title or plot_title, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"]
)
ax_lower.set_ylabel(ylabel, fontsize=12)
ax_lower.set_xlabel(xlabel, fontsize=12, labelpad=10)
for ax in (ax_upper, ax_lower):
_style_axes(ax, general)
plt.tight_layout()
return fig
fig = _build(colors)
fig_grey = _build("grey")
if save_name is not None:
save_figure(fig, output_dir, "sample_cell_counts", f"{save_name}_{plot_type}_break_axis.png", config)
save_figure(fig_grey, output_dir, "sample_cell_counts", f"{save_name}_{plot_type}_break_axis_grey.png", config)
return fig, fig_grey
[docs]
def cells_per_patient_boxplot(
adata,
level_column: str,
sample_column: str = "Sample_ID",
figsize: tuple = (10, 6),
title: str | None = None,
save_name: str | None = None,
palette: dict | None = None,
xlabel: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Boxplot of cells-per-patient distribution, grouped by ``level_column``.
Returns ``(fig, fig_grey)`` -- a colored version and a grayscale twin.
"""
config = PlotConfig.load(config)
require_columns(adata, level_column, sample_column)
general = config.general
cells_per_patient = (
adata.obs.groupby(sample_column, observed=True)
.agg({level_column: "first", sample_column: "size"})
.rename(columns={sample_column: "cell_count"})
)
avg_cells = cells_per_patient.groupby(level_column, observed=True)["cell_count"].mean()
sorted_categories = avg_cells.sort_values(ascending=False).index.tolist()
color_dict = _category_colors(config, level_column, sorted_categories, palette)
colors = [color_dict[c] for c in sorted_categories]
xlabel = xlabel or level_column.replace("_", " ").title()
plot_title = title or "Cells per Patient by Category"
data_for_boxplot = [
cells_per_patient[cells_per_patient[level_column] == c]["cell_count"].values for c in sorted_categories
]
def _build(box_colors):
fig, ax = plt.subplots(figsize=figsize)
bp = ax.boxplot(
data_for_boxplot, patch_artist=True,
boxprops=dict(linewidth=1.8), whiskerprops=dict(linewidth=1.8),
capprops=dict(linewidth=1.8), medianprops=dict(linewidth=1.8, color="black"),
)
ax.set_xticks(range(1, len(sorted_categories) + 1))
ax.set_xticklabels(sorted_categories)
for patch, color in zip(bp["boxes"], box_colors):
patch.set_facecolor(color)
patch.set_alpha(0.7)
ax.set_title(plot_title, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
ax.set_ylabel("Number of Cells per Patient", fontsize=12)
ax.set_xlabel(xlabel, fontsize=12, labelpad=10)
ax.tick_params(axis="x", rotation=45)
ax.tick_params(axis="both", width=1.8)
for sp in ax.spines.values():
sp.set_linewidth(1.8)
ax.grid(False)
plt.tight_layout()
return fig
fig = _build(colors)
fig_grey = _build(["grey"] * len(sorted_categories))
if save_name is not None:
save_figure(fig, output_dir, "cells_per_patient_boxplot", f"{save_name}.png", config)
save_figure(fig_grey, output_dir, "cells_per_patient_boxplot", f"{save_name}_all_grey.png", config)
return fig, fig_grey
[docs]
def cell_abundance_barplot(
adata,
cell_type_column: str,
figsize: tuple = (8, 5),
title_suffix: str | None = None,
save_name: str | None = None,
color: str = "grey",
xlabel: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Absolute and relative cell-count bar plots for one column.
Returns ``(fig_counts, fig_relative)``.
"""
config = PlotConfig.load(config)
require_columns(adata, cell_type_column)
general = config.general
cell_counts = adata.obs[cell_type_column].value_counts()
relative_abundance = (cell_counts / cell_counts.sum()) * 100
sorted_categories = cell_counts.sort_values(ascending=False).index.tolist()
cell_counts = cell_counts.reindex(sorted_categories)
relative_abundance = relative_abundance.reindex(sorted_categories)
x_label = xlabel or cell_type_column.replace("_", " ").title()
title_suffix = title_suffix or ""
fig_counts, ax_counts = plt.subplots(1, 1, figsize=figsize)
ax_counts.bar(cell_counts.index, cell_counts.values, color=color)
ax_counts.set_title(
f"Absolute Cell Abundance{title_suffix}",
fontsize=general["title_fontsize"], fontweight=general["title_fontweight"],
)
ax_counts.set_ylabel("Number of Cells", fontsize=12)
ax_counts.set_xlabel(x_label, fontsize=12, labelpad=10)
_style_axes(ax_counts, general)
plt.tight_layout()
fig_relative, ax_relative = plt.subplots(1, 1, figsize=figsize)
ax_relative.bar(relative_abundance.index, relative_abundance.values, color=color)
ax_relative.set_title(
f"Relative Cell Abundance{title_suffix}",
fontsize=general["title_fontsize"], fontweight=general["title_fontweight"],
)
ax_relative.set_ylabel("Percentage of Total Cells (%)", fontsize=12)
ax_relative.set_xlabel(x_label, fontsize=12, labelpad=10)
_style_axes(ax_relative, general)
plt.tight_layout()
if save_name is not None:
save_figure(fig_counts, output_dir, "cell_abundance_barplot", f"{save_name}_absolute_counts.png", config)
save_figure(fig_relative, output_dir, "cell_abundance_barplot", f"{save_name}_relative_abundance.png", config)
return fig_counts, fig_relative
[docs]
def compare_cell_abundance_barplot(
adata1,
adata2,
cell_type_column: str,
label1: str = "Dataset 1",
label2: str = "Dataset 2",
color1: str = "lightgray",
color2: str = "dimgray",
figsize: tuple = (12, 6),
title_suffix: str | None = None,
save_name: str | None = None,
xlabel: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Grouped bar plots comparing cell-type abundance between two datasets.
Returns ``(fig_counts, fig_relative)``.
"""
config = PlotConfig.load(config)
require_columns(adata1, cell_type_column)
require_columns(adata2, cell_type_column)
general = config.general
counts1 = adata1.obs[cell_type_column].value_counts()
counts2 = adata2.obs[cell_type_column].value_counts()
relative1 = (counts1 / counts1.sum()) * 100
relative2 = (counts2 / counts2.sum()) * 100
all_cell_types = counts1.index.union(counts2.index)
counts1 = counts1.reindex(all_cell_types, fill_value=0)
counts2 = counts2.reindex(all_cell_types, fill_value=0)
relative1 = relative1.reindex(all_cell_types, fill_value=0)
relative2 = relative2.reindex(all_cell_types, fill_value=0)
sorted_categories = (counts1 + counts2).sort_values(ascending=False).index.tolist()
counts1, counts2 = counts1.reindex(sorted_categories), counts2.reindex(sorted_categories)
relative1, relative2 = relative1.reindex(sorted_categories), relative2.reindex(sorted_categories)
x_label = xlabel or cell_type_column.replace("_", " ").title()
title_suffix = title_suffix or ""
x = np.arange(len(sorted_categories))
width = 0.35
def _grouped(ax, v1, v2, ylabel, title):
ax.bar(x - width / 2, v1, width, label=label1, color=color1)
ax.bar(x + width / 2, v2, width, label=label2, color=color2)
ax.set_title(title, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
ax.set_ylabel(ylabel, fontsize=12)
ax.set_xlabel(x_label, fontsize=12, labelpad=10)
ax.set_xticks(x, sorted_categories, rotation=90)
ax.tick_params(axis="both", width=1.8)
for sp in ax.spines.values():
sp.set_linewidth(1.8)
ax.legend(loc="upper right")
ax.grid(False)
fig_counts, ax_counts = plt.subplots(1, 1, figsize=figsize)
_grouped(
ax_counts, counts1.values, counts2.values,
"Number of Cells", f"Absolute Cell Abundance Comparison{title_suffix}",
)
plt.tight_layout()
fig_relative, ax_relative = plt.subplots(1, 1, figsize=figsize)
_grouped(
ax_relative, relative1.values, relative2.values,
"Percentage of Total Cells (%)", f"Relative Cell Abundance Comparison{title_suffix}",
)
plt.tight_layout()
if save_name is not None:
save_figure(
fig_counts, output_dir, "cell_abundance_barplot", f"{save_name}_absolute_counts_comparison.png", config
)
save_figure(
fig_relative, output_dir, "cell_abundance_barplot", f"{save_name}_relative_abundance_comparison.png", config
)
return fig_counts, fig_relative
def _draw_box(ax, x, vals, color, box_w, point_markers, rng):
if point_markers is None:
point_markers = ["o"] * len(vals)
if len(vals) < 2:
ax.scatter([x], vals, marker=point_markers[0], color=color, edgecolors="black", s=40, linewidths=0.8, zorder=5)
return
q1, med, q3 = np.percentile(vals, [25, 50, 75])
iqr = q3 - q1
lo_fence, hi_fence = q1 - 1.5 * iqr, q3 + 1.5 * iqr
in_lo, in_hi = vals[vals >= lo_fence], vals[vals <= hi_fence]
wlo = in_lo.min() if len(in_lo) > 0 else vals.min()
whi = in_hi.max() if len(in_hi) > 0 else vals.max()
cap_w = box_w * 0.45
ax.add_patch(
mpatches.Rectangle((x - box_w / 2, q1), box_w, q3 - q1, linewidth=1.5, edgecolor="black", facecolor=color,
alpha=0.75, zorder=2)
)
ax.hlines(med, x - box_w / 2, x + box_w / 2, color="black", linewidth=2, zorder=4)
ax.vlines(x, wlo, q1, color="black", linewidth=1.5, zorder=2)
ax.vlines(x, q3, whi, color="black", linewidth=1.5, zorder=2)
ax.hlines(wlo, x - cap_w / 2, x + cap_w / 2, color="black", linewidth=1.5, zorder=2)
ax.hlines(whi, x - cap_w / 2, x + cap_w / 2, color="black", linewidth=1.5, zorder=2)
jitter = rng.uniform(-0.08, 0.08, size=len(vals))
for m in dict.fromkeys(point_markers):
idx = [j for j, pm in enumerate(point_markers) if pm == m]
ax.scatter(x + jitter[idx], vals[idx], marker=m, color=color, edgecolors="black", s=35, linewidths=0.7,
zorder=5, alpha=0.9)
[docs]
def compare_cell_abundance_boxplot(
adata1,
adata2,
cell_type_column: str,
sample_column: str = "Sample_ID",
label1: str = "Dataset 1",
label2: str = "Dataset 2",
colored: bool = False,
color1: str | None = None,
color2: str | None = None,
figsize: tuple = (12, 5),
title: str | None = None,
save_name: str | None = None,
high_dif: int = 0,
marker_column: str | None = None,
groups_of_interest: Sequence | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Paired per-sample boxplots comparing relative cell-type abundance between two datasets.
For each cell type, computes per-sample relative abundance (%) in each
dataset and draws paired boxplots with individual sample dots overlaid --
unlike :func:`compare_cell_abundance_barplot`, this shows the full
per-sample distribution rather than just a single dataset-wide bar.
Parameters
----------
adata1, adata2:
Datasets to compare.
cell_type_column, sample_column:
Columns in both datasets' ``.obs``.
colored:
Use two distinct auto-generated colors for ``label1``/``label2``
instead of the default grayscale (``"lightgray"``/``"dimgray"``).
Ignored if ``color1``/``color2`` are given explicitly.
high_dif:
If > 0, show only the top N cell types by absolute difference in mean
relative abundance. ``0`` shows all.
marker_column:
Optional ``.obs`` column (present in both datasets) whose unique
values map to different point marker shapes.
groups_of_interest:
If given, restrict (and order) the plot to these cell types instead
of using ``high_dif``.
"""
config = PlotConfig.load(config)
require_columns(adata1, cell_type_column, sample_column, marker_column)
require_columns(adata2, cell_type_column, sample_column, marker_column)
general = config.general
if color1 is None or color2 is None:
if colored:
auto_colors = auto_palette([label1, label2])
color1 = color1 or auto_colors[label1]
color2 = color2 or auto_colors[label2]
else:
color1 = color1 or "lightgray"
color2 = color2 or "dimgray"
def _sample_props(adata):
return adata.obs.groupby(sample_column, observed=True)[cell_type_column].value_counts(normalize=True).mul(
100
).unstack(fill_value=0)
props1, props2 = _sample_props(adata1), _sample_props(adata2)
all_cats = props1.columns.union(props2.columns)
props1 = props1.reindex(columns=all_cats, fill_value=0)
props2 = props2.reindex(columns=all_cats, fill_value=0)
sorted_cats = (props1.mean() + props2.mean()).sort_values(ascending=False).index.tolist()
if groups_of_interest is not None:
missing = [g for g in groups_of_interest if g not in all_cats]
if missing:
print(f"Warning: the following groups were not found and will be skipped: {missing}")
sorted_cats = [g for g in groups_of_interest if g in all_cats]
elif high_dif > 0:
abs_diff = (props1.mean() - props2.mean()).abs().reindex(sorted_cats)
top_cats = set(abs_diff.nlargest(high_dif).index)
sorted_cats = [c for c in sorted_cats if c in top_cats]
if not sorted_cats:
raise ValueError("No cell types available to plot after applying groups_of_interest/high_dif filtering.")
props1, props2 = props1[sorted_cats], props2[sorted_cats]
marker_val_map = {}
markers1 = markers2 = None
if marker_column is not None:
sm1 = adata1.obs.groupby(sample_column, observed=True)[marker_column].first()
sm2 = adata2.obs.groupby(sample_column, observed=True)[marker_column].first()
all_marker_vals = sorted(set(sm1.tolist()) | set(sm2.tolist()), key=str)
marker_val_map = {v: _POINT_MARKERS[i % len(_POINT_MARKERS)] for i, v in enumerate(all_marker_vals)}
markers1 = [marker_val_map.get(sm1.get(s), "o") for s in props1.index]
markers2 = [marker_val_map.get(sm2.get(s), "o") for s in props2.index]
fig, ax = plt.subplots(figsize=figsize)
n_cats = len(sorted_cats)
spacing, offset, box_w = 1.0, 0.22, 0.35
rng = np.random.default_rng(seed=42)
group_centers = []
for i, cat in enumerate(sorted_cats):
cx = i * spacing
group_centers.append(cx)
_draw_box(ax, cx - offset, props1[cat].values, color1, box_w, markers1, rng)
_draw_box(ax, cx + offset, props2[cat].values, color2, box_w, markers2, rng)
ax.set_xticks(group_centers)
ax.set_xticklabels([str(c).replace("_", " ") for c in sorted_cats], rotation=45, ha="right",
fontsize=general["legend_fontsize"])
ax.set_ylabel("Relative abundance (%)", fontsize=12)
ax.set_title(title or f"Cell type abundance: {label1} vs {label2}", fontsize=general["title_fontsize"],
fontweight=general["title_fontweight"], pad=10)
ax.set_xlim(-0.6, (n_cats - 1) * spacing + 0.6)
ax.set_ylim(bottom=0)
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
for sp in ("left", "bottom"):
ax.spines[sp].set_linewidth(1.8)
ax.tick_params(axis="both", width=1.8, length=5)
ax.grid(False)
legend_handles = [
mpatches.Patch(facecolor=color1, edgecolor="black", linewidth=1.2, label=label1),
mpatches.Patch(facecolor=color2, edgecolor="black", linewidth=1.2, label=label2),
]
if marker_column and marker_val_map:
from matplotlib.lines import Line2D
legend_handles.append(Line2D([], [], linestyle="none", label=""))
legend_handles += [
Line2D([0], [0], marker=msym, linestyle="none", color="grey", markeredgecolor="black", markersize=5,
markeredgewidth=0.7, label=str(val).replace("_", " "))
for val, msym in marker_val_map.items()
]
ax.legend(handles=legend_handles, frameon=False, fontsize=general["legend_fontsize"], loc="upper right")
plt.tight_layout()
if save_name is not None:
save_figure(fig, output_dir, "cell_abundance_barplot", f"{save_name}.png", config)
return fig
