Source code for figures.dense_features
"""The dense feature matrix: every significant feature's signed drift, per class and axis.
The top-five lollipops of the {py:mod}`~figures.category_decomposition` figure name the leading
features; this figure shows them all. Each row is a feature that clears false-discovery-rate control
in at least one class on at least one axis (227 of the 238 do), each column is a class on an axis,
and the cell colour is the signed separation-standardised displacement, red for a rise and blue for
a fall, with the non-significant cells left faint so the significance pattern reads at a glance. The
rows are grouped by a chosen key, the author symptom category for $H_0^F$ or the instrument referent
for $H_0^G$, with a colour sidebar and a divider between groups, so the concentration the summary
figures report is visible feature by feature.
"""
from __future__ import annotations
import numpy as np
import pandas as pd
from matplotlib.colors import ListedColormap
from matplotlib.figure import Figure
from matplotlib.gridspec import GridSpec
from figures import style
from figures.category_decomposition import _CATEGORY_ORDER, _category_of
_NICE_AXIS = {"era": "era", "age_at_diagnosis": "age"}
def _feature_groups(
features: dict[str, pd.DataFrame], group_by: str
) -> tuple[dict[str, str], set[str]]:
"""Return the group of every feature and the set significant in at least one cell."""
group_of: dict[str, str] = {}
significant: set[str] = set()
for frame in features.values():
for feature, group, reject in zip(
frame["feature"], frame[group_by], frame["reject"], strict=True
):
group_of.setdefault(str(feature), _category_of(group))
if bool(reject):
significant.add(str(feature))
return group_of, significant
def _order_features(
features: list[str],
group_of: dict[str, str],
peak: dict[str, float],
group_order: tuple[str, ...],
) -> list[str]:
"""Order features by group (in ``group_order``), then by descending peak magnitude within it."""
rank = {name: i for i, name in enumerate(group_order)}
return sorted(
features,
key=lambda f: (rank.get(group_of[f], len(rank)), group_of[f], -peak.get(f, 0.0)),
)
[docs]
def dense_feature_figure(
features: dict[str, pd.DataFrame],
meta: dict,
*,
group_by: str = "category",
group_order: tuple[str, ...] = _CATEGORY_ORDER,
group_label: str = "category",
) -> Figure:
"""Build the dense signed-displacement matrix over every significant feature.
Parameters
----------
features : dict of str to pandas.DataFrame
The ``feature_displacement_<axis>.parquet`` frame per axis, each carrying ``feature``,
``ref_class``, ``class_name``, ``displacement``, ``reject``, and the grouping column.
meta : dict
Figure metadata; unused beyond documenting provenance.
group_by : str, optional
The column the rows are grouped by (``"category"`` for $H_0^F$, ``"referent"`` for $H_0^G$).
group_order : tuple of str, optional
The order the groups are shown in, top to bottom.
group_label : str, optional
The word for the grouping, used in the sidebar title.
Returns
-------
matplotlib.figure.Figure
The dense feature-by-class heatmap with a group colour sidebar.
"""
import matplotlib.pyplot as plt
axes_order = [a for a in ("era", "age_at_diagnosis") if a in features]
disp = {
a: features[a].pivot(index="feature", columns="ref_class", values="displacement")
for a in axes_order
}
reject = {
a: features[a].pivot(index="feature", columns="ref_class", values="reject")
for a in axes_order
}
classes = sorted(int(c) for c in features[axes_order[0]]["ref_class"].unique())
names = (
features[axes_order[0]]
.drop_duplicates("ref_class")
.set_index("ref_class")["class_name"]
.to_dict()
)
group_of, significant = _feature_groups(features, group_by)
present_groups = [g for g in group_order if g in {group_of[f] for f in significant}]
present_groups += sorted({group_of[f] for f in significant} - set(present_groups))
colours = {g: style.PALETTE[i % len(style.PALETTE)] for i, g in enumerate(present_groups)}
peak = {
f: float(
np.nanmax([np.nanmax(np.abs(disp[a].reindex([f]).to_numpy())) for a in axes_order])
)
for f in significant
}
ordered = _order_features(sorted(significant), group_of, peak, tuple(present_groups))
# The signed-displacement matrix (features by axis-class), with non-significant cells masked so
# only the significant drift shows in colour.
columns: list[tuple[str, int]] = [(a, c) for a in axes_order for c in classes]
matrix = np.full((len(ordered), len(columns)), np.nan)
for j, (axis, cls) in enumerate(columns):
col = disp[axis].reindex(ordered).get(cls)
rej = reject[axis].reindex(ordered).get(cls)
if col is not None:
values = col.to_numpy(dtype=float)
mask = np.asarray(rej.to_numpy() if rej is not None else np.zeros(len(ordered)), bool)
matrix[mask, j] = values[mask]
cap = float(np.nanpercentile(np.abs(matrix), 96)) if np.isfinite(matrix).any() else 1.0
cap = max(cap, 0.1)
with style.house_style():
height = max(4.0, min(20.0, 0.045 * len(ordered) + 1.5))
fig = plt.figure(figsize=(7.6, height))
grid = GridSpec(1, 2, width_ratios=(0.05, 1.0), wspace=0.02, figure=fig)
# The group colour sidebar.
ax_side = fig.add_subplot(grid[0, 0])
band = np.array([[present_groups.index(group_of[f])] for f in ordered])
ax_side.imshow(
band,
aspect="auto",
cmap=ListedColormap([colours[g] for g in present_groups]),
vmin=0,
vmax=max(len(present_groups) - 1, 1),
)
ax_side.set_xticks([])
ax_side.set_yticks([])
ax_side.set_title(f"n = {len(ordered)}", fontsize=7)
boundaries = _group_boundaries([group_of[f] for f in ordered])
for _, lo, _hi in boundaries[1:]:
ax_side.axhline(lo - 0.5, color="white", linewidth=0.8)
ax = fig.add_subplot(grid[0, 1])
image = ax.imshow(matrix, aspect="auto", cmap="RdBu_r", vmin=-cap, vmax=cap)
image.cmap.set_bad("#f2f2f2")
ax.set_yticks([])
for _, lo, _hi in boundaries[1:]:
ax.axhline(lo - 0.5, color="white", linewidth=0.8)
ax.set_xticks(range(len(columns)))
ax.set_xticklabels(
[f"{names.get(c, c)}" for _a, c in columns], rotation=45, ha="right", fontsize=6.5
)
# A divider and axis super-labels between the era block and the age block.
if len(axes_order) == 2:
ax.axvline(len(classes) - 0.5, color=style.REFERENCE_COLOUR, linewidth=1.2)
for i, axis in enumerate(axes_order):
ax.text(
(i + 0.5) * len(classes) - 0.5,
-0.02 * len(ordered) - 1.0,
_NICE_AXIS[axis],
ha="center",
va="bottom",
fontsize=8,
fontweight="bold",
)
bar = fig.colorbar(image, ax=ax, fraction=0.03, pad=0.02)
bar.set_label("signed displacement (sep-standardised)", fontsize=7.5)
bar.ax.tick_params(labelsize=7)
handles = [
plt.Line2D(
[],
[],
marker="s",
linestyle="none",
markersize=7,
color=colours[g],
label=style.CATEGORY_LABELS.get(g, g),
)
for g in present_groups
]
fig.legend(
handles=handles,
title=group_label,
loc="lower center",
ncol=min(len(handles), 5),
fontsize=7,
title_fontsize=7,
bbox_to_anchor=(0.5, -0.04),
)
return fig
def _group_boundaries(sequence: list[str]) -> list[tuple[str, int, int]]:
"""Return each contiguous run in a sequence as ``(value, start, stop)`` (stop exclusive)."""
runs: list[tuple[str, int, int]] = []
start = 0
for i in range(1, len(sequence) + 1):
if i == len(sequence) or sequence[i] != sequence[start]:
runs.append((sequence[start], start, i))
start = i
return runs