From bc1dea08add011c791981093cce7dce51984e9a2 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Fri, 24 Jul 2026 09:03:26 +0900 Subject: [PATCH 1/2] feat(functional): add plot_lineage_tree for matplotlib lineage trees Add plot_lineage_tree in tracksdata.functional._plot: render lineage trees with matplotlib, supporting attribute-bound node colors and sizes, time windows, and exact timestamps. Export from tracksdata.functional and add an optional `plot` extra (matplotlib). Co-Authored-By: Claude Opus 4.8 --- pyproject.toml | 2 + src/tracksdata/functional/__init__.py | 2 + src/tracksdata/functional/_plot.py | 430 +++++++++++++++++++ src/tracksdata/functional/_test/test_plot.py | 265 ++++++++++++ 4 files changed, 699 insertions(+) create mode 100644 src/tracksdata/functional/_plot.py create mode 100644 src/tracksdata/functional/_test/test_plot.py diff --git a/pyproject.toml b/pyproject.toml index f94aabd6..b413abbb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -62,9 +62,11 @@ dependencies = [ [project.optional-dependencies] spatial = ["spatial-graph"] motile = ["motile"] +plot = ["matplotlib"] test = [ "spatial-graph", "motile", + "matplotlib", "pytest>=7.0", "pytest-cov", "pytest-html", diff --git a/src/tracksdata/functional/__init__.py b/src/tracksdata/functional/__init__.py index eb9da557..3fae4368 100644 --- a/src/tracksdata/functional/__init__.py +++ b/src/tracksdata/functional/__init__.py @@ -6,12 +6,14 @@ from tracksdata.functional._labeling import ancestral_connected_edges from tracksdata.functional._motile import to_motile_graph from tracksdata.functional._napari import rx_digraph_to_napari_dict, to_napari_format +from tracksdata.functional._plot import plot_lineage_tree __all__ = [ "TilingScheme", "ancestral_connected_edges", "apply_tiled", "join_node_attrs_to_edges", + "plot_lineage_tree", "rx_digraph_to_napari_dict", "shift_division", "to_motile_graph", diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py new file mode 100644 index 00000000..db255149 --- /dev/null +++ b/src/tracksdata/functional/_plot.py @@ -0,0 +1,430 @@ +"""Matplotlib-based plotting utilities for lineage trees.""" + +from collections.abc import Mapping, Sequence +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import rustworkx as rx +from numpy.typing import ArrayLike + +from tracksdata.constants import DEFAULT_ATTR_KEYS +from tracksdata.graph._base_graph import BaseGraph + +if TYPE_CHECKING: + from matplotlib.axes import Axes + from matplotlib.colors import Colormap, Normalize + +__all__ = ["plot_lineage_tree"] + + +def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: + """ + Assign a tree-axis coordinate to each tracklet of a tracklet graph. + + Leaf tracklets receive consecutive integer coordinates and each parent + tracklet is centered at the mean coordinate of its children, resulting + in the classic dendrogram-like lineage tree layout. + + Parameters + ---------- + tracklet_graph : rx.PyDiGraph + Compressed tracklet graph as returned by + [BaseGraph.tracklet_graph][tracksdata.graph.BaseGraph.tracklet_graph], + where node values are tracklet ids and edges point from parent to child. + + Returns + ------- + dict[int, float] + Mapping of tracklet id to tree-axis coordinate. + """ + positions: dict[int, float] = {} + visited: set[int] = set() + next_leaf = 0.0 + + roots = sorted( + (rx_id for rx_id in tracklet_graph.node_indices() if tracklet_graph.in_degree(rx_id) == 0), + key=tracklet_graph.__getitem__, + ) + + for root in roots: + # iterative post-order traversal: children are positioned before parents + stack: list[tuple[int, bool]] = [(root, False)] + while stack: + rx_id, expanded = stack.pop() + if expanded: + children_pos = [ + positions[tracklet_graph[child]] + for child in tracklet_graph.successor_indices(rx_id) + if tracklet_graph[child] in positions + ] + if children_pos: + positions[tracklet_graph[rx_id]] = float(np.mean(children_pos)) + else: + positions[tracklet_graph[rx_id]] = next_leaf + next_leaf += 1.0 + elif rx_id not in visited: + visited.add(rx_id) + stack.append((rx_id, True)) + for child in sorted( + tracklet_graph.successor_indices(rx_id), + key=tracklet_graph.__getitem__, + reverse=True, + ): + if child not in visited: + stack.append((child, False)) + + return positions + + +def _time_axis_positions( + time_points: list[int], + time_positions: "Mapping[int, float] | ArrayLike | None", +) -> dict[int, float]: + """ + Map each time point to its coordinate along the time axis. + + Parameters + ---------- + time_points : list[int] + Sorted unique time points to be displayed. + time_positions : Mapping[int, float] | ArrayLike | None + Exact time-axis coordinates (e.g. timestamps). Either a mapping of + time point to coordinate or a sequence indexed by time point. + If None, time points are evenly separated in their sorted order. + + Returns + ------- + dict[int, float] + Mapping of time point to time-axis coordinate. + """ + if time_positions is None: + return {t: float(i) for i, t in enumerate(time_points)} + + if isinstance(time_positions, Mapping): + missing = [t for t in time_points if t not in time_positions] + if missing: + raise ValueError(f"`time_positions` is missing positions for time points {missing}") + return {t: float(time_positions[t]) for t in time_points} + + time_positions = np.asarray(time_positions) + if time_positions.ndim != 1: + raise ValueError(f"`time_positions` must be 1-dimensional, got {time_positions.ndim} dimensions.") + if time_points[-1] >= len(time_positions): + raise ValueError( + f"`time_positions` of length {len(time_positions)} cannot be indexed " + f"by the maximum time point {time_points[-1]}." + ) + return {t: float(time_positions[t]) for t in time_points} + + +def _map_to_size_range( + values: np.ndarray, + size_norm: tuple[float, float] | None, + size_range: tuple[float, float], +) -> np.ndarray: + """ + Linearly map attribute values to marker sizes within `size_range`. + + Parameters + ---------- + values : np.ndarray + Attribute values to map. + size_norm : tuple[float, float] | None + The (vmin, vmax) values mapped to the limits of `size_range`. + If None, the minimum and maximum of `values` are used. + size_range : tuple[float, float] + The (smallest, largest) marker sizes in points**2. + + Returns + ------- + np.ndarray + Marker sizes, one per value. + """ + values = np.asarray(values, dtype=float) + if size_norm is None: + vmin, vmax = np.nanmin(values), np.nanmax(values) + else: + vmin, vmax = size_norm + + smin, smax = size_range + if vmax <= vmin: + return np.full(values.shape, (smin + smax) / 2) + + fraction = np.clip((values - vmin) / (vmax - vmin), 0.0, 1.0) + return smin + fraction * (smax - smin) + + +def _bridged_edge_segments( + successors: dict[int, list[int]], + node_coords: dict[int, tuple[float, float]], +) -> list[tuple[tuple[float, float], tuple[float, float]]]: + """ + Build edge segments between displayed nodes, bridging across hidden ones. + + Each displayed node is connected to its nearest displayed descendants by + walking forward through the tracking graph and skipping over nodes that are + not displayed. This keeps the lineage structure visible when only a subset + of time points is shown. When all nodes are displayed it reduces to the + direct edges of the graph. + + Parameters + ---------- + successors : dict[int, list[int]] + Forward adjacency of the full (sub)graph, mapping each source node id + to the list of its target node ids. + node_coords : dict[int, tuple[float, float]] + Plot coordinates of the displayed nodes, keyed by node id. + + Returns + ------- + list[tuple[tuple[float, float], tuple[float, float]]] + Line segments connecting the coordinates of displayed nodes. + """ + segments = [] + for source in node_coords: + # walk forward to the nearest displayed descendants, skipping hidden nodes + stack = list(successors.get(source, ())) + seen: set[int] = set() + while stack: + node = stack.pop() + if node in seen: + continue + seen.add(node) + if node in node_coords: + segments.append((node_coords[source], node_coords[node])) + else: + stack.extend(successors.get(node, ())) + return segments + + +def plot_lineage_tree( + graph: BaseGraph, + *, + ax: "Axes | None" = None, + tracklet_id_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, + color_attr: str | None = None, + cmap: "str | Colormap" = "viridis", + color_norm: "Normalize | tuple[float, float] | None" = None, + size_attr: str | None = None, + size_norm: tuple[float, float] | None = None, + size_range: tuple[float, float] = (10.0, 100.0), + node_size: float = 30.0, + time_range: tuple[int, int] | None = None, + time_points: Sequence[int] | None = None, + time_positions: "Mapping[int, float] | ArrayLike | None" = None, + orientation: Literal["vertical", "horizontal"] = "vertical", + scatter_kwargs: dict[str, Any] | None = None, + line_kwargs: dict[str, Any] | None = None, +) -> "Axes": + """ + Plot a graph as a lineage tree with matplotlib. + + Nodes are drawn as points aligned in time and grouped by tracklet, + with parent tracklets centered above their children. Edges are drawn + as line segments, so divisions appear as forks in the tree. When only a + subset of time points is shown, each node is connected to its nearest + displayed descendants, bridging over the hidden time points so the lineage + stays connected. + + Requires `matplotlib`, which is an optional dependency + (`pip install "tracksdata[plot]"`). + + IMPORTANT: If `tracklet_id_key` is not an existing node attribute, + tracklet ids are assigned on the fly, modifying the graph. + To plot only solution nodes, pass the solution subgraph, e.g. + `graph.filter(NodeAttr("solution") == True, EdgeAttr("solution") == True).subgraph()`. + + Parameters + ---------- + graph : BaseGraph + The graph to plot. + ax : Axes | None, optional + The matplotlib axes to plot into. If None, a new figure and axes + are created. + tracklet_id_key : str, optional + The key of the tracklet id node attribute. If the key does not exist, + [BaseGraph.assign_tracklet_ids][tracksdata.graph.BaseGraph.assign_tracklet_ids] + is called first. + color_attr : str | None, optional + Node attribute key bound to the marker colors. Must be numeric. + cmap : str | Colormap, optional + Colormap used with `color_attr`. + color_norm : Normalize | tuple[float, float] | None, optional + Normalization for the colors, either a matplotlib `Normalize` + instance or a `(vmin, vmax)` tuple. If None, the data range is used. + size_attr : str | None, optional + Node attribute key bound to the marker sizes. Must be numeric. + size_norm : tuple[float, float] | None, optional + The `(vmin, vmax)` attribute values mapped to the limits of + `size_range`. If None, the data range is used. + size_range : tuple[float, float], optional + The marker sizes in points**2 assigned to the smallest and largest + values of `size_attr`. + node_size : float, optional + Marker size in points**2 used when `size_attr` is None. + time_range : tuple[int, int] | None, optional + Inclusive `(start, end)` range of time points to display. + If None, all time points are displayed. Mutually exclusive with + `time_points`. + time_points : Sequence[int] | None, optional + Explicit subset of time points to display, which need not be + contiguous (e.g. `[0, 5, 10]`). Edges bridge over the hidden time + points, connecting each displayed node to its nearest displayed + descendants. Mutually exclusive with `time_range`. + time_positions : Mapping[int, float] | ArrayLike | None, optional + Exact positions of the time points along the time axis + (e.g. acquisition timestamps). Either a mapping of time point to + position or a sequence indexed by time point. If None, the displayed + time points are evenly separated and labeled with their values. + orientation : {"vertical", "horizontal"}, optional + If "vertical", time runs downward along the y-axis. + If "horizontal", time runs rightward along the x-axis. + scatter_kwargs : dict[str, Any] | None, optional + Additional keyword arguments forwarded to `Axes.scatter`, + e.g. `edgecolors` and `linewidths` to style the marker borders. + line_kwargs : dict[str, Any] | None, optional + Additional keyword arguments forwarded to the edge + `LineCollection` (e.g. `color`, `linewidth`). + + Returns + ------- + Axes + The matplotlib axes containing the lineage tree. The node + `PathCollection` is the last entry of `Axes.collections`, which + can be used to add a colorbar. + + Examples + -------- + ```python + from tracksdata.functional import plot_lineage_tree + + ax = plot_lineage_tree(graph, color_attr="area", cmap="magma", size_attr="area") + ax.figure.colorbar(ax.collections[-1], ax=ax, label="area") + ``` + + Display only a time window with timestamps in seconds: + + ```python + ax = plot_lineage_tree( + graph, + time_range=(10, 20), + time_positions={t: t * 30.0 for t in range(50)}, + ) + ``` + + Display an arbitrary subset of time points with styled marker borders: + + ```python + ax = plot_lineage_tree( + graph, + time_points=[0, 5, 10, 15], + scatter_kwargs={"edgecolors": "black", "linewidths": 0.5}, + ) + ``` + """ + try: + import matplotlib.pyplot as plt + from matplotlib.collections import LineCollection + from matplotlib.colors import Normalize + except ImportError as e: + raise ImportError( + "matplotlib is required for `plot_lineage_tree`. " + "Install it with `pip install matplotlib` or `pip install 'tracksdata[plot]'`." + ) from e + + if orientation not in ("vertical", "horizontal"): + raise ValueError(f"`orientation` must be 'vertical' or 'horizontal', got '{orientation}'.") + + if tracklet_id_key not in graph.node_attr_keys(): + graph.assign_tracklet_ids(tracklet_id_key) + + attr_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key] + for key in (color_attr, size_attr): + if key is None or key in attr_keys: + continue + if key not in graph.node_attr_keys(): + raise ValueError(f"Attribute '{key}' not found in graph. Expected one of {graph.node_attr_keys()}") + attr_keys.append(key) + + if time_range is not None and time_points is not None: + raise ValueError("`time_range` and `time_points` are mutually exclusive, provide at most one.") + + nodes_df = graph.node_attrs(attr_keys=attr_keys) + + if time_range is not None: + start, end = time_range + nodes_df = nodes_df.filter((nodes_df[DEFAULT_ATTR_KEYS.T] >= start) & (nodes_df[DEFAULT_ATTR_KEYS.T] <= end)) + elif time_points is not None: + nodes_df = nodes_df.filter(nodes_df[DEFAULT_ATTR_KEYS.T].is_in(list(time_points))) + + if len(nodes_df) == 0: + raise ValueError("No nodes to plot. The graph is empty or `time_range`/`time_points` excluded all nodes.") + + # tree-axis coordinate per tracklet, computed on the full graph so the + # layout is independent of the displayed time range + tracklet_positions = _tracklet_tree_layout(graph.tracklet_graph(tracklet_id_key=tracklet_id_key)) + + time_points = nodes_df[DEFAULT_ATTR_KEYS.T].unique().sort().to_list() + time_axis_positions = _time_axis_positions(time_points, time_positions) + + tree_coords = np.asarray([tracklet_positions[tid] for tid in nodes_df[tracklet_id_key]]) + time_coords = np.asarray([time_axis_positions[t] for t in nodes_df[DEFAULT_ATTR_KEYS.T]]) + + if orientation == "vertical": + x_coords, y_coords = tree_coords, time_coords + else: + x_coords, y_coords = time_coords, tree_coords + + node_coords = { + node_id: (x, y) for node_id, x, y in zip(nodes_df[DEFAULT_ATTR_KEYS.NODE_ID], x_coords, y_coords, strict=True) + } + + edges_df = graph.edge_attrs(attr_keys=[]) + successors: dict[int, list[int]] = {} + for source, target in zip( + edges_df[DEFAULT_ATTR_KEYS.EDGE_SOURCE].to_list(), + edges_df[DEFAULT_ATTR_KEYS.EDGE_TARGET].to_list(), + strict=True, + ): + successors.setdefault(source, []).append(target) + + segments = _bridged_edge_segments(successors, node_coords) + + if ax is None: + _, ax = plt.subplots() + + line_kwargs = {"color": "0.6", "linewidth": 1.0, "zorder": 1, **(line_kwargs or {})} + ax.add_collection(LineCollection(segments, **line_kwargs)) + + scatter_kwargs = {"zorder": 2, **(scatter_kwargs or {})} + if color_attr is not None: + if isinstance(color_norm, tuple): + color_norm = Normalize(*color_norm) + scatter_kwargs["c"] = nodes_df[color_attr].to_numpy() + scatter_kwargs["cmap"] = cmap + scatter_kwargs["norm"] = color_norm + if size_attr is not None: + scatter_kwargs["s"] = _map_to_size_range(nodes_df[size_attr].to_numpy(), size_norm, size_range) + else: + scatter_kwargs.setdefault("s", node_size) + + ax.scatter(x_coords, y_coords, **scatter_kwargs) + + if orientation == "vertical": + time_axis, tree_axis = ax.yaxis, ax.xaxis + ax.set_ylabel("time") + if not ax.yaxis_inverted(): + ax.invert_yaxis() + else: + time_axis, tree_axis = ax.xaxis, ax.yaxis + ax.set_xlabel("time") + + tree_axis.set_ticks([]) + + if time_positions is None: + # evenly separated positions: label the ticks with the time point values + stride = max(1, len(time_points) // 10) + ticks = time_points[::stride] + time_axis.set_ticks([time_axis_positions[t] for t in ticks], labels=[str(t) for t in ticks]) + + return ax diff --git a/src/tracksdata/functional/_test/test_plot.py b/src/tracksdata/functional/_test/test_plot.py new file mode 100644 index 00000000..853d25c6 --- /dev/null +++ b/src/tracksdata/functional/_test/test_plot.py @@ -0,0 +1,265 @@ +import numpy as np +import polars as pl +import pytest + +matplotlib = pytest.importorskip("matplotlib") +matplotlib.use("Agg") + +import matplotlib.pyplot as plt # noqa: E402 +from matplotlib.axes import Axes # noqa: E402 + +from tracksdata.constants import DEFAULT_ATTR_KEYS # noqa: E402 +from tracksdata.functional import plot_lineage_tree # noqa: E402 +from tracksdata.graph import RustWorkXGraph # noqa: E402 + + +@pytest.fixture(autouse=True) +def _close_figures() -> None: + yield + plt.close("all") + + +def _dividing_graph() -> RustWorkXGraph: + """Build a graph with a single lineage: tracklet 1 divides into tracklets 2 and 3.""" + positions = np.asarray( + [ + [0, 0, 0], # t=0, tracklet 1 + [1, 0, 0], # t=1, tracklet 1 + [2, 0, 0], # t=2, tracklet 2 + [3, 0, 0], # t=3, tracklet 2 + [2, 1, 1], # t=2, tracklet 3 + [3, 1, 1], # t=3, tracklet 3 + ] + ) + tracklet_ids = np.asarray([1, 1, 2, 2, 3, 3]) + graph = RustWorkXGraph.from_array( + positions, + tracklet_ids=tracklet_ids, + tracklet_id_graph={2: 1, 3: 1}, + ) + graph.add_node_attr_key("feature", pl.Float64) + graph.update_node_attrs( + node_ids=graph.node_ids(), + attrs={"feature": [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]}, + ) + return graph + + +def test_plot_lineage_tree_basic() -> None: + """Test the default lineage tree layout and edge drawing.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph) + + assert isinstance(ax, Axes) + + lines, scatter = ax.collections + offsets = np.asarray(scatter.get_offsets()) + assert offsets.shape == (graph.num_nodes(), 2) + assert len(lines.get_segments()) == graph.num_edges() + + # vertical orientation: time on the (inverted) y-axis + assert ax.yaxis_inverted() + assert ax.get_ylabel() == "time" + np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 1])), [0.0, 1.0, 2.0, 3.0]) + + # the parent tracklet is centered between its two children + nodes_df = graph.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.TRACKLET_ID]) + tracklet_ids = nodes_df[DEFAULT_ATTR_KEYS.TRACKLET_ID].to_numpy() + tree_coords = {tid: set(offsets[tracklet_ids == tid, 0]) for tid in (1, 2, 3)} + for tid in (1, 2, 3): + assert len(tree_coords[tid]) == 1 # all nodes of a tracklet share the same coordinate + (parent_x,) = tree_coords[1] + (child_a_x,) = tree_coords[2] + (child_b_x,) = tree_coords[3] + assert child_a_x != child_b_x + assert parent_x == pytest.approx((child_a_x + child_b_x) / 2) + + +def test_plot_lineage_tree_color_and_size() -> None: + """Test binding attributes to marker colors and sizes.""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + color_attr="feature", + cmap="magma", + color_norm=(0.0, 10.0), + size_attr="feature", + size_range=(10.0, 50.0), + ) + + scatter = ax.collections[-1] + + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + np.testing.assert_array_equal(np.asarray(scatter.get_array()), feature) + assert scatter.get_cmap().name == "magma" + assert scatter.norm.vmin == 0.0 + assert scatter.norm.vmax == 10.0 + + sizes = np.asarray(scatter.get_sizes()) + expected = 10.0 + (feature - feature.min()) / (feature.max() - feature.min()) * 40.0 + np.testing.assert_allclose(sizes, expected) + + +def test_plot_lineage_tree_size_norm() -> None: + """Test explicit size normalization limits with clipping.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, size_attr="feature", size_norm=(0.0, 2.0), size_range=(10.0, 50.0)) + + sizes = np.asarray(ax.collections[-1].get_sizes()) + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + expected = 10.0 + np.clip(feature / 2.0, 0.0, 1.0) * 40.0 + np.testing.assert_allclose(sizes, expected) + + +def test_plot_lineage_tree_time_range() -> None: + """Test that time_range limits the displayed nodes and edges.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, time_range=(1, 2)) + + lines, scatter = ax.collections + # nodes: t=1 (tracklet 1) and t=2 (tracklets 2 and 3) + assert len(scatter.get_offsets()) == 3 + # edges: only the two division edges are fully within the range + assert len(lines.get_segments()) == 2 + + # evenly separated positions labeled with the actual time points + labels = [tick.get_text() for tick in ax.get_yticklabels()] + assert labels == ["1", "2"] + + +def test_plot_lineage_tree_time_points() -> None: + """Test selecting an arbitrary, non-contiguous subset of time points.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, time_points=[0, 3]) + + lines, scatter = ax.collections + # nodes: t=0 (tracklet 1) and t=3 (tracklets 2 and 3) + assert len(scatter.get_offsets()) == 3 + + # edges bridge over the hidden frames: the single t=0 node connects to each + # of the two t=3 nodes through the (hidden) division at t=2 + segments = lines.get_segments() + assert len(segments) == 2 + # both bridged segments start at the same point: the single displayed t=0 node, + # which sits at the minimum (topmost) time coordinate + starts = np.asarray([seg[0] for seg in segments]) + np.testing.assert_array_equal(starts[0], starts[1]) + assert starts[0, 1] == 0.0 # t=0 evenly-separated position + # the two endpoints are the two distinct t=3 nodes + ends = np.asarray([seg[1] for seg in segments]) + assert ends[0, 0] != ends[1, 0] + np.testing.assert_array_equal(ends[:, 1], [1.0, 1.0]) # both at t=3 position + + # the two displayed time points are evenly separated and labeled with their values + offsets = np.asarray(scatter.get_offsets()) + np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 1])), [0.0, 1.0]) + labels = [tick.get_text() for tick in ax.get_yticklabels()] + assert labels == ["0", "3"] + + +def test_plot_lineage_tree_time_points_mutually_exclusive() -> None: + """Test that time_range and time_points cannot be combined.""" + graph = _dividing_graph() + + with pytest.raises(ValueError, match="mutually exclusive"): + plot_lineage_tree(graph, time_range=(0, 2), time_points=[0, 1]) + + +def test_plot_lineage_tree_edge_colors() -> None: + """Test styling marker borders via scatter_kwargs (edgecolors/linewidths).""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + color_attr="feature", + scatter_kwargs={"edgecolors": "red", "linewidths": 1.5}, + ) + + scatter = ax.collections[-1] + np.testing.assert_allclose(scatter.get_edgecolors()[0], [1.0, 0.0, 0.0, 1.0]) + np.testing.assert_allclose(scatter.get_linewidths(), [1.5]) + # face colors still come from the colormap, independent of the edge color + np.testing.assert_array_equal(np.asarray(scatter.get_array()), graph.node_attrs(attr_keys=["feature"])["feature"]) + + +def test_plot_lineage_tree_time_positions() -> None: + """Test exact time positions given as a mapping and as a sequence.""" + graph = _dividing_graph() + + timestamps = {t: 100.0 + 10.0 * t for t in range(4)} + ax = plot_lineage_tree(graph, time_positions=timestamps) + offsets = np.asarray(ax.collections[-1].get_offsets()) + np.testing.assert_array_equal( + np.sort(np.unique(offsets[:, 1])), + [100.0, 110.0, 120.0, 130.0], + ) + + ax = plot_lineage_tree(graph, time_positions=np.asarray([0.0, 1.0, 2.0, 10.0])) + offsets = np.asarray(ax.collections[-1].get_offsets()) + np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 1])), [0.0, 1.0, 2.0, 10.0]) + + with pytest.raises(ValueError, match="missing positions"): + plot_lineage_tree(graph, time_positions={0: 0.0}) + + with pytest.raises(ValueError, match="cannot be indexed"): + plot_lineage_tree(graph, time_positions=np.asarray([0.0, 1.0])) + + +def test_plot_lineage_tree_horizontal() -> None: + """Test horizontal orientation with time on the x-axis.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, orientation="horizontal") + + offsets = np.asarray(ax.collections[-1].get_offsets()) + np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 0])), [0.0, 1.0, 2.0, 3.0]) + assert ax.get_xlabel() == "time" + assert not ax.yaxis_inverted() + + with pytest.raises(ValueError, match="`orientation` must be"): + plot_lineage_tree(graph, orientation="diagonal") + + +def test_plot_lineage_tree_assigns_tracklet_ids() -> None: + """Test that tracklet ids are assigned when the key is missing.""" + positions = np.asarray([[0, 0, 0], [1, 5, 5]]) + graph = RustWorkXGraph.from_array(positions) + + assert "my_tracklet_id" not in graph.node_attr_keys() + + ax = plot_lineage_tree(graph, tracklet_id_key="my_tracklet_id") + + assert "my_tracklet_id" in graph.node_attr_keys() + assert len(ax.collections[-1].get_offsets()) == 2 + + +def test_plot_lineage_tree_existing_axes_and_kwargs() -> None: + """Test plotting into an existing axes with custom artist kwargs.""" + graph = _dividing_graph() + + _, ax = plt.subplots() + returned_ax = plot_lineage_tree( + graph, + ax=ax, + scatter_kwargs={"alpha": 0.5}, + line_kwargs={"color": "red"}, + ) + + assert returned_ax is ax + assert ax.collections[-1].get_alpha() == 0.5 + + +def test_plot_lineage_tree_errors() -> None: + """Test error handling for empty selections and missing attributes.""" + graph = _dividing_graph() + + with pytest.raises(ValueError, match="No nodes to plot"): + plot_lineage_tree(graph, time_range=(10, 20)) + + with pytest.raises(ValueError, match="not found in graph"): + plot_lineage_tree(graph, color_attr="does_not_exist") From 78930a89f4e35cc072ca7a84379d7ffcac3bcf35 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Fri, 24 Jul 2026 16:09:26 +0900 Subject: [PATCH 2/2] feat(functional): callable color/size/marker/text aesthetics in plot_lineage_tree Rename color_attr/size_attr to color/size and let color, size, marker, and text each accept either a fixed value or a callable resolving a per-node value from the node's attribute row: - color: numeric attr or callable returning numbers -> cmap + colorbar; callable returning literal colors (names/hex/RGBA) -> used verbatim. - size: attr mapped to size_range, callable returning raw sizes, or constant. - marker: single glyph, or callable returning a per-node glyph (nodes are grouped by glyph, one scatter call per group, sharing one normalization). - text: attr name or callable -> per-node annotation, styled via text_kwargs. Add `attrs` to declare the node keys callables read. If a callable is given without `attrs`, warn and load all node attributes (may pull mask blobs). Co-Authored-By: Claude Opus 4.8 --- src/tracksdata/functional/_plot.py | 261 ++++++++++++++++--- src/tracksdata/functional/_test/test_plot.py | 143 +++++++++- 2 files changed, 360 insertions(+), 44 deletions(-) diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py index db255149..0c7ad8c6 100644 --- a/src/tracksdata/functional/_plot.py +++ b/src/tracksdata/functional/_plot.py @@ -1,6 +1,7 @@ """Matplotlib-based plotting utilities for lineage trees.""" -from collections.abc import Mapping, Sequence +import warnings +from collections.abc import Callable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Literal import numpy as np @@ -17,6 +18,34 @@ __all__ = ["plot_lineage_tree"] +def _resolve_color_values(raw: list) -> tuple[Any, bool]: + """ + Interpret per-node color-callable outputs as either scalars or literal colors. + + Parameters + ---------- + raw : list + One value per node, as returned by a `color` callable. + + Returns + ------- + tuple[Any, bool] + `(values, is_scalar)`. If the outputs form a 1-D numeric array, + `values` is that array and `is_scalar` is True, so they are mapped + through a colormap (and support a colorbar). Otherwise `values` is + the original list of literal colors (names, hex, or RGB(A) tuples) + and `is_scalar` is False. + """ + try: + arr = np.asarray(raw, dtype=float) + except (ValueError, TypeError): + return list(raw), False + if arr.ndim == 1: + return arr, True + # (N, 3) or (N, 4): literal RGB(A) colors, not colormap-able scalars + return list(raw), False + + def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: """ Assign a tree-axis coordinate to each tracklet of a tracklet graph. @@ -202,13 +231,17 @@ def plot_lineage_tree( *, ax: "Axes | None" = None, tracklet_id_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, - color_attr: str | None = None, + color: "str | Callable[[Mapping[str, Any]], Any] | None" = None, cmap: "str | Colormap" = "viridis", color_norm: "Normalize | tuple[float, float] | None" = None, - size_attr: str | None = None, + size: "str | Callable[[Mapping[str, Any]], float] | float | None" = None, size_norm: tuple[float, float] | None = None, size_range: tuple[float, float] = (10.0, 100.0), node_size: float = 30.0, + marker: "str | Callable[[Mapping[str, Any]], str] | None" = None, + text: "str | Callable[[Mapping[str, Any]], Any] | None" = None, + text_kwargs: dict[str, Any] | None = None, + attrs: Sequence[str] | None = None, time_range: tuple[int, int] | None = None, time_points: Sequence[int] | None = None, time_positions: "Mapping[int, float] | ArrayLike | None" = None, @@ -226,6 +259,21 @@ def plot_lineage_tree( displayed descendants, bridging over the hidden time points so the lineage stays connected. + The `color`, `size`, `marker`, and `text` aesthetics each accept either a + fixed value or a callable, which is the main way to customize the markers: + + - As a string, `color`/`size`/`text` name a numeric node attribute, and + `marker` is a single matplotlib marker glyph applied to every node. + - As a callable, they receive each node's attribute row (a mapping of + attribute key to value) and return that node's color, size, marker glyph, + or text label. This allows categorical colors, per-node marker shapes, + and colors derived from a computed quantity (e.g. `np.log1p(row["area"])`). + + A colorbar-compatible mapping is available whenever `color` produces numeric + values (a numeric attribute name, or a callable returning numbers) together + with `cmap`. If a callable returns literal colors (names, hex, or RGB(A)), + those colors are used verbatim and no colorbar mapping exists. + Requires `matplotlib`, which is an optional dependency (`pip install "tracksdata[plot]"`). @@ -245,23 +293,51 @@ def plot_lineage_tree( The key of the tracklet id node attribute. If the key does not exist, [BaseGraph.assign_tracklet_ids][tracksdata.graph.BaseGraph.assign_tracklet_ids] is called first. - color_attr : str | None, optional - Node attribute key bound to the marker colors. Must be numeric. + color : str | Callable | None, optional + Marker color. A string names a numeric node attribute mapped through + `cmap`/`color_norm` (a colorbar mapping is available). A callable + receives each node's attribute row and returns either a number (mapped + through `cmap`, colorbar available) or a literal color (used as-is, no + colorbar). If None, matplotlib's default color is used. cmap : str | Colormap, optional - Colormap used with `color_attr`. + Colormap used when `color` yields numeric values. color_norm : Normalize | tuple[float, float] | None, optional - Normalization for the colors, either a matplotlib `Normalize` + Normalization for numeric colors, either a matplotlib `Normalize` instance or a `(vmin, vmax)` tuple. If None, the data range is used. - size_attr : str | None, optional - Node attribute key bound to the marker sizes. Must be numeric. + A single shared normalization is applied across all marker groups. + size : str | Callable | float | None, optional + Marker size. A string names a numeric node attribute mapped into + `size_range`. A callable receives each node's attribute row and returns + the marker size in points**2 directly. A number sets a constant size. + If None, `node_size` is used. size_norm : tuple[float, float] | None, optional The `(vmin, vmax)` attribute values mapped to the limits of - `size_range`. If None, the data range is used. + `size_range`, used when `size` is an attribute name. If None, the data + range is used. size_range : tuple[float, float], optional The marker sizes in points**2 assigned to the smallest and largest - values of `size_attr`. + values when `size` is an attribute name. node_size : float, optional - Marker size in points**2 used when `size_attr` is None. + Marker size in points**2 used when `size` is None. + marker : str | Callable | None, optional + Marker shape. A string is a single matplotlib marker glyph (e.g. "s") + applied to every node. A callable receives each node's attribute row + and returns the marker glyph for that node; nodes are grouped by glyph + and drawn with one `Axes.scatter` call per group. If None, "o" is used. + text : str | Callable | None, optional + Per-node text label. A string names a node attribute whose value is + annotated at each node. A callable receives each node's attribute row + and returns the label. If None, no labels are drawn. Labels are drawn + per node and can clutter large trees. + text_kwargs : dict[str, Any] | None, optional + Additional keyword arguments forwarded to `Axes.annotate` for the text + labels (e.g. `fontsize`, `color`, `xytext`). + attrs : Sequence[str] | None, optional + Extra node attribute keys to load so the `color`/`size`/`marker`/`text` + callables can read them. If a callable is passed but `attrs` is None, a + warning is emitted and all node attributes are loaded, which may be slow + or memory-heavy (e.g. mask attributes). Ignored keys already loaded for + other reasons are harmless. time_range : tuple[int, int] | None, optional Inclusive `(start, end)` range of time points to display. If None, all time points are displayed. Mutually exclusive with @@ -289,36 +365,55 @@ def plot_lineage_tree( Returns ------- Axes - The matplotlib axes containing the lineage tree. The node - `PathCollection` is the last entry of `Axes.collections`, which - can be used to add a colorbar. + The matplotlib axes containing the lineage tree. When `color` yields + numeric values, the last node `PathCollection` in `Axes.collections` + is a colorbar-compatible mapping (all marker groups share the same + normalization and colormap). Examples -------- + Continuous color and size from an attribute, with a colorbar: + ```python from tracksdata.functional import plot_lineage_tree - ax = plot_lineage_tree(graph, color_attr="area", cmap="magma", size_attr="area") + ax = plot_lineage_tree(graph, color="area", cmap="magma", size="area") ax.figure.colorbar(ax.collections[-1], ax=ax, label="area") ``` - Display only a time window with timestamps in seconds: + Color by a computed quantity (still colorbar-compatible) and shape markers + by a categorical attribute: ```python + import numpy as np + ax = plot_lineage_tree( graph, - time_range=(10, 20), - time_positions={t: t * 30.0 for t in range(50)}, + color=lambda row: np.log1p(row["area"]), + marker=lambda row: "s" if row["is_dividing"] else "o", + attrs=["area", "is_dividing"], ) ``` - Display an arbitrary subset of time points with styled marker borders: + Categorical colors and per-node text labels: ```python + palette = {"A": "tab:red", "B": "tab:blue"} ax = plot_lineage_tree( graph, - time_points=[0, 5, 10, 15], - scatter_kwargs={"edgecolors": "black", "linewidths": 0.5}, + color=lambda row: palette[row["class"]], + text=lambda row: row["class"], + attrs=["class"], + ) + ``` + + Display only a time window with timestamps in seconds: + + ```python + ax = plot_lineage_tree( + graph, + time_range=(10, 20), + time_positions={t: t * 30.0 for t in range(50)}, ) ``` """ @@ -338,13 +433,32 @@ def plot_lineage_tree( if tracklet_id_key not in graph.node_attr_keys(): graph.assign_tracklet_ids(tracklet_id_key) + has_callable = any(callable(spec) for spec in (color, size, marker, text)) + attr_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key] - for key in (color_attr, size_attr): - if key is None or key in attr_keys: - continue - if key not in graph.node_attr_keys(): - raise ValueError(f"Attribute '{key}' not found in graph. Expected one of {graph.node_attr_keys()}") - attr_keys.append(key) + if has_callable and attrs is None: + warnings.warn( + "A `color`/`size`/`marker`/`text` callable was given without `attrs`; " + "loading all node attributes, which may be slow or memory-heavy " + "(e.g. mask attributes). Pass `attrs=[...]` to load only the keys the callables need.", + stacklevel=2, + ) + for key in graph.node_attr_keys(): + if key not in attr_keys: + attr_keys.append(key) + else: + # attribute names referenced directly (string aesthetics) plus any + # extra keys the callables need. `marker` as a string is a matplotlib + # glyph, not an attribute name, so it is not loaded. + requested = [spec for spec in (color, size, text) if isinstance(spec, str)] + if attrs is not None: + requested.extend(attrs) + for key in requested: + if key in attr_keys: + continue + if key not in graph.node_attr_keys(): + raise ValueError(f"Attribute '{key}' not found in graph. Expected one of {graph.node_attr_keys()}") + attr_keys.append(key) if time_range is not None and time_points is not None: raise ValueError("`time_range` and `time_points` are mutually exclusive, provide at most one.") @@ -396,19 +510,88 @@ def plot_lineage_tree( line_kwargs = {"color": "0.6", "linewidth": 1.0, "zorder": 1, **(line_kwargs or {})} ax.add_collection(LineCollection(segments, **line_kwargs)) - scatter_kwargs = {"zorder": 2, **(scatter_kwargs or {})} - if color_attr is not None: - if isinstance(color_norm, tuple): - color_norm = Normalize(*color_norm) - scatter_kwargs["c"] = nodes_df[color_attr].to_numpy() - scatter_kwargs["cmap"] = cmap - scatter_kwargs["norm"] = color_norm - if size_attr is not None: - scatter_kwargs["s"] = _map_to_size_range(nodes_df[size_attr].to_numpy(), size_norm, size_range) + # per-node attribute rows, only materialized when a callable needs them + rows = list(nodes_df.iter_rows(named=True)) if has_callable else [] + + # resolve the color channel to values passed to scatter's `c` + color_values: Any = None + color_is_scalar = False + if color is not None: + if callable(color): + color_values, color_is_scalar = _resolve_color_values([color(row) for row in rows]) + else: + color_values = nodes_df[color].to_numpy() + color_is_scalar = True + + # a single shared normalization so colors are consistent across marker groups + norm: Normalize | None = None + if color_is_scalar: + if color_norm is None: + norm = Normalize(vmin=float(np.nanmin(color_values)), vmax=float(np.nanmax(color_values))) + elif isinstance(color_norm, tuple): + norm = Normalize(*color_norm) + else: + norm = color_norm + + # resolve the size channel: attribute name -> mapped range, callable -> raw + # sizes, number -> constant, None -> node_size default + if size is None: + size_values: Any = None + elif callable(size): + size_values = np.asarray([size(row) for row in rows], dtype=float) + elif isinstance(size, str): + size_values = _map_to_size_range(nodes_df[size].to_numpy(), size_norm, size_range) else: - scatter_kwargs.setdefault("s", node_size) + size_values = float(size) - ax.scatter(x_coords, y_coords, **scatter_kwargs) + # resolve the marker channel: callable -> per-node glyphs (grouped), string + # -> single glyph, None -> "o" + if callable(marker): + marker_values = [marker(row) for row in rows] + else: + marker_values = None + single_marker = marker if isinstance(marker, str) else "o" + + scatter_kwargs = {"zorder": 2, **(scatter_kwargs or {})} + + def _scatter_group(idx: np.ndarray, marker_glyph: str) -> Any: + kwargs = dict(scatter_kwargs) + if color_is_scalar: + kwargs["c"] = color_values[idx] + kwargs["cmap"] = cmap + kwargs["norm"] = norm + elif color_values is not None: + kwargs["c"] = [color_values[i] for i in idx] + if size_values is None: + kwargs.setdefault("s", node_size) + elif np.isscalar(size_values): + kwargs.setdefault("s", size_values) + else: + kwargs["s"] = size_values[idx] + return ax.scatter(x_coords[idx], y_coords[idx], marker=marker_glyph, **kwargs) + + if marker_values is None: + _scatter_group(np.arange(len(nodes_df)), single_marker) + else: + marker_arr = np.asarray(marker_values, dtype=object) + # one scatter call per distinct glyph (scatter accepts a single marker) + for glyph in dict.fromkeys(marker_values): + idx = np.nonzero(marker_arr == glyph)[0] + _scatter_group(idx, glyph) + + if text is not None: + if callable(text): + labels = [text(row) for row in rows] + else: + labels = nodes_df[text].to_list() + annotate_kwargs = { + "fontsize": 8, + "xytext": (3.0, 0.0), + "textcoords": "offset points", + **(text_kwargs or {}), + } + for x, y, label in zip(x_coords, y_coords, labels, strict=True): + ax.annotate(str(label), (x, y), **annotate_kwargs) if orientation == "vertical": time_axis, tree_axis = ax.yaxis, ax.xaxis diff --git a/src/tracksdata/functional/_test/test_plot.py b/src/tracksdata/functional/_test/test_plot.py index 853d25c6..4666d46b 100644 --- a/src/tracksdata/functional/_test/test_plot.py +++ b/src/tracksdata/functional/_test/test_plot.py @@ -82,10 +82,10 @@ def test_plot_lineage_tree_color_and_size() -> None: ax = plot_lineage_tree( graph, - color_attr="feature", + color="feature", cmap="magma", color_norm=(0.0, 10.0), - size_attr="feature", + size="feature", size_range=(10.0, 50.0), ) @@ -106,7 +106,7 @@ def test_plot_lineage_tree_size_norm() -> None: """Test explicit size normalization limits with clipping.""" graph = _dividing_graph() - ax = plot_lineage_tree(graph, size_attr="feature", size_norm=(0.0, 2.0), size_range=(10.0, 50.0)) + ax = plot_lineage_tree(graph, size="feature", size_norm=(0.0, 2.0), size_range=(10.0, 50.0)) sizes = np.asarray(ax.collections[-1].get_sizes()) feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() @@ -176,7 +176,7 @@ def test_plot_lineage_tree_edge_colors() -> None: ax = plot_lineage_tree( graph, - color_attr="feature", + color="feature", scatter_kwargs={"edgecolors": "red", "linewidths": 1.5}, ) @@ -262,4 +262,137 @@ def test_plot_lineage_tree_errors() -> None: plot_lineage_tree(graph, time_range=(10, 20)) with pytest.raises(ValueError, match="not found in graph"): - plot_lineage_tree(graph, color_attr="does_not_exist") + plot_lineage_tree(graph, color="does_not_exist") + + +def test_plot_lineage_tree_color_callable_scalar() -> None: + """A color callable returning numbers is colormap-mapped and colorbar-compatible.""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + color=lambda row: 2.0 * row["feature"], + cmap="magma", + attrs=["feature"], + ) + + scatter = ax.collections[-1] + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + # numeric callable output feeds scatter's color array -> colorbar mapping exists + assert scatter.get_array() is not None + np.testing.assert_array_equal(np.asarray(scatter.get_array()), 2.0 * feature) + assert scatter.get_cmap().name == "magma" + + +def test_plot_lineage_tree_color_callable_categorical() -> None: + """A color callable returning literal colors is used verbatim (no colorbar).""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + color=lambda row: "red" if row["feature"] < 3.0 else "blue", + attrs=["feature"], + ) + + scatter = ax.collections[-1] + # literal colors -> no scalar mapping + assert scatter.get_array() is None + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + facecolors = scatter.get_facecolors() + red = np.array([1.0, 0.0, 0.0, 1.0]) + blue = np.array([0.0, 0.0, 1.0, 1.0]) + expected = np.where((feature < 3.0)[:, None], red, blue) + np.testing.assert_allclose(facecolors, expected) + + +def test_plot_lineage_tree_size_callable() -> None: + """A size callable returns marker sizes directly.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, size=lambda row: 5.0 + row["feature"], attrs=["feature"]) + + sizes = np.asarray(ax.collections[-1].get_sizes()) + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + np.testing.assert_allclose(sizes, 5.0 + feature) + + +def test_plot_lineage_tree_size_constant() -> None: + """A numeric size sets a constant marker size.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, size=42.0) + + sizes = np.asarray(ax.collections[-1].get_sizes()) + np.testing.assert_allclose(sizes, [42.0]) + + +def test_plot_lineage_tree_marker_callable_groups() -> None: + """A marker callable groups nodes by glyph into one scatter call each.""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + marker=lambda row: "s" if row["feature"] < 3.0 else "^", + attrs=["feature"], + ) + + # one LineCollection for edges + one PathCollection per distinct glyph + scatters = ax.collections[1:] + assert len(scatters) == 2 + total = sum(len(s.get_offsets()) for s in scatters) + assert total == graph.num_nodes() + + +def test_plot_lineage_tree_marker_single_glyph() -> None: + """A marker string applies a single glyph to every node in one scatter call.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, marker="s") + + _lines, scatter = ax.collections + assert len(scatter.get_offsets()) == graph.num_nodes() + + +def test_plot_lineage_tree_marker_and_color_share_norm() -> None: + """Marker groups share one normalization so colors stay consistent and colorbar-compatible.""" + graph = _dividing_graph() + + ax = plot_lineage_tree( + graph, + color="feature", + color_norm=(0.0, 10.0), + marker=lambda row: "s" if row["feature"] < 3.0 else "^", + attrs=["feature"], + ) + + scatters = ax.collections[1:] + assert len(scatters) == 2 + for scatter in scatters: + assert scatter.norm.vmin == 0.0 + assert scatter.norm.vmax == 10.0 + assert scatter.get_cmap().name == "viridis" + + +def test_plot_lineage_tree_text() -> None: + """Text labels are annotated per node, from an attribute name or a callable.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, text="feature") + texts = {t.get_text() for t in ax.texts} + assert len(ax.texts) == graph.num_nodes() + assert "0.0" in texts + + _, ax2 = plt.subplots() + plot_lineage_tree(graph, ax=ax2, text=lambda row: f"n{row['feature']:.0f}", attrs=["feature"]) + labels = {t.get_text() for t in ax2.texts} + assert "n0" in labels and "n5" in labels + + +def test_plot_lineage_tree_callable_without_attrs_warns() -> None: + """A callable without `attrs` warns and still plots by loading all attributes.""" + graph = _dividing_graph() + + with pytest.warns(UserWarning, match="loading all node attributes"): + ax = plot_lineage_tree(graph, color=lambda row: row["feature"]) + + assert len(ax.collections[-1].get_offsets()) == graph.num_nodes()