"""Sankey diagrams, sunbursts and treemaps showing how cells flow through /
nest within a hierarchy of annotation levels (e.g. Level_1 -> Level_4)."""
from __future__ import annotations
from collections import defaultdict
from pathlib import Path
from matplotlib.colors import to_hex, to_rgb
from ._config import PlotConfig
from ._palettes import resolve_palette
from ._utils import require_columns
__all__ = [
"hierarchical_counts",
"hierarchical_counts_string",
"sankey_plot",
"sunburst_plot",
"treemap_plot",
]
[docs]
def hierarchical_counts(adata, levels: list[str]) -> list[str]:
"""Return indented ``"label [count]"`` lines describing cell counts nested
across 1-3 annotation ``levels`` (coarse to fine)."""
if not levels or len(levels) > 3:
raise ValueError("Provide between 1 and 3 levels.")
require_columns(adata, *levels)
obs_df = adata.obs
if len(levels) == 1:
hierarchy = defaultdict(int)
elif len(levels) == 2:
hierarchy = defaultdict(lambda: defaultdict(int))
else:
hierarchy = defaultdict(lambda: defaultdict(lambda: defaultdict(int)))
for _, row in obs_df.iterrows():
if len(levels) == 1:
hierarchy[row[levels[0]]] += 1
elif len(levels) == 2:
hierarchy[row[levels[0]]][row[levels[1]]] += 1
else:
hierarchy[row[levels[0]]][row[levels[1]]][row[levels[2]]] += 1
result_list = []
if len(levels) == 1:
for level1, count in hierarchy.items():
result_list.append(f"{level1} [{count}]")
elif len(levels) == 2:
for level1, level2_dict in hierarchy.items():
result_list.append(f"{level1} [{sum(level2_dict.values())}]")
for level2, count in level2_dict.items():
result_list.append(f" {level2} [{count}]")
else:
for level1, level2_dict in hierarchy.items():
level1_count = sum(sum(d.values()) for d in level2_dict.values())
result_list.append(f"{level1} [{level1_count}]")
for level2, level3_dict in level2_dict.items():
result_list.append(f" {level2} [{sum(level3_dict.values())}]")
for level3, count in level3_dict.items():
result_list.append(f" {level3} [{count}]")
return result_list
[docs]
def hierarchical_counts_string(adata, levels: list[str]) -> str:
"""Newline-joined version of :func:`hierarchical_counts`."""
return "\n".join(hierarchical_counts(adata, levels))
[docs]
def sankey_plot(
adata,
levels: list[str],
save_name: str | None = None,
width: int = 1100,
height: int = 800,
show_labels: bool = True,
node_color: str = "grey",
link_color: str = "lightgrey",
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Sankey diagram of cell flow across 1-3 nested annotation ``levels``.
Requires the optional ``sankey`` extra (``pip install scplotkit[sankey]``)
for ``plotly``/``kaleido``.
"""
try:
import plotly.graph_objects as go
except ImportError as e:
raise ImportError(
"sankey_plot requires plotly (and kaleido to save images). "
"Install with `pip install scplotkit[sankey]`."
) from e
config = PlotConfig.load(config)
hierarchy_string = hierarchical_counts_string(adata, levels)
labels, sources, targets, values = [], [], [], []
label_map: dict[str, int] = {}
def label_index(label: str) -> int:
if label not in label_map:
label_map[label] = len(labels)
labels.append(label)
return label_map[label]
parent_stack: list[str] = []
for line in hierarchy_string.strip().split("\n"):
indent = len(line) - len(line.lstrip())
name, count = line.strip().rsplit(" [", 1)
value = int(count[:-1])
level = indent // 2
parent_stack = parent_stack[:level]
parent_stack.append(name)
if level > 0:
sources.append(label_index(parent_stack[level - 1]))
targets.append(label_index(parent_stack[level]))
values.append(value)
else:
label_index(name)
node_kwargs = dict(pad=15, thickness=20, color=node_color)
if show_labels:
node_kwargs["label"] = labels
fig = go.Figure(
data=[go.Sankey(node=node_kwargs, link=dict(source=sources, target=targets, value=values, color=link_color))]
)
fig.update_layout(width=width, height=height, font_size=12, margin=dict(l=50, r=50, t=50, b=50))
fig.data[0].arrangement = "snap"
if save_name is not None:
save_dir = Path(output_dir) / "sankey_plot"
save_dir.mkdir(parents=True, exist_ok=True)
save_path = save_dir / f"{save_name}.png"
fig.write_image(str(save_path), width=width, height=height, scale=2)
print(f"Figure saved to: {save_path}")
return fig
def _hierarchy_nodes(adata, levels: list[str]) -> tuple[list[str], list[str], list[str], list[int], list[int]]:
"""Build ``(ids, labels, parents, values, depths)`` for every node across nested ``levels``.
Each node id is the full ``level1/level2/...`` path rather than the bare
category name, so two branches that happen to share a child label (e.g.
an "Other" under both "Myeloid" and "Lymphoid") stay distinct nodes
instead of silently merging.
"""
require_columns(adata, *levels)
counts = adata.obs.groupby(list(levels), observed=True).size()
value_by_id: dict[str, int] = {}
parent_by_id: dict[str, str] = {}
label_by_id: dict[str, str] = {}
depth_by_id: dict[str, int] = {}
order: list[str] = []
for key, count in counts.items():
key_tuple = key if isinstance(key, tuple) else (key,)
parent_id = ""
for depth, name in enumerate(key_tuple):
node_id = f"{parent_id}/{name}" if parent_id else str(name)
if node_id not in value_by_id:
value_by_id[node_id] = 0
parent_by_id[node_id] = parent_id
label_by_id[node_id] = str(name)
depth_by_id[node_id] = depth
order.append(node_id)
value_by_id[node_id] += count
parent_id = node_id
labels = [label_by_id[i] for i in order]
parents = [parent_by_id[i] for i in order]
values = [value_by_id[i] for i in order]
depths = [depth_by_id[i] for i in order]
return order, labels, parents, values, depths
def _lighten(color, amount: float) -> str:
"""Blend ``color`` toward white by ``amount`` (0 = unchanged, 1 = white)."""
r, g, b = to_rgb(color)
return to_hex((r + (1 - r) * amount, g + (1 - g) * amount, b + (1 - b) * amount))
def _hierarchy_colors(
ids: list[str], depths: list[int], config: PlotConfig, level0_column: str, palette: dict | None
) -> list[str]:
"""One color per node: a categorical color per top-level branch, lightened with depth."""
top_level_names = [ids[i].split("/", 1)[0] for i, d in enumerate(depths) if d == 0]
base_colors = resolve_palette(config, level0_column, top_level_names, palette)
return [
_lighten(base_colors[node_id.split("/", 1)[0]], min(depth, 4) * 0.16) for node_id, depth in zip(ids, depths)
]
def _hierarchy_figure(
build_figure,
adata,
levels: list[str],
palette: dict | None,
title: str | None,
width: int,
height: int,
config: PlotConfig | None,
extra_trace_kwargs: dict,
):
config = PlotConfig.load(config)
if not levels or len(levels) > 4:
raise ValueError("Provide between 1 and 4 levels.")
ids, labels, parents, values, depths = _hierarchy_nodes(adata, levels)
colors = _hierarchy_colors(ids, depths, config, levels[0], palette)
fig = build_figure(
ids=ids,
labels=labels,
parents=parents,
values=values,
branchvalues="total",
marker=dict(colors=colors, line=dict(color="white", width=1)),
**extra_trace_kwargs,
)
general = config.general
fig.update_layout(
width=width,
height=height,
title=dict(text=title, font=dict(size=general["title_fontsize"], family=general["font_family"])),
font=dict(size=general["legend_fontsize"], family=general["font_family"]),
margin=dict(l=10, r=10, t=50 if title else 10, b=10),
)
return fig
[docs]
def sunburst_plot(
adata,
levels: list[str],
title: str | None = None,
save_name: str | None = None,
width: int = 900,
height: int = 900,
palette: dict | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Sunburst of a nested cell-type taxonomy (e.g. ``["Level_1", ..., "Level_4"]``),
an alternative to :func:`sankey_plot` for seeing the whole tree at a glance.
Ring 1 (innermost) is colored by ``levels[0]`` using the config's
registered palette for that column (falling back to an auto-generated
one); deeper rings reuse their top-level ancestor's color, progressively
lightened by depth. Purely descriptive -- segment size is raw cell count,
no statistical framing.
Requires the optional ``sankey`` extra (``pip install scplotkit[sankey]``)
for ``plotly``/``kaleido``.
"""
try:
import plotly.graph_objects as go
except ImportError as e:
raise ImportError(
"sunburst_plot requires plotly (and kaleido to save images). "
"Install with `pip install scplotkit[sankey]`."
) from e
fig = _hierarchy_figure(
lambda **kw: go.Figure(data=[go.Sunburst(**kw)]),
adata, levels, palette, title, width, height, config,
extra_trace_kwargs=dict(textinfo="label+percent parent"),
)
if save_name is not None:
save_dir = Path(output_dir) / "sunburst_plot"
save_dir.mkdir(parents=True, exist_ok=True)
save_path = save_dir / f"{save_name}.png"
fig.write_image(str(save_path), width=width, height=height, scale=2)
print(f"Figure saved to: {save_path}")
return fig
[docs]
def treemap_plot(
adata,
levels: list[str],
title: str | None = None,
save_name: str | None = None,
width: int = 1100,
height: int = 700,
palette: dict | None = None,
config: PlotConfig | None = None,
output_dir: str | Path | None = "figures",
):
"""Treemap of a nested cell-type taxonomy (e.g. ``["Level_1", ..., "Level_4"]``),
an alternative to :func:`sankey_plot` -- rectangle area makes relative size
easier to compare at a glance than a sunburst's arc angles, especially
with many leaf categories.
Coloring follows the same rule as :func:`sunburst_plot`: top-level
branches take the config's palette for ``levels[0]``, descendants reuse
their ancestor's color lightened by depth. Purely descriptive -- tile
area is raw cell count, no statistical framing.
Requires the optional ``sankey`` extra (``pip install scplotkit[sankey]``)
for ``plotly``/``kaleido``.
"""
try:
import plotly.graph_objects as go
except ImportError as e:
raise ImportError(
"treemap_plot requires plotly (and kaleido to save images). "
"Install with `pip install scplotkit[sankey]`."
) from e
fig = _hierarchy_figure(
lambda **kw: go.Figure(data=[go.Treemap(**kw)]),
adata, levels, palette, title, width, height, config,
extra_trace_kwargs=dict(textinfo="label+value", tiling=dict(packing="squarify")),
)
if save_name is not None:
save_dir = Path(output_dir) / "treemap_plot"
save_dir.mkdir(parents=True, exist_ok=True)
save_path = save_dir / f"{save_name}.png"
fig.write_image(str(save_path), width=width, height=height, scale=2)
print(f"Figure saved to: {save_path}")
return fig
