From bc1dea08add011c791981093cce7dce51984e9a2 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Fri, 24 Jul 2026 09:03:26 +0900 Subject: [PATCH 1/6] 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/6] 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() From 1dd825839567de63e8a2e3a714853b22668a5fb7 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Thu, 10 Sep 2026 09:22:15 +0900 Subject: [PATCH 3/6] update --- src/tracksdata/functional/_plot.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py index 0c7ad8c6..d2f27083 100644 --- a/src/tracksdata/functional/_plot.py +++ b/src/tracksdata/functional/_plot.py @@ -76,11 +76,19 @@ def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: ) for root in roots: - # iterative post-order traversal: children are positioned before parents + # Iterative post-order DFS: children are positioned before their parents. + # DFS visits a subtree contiguously, so its leaves take a consecutive block of + # slots and its parents, being means of their children, land inside that block. + # Sibling blocks are therefore disjoint and no edges cross. + # + # For 1 -> (2, 3) and 3 -> (4, 5), tracklets are positioned in the order + # 2, 4, 5, 3, 1, giving 2: 0.0, 4: 1.0, 5: 2.0, 3: 1.5, 1: 0.75. stack: list[tuple[int, bool]] = [(root, False)] while stack: + # `rx_id` is a rustworkx node index, `tracklet_graph[rx_id]` the tracklet id rx_id, expanded = stack.pop() if expanded: + # second visit: every child has been positioned already children_pos = [ positions[tracklet_graph[child]] for child in tracklet_graph.successor_indices(rx_id) @@ -89,16 +97,20 @@ def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: if children_pos: positions[tracklet_graph[rx_id]] = float(np.mean(children_pos)) else: + # leaf tracklet: take the next free slot positions[tracklet_graph[rx_id]] = next_leaf next_leaf += 1.0 elif rx_id not in visited: + # first visit: re-push self as expanded, then push children on top visited.add(rx_id) stack.append((rx_id, True)) + # reversed order, so LIFO pops children by ascending tracklet id for child in sorted( tracklet_graph.successor_indices(rx_id), key=tracklet_graph.__getitem__, reverse=True, ): + # skip children already reached via another parent (merges) if child not in visited: stack.append((child, False)) From 4b5d1093605d24e1f5b6270bd00865a43ed4bdd6 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Thu, 10 Sep 2026 14:44:57 +0900 Subject: [PATCH 4/6] feat(functional): add edge_color and slim down plot_lineage_tree signature - add `edge_color`, resolved like `color` (attribute name, callable, or literal color); numeric edge values share `cmap`/`color_norm` with the face color and are mapped manually since scatter only colormaps `c` - a `color`/`edge_color` string that is not a node attribute is taken as a literal matplotlib color (previously a ValueError) - drop `node_size`: `size` now defaults to 30.0 and accepts a float - drop `time_range`: `time_points=range(a, b + 1)` covers it Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WwHX5ciJfQPm9z6zmBsb27 --- src/tracksdata/functional/_plot.py | 192 ++++++++++++------- src/tracksdata/functional/_test/test_plot.py | 94 +++++++-- 2 files changed, 210 insertions(+), 76 deletions(-) diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py index d2f27083..59d6201e 100644 --- a/src/tracksdata/functional/_plot.py +++ b/src/tracksdata/functional/_plot.py @@ -25,7 +25,7 @@ def _resolve_color_values(raw: list) -> tuple[Any, bool]: Parameters ---------- raw : list - One value per node, as returned by a `color` callable. + One value per node, as returned by a `color`/`edge_color` callable. Returns ------- @@ -46,6 +46,40 @@ def _resolve_color_values(raw: list) -> tuple[Any, bool]: return list(raw), False +def _resolve_color_channel( + spec: "str | Callable[[Mapping[str, Any]], Any] | None", + nodes_df: Any, + rows: list[Mapping[str, Any]], +) -> tuple[Any, bool]: + """ + Resolve a `color`/`edge_color` specification into per-node color values. + + Parameters + ---------- + spec : str | Callable | None + An attribute name (loaded in `nodes_df`), a literal matplotlib color, + a callable of the node attribute row, or None. + nodes_df : pl.DataFrame + Displayed nodes with the attributes referenced by string specs loaded. + rows : list[Mapping[str, Any]] + Per-node attribute rows, fed to callables (unused otherwise). + + Returns + ------- + tuple[Any, bool] + `(values, is_scalar)` as in `_resolve_color_values`, or `(None, False)` + when `spec` is None. A literal color string becomes one entry per node. + """ + if spec is None: + return None, False + if callable(spec): + return _resolve_color_values([spec(row) for row in rows]) + if spec in nodes_df.columns: + return nodes_df[spec].to_numpy(), True + # validated upfront as a matplotlib color: same literal for every node + return [spec] * len(nodes_df), False + + def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: """ Assign a tree-axis coordinate to each tracklet of a tracklet graph. @@ -244,17 +278,16 @@ def plot_lineage_tree( ax: "Axes | None" = None, tracklet_id_key: str = DEFAULT_ATTR_KEYS.TRACKLET_ID, color: "str | Callable[[Mapping[str, Any]], Any] | None" = None, + edge_color: "str | Callable[[Mapping[str, Any]], Any] | None" = None, cmap: "str | Colormap" = "viridis", color_norm: "Normalize | tuple[float, float] | None" = None, - size: "str | Callable[[Mapping[str, Any]], float] | float | None" = None, + size: "str | Callable[[Mapping[str, Any]], float] | float" = 30.0, 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, orientation: Literal["vertical", "horizontal"] = "vertical", @@ -271,11 +304,15 @@ 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: + The `color`, `edge_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 string, `color`/`edge_color`/`size`/`text` name a numeric node + attribute, and `marker` is a single matplotlib marker glyph applied to + every node. A `color`/`edge_color` string that is not an attribute name + is taken as a literal matplotlib color (e.g. `"tab:red"`, `"none"`) + 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, @@ -284,7 +321,9 @@ def plot_lineage_tree( 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. + those colors are used verbatim and no colorbar mapping exists. Numeric + `edge_color` values are mapped through the same `cmap`/`color_norm` as + `color`, so face and edge colors are directly comparable. Requires `matplotlib`, which is an optional dependency (`pip install "tracksdata[plot]"`). @@ -306,22 +345,30 @@ def plot_lineage_tree( [BaseGraph.assign_tracklet_ids][tracksdata.graph.BaseGraph.assign_tracklet_ids] is called first. 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. + Marker face color. A string names a numeric node attribute mapped + through `cmap`/`color_norm` (a colorbar mapping is available); if it + is not an attribute name, it is a literal matplotlib color applied to + every node. 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. + edge_color : str | Callable | None, optional + Marker edge (border) color, resolved exactly like `color`: an attribute + name or numeric callable output is mapped through the shared + `cmap`/`color_norm`, a literal color or a callable returning literal + colors is used as-is. If None, edges take matplotlib's default (the + face color). Set the border width with `scatter_kwargs={"linewidths": ...}`. cmap : str | Colormap, optional Colormap used when `color` yields numeric values. color_norm : Normalize | tuple[float, float] | None, optional Normalization for numeric colors, either a matplotlib `Normalize` instance or a `(vmin, vmax)` tuple. If None, the data range is used. A single shared normalization is applied across all marker groups. - size : str | Callable | float | None, optional + size : str | Callable | float, 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. + the marker size in points**2 directly. A number sets a constant size + in points**2 for every node. size_norm : tuple[float, float] | None, optional The `(vmin, vmax)` attribute values mapped to the limits of `size_range`, used when `size` is an attribute name. If None, the data @@ -329,8 +376,6 @@ def plot_lineage_tree( size_range : tuple[float, float], optional The marker sizes in points**2 assigned to the smallest and largest values when `size` is an attribute name. - node_size : float, optional - 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 @@ -350,15 +395,11 @@ def plot_lineage_tree( 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 - `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 + Time points to display, e.g. `range(10, 21)` for a contiguous window + or `[0, 5, 10]` for a sparse subset. Edges bridge over the hidden time points, connecting each displayed node to its nearest displayed - descendants. Mutually exclusive with `time_range`. + descendants. If None, all time points are displayed. 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 @@ -369,7 +410,7 @@ def plot_lineage_tree( 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. + e.g. `linewidths` to set the marker border width or `alpha`. line_kwargs : dict[str, Any] | None, optional Additional keyword arguments forwarded to the edge `LineCollection` (e.g. `color`, `linewidth`). @@ -407,6 +448,18 @@ def plot_lineage_tree( ) ``` + Outline dividing cells on top of a continuous face color: + + ```python + ax = plot_lineage_tree( + graph, + color="area", + edge_color=lambda row: "black" if row["is_dividing"] else "none", + attrs=["is_dividing"], + scatter_kwargs={"linewidths": 1.5}, + ) + ``` + Categorical colors and per-node text labels: ```python @@ -424,7 +477,7 @@ def plot_lineage_tree( ```python ax = plot_lineage_tree( graph, - time_range=(10, 20), + time_points=range(10, 21), time_positions={t: t * 30.0 for t in range(50)}, ) ``` @@ -432,7 +485,7 @@ def plot_lineage_tree( try: import matplotlib.pyplot as plt from matplotlib.collections import LineCollection - from matplotlib.colors import Normalize + from matplotlib.colors import Normalize, is_color_like except ImportError as e: raise ImportError( "matplotlib is required for `plot_lineage_tree`. " @@ -442,49 +495,58 @@ def plot_lineage_tree( 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(): + node_attr_keys = graph.node_attr_keys() + if tracklet_id_key not in node_attr_keys: graph.assign_tracklet_ids(tracklet_id_key) - has_callable = any(callable(spec) for spec in (color, size, marker, text)) + has_callable = any(callable(spec) for spec in (color, edge_color, size, marker, text)) + + # a `color`/`edge_color` string is an attribute name when one exists, + # otherwise it must be a literal matplotlib color + color_attr_keys = [] + for spec in (color, edge_color): + if not isinstance(spec, str): + continue + if spec in node_attr_keys: + color_attr_keys.append(spec) + elif not is_color_like(spec): + raise ValueError( + f"Color '{spec}' not found in graph attributes and is not a valid matplotlib color. " + f"Expected a color or one of {node_attr_keys}" + ) attr_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key] if has_callable and attrs is None: warnings.warn( - "A `color`/`size`/`marker`/`text` callable was given without `attrs`; " + "A `color`/`edge_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(): + for key in 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)] + requested = color_attr_keys + [spec for spec in (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()}") + if key not in node_attr_keys: + raise ValueError(f"Attribute '{key}' not found in graph. Expected one of {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: + if 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.") + raise ValueError("No nodes to plot. The graph is empty or `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 @@ -525,32 +587,32 @@ def plot_lineage_tree( # 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 + # resolve the face color channel (scatter's `c`) and the edge color channel + color_values, color_is_scalar = _resolve_color_channel(color, nodes_df, rows) + edge_values, edge_is_scalar = _resolve_color_channel(edge_color, nodes_df, rows) - # a single shared normalization so colors are consistent across marker groups + # a single shared normalization so colors are consistent across marker + # groups and between face and edge colors norm: Normalize | None = None - if color_is_scalar: + scalar_channels = [ + values for values, is_scalar in ((color_values, color_is_scalar), (edge_values, edge_is_scalar)) if is_scalar + ] + if scalar_channels: if color_norm is None: - norm = Normalize(vmin=float(np.nanmin(color_values)), vmax=float(np.nanmax(color_values))) + stacked = np.concatenate(scalar_channels) + norm = Normalize(vmin=float(np.nanmin(stacked)), vmax=float(np.nanmax(stacked))) elif isinstance(color_norm, tuple): norm = Normalize(*color_norm) else: norm = color_norm + # scatter only colormaps `c`, so numeric edge colors are mapped here + edge_cmap = plt.get_cmap(cmap) if edge_is_scalar else None + # 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) + # sizes, number -> constant + if callable(size): + size_values: Any = 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: @@ -574,9 +636,11 @@ def _scatter_group(idx: np.ndarray, marker_glyph: str) -> Any: 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): + if edge_is_scalar: + kwargs["edgecolors"] = edge_cmap(norm(edge_values[idx])) + elif edge_values is not None: + kwargs["edgecolors"] = [edge_values[i] for i in idx] + if np.isscalar(size_values): kwargs.setdefault("s", size_values) else: kwargs["s"] = size_values[idx] diff --git a/src/tracksdata/functional/_test/test_plot.py b/src/tracksdata/functional/_test/test_plot.py index 4666d46b..1c6857f2 100644 --- a/src/tracksdata/functional/_test/test_plot.py +++ b/src/tracksdata/functional/_test/test_plot.py @@ -114,11 +114,11 @@ def test_plot_lineage_tree_size_norm() -> None: np.testing.assert_allclose(sizes, expected) -def test_plot_lineage_tree_time_range() -> None: - """Test that time_range limits the displayed nodes and edges.""" +def test_plot_lineage_tree_time_points_window() -> None: + """A contiguous `time_points` window (a range) limits the displayed nodes and edges.""" graph = _dividing_graph() - ax = plot_lineage_tree(graph, time_range=(1, 2)) + ax = plot_lineage_tree(graph, time_points=range(1, 3)) lines, scatter = ax.collections # nodes: t=1 (tracklet 1) and t=2 (tracklets 2 and 3) @@ -162,14 +162,6 @@ def test_plot_lineage_tree_time_points() -> None: 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() @@ -259,7 +251,7 @@ def test_plot_lineage_tree_errors() -> None: graph = _dividing_graph() with pytest.raises(ValueError, match="No nodes to plot"): - plot_lineage_tree(graph, time_range=(10, 20)) + plot_lineage_tree(graph, time_points=range(10, 21)) with pytest.raises(ValueError, match="not found in graph"): plot_lineage_tree(graph, color="does_not_exist") @@ -396,3 +388,81 @@ def test_plot_lineage_tree_callable_without_attrs_warns() -> None: ax = plot_lineage_tree(graph, color=lambda row: row["feature"]) assert len(ax.collections[-1].get_offsets()) == graph.num_nodes() + + +def test_plot_lineage_tree_color_literal() -> None: + """A `color` string that is not an attribute name is a literal color for every node.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, color="tab:blue") + + scatter = ax.collections[-1] + assert scatter.get_array() is None + expected = np.tile(matplotlib.colors.to_rgba("tab:blue"), (graph.num_nodes(), 1)) + np.testing.assert_allclose(scatter.get_facecolors(), expected) + + +def test_plot_lineage_tree_edge_color_literal() -> None: + """A literal `edge_color` outlines every node without touching the face colors.""" + graph = _dividing_graph() + + ax = plot_lineage_tree(graph, color="feature", edge_color="red") + + scatter = ax.collections[-1] + expected = np.tile([1.0, 0.0, 0.0, 1.0], (graph.num_nodes(), 1)) + np.testing.assert_allclose(scatter.get_edgecolors(), expected) + np.testing.assert_array_equal(np.asarray(scatter.get_array()), graph.node_attrs(attr_keys=["feature"])["feature"]) + + +def test_plot_lineage_tree_edge_color_attribute_shares_cmap_and_norm() -> None: + """A numeric `edge_color` is mapped through the same colormap and normalization as `color`.""" + graph = _dividing_graph() + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + + ax = plot_lineage_tree(graph, color="feature", edge_color="feature", cmap="magma") + + scatter = ax.collections[-1] + expected = plt.get_cmap("magma")(matplotlib.colors.Normalize(0.0, 5.0)(feature)) + np.testing.assert_allclose(scatter.get_edgecolors(), expected) + # colormapped face colors are only resolved when drawn + ax.figure.canvas.draw() + np.testing.assert_allclose(scatter.get_facecolors(), expected) + + # the shared normalization spans the union of face and edge value ranges + _, ax2 = plt.subplots() + plot_lineage_tree( + graph, + ax=ax2, + color="feature", + edge_color=lambda row: 2.0 * row["feature"], + attrs=["feature"], + ) + assert ax2.collections[-1].norm.vmin == 0.0 + assert ax2.collections[-1].norm.vmax == 10.0 + + +def test_plot_lineage_tree_edge_color_callable_categorical() -> None: + """An `edge_color` callable returning literal colors is used per node, also across marker groups.""" + graph = _dividing_graph() + feature = graph.node_attrs(attr_keys=["feature"])["feature"].to_numpy() + + ax = plot_lineage_tree( + graph, + edge_color=lambda row: "red" if row["feature"] < 3.0 else "blue", + marker=lambda row: "s" if row["feature"] < 3.0 else "^", + attrs=["feature"], + ) + + red = np.array([1.0, 0.0, 0.0, 1.0]) + blue = np.array([0.0, 0.0, 1.0, 1.0]) + squares, triangles = ax.collections[1:] + np.testing.assert_allclose(squares.get_edgecolors(), np.tile(red, (int((feature < 3.0).sum()), 1))) + np.testing.assert_allclose(triangles.get_edgecolors(), np.tile(blue, (int((feature >= 3.0).sum()), 1))) + + +def test_plot_lineage_tree_edge_color_invalid() -> None: + """An `edge_color` string that is neither an attribute nor a color is rejected.""" + graph = _dividing_graph() + + with pytest.raises(ValueError, match="not found in graph"): + plot_lineage_tree(graph, edge_color="does_not_exist") From 88db0631b540df725768eca0d47de90beb235b43 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Thu, 10 Sep 2026 15:11:31 +0900 Subject: [PATCH 5/6] refactor(functional): simplify plot_lineage_tree internals - resolve requested attribute keys through a single validation path and dedupe with dict.fromkeys; typos in `size`/`text` now raise ValueError even when a callable triggers loading all attributes - fold `_resolve_color_values` into `_resolve_color_channel` - draw every marker glyph through one grouped scatter loop, inlining the former `_scatter_group` closure; scalar sizes are expanded to an array so `size` consistently overrides `scatter_kwargs["s"]` - build the successor adjacency with a polars group_by and look up node coordinates with `Series.replace_strict` - compute the orientation flag once and label the time axis via `Axis.set_label_text` - rename the shadowed `time_points` local to `shown_time_points` Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WwHX5ciJfQPm9z6zmBsb27 --- src/tracksdata/functional/_plot.py | 165 +++++++------------ src/tracksdata/functional/_test/test_plot.py | 2 +- 2 files changed, 61 insertions(+), 106 deletions(-) diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py index 59d6201e..498b96ee 100644 --- a/src/tracksdata/functional/_plot.py +++ b/src/tracksdata/functional/_plot.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, Literal import numpy as np +import polars as pl import rustworkx as rx from numpy.typing import ArrayLike @@ -18,37 +19,9 @@ __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`/`edge_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 _resolve_color_channel( spec: "str | Callable[[Mapping[str, Any]], Any] | None", - nodes_df: Any, + nodes_df: pl.DataFrame, rows: list[Mapping[str, Any]], ) -> tuple[Any, bool]: """ @@ -67,17 +40,29 @@ def _resolve_color_channel( Returns ------- tuple[Any, bool] - `(values, is_scalar)` as in `_resolve_color_values`, or `(None, False)` - when `spec` is None. A literal color string becomes one entry per node. + `(values, is_scalar)`. If the values 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 a + list of literal colors (names, hex, or RGB(A) tuples), one per node, + and `is_scalar` is False. `(None, False)` when `spec` is None. """ if spec is None: return None, False - if callable(spec): - return _resolve_color_values([spec(row) for row in rows]) - if spec in nodes_df.columns: - return nodes_df[spec].to_numpy(), True - # validated upfront as a matplotlib color: same literal for every node - return [spec] * len(nodes_df), False + if not callable(spec): + if spec in nodes_df.columns: + return nodes_df[spec].to_numpy(), True + # validated upfront as a matplotlib color: same literal for every node + return [spec] * len(nodes_df), False + + raw = [spec(row) for row in rows] + try: + arr = np.asarray(raw, dtype=float) + except (ValueError, TypeError): + return raw, False + if arr.ndim == 1: + return arr, True + # (N, 3) or (N, 4): literal RGB(A) colors, not colormap-able scalars + return raw, False def _tracklet_tree_layout(tracklet_graph: rx.PyDiGraph) -> dict[int, float]: @@ -411,6 +396,8 @@ def plot_lineage_tree( scatter_kwargs : dict[str, Any] | None, optional Additional keyword arguments forwarded to `Axes.scatter`, e.g. `linewidths` to set the marker border width or `alpha`. + `c`, `edgecolors`, `s`, and `marker` are set from the aesthetics + above and take precedence. line_kwargs : dict[str, Any] | None, optional Additional keyword arguments forwarded to the edge `LineCollection` (e.g. `color`, `linewidth`). @@ -494,6 +481,7 @@ def plot_lineage_tree( if orientation not in ("vertical", "horizontal"): raise ValueError(f"`orientation` must be 'vertical' or 'horizontal', got '{orientation}'.") + vertical = orientation == "vertical" node_attr_keys = graph.node_attr_keys() if tracklet_id_key not in node_attr_keys: @@ -515,7 +503,13 @@ def plot_lineage_tree( f"Expected a color or one of {node_attr_keys}" ) - attr_keys = [DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key] + # 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 = color_attr_keys + [spec for spec in (size, text) if isinstance(spec, str)] + list(attrs or []) + missing = [key for key in requested if key not in node_attr_keys] + if missing: + raise ValueError(f"Attributes {missing} not found in graph. Expected one of {node_attr_keys}") if has_callable and attrs is None: warnings.warn( "A `color`/`edge_color`/`size`/`marker`/`text` callable was given without `attrs`; " @@ -523,22 +517,8 @@ def plot_lineage_tree( "(e.g. mask attributes). Pass `attrs=[...]` to load only the keys the callables need.", stacklevel=2, ) - for key in 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 = color_attr_keys + [spec for spec in (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 node_attr_keys: - raise ValueError(f"Attribute '{key}' not found in graph. Expected one of {node_attr_keys}") - attr_keys.append(key) + requested = node_attr_keys + attr_keys = list(dict.fromkeys([DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key, *requested])) nodes_df = graph.node_attrs(attr_keys=attr_keys) @@ -552,30 +532,19 @@ def plot_lineage_tree( # 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]]) + shown_time_points = nodes_df[DEFAULT_ATTR_KEYS.T].unique().sort().to_list() + time_axis_positions = _time_axis_positions(shown_time_points, time_positions) - if orientation == "vertical": - x_coords, y_coords = tree_coords, time_coords - else: - x_coords, y_coords = time_coords, tree_coords + tree_coords = nodes_df[tracklet_id_key].replace_strict(tracklet_positions, return_dtype=pl.Float64).to_numpy() + time_coords = nodes_df[DEFAULT_ATTR_KEYS.T].replace_strict(time_axis_positions, return_dtype=pl.Float64).to_numpy() + x_coords, y_coords = (tree_coords, time_coords) if vertical else (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) - + successors = dict(edges_df.group_by(DEFAULT_ATTR_KEYS.EDGE_SOURCE).agg(DEFAULT_ATTR_KEYS.EDGE_TARGET).iter_rows()) segments = _bridged_edge_segments(successors, node_coords) if ax is None: @@ -609,26 +578,28 @@ def plot_lineage_tree( # scatter only colormaps `c`, so numeric edge colors are mapped here edge_cmap = plt.get_cmap(cmap) if edge_is_scalar else None - # resolve the size channel: attribute name -> mapped range, callable -> raw - # sizes, number -> constant + # resolve the size channel to one size per node: attribute name -> mapped + # range, callable -> raw sizes, number -> constant if callable(size): - size_values: Any = np.asarray([size(row) for row in rows], dtype=float) + 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: - size_values = float(size) + size_values = np.full(len(nodes_df), float(size)) - # resolve the marker channel: callable -> per-node glyphs (grouped), string - # -> single glyph, None -> "o" + # resolve the marker channel: callable -> per-node glyphs, 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" + marker_values = [marker or "o"] * len(nodes_df) + marker_arr = np.asarray(marker_values, dtype=object) scatter_kwargs = {"zorder": 2, **(scatter_kwargs or {})} - def _scatter_group(idx: np.ndarray, marker_glyph: str) -> Any: + # 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] kwargs = dict(scatter_kwargs) if color_is_scalar: kwargs["c"] = color_values[idx] @@ -640,20 +611,8 @@ def _scatter_group(idx: np.ndarray, marker_glyph: str) -> Any: kwargs["edgecolors"] = edge_cmap(norm(edge_values[idx])) elif edge_values is not None: kwargs["edgecolors"] = [edge_values[i] for i in idx] - if 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) + kwargs["s"] = size_values[idx] + ax.scatter(x_coords[idx], y_coords[idx], marker=glyph, **kwargs) if text is not None: if callable(text): @@ -669,21 +628,17 @@ def _scatter_group(idx: np.ndarray, marker_glyph: str) -> Any: 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 - 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") - + time_axis, tree_axis = (ax.yaxis, ax.xaxis) if vertical else (ax.xaxis, ax.yaxis) + time_axis.set_label_text("time") tree_axis.set_ticks([]) + # time runs downward in the vertical layout + if vertical and not ax.yaxis_inverted(): + ax.invert_yaxis() 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] + stride = max(1, len(shown_time_points) // 10) + ticks = shown_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 index 1c6857f2..d91f0d0a 100644 --- a/src/tracksdata/functional/_test/test_plot.py +++ b/src/tracksdata/functional/_test/test_plot.py @@ -315,7 +315,7 @@ def test_plot_lineage_tree_size_constant() -> None: ax = plot_lineage_tree(graph, size=42.0) sizes = np.asarray(ax.collections[-1].get_sizes()) - np.testing.assert_allclose(sizes, [42.0]) + np.testing.assert_allclose(sizes, np.full(graph.num_nodes(), 42.0)) def test_plot_lineage_tree_marker_callable_groups() -> None: From de92725a56d04afdc9e468aa6cfe23ee3bdd2317 Mon Sep 17 00:00:00 2001 From: Yohsuke Fukai Date: Thu, 10 Sep 2026 15:30:24 +0900 Subject: [PATCH 6/6] feat(functional): draw lineage edges through hidden time points - `time_points` now only limits where markers are drawn; edges are drawn for every node between the first and last displayed time point, so divisions stay at their true time instead of being bridged - with `time_positions=None`, all time points in the displayed range are evenly separated (hidden ones included) and only the displayed ones are labeled - build edge segments with two polars joins on node coordinates and drop the `_bridged_edge_segments` walk - map attribute sizes with `matplotlib.colors.Normalize(clip=True)` instead of `_map_to_size_range`; `size_norm` also accepts a `Normalize` instance, mirroring `color_norm` Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01WwHX5ciJfQPm9z6zmBsb27 --- src/tracksdata/functional/_plot.py | 170 ++++++------------- src/tracksdata/functional/_test/test_plot.py | 40 ++--- 2 files changed, 77 insertions(+), 133 deletions(-) diff --git a/src/tracksdata/functional/_plot.py b/src/tracksdata/functional/_plot.py index 498b96ee..f5ba49b9 100644 --- a/src/tracksdata/functional/_plot.py +++ b/src/tracksdata/functional/_plot.py @@ -146,7 +146,8 @@ def _time_axis_positions( Parameters ---------- time_points : list[int] - Sorted unique time points to be displayed. + Sorted unique time points needing a coordinate: the displayed ones + and the hidden ones in between that edges pass through. 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. @@ -177,86 +178,6 @@ def _time_axis_positions( 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, *, @@ -267,7 +188,7 @@ def plot_lineage_tree( cmap: "str | Colormap" = "viridis", color_norm: "Normalize | tuple[float, float] | None" = None, size: "str | Callable[[Mapping[str, Any]], float] | float" = 30.0, - size_norm: tuple[float, float] | None = None, + size_norm: "Normalize | tuple[float, float] | None" = None, size_range: tuple[float, float] = (10.0, 100.0), marker: "str | Callable[[Mapping[str, Any]], str] | None" = None, text: "str | Callable[[Mapping[str, Any]], Any] | None" = None, @@ -285,9 +206,9 @@ def plot_lineage_tree( 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. + subset of time points is shown, markers are limited to those time points + while edges still run through the hidden nodes in between, so the tree + keeps its shape and divisions stay at their true time. The `color`, `edge_color`, `size`, `marker`, and `text` aesthetics each accept either a fixed value or a callable, which is the main way to @@ -354,9 +275,10 @@ def plot_lineage_tree( `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 in points**2 for every node. - size_norm : tuple[float, float] | None, optional - The `(vmin, vmax)` attribute values mapped to the limits of - `size_range`, used when `size` is an attribute name. If None, the data + size_norm : Normalize | tuple[float, float] | None, optional + Normalization of attribute values onto `size_range`, used when `size` + is an attribute name: a matplotlib `Normalize` instance or a + `(vmin, vmax)` tuple (values outside are clipped). 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 @@ -381,15 +303,17 @@ def plot_lineage_tree( or memory-heavy (e.g. mask attributes). Ignored keys already loaded for other reasons are harmless. time_points : Sequence[int] | None, optional - Time points to display, e.g. `range(10, 21)` for a contiguous window - or `[0, 5, 10]` for a sparse subset. Edges bridge over the hidden time - points, connecting each displayed node to its nearest displayed - descendants. If None, all time points are displayed. + Time points at which markers are drawn, e.g. `range(10, 21)` for a + contiguous window or `[0, 5, 10]` for a sparse subset. Edges are drawn + for every node between the first and last displayed time point, hidden + ones included. If None, all time points are displayed. 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. + position or a sequence indexed by time point, covering every time + point within the displayed range. If None, every time point within + the displayed range is evenly separated, hidden ones included, and + the displayed ones are 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. @@ -520,32 +444,44 @@ def plot_lineage_tree( requested = node_attr_keys attr_keys = list(dict.fromkeys([DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.T, tracklet_id_key, *requested])) - nodes_df = graph.node_attrs(attr_keys=attr_keys) + all_nodes_df = graph.node_attrs(attr_keys=attr_keys) + shown_time_points = all_nodes_df[DEFAULT_ATTR_KEYS.T].unique().sort() if 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: + shown_time_points = shown_time_points.filter(shown_time_points.is_in(list(time_points))) + shown_time_points = shown_time_points.to_list() + if not shown_time_points: raise ValueError("No nodes to plot. The graph is empty or `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)) - shown_time_points = nodes_df[DEFAULT_ATTR_KEYS.T].unique().sort().to_list() - time_axis_positions = _time_axis_positions(shown_time_points, time_positions) - - tree_coords = nodes_df[tracklet_id_key].replace_strict(tracklet_positions, return_dtype=pl.Float64).to_numpy() - time_coords = nodes_df[DEFAULT_ATTR_KEYS.T].replace_strict(time_axis_positions, return_dtype=pl.Float64).to_numpy() - x_coords, y_coords = (tree_coords, time_coords) if vertical else (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 run through every node inside the displayed time span, hidden time + # points included, so divisions appear at their true time + span_df = all_nodes_df.filter(pl.col(DEFAULT_ATTR_KEYS.T).is_between(shown_time_points[0], shown_time_points[-1])) + span_time_points = span_df[DEFAULT_ATTR_KEYS.T].unique().sort().to_list() + time_axis_positions = _time_axis_positions(span_time_points, time_positions) + + tree_coords = span_df[tracklet_id_key].replace_strict(tracklet_positions, return_dtype=pl.Float64).to_numpy() + time_coords = span_df[DEFAULT_ATTR_KEYS.T].replace_strict(time_axis_positions, return_dtype=pl.Float64).to_numpy() + span_x, span_y = (tree_coords, time_coords) if vertical else (time_coords, tree_coords) + + # edge segments as (E, 2, 2) start/end points; the inner joins drop edges + # with an endpoint outside the span + node_id, source, target = DEFAULT_ATTR_KEYS.NODE_ID, DEFAULT_ATTR_KEYS.EDGE_SOURCE, DEFAULT_ATTR_KEYS.EDGE_TARGET + coords_df = pl.DataFrame({node_id: span_df[node_id], "x": span_x, "y": span_y}) + segments_df = ( + graph.edge_attrs(attr_keys=[]) + .join(coords_df.rename({node_id: source, "x": "x0", "y": "y0"}), on=source) + .join(coords_df.rename({node_id: target, "x": "x1", "y": "y1"}), on=target) + ) + segments = segments_df.select("x0", "y0", "x1", "y1").to_numpy().reshape(-1, 2, 2) - edges_df = graph.edge_attrs(attr_keys=[]) - successors = dict(edges_df.group_by(DEFAULT_ATTR_KEYS.EDGE_SOURCE).agg(DEFAULT_ATTR_KEYS.EDGE_TARGET).iter_rows()) - segments = _bridged_edge_segments(successors, node_coords) + # markers only at the displayed time points + is_shown = span_df[DEFAULT_ATTR_KEYS.T].is_in(shown_time_points).to_numpy() + nodes_df = span_df.filter(is_shown) + x_coords, y_coords = span_x[is_shown], span_y[is_shown] if ax is None: _, ax = plt.subplots() @@ -578,12 +514,18 @@ def plot_lineage_tree( # scatter only colormaps `c`, so numeric edge colors are mapped here edge_cmap = plt.get_cmap(cmap) if edge_is_scalar else None - # resolve the size channel to one size per node: attribute name -> mapped - # range, callable -> raw sizes, number -> constant + # resolve the size channel to one size per node: attribute name -> linearly + # mapped into `size_range`, callable -> raw sizes, number -> constant if 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) + if isinstance(size_norm, Normalize): + size_scale = size_norm + else: + # autoscales to the data range when no limits are given + size_scale = Normalize(*(size_norm or (None, None)), clip=True) + fraction = size_scale(np.ma.masked_invalid(nodes_df[size].to_numpy().astype(float))) + size_values = np.ma.filled(size_range[0] + fraction * (size_range[1] - size_range[0]), np.nan) else: size_values = np.full(len(nodes_df), float(size)) diff --git a/src/tracksdata/functional/_test/test_plot.py b/src/tracksdata/functional/_test/test_plot.py index d91f0d0a..7c3d622b 100644 --- a/src/tracksdata/functional/_test/test_plot.py +++ b/src/tracksdata/functional/_test/test_plot.py @@ -113,6 +113,12 @@ def test_plot_lineage_tree_size_norm() -> None: expected = 10.0 + np.clip(feature / 2.0, 0.0, 1.0) * 40.0 np.testing.assert_allclose(sizes, expected) + # a matplotlib Normalize instance is accepted as well + _, ax2 = plt.subplots() + norm = matplotlib.colors.Normalize(0.0, 2.0, clip=True) + plot_lineage_tree(graph, ax=ax2, size="feature", size_norm=norm, size_range=(10.0, 50.0)) + np.testing.assert_allclose(np.asarray(ax2.collections[-1].get_sizes()), expected) + def test_plot_lineage_tree_time_points_window() -> None: """A contiguous `time_points` window (a range) limits the displayed nodes and edges.""" @@ -138,29 +144,25 @@ def test_plot_lineage_tree_time_points() -> None: 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 + # markers: t=0 (tracklet 1) and t=3 (tracklets 2 and 3) offsets = np.asarray(scatter.get_offsets()) - np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 1])), [0.0, 1.0]) + assert len(offsets) == 3 + # every time point in the displayed range is evenly separated, hidden ones + # included, and only the displayed ones are labeled + np.testing.assert_array_equal(np.sort(np.unique(offsets[:, 1])), [0.0, 3.0]) labels = [tick.get_text() for tick in ax.get_yticklabels()] assert labels == ["0", "3"] + # edges still run through the hidden t=1 and t=2 nodes + segments = np.asarray(lines.get_segments()) + assert len(segments) == graph.num_edges() + np.testing.assert_array_equal(np.sort(np.unique(segments[:, :, 1])), [0.0, 1.0, 2.0, 3.0]) + # the division fork starts at the hidden t=1 node and ends at the two distinct t=2 nodes + forks = segments[segments[:, 0, 1] == 1.0] + assert len(forks) == 2 + np.testing.assert_array_equal(forks[:, 1, 1], [2.0, 2.0]) + assert forks[0, 1, 0] != forks[1, 1, 0] + def test_plot_lineage_tree_edge_colors() -> None: """Test styling marker borders via scatter_kwargs (edgecolors/linewidths)."""