Skip to content

fix: robust annotation legends #13

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Sergio-ote merged 2 commits into from
Sep 9, 2026
Merged

fix: robust annotation legends #13

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3094a4d
fix: relocate legends in UMAP and spatial plots to avoid overlapping …
Sergio-ote Sep 9, 2026
f61c6c7
fix: clustree and silhouette display
Sergio-ote Sep 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions docs/rules.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ Generates integrated (Harmony) and unintegrated results, with UMAPs coloured by
#### `sample_report`
Assembles a compact, multi-page PDF with graphical outputs per sample, ideal to share with internal and external collaborators.

Each page contains the most important data-descriptive information: cluster / annotation / region / niche UMAPs and matching spatial maps, bar plots for absolute and relative composition, and a dot plot with the top-10 expressed markers per cell type. Only annotation sources containing real labels are shown; in runs without a manual TSV annotation, the external or ingest annotation also drives cell-type composition.
Each page contains the most important data-descriptive information: cluster / annotation / region / niche UMAPs and matching spatial maps, bar plots for absolute and relative composition, and a dot plot with the top-10 expressed markers per cell type. Only annotation sources containing real labels are shown; in runs without a manual TSV annotation, the external or ingest annotation also drives cell-type composition. Categorical legends use reserved bands below their plots, with adaptive rows and wrapped long labels, and always remain on the same PDF page as the corresponding panels.

---

Expand All @@ -202,12 +202,12 @@ Integrated exploration scales are standardized to ease interpretation of the res
#### `explore_genes_integrated`
Exploration of the integrated samples.

Loads the Harmony-integrated object once and computes AUCell scores for the input gene or gene set. It then writes shared expression ranges and produces integrated UMAP/dotplot PNGs per gene or gene set.
Loads the Harmony-integrated object once and computes AUCell scores for the input gene or gene set. It then writes shared expression ranges and produces integrated UMAP/dotplot PNGs per gene or gene set. Multi-panel categorical legends are kept in dedicated bands within each composite.

#### `explore_genes_sample`
Exploration of individual samples.

Using the shared expression ranges from the integrated exploration, produces per-sample spatial composites (expression + cell type + region) and stacked dotplot composites for each query entry.
Using the shared expression ranges from the integrated exploration, produces per-sample spatial composites (expression + cell type + region) and stacked dotplot composites for each query entry. Multi-panel categorical legends are kept in dedicated bands within each composite.

---

Expand Down
45 changes: 35 additions & 10 deletions workflow/scripts/annotate_cells.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,11 @@
record_active_annotation_columns,
)
from external_metadata import optional_input_path, read_cell_metadata
from plotting_legends import (
category_labels,
move_legend_to_axis,
panel_figure,
)


# ── Logging ──────────────────────────────────────────────────────────────────
Expand Down Expand Up @@ -433,26 +438,46 @@ def generate_annotation_plots(adata, annot_key, label, plots_dir, sample_id,
annot_cols_present = active_annotations
n_panels = len(annot_cols_present)
if n_panels >= 2:
fig, axes = plt.subplots(1, n_panels, figsize=(8 * n_panels, 6),
gridspec_kw={"wspace": 0.5})
for ax, col in zip(axes, annot_cols_present):
labels_by_panel = [category_labels(adata, col) for col in annot_cols_present]
fig, plot_axes, legend_axes = panel_figure(
n_panels,
labels_by_panel,
panel_width=8,
plot_height=6,
wspace=0.18,
)
axes = plot_axes[0]
for ax, legend_ax, col in zip(axes, legend_axes, annot_cols_present):
title = col.replace("cell_type_", "")
sc.pl.umap(adata, color=col, size=UMAP_POINT_SIZE, frameon=False,
title=col.replace("cell_type_", ""), ax=ax, show=False,
legend_fontsize=6, na_in_legend=False)
title=title, ax=ax, show=False,
legend_fontsize=12, na_in_legend=False)
move_legend_to_axis(ax, legend_ax, title=title)
plt.savefig(os.path.join(plots_dir, f"UMAP_comparison_{sample_id}.png"),
dpi=DPI, bbox_inches="tight")
plt.close()

# Side-by-side SPATIAL comparison (only if >1 annotation method)
if n_panels >= 2:
try:
fig, axes = plt.subplots(1, n_panels, figsize=(8 * n_panels, 6),
gridspec_kw={"wspace": 0.5})
for ax, col in zip(axes, annot_cols_present):
labels_by_panel = [
category_labels(adata, col) for col in annot_cols_present
]
fig, plot_axes, legend_axes = panel_figure(
n_panels,
labels_by_panel,
panel_width=8,
plot_height=6,
wspace=0.18,
)
axes = plot_axes[0]
for ax, legend_ax, col in zip(axes, legend_axes, annot_cols_present):
title = col.replace("cell_type_", "")
sc.pl.spatial(adata, color=col, spot_size=SPATIAL_POINT_SIZE, frameon=False,
title=col.replace("cell_type_", ""),
title=title,
library_id=library_id, ax=ax, show=False,
legend_fontsize=6, na_in_legend=False)
legend_fontsize=12, na_in_legend=False)
move_legend_to_axis(ax, legend_ax, title=title)
plt.savefig(os.path.join(plots_dir, f"spatial_comparison_{sample_id}.png"),
dpi=DPI, bbox_inches="tight")
plt.close()
Expand Down
79 changes: 39 additions & 40 deletions workflow/scripts/explore_genes_common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
import matplotlib.pyplot as plt
from matplotlib.patches import Patch
from PIL import Image
from plotting_legends import legend_layout

log = logging.getLogger(__name__)

Expand Down Expand Up @@ -75,24 +76,6 @@ def make_score_adata(adata, score_col):
)


def _compact_legend(ax, title="", per_col=20, fontsize=5):
"""Recast an axis legend into multiple narrow columns so a high niche count
does not blow up the panel (~per_col entries per column)."""
leg = ax.get_legend()
if leg is None:
return
handles = (leg.legend_handles if hasattr(leg, "legend_handles")
else getattr(leg, "legendHandles", []))
labels = [t.get_text() for t in leg.get_texts()]
if not handles:
return
ncol = max(1, (len(labels) + per_col - 1) // per_col)
ax.legend(handles, labels, loc="center left", bbox_to_anchor=(1.0, 0.5),
ncol=ncol, fontsize=fontsize, frameon=False, title=title,
title_fontsize=fontsize + 1, handletextpad=0.3,
columnspacing=0.6, labelspacing=0.25, borderaxespad=0.2)


def annotate_ct_region_dotplot(annot_key, annotation_colors, region_colors):
"""Add cell type and region colour bars BETWEEN the row labels and the
dot grid, with colour legends below. tick_params(pad) creates space
Expand Down Expand Up @@ -181,8 +164,12 @@ def annotate_ct_region_dotplot(annot_key, annotation_colors, region_colors):

def create_annotation_legend(annot_key, annotation_colors, region_colors,
out_path, dpi):
"""Create a standalone legend image for cell type + region colour bars.
Legends are stacked vertically so they never overlap."""
"""Create a temporary legend band composited beneath its dotplots.

Cell-type and region legends use dedicated rows with adaptive height, so long
labels cannot overlap. The caller embeds this image in the same output PNG and
deletes the temporary file; it is never a separate report page.
"""
try:
from matplotlib.patches import Patch

Expand Down Expand Up @@ -213,29 +200,41 @@ def create_annotation_legend(annot_key, annotation_colors, region_colors,
if not unique_cts and not unique_regs:
return

# Two axes stacked vertically — one legend per row, no overlap
n_rows = (1 if unique_cts else 0) + (1 if unique_regs else 0)
fig_leg, axes_leg = plt.subplots(n_rows, 1, figsize=(12, 1.5 * n_rows),
squeeze=False)
row = 0
groups = []
if unique_cts:
ax = axes_leg[row, 0]
ax.set_axis_off()
ct_patches = [Patch(facecolor=c, label=n) for n, c in unique_cts.items()]
ax.legend(handles=ct_patches, title="Cell type",
loc="center", fontsize=12, title_fontsize=13,
frameon=True, edgecolor="lightgray",
ncol=max(1, len(unique_cts) // 4 + 1))
row += 1

groups.append(("Cell type", unique_cts))
if unique_regs:
ax = axes_leg[row, 0]
groups.append(("Region", unique_regs))

layouts = [legend_layout(list(entries), preferred_rows=4)
for _, entries in groups]
heights = [height for _, _, height in layouts]
fig_leg = plt.figure(figsize=(12, sum(heights)))
grid = fig_leg.add_gridspec(len(groups), 1, height_ratios=heights, hspace=0.15)
for row, ((title, entries), (labels, ncol, _)) in enumerate(
zip(groups, layouts)
):
ax = fig_leg.add_subplot(grid[row, 0])
ax.set_axis_off()
reg_patches = [Patch(facecolor=c, label=n) for n, c in unique_regs.items()]
ax.legend(handles=reg_patches, title="Region",
loc="center", fontsize=12, title_fontsize=13,
frameon=True, edgecolor="lightgray",
ncol=max(1, len(unique_regs) // 3 + 1))
patches = [
Patch(facecolor=color, label=label)
for label, color in entries.items()
]
ax.legend(
handles=patches,
labels=labels,
title=title,
loc="upper center",
fontsize=12,
title_fontsize=13,
frameon=True,
edgecolor="lightgray",
ncol=ncol,
handletextpad=0.35,
columnspacing=0.8,
labelspacing=0.35,
borderaxespad=0.0,
)

fig_leg.savefig(out_path, dpi=dpi, bbox_inches="tight")
plt.close(fig_leg)
Expand Down
38 changes: 30 additions & 8 deletions workflow/scripts/explore_genes_integrated.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,12 @@
from explore_genes_common import (
read_tsv_to_dict, classify_entries, apply_annotation_palette, apply_region_palette,
make_score_adata, annotate_ct_region_dotplot, create_annotation_legend,
composite_vertical, _compact_legend,
composite_vertical,
)
from plotting_legends import (
category_labels,
move_legend_to_axis,
panel_figure,
)

# Large cell type × region dotplots can exceed PIL's default pixel limit
Expand Down Expand Up @@ -165,7 +170,8 @@ def generate_dotplots(adata, var_names, annot_key, has_regions, out_dir,
plt.savefig(tmp, dpi=dpi, bbox_inches="tight")
tmp_files.append(tmp)

# Standalone annotation legend (before plt.close so gcf() works)
# Temporary annotation-legend band (before plt.close so gcf() works).
# It is composited into the same output PNG below and then deleted.
if annotation_colors or region_colors:
tmp_leg = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
create_annotation_legend(annot_key,
Expand Down Expand Up @@ -202,8 +208,20 @@ def generate_umap_composite(adata, color_col, annot_key, has_regions, out_path,
multi-column legend for high niche counts."""
has_niche = bool(niche_col and niche_col in adata.obs.columns)
n_panels = 2 + (1 if has_regions else 0) + (1 if has_niche else 0)
fig, axes = plt.subplots(1, n_panels, figsize=(8 * n_panels, 7),
gridspec_kw={"wspace": 0.5})
panel_keys = [None, annot_key]
if has_regions:
panel_keys.append("region_annotation")
if has_niche:
panel_keys.append(niche_col)
labels_by_panel = [category_labels(adata, key) for key in panel_keys]
fig, plot_axes, legend_axes = panel_figure(
n_panels,
labels_by_panel,
panel_width=8,
plot_height=6,
wspace=0.18,
)
axes = plot_axes[0]

try:
sc.pl.umap(adata, color=color_col, size=UMAP_POINT_SIZE, frameon=False,
Expand All @@ -214,19 +232,21 @@ def generate_umap_composite(adata, color_col, annot_key, has_regions, out_path,

try:
sc.pl.umap(adata, color=annot_key, size=UMAP_POINT_SIZE, frameon=False,
title="Cell types", legend_fontsize=6,
title="Cell types", legend_fontsize=12,
na_in_legend=False,
ax=axes[1], show=False)
move_legend_to_axis(axes[1], legend_axes[1], title="Cell types")
except Exception as e:
log.warning(" UMAP celltype failed: %s", e)

idx = 2
if has_regions:
try:
sc.pl.umap(adata, color="region_annotation", size=UMAP_POINT_SIZE, frameon=False,
title="Regions", legend_fontsize=6,
title="Regions", legend_fontsize=12,
na_in_legend=False,
ax=axes[idx], show=False)
move_legend_to_axis(axes[idx], legend_axes[idx], title="Regions")
except Exception as e:
log.warning(" UMAP region failed: %s", e)
idx += 1
Expand All @@ -240,9 +260,11 @@ def generate_umap_composite(adata, color_col, annot_key, has_regions, out_path,
adata.uns[f"{niche_col}_colors"] = [
niche_palette.get(str(c), "#cccccc") for c in cats]
sc.pl.umap(adata, color=niche_col, size=UMAP_POINT_SIZE, frameon=False,
title="Spatial niches", legend_fontsize=5,
title="Spatial niches", legend_fontsize=12,
na_in_legend=False, ax=axes[idx], show=False)
_compact_legend(axes[idx], title="Niche")
move_legend_to_axis(
axes[idx], legend_axes[idx], title="Spatial niches", fontsize=12
)
except Exception as e:
log.warning(" UMAP niche failed: %s", e)
idx += 1
Expand Down
41 changes: 32 additions & 9 deletions workflow/scripts/explore_genes_sample.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,12 @@
from explore_genes_common import (
read_tsv_to_dict, classify_entries, apply_annotation_palette, apply_region_palette,
make_score_adata, annotate_ct_region_dotplot, create_annotation_legend,
composite_vertical, _compact_legend,
composite_vertical,
)
from plotting_legends import (
category_labels,
move_legend_to_axis,
panel_figure,
)

# Large cell type × region dotplots can exceed PIL's default pixel limit
Expand Down Expand Up @@ -72,8 +77,21 @@ def generate_spatial_composite(adata, color_col, annot_key, has_regions,
multi-column legend so a high niche count does not distort the composite."""
has_niche = bool(niche_col and niche_col in adata.obs.columns)
n_panels = 2 + (1 if has_regions else 0) + (1 if has_niche else 0)
fig, axes = plt.subplots(1, n_panels, figsize=(8 * n_panels, 7),
gridspec_kw={"wspace": 0.5})
panel_keys = [None, annot_key]
if has_regions:
panel_keys.append("region_annotation")
if has_niche:
panel_keys.append(niche_col)
labels_by_panel = [category_labels(adata, key) for key in panel_keys]
fig, plot_axes, legend_axes = panel_figure(
n_panels,
labels_by_panel,
panel_width=8,
plot_height=6,
wspace=0.18,
top_margin=0.8,
)
axes = plot_axes[0]
fig.suptitle(sample_id, fontsize=16, fontweight="bold", y=1.02)

# (1) Expression / AUCell
Expand All @@ -90,8 +108,9 @@ def generate_spatial_composite(adata, color_col, annot_key, has_regions,
try:
sc.pl.spatial(adata, color=annot_key, spot_size=SPATIAL_POINT_SIZE, frameon=False,
title="Cell types", library_id=library_id,
legend_fontsize=6, na_in_legend=False,
legend_fontsize=12, na_in_legend=False,
ax=axes[1], show=False)
move_legend_to_axis(axes[1], legend_axes[1], title="Cell types")
except Exception as e:
log.warning(" Spatial celltype failed: %s", e)
axes[1].set_title("Cell types (failed)")
Expand All @@ -103,8 +122,9 @@ def generate_spatial_composite(adata, color_col, annot_key, has_regions,
sc.pl.spatial(adata, color="region_annotation", spot_size=SPATIAL_POINT_SIZE,
frameon=False, title="Regions",
library_id=library_id,
legend_fontsize=6, na_in_legend=False,
legend_fontsize=12, na_in_legend=False,
ax=axes[idx], show=False)
move_legend_to_axis(axes[idx], legend_axes[idx], title="Regions")
except Exception as e:
log.warning(" Spatial region failed: %s", e)
axes[idx].set_title("Regions (failed)")
Expand All @@ -122,9 +142,11 @@ def generate_spatial_composite(adata, color_col, annot_key, has_regions,
niche_palette.get(str(c), "#cccccc") for c in cats]
sc.pl.spatial(adata, color=niche_col, spot_size=SPATIAL_POINT_SIZE, frameon=False,
title="Spatial niches", library_id=library_id,
legend_fontsize=5, na_in_legend=False,
legend_fontsize=12, na_in_legend=False,
ax=axes[idx], show=False)
_compact_legend(axes[idx], title="Niche")
move_legend_to_axis(
axes[idx], legend_axes[idx], title="Spatial niches", fontsize=12
)
except Exception as e:
log.warning(" Spatial niche failed: %s", e)
axes[idx].set_title("Spatial niches (failed)")
Expand Down Expand Up @@ -197,8 +219,9 @@ def generate_dotplot_composite(adata, var_names, annot_key, has_regions,
plt.savefig(tmp, dpi=dpi, bbox_inches="tight")
tmp_files.append(tmp)

# Standalone annotation legend (before plt.close so gcf() can
# still read the dotplot's ytick labels)
# Temporary annotation-legend band (before plt.close so gcf() can
# still read the dotplot's ytick labels). It is composited into the
# same output PNG below and then deleted.
if annotation_colors or region_colors:
tmp_leg = tempfile.NamedTemporaryFile(suffix=".png", delete=False).name
create_annotation_legend(annot_key,
Expand Down
Loading
Loading