"""Pseudobulk plots: per-sample aggregated gene expression across groups or genes.
Each sample's expression of a gene is first collapsed to a single mean value
(the "pseudobulk"), then compared across groups with boxplots and individual
sample points overlaid -- the standard way to avoid pseudo-replication when
comparing conditions across samples with many cells each.
"""
from __future__ import annotations
from collections.abc import Sequence
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import patches as mpatches
from matplotlib.lines import Line2D
from scipy.sparse import issparse
from ._config import PlotConfig
from ._palettes import resolve_palette
from ._style import save_figure
from ._utils import require_columns
__all__ = ["pseudobulk_boxplot", "pseudobulk_multigene_boxplot"]
_POINT_MARKERS = ["o", "X", "s", "^", "D", "P", "*"]
def _expr_values(adata, gene: str, layer: str) -> np.ndarray:
if gene not in adata.var_names:
raise ValueError(f"{gene} not in adata.var_names")
if layer not in adata.layers:
raise ValueError(f"Layer '{layer}' not in adata.layers")
expr = adata[:, gene].layers[layer]
return expr.toarray().flatten() if issparse(expr) else np.asarray(expr).flatten()
def _subset_celltype(adata, celltype_column, celltype_of_interest):
if celltype_of_interest is None:
return adata
require_columns(adata, celltype_column)
values = [celltype_of_interest] if isinstance(celltype_of_interest, str) else celltype_of_interest
subset = adata[adata.obs[celltype_column].isin(values)]
if subset.n_obs == 0:
raise ValueError("No cells left after subsetting on celltype_of_interest.")
return subset
[docs]
def pseudobulk_boxplot(
adata,
gene: str,
layer: str,
sample_column: str = "Sample_ID",
group_by: str | None = None,
celltype_column: str | None = None,
celltype_of_interest: str | Sequence | None = None,
palette: dict | None = None,
marker_column: str | None = None,
figsize: tuple = (4, 5),
title: str | None = None,
save_name: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Boxplot of per-sample pseudobulk (mean) expression of one gene.
Computes each sample's mean ``gene`` expression from ``layer``, then draws
one box per ``group_by`` value (or a single box if omitted) with
individual sample points jittered on top. Point shape can additionally be
mapped to another column via ``marker_column``.
"""
config = PlotConfig.load(config)
require_columns(adata, sample_column, group_by, marker_column)
general = config.general
adata = _subset_celltype(adata, celltype_column, celltype_of_interest)
expr = _expr_values(adata, gene, layer)
marker_val_map = {}
sample_marker = None
if marker_column is not None:
sample_marker = adata.obs.groupby(sample_column, observed=True)[marker_column].first()
marker_val_map = {
v: _POINT_MARKERS[i % len(_POINT_MARKERS)]
for i, v in enumerate(sorted(sample_marker.unique().tolist(), key=str))
}
obs_cols = [sample_column] + ([group_by] if group_by else [])
df = adata.obs[obs_cols].copy()
df["_expr"] = expr
pb = df.groupby(obs_cols, observed=True)["_expr"].mean().reset_index().rename(columns={"_expr": "mean_expr"})
if sample_marker is not None:
pb["_marker_val"] = pb[sample_column].map(sample_marker)
if group_by:
groups = pb[group_by].unique().tolist()
pb["group"] = pb[group_by]
else:
groups = [gene]
pb["group"] = gene
color_map = resolve_palette(config, group_by or gene, groups, palette)
fig, ax = plt.subplots(figsize=figsize)
sns.boxplot(
data=pb, x="group", y="mean_expr", order=groups, hue="group", palette=color_map,
width=0.5, linewidth=1.5, fliersize=0, legend=False, ax=ax,
)
if sample_marker is not None:
for mval, mgroup in pb.groupby("_marker_val", observed=True):
sns.stripplot(
data=mgroup, x="group", y="mean_expr", order=groups, hue="group", palette=color_map,
jitter=0.12, size=6, edgecolor="black", linewidth=0.8, ax=ax, zorder=3, legend=False,
marker=marker_val_map[mval],
)
else:
sns.stripplot(
data=pb, x="group", y="mean_expr", order=groups, hue="group", palette=color_map,
jitter=0.12, size=6, edgecolor="black", linewidth=0.8, ax=ax, zorder=3, legend=False,
)
ax.set_xticks(range(len(groups)))
ax.set_xticklabels(
[str(g).replace("_", " ") for g in groups],
rotation=45 if len(groups) > 1 else 0, ha="right" if len(groups) > 1 else "center",
fontsize=general["legend_fontsize"],
)
ax.set_xlabel(group_by.replace("_", " ").title() if group_by else "", fontsize=12)
ax.set_ylabel(f"{gene} mean expression\n({layer})", fontsize=12)
ax.set_title(title or gene, fontsize=general["title_fontsize"], fontweight=general["title_fontweight"])
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(width=1.8, length=5)
ax.grid(False)
if marker_column and marker_val_map:
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, title=marker_column.replace("_", " ").title(),
frameon=False, fontsize=general["legend_fontsize"],
)
fig.tight_layout()
if save_name is not None:
save_figure(fig, output_dir, "pseudobulk", f"{save_name}_box.png", config)
return fig
def _draw_manual_box(ax, x, vals, color, box_w, rng):
if len(vals) < 2:
ax.scatter([x], vals, 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
# Candidates for the whisker ends must also lie within [q1, q3], otherwise a value
# excluded as a fence outlier on one side can leave the *nearest remaining* point on
# the wrong side of q1/q3 (since q1/q3 are interpolated, not necessarily actual data
# points) -- producing a whisker that runs backwards into the box instead of away from it.
in_lo = vals[(vals >= lo_fence) & (vals <= q1)]
in_hi = vals[(vals <= hi_fence) & (vals >= q3)]
wlo = in_lo.min() if len(in_lo) > 0 else q1
whi = in_hi.max() if len(in_hi) > 0 else q3
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.12, 0.12, size=len(vals))
ax.scatter(x + jitter, vals, color=color, edgecolors="black", s=40, linewidths=0.8, alpha=0.9, zorder=3)
[docs]
def pseudobulk_multigene_boxplot(
adata,
genes: Sequence[str],
layer: str,
color_column: str,
sample_column: str = "Sample_ID",
marker_column: str | None = None,
celltype_column: str | None = None,
celltype_of_interest: str | Sequence | None = None,
palette: dict | None = None,
groups_order: Sequence | None = None,
min_max_scale: bool = True,
figsize: tuple = (12, 5),
title: str | None = None,
save_name: str | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Pseudobulk expression boxplot for multiple genes, optionally min-max scaled per gene.
For each gene, computes per-sample mean expression, scaled to ``[0, 1]``
by default (set ``min_max_scale=False`` to plot raw mean expression
instead). Draws one boxplot per ``color_column`` group per gene (side by
side), with individual sample points overlaid; point shape can be mapped
to ``marker_column``.
If ``save_name`` is given, also exports ``{save_name}_raw.csv`` (unscaled
per-sample mean expression per gene) and, when ``min_max_scale`` is
``True``, ``{save_name}_minmax_scaled.csv`` (the scaled values used for
the plot), both wide-format with one row per sample.
"""
config = PlotConfig.load(config)
require_columns(adata, sample_column, color_column, marker_column)
missing_genes = [g for g in genes if g not in adata.var_names]
if missing_genes:
raise ValueError(f"Genes not in adata.var_names: {missing_genes}")
if layer not in adata.layers:
raise ValueError(f"Layer '{layer}' not in adata.layers.")
general = config.general
adata = _subset_celltype(adata, celltype_column, celltype_of_interest)
meta_cols = [sample_column, color_column] + ([marker_column] if marker_column else [])
sample_meta = adata.obs[meta_cols].drop_duplicates(subset=[sample_column]).set_index(sample_column)
groups = list(groups_order) if groups_order is not None else list(dict.fromkeys(sample_meta[color_column]))
color_map = resolve_palette(config, color_column, groups, palette)
marker_val_map = {}
if marker_column:
unique_marker_vals = sorted(sample_meta[marker_column].unique().tolist(), key=str)
marker_val_map = {v: _POINT_MARKERS[i % len(_POINT_MARKERS)] for i, v in enumerate(unique_marker_vals)}
pb_plot, pb_raw = {}, {}
for gene in genes:
expr = _expr_values(adata, gene, layer)
df_gene = adata.obs[[sample_column]].copy()
df_gene["_expr"] = expr
mean_per_sample = df_gene.groupby(sample_column, observed=True)["_expr"].mean()
pb_raw[gene] = mean_per_sample
if min_max_scale:
mn, mx = mean_per_sample.min(), mean_per_sample.max()
pb_plot[gene] = (mean_per_sample - mn) / (mx - mn) if mx > mn else mean_per_sample * 0.0
else:
pb_plot[gene] = mean_per_sample
fig, ax = plt.subplots(figsize=figsize)
n_groups = len(groups)
spacing, box_w = 1.0, 0.3
total_width = box_w * n_groups + 0.05 * (n_groups - 1)
offsets = np.linspace(-total_width / 2 + box_w / 2, total_width / 2 - box_w / 2, n_groups)
rng = np.random.default_rng(seed=42)
group_centers = []
for i, gene in enumerate(genes):
cx = i * spacing
group_centers.append(cx)
plot_vals = pb_plot[gene]
for gi, grp in enumerate(groups):
grp_samples = sample_meta[sample_meta[color_column] == grp].index
vals = plot_vals.reindex(grp_samples).dropna().values
if len(vals) == 0:
continue
if marker_column:
point_markers = [
marker_val_map.get(sample_meta.loc[s, marker_column], "o")
for s in grp_samples if s in plot_vals.index and not np.isnan(plot_vals[s])
]
jitter = rng.uniform(-0.06, 0.06, size=len(vals))
for m in dict.fromkeys(point_markers):
idx = [j for j, pm in enumerate(point_markers) if pm == m]
ax.scatter(
cx + offsets[gi] + jitter[idx], vals[idx], marker=m, color=color_map[grp],
edgecolors="black", s=35, linewidths=0.7, zorder=5, alpha=0.9,
)
_draw_box_frame(ax, cx + offsets[gi], vals, color_map[grp], box_w)
else:
_draw_manual_box(ax, cx + offsets[gi], vals, color_map[grp], box_w, rng)
ax.set_xticks(group_centers)
ax.set_xticklabels([g.replace("_", " ") for g in genes], rotation=45, ha="right",
fontsize=general["legend_fontsize"])
if min_max_scale:
ax.set_ylabel("Scaled expression (min-max)", fontsize=12)
ax.set_ylim(-0.05, 1.1)
else:
ax.set_ylabel(f"Mean expression\n({layer})", fontsize=12)
all_vals = np.concatenate([pb_plot[g].values for g in genes]) if genes else np.array([0.0])
lo, hi = np.nanmin(all_vals), np.nanmax(all_vals)
pad = (hi - lo) * 0.05 or 1.0
ax.set_ylim(lo - pad, hi + pad)
ax.set_xlim(-0.6, (len(genes) - 1) * spacing + 0.6)
ax.set_title(title or "Pseudobulk gene expression", fontsize=general["title_fontsize"],
fontweight=general["title_fontweight"], pad=10)
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)
# Fit labels/ticks/title first, then shrink the axes to leave room on the
# right for the legend(s) -- they're anchored outside the axes via
# bbox_to_anchor, so a later tight_layout() would not account for them
# and a plain (non-tight-bbox) render, e.g. inline in a notebook, would
# clip them off the canvas.
fig.tight_layout()
fig.subplots_adjust(right=0.78)
color_handles = [
mpatches.Patch(facecolor=color_map[v], edgecolor="black", linewidth=1.0, label=str(v).replace("_", " "))
for v in groups
]
legend1 = ax.legend(
handles=color_handles, title=color_column.replace("_", " ").title(),
frameon=False, fontsize=general["legend_fontsize"], bbox_to_anchor=(1.02, 1), loc="upper left",
)
legends = [legend1]
if marker_column and marker_val_map:
ax.add_artist(legend1)
marker_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()
]
legend2 = ax.legend(
handles=marker_handles, title=marker_column.replace("_", " ").title(),
frameon=False, fontsize=general["legend_fontsize"], bbox_to_anchor=(1.02, 0), loc="lower left",
)
legends.append(legend2)
# The 0.78 above is just a starting guess -- how much width a legend actually
# needs depends on label length and font size, neither of which tight_layout()
# accounts for (it ignores artists anchored outside the axes). Measure the
# widest legend after a real draw and reserve exactly that much figure space,
# so labels never get clipped regardless of figsize or how long they are.
fig.canvas.draw()
renderer = fig.canvas.get_renderer()
max_legend_width_in = max(leg.get_window_extent(renderer).width / fig.dpi for leg in legends)
right_frac = max(0.4, 1 - max_legend_width_in / fig.get_size_inches()[0] - 0.03)
fig.subplots_adjust(right=right_frac)
if save_name is not None:
save_figure(fig, output_dir, "pseudobulk", f"{save_name}.png", config, include_legends=True)
save_dir = Path(output_dir) / "pseudobulk"
save_dir.mkdir(parents=True, exist_ok=True)
export_meta = sample_meta[[color_column] + ([marker_column] if marker_column else [])]
raw_df = export_meta.join(pd.DataFrame(pb_raw), how="right")
raw_df.index.name = sample_column
raw_df.reset_index().to_csv(save_dir / f"{save_name}_raw.csv", index=False)
if min_max_scale:
scaled_df = export_meta.join(pd.DataFrame(pb_plot), how="right")
scaled_df.index.name = sample_column
scaled_df.reset_index().to_csv(save_dir / f"{save_name}_minmax_scaled.csv", index=False)
return fig
def _draw_box_frame(ax, x, vals, color, box_w):
"""Draw the box/whisker frame filled with ``color`` (points are drawn separately per-marker)."""
if len(vals) < 2:
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 = vals[(vals >= lo_fence) & (vals <= q1)]
in_hi = vals[(vals <= hi_fence) & (vals >= q3)]
wlo = in_lo.min() if len(in_lo) > 0 else q1
whi = in_hi.max() if len(in_hi) > 0 else q3
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=1)
)
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=1)
ax.vlines(x, q3, whi, color="black", linewidth=1.5, zorder=1)
ax.hlines(wlo, x - cap_w / 2, x + cap_w / 2, color="black", linewidth=1.5, zorder=1)
ax.hlines(whi, x - cap_w / 2, x + cap_w / 2, color="black", linewidth=1.5, zorder=1)
