Action benchmarks

.github/workflows/benchmark-report.yml

@@ -29,7 +29,7 @@ jobs:
         run: uv sync --no-dev --group testing
 
       - name: Run benchmarks
-        run: uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json output.json
+        run: uv run pytest tests/benchmarks/ --benchmark-json output.json
 
       - name: Store benchmark results
         uses: benchmark-action/github-action-benchmark@v1

.github/workflows/benchmark_pr.py

@@ -31,15 +31,18 @@ def make_report(old_path, new_path, out_file, header=None):
     old = load_stats(old_path)
     new = load_stats(new_path)
 
-    # Merge on benchmark name
-    df = old[-1].merge(new[-1], on="Benchmark", suffixes=("_old", "_new"))
+    # Merge on benchmark name. Outer join so benchmarks that exist on only one
+    # side (added or removed by the PR) still show up in the report.
+    df = old[-1].merge(new[-1], on="Benchmark", suffixes=("_old", "_new"), how="outer")
 
     pct_change = 100 * (df["mean_new"] - df["mean_old"]) / df["mean_old"]
-    df["Percent Change"] = pct_change.map("{:+.2f}".format)
+    df["Percent Change"] = pct_change.map("{:+.2f}".format).where(
+        pct_change.notna(), "n/a"
+    )
 
     # Format runtimes
-    df["mean_old"] = df["mean_old"].map("{:.5f}".format)
-    df["mean_new"] = df["mean_new"].map("{:.5f}".format)
+    for col in ("mean_old", "mean_new"):
+        df[col] = df[col].map("{:.5f}".format).where(df[col].notna(), "-")
 
     # Change column names to commit ids
     df = df.rename(

.github/workflows/test.yml

@@ -85,7 +85,9 @@ jobs:
         if: steps.cache_baseline.outputs.cache-hit != 'true'
         run: |
           git checkout ${{ github.event.pull_request.base.sha }}
-          uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json baseline.json
+          # -o python_files: the base commit's pyproject.toml may predate the
+          # bench_*.py collection config, so set it on the command line instead.
+          uv run pytest tests/benchmarks/ -o "python_files=test_*.py bench_*.py" --benchmark-json baseline.json
 
       - name: Cache baseline results
         uses: actions/cache/save@v5
@@ -97,7 +99,7 @@ jobs:
       - name: Run benchmark on PR head
         run: |
           git checkout ${{ github.event.pull_request.head.sha }}
-          uv run pytest tests/benchmarks/bench_candidate_graph.py --benchmark-json pr.json
+          uv run pytest tests/benchmarks/ -o "python_files=test_*.py bench_*.py" --benchmark-json pr.json
 
       - name: Generate report
         id: report

pyproject.toml

@@ -74,6 +74,10 @@ dev = [
 
 [tool.setuptools_scm]
 
+[tool.pytest.ini_options]
+python_files = ["test_*.py", "bench_*.py"]
+norecursedirs = ["tests/benchmarks"]
+
 [tool.ruff]
 line-length = 90
 target-version = "py310"

tests/benchmarks/_graph_builders.py

@@ -0,0 +1,127 @@
+"""Builders for the synthetic graphs the benchmarks run against.
+
+A plain module rather than conftest.py, because benchmark modules import these builders
+directly and a conftest is not importable. The conftest next door holds only fixtures.
+"""
+
+import numpy as np
+import polars as pl
+import tracksdata as td
+from skimage.draw import disk
+from tracksdata.nodes import Mask
+
+from funtracks.data_model import Tracks
+from funtracks.utils.tracksdata_utils import create_empty_graph
+
+
+def make_tracks(
+    n_frames: int,
+    cells_per_frame: int,
+    frame_shape: tuple[int, int],
+    seed: int = 42,
+) -> Tracks:
+    """Build a synthetic Tracks with segmentation masks.
+
+    Produces a *solution* graph, not a candidate graph: cells are linked 1-to-1 between
+    adjacent frames so every node has out-degree <= 1. This matters because the user
+    actions assume a tracking solution -- UserDeleteEdge rejects a removal that would
+    leave out-degree > 1, and UserAddEdge rejects merges -- so they raise
+    InvalidActionError on the multi-successor candidate graphs that
+    compute_graph_from_seg produces.
+
+    Args:
+        n_frames: Number of time points.
+        cells_per_frame: Cells generated per time point.
+        frame_shape: Spatial (y, x) shape of each frame.
+        seed: Seed for the random generator, for reproducible geometry.
+
+    Returns:
+        A Tracks with pos, area, mask and bbox node attrs, iou edge attrs, and a
+        writable unmanaged "score" node attr (populated on every node) for
+        attribute-update benchmarks.
+    """
+    rng = np.random.default_rng(seed)
+    graph = create_empty_graph(
+        node_attributes=[
+            "pos",
+            "area",
+            td.DEFAULT_ATTR_KEYS.MASK,
+            td.DEFAULT_ATTR_KEYS.BBOX,
+        ],
+        node_default_values=[0.0, 0.0, 0.0, 0.0],
+        edge_attributes=["iou"],
+        position_attrs=["pos"],
+        ndim=3,
+    )
+
+    nodes: list[dict] = []
+    node_ids: list[int] = []
+    edges: list[dict] = []
+    node_id = 1
+    prev_frame_ids: list[int] = []
+
+    for t in range(n_frames):
+        frame_ids = []
+        for _ in range(cells_per_frame):
+            cy = int(rng.integers(25, frame_shape[0] - 25))
+            cx = int(rng.integers(25, frame_shape[1] - 25))
+            radius = int(rng.integers(8, 15))
+
+            # Mask is stored in its own local frame with an offset bbox, so the
+            # disk is drawn centered in a (2r-1, 2r-1) box.
+            mask_shape = (2 * radius - 1, 2 * radius - 1)
+            rr, cc = disk(
+                center=(radius - 1, radius - 1), radius=radius, shape=mask_shape
+            )
+            mask_arr = np.zeros(mask_shape, dtype=bool)
+            mask_arr[rr, cc] = True
+            bbox = np.array([cy - radius + 1, cx - radius + 1, cy + radius, cx + radius])
+
+            nodes.append(
+                {
+                    "t": t,
+                    "pos": [float(cy), float(cx)],
+                    "area": float(mask_arr.sum()),
+                    "solution": True,
+                    td.DEFAULT_ATTR_KEYS.MASK: Mask(mask_arr, bbox=bbox),
+                    td.DEFAULT_ATTR_KEYS.BBOX: bbox,
+                }
+            )
+            node_ids.append(node_id)
+            frame_ids.append(node_id)
+            node_id += 1
+
+        # Link 1-to-1 with the previous frame to keep out-degree <= 1.
+        for prev_id, cur_id in zip(prev_frame_ids, frame_ids, strict=False):
+            edges.append(
+                {
+                    "source_id": prev_id,
+                    "target_id": cur_id,
+                    "solution": True,
+                    "iou": float(rng.uniform(0.1, 0.9)),
+                }
+            )
+        prev_frame_ids = frame_ids
+
+    # Register "score" before adding the nodes so every node gets a stored value. Adding
+    # the key afterwards leaves existing nodes null, and tracksdata applies the default
+    # lazily on read (_maybe_fill_null in graph/_rustworkx_graph.py, guarded by
+    # s.has_nulls()). That backfill costs two extra polars collects per read until the
+    # node is first written, which made the first update of any node ~2x slower than
+    # subsequent ones and turned the attribute benchmarks into a measure of null
+    # backfill rather than of the update path.
+    graph.add_node_attr_key("score", dtype=pl.Float64, default_value=0.0)
+    for node in nodes:
+        node["score"] = 0.0
+
+    graph.bulk_add_nodes(nodes=nodes, indices=node_ids)
+    graph.bulk_add_edges(edges)
+    graph._update_metadata(segmentation_shape=(n_frames, *frame_shape))
+
+    return Tracks(
+        graph,
+        time_attr="t",
+        pos_attr="pos",
+        tracklet_attr="tracklet_id",
+        lineage_attr="lineage_id",
+    )

tests/benchmarks/bench_actions.py

@@ -0,0 +1,180 @@
+import pytest
+
+from funtracks.user_actions import (
+    UserAddEdge,
+    UserDeleteEdge,
+    UserDeleteNode,
+    UserUpdateNodeAttrs,
+    UserUpdateNodesAttrs,
+    UserUpdateSegmentation,
+)
+from funtracks.utils.tracksdata_utils import td_mask_to_pixels
+
+from ._graph_builders import make_tracks
+
+NUM_FRAMES = 20
+FRAME_SHAPE = (500, 500)
+CELLS_PER_FRAME = 50
+
+# Operations per timed call. Even, so the self-inverting benchmarks return the graph to
+# its starting state. Raise this (rather than ROUNDS) if a benchmark reads as noisy.
+N_OPS = 50
+ROUNDS = 3
+
+# Fraction of a node's mask painted in the segmentation benchmark. Must stay < 1: a
+# patch that fully covers the mask makes UserUpdateSegmentation delete the node instead
+# of updating its mask, which is a different (and much more expensive) code path.
+PATCH_FRACTION = 3
+
+
+@pytest.fixture(scope="module")
+def tracks(_warm_jit):
+    """One Tracks object shared by every benchmark in this module."""
+    return make_tracks(
+        n_frames=NUM_FRAMES,
+        cells_per_frame=CELLS_PER_FRAME,
+        frame_shape=FRAME_SHAPE,
+    )
+
+
+@pytest.fixture(scope="module")
+def solution_edges(tracks):
+    """Every existing (source, target) edge, in node id order.
+
+    Node ids are assigned frame by frame, so this is ordered by time: consecutive entries
+    are cells in the same frame, and the benchmarks' per-round slices therefore walk
+    forward through the movie rather than repeatedly hitting one frame.
+    """
+    edges = []
+    for node in sorted(int(n) for n in tracks.graph_solution.node_ids()):
+        successors = tracks.successors(node)
+        if successors:
+            edges.append((node, int(successors[0])))
+    return edges
+
+
+def _node_batches(solution_edges, n_rounds=ROUNDS, per_round=N_OPS):
+    """Split source nodes into one distinct batch per round.
+
+    Each round gets its own nodes, and callers offset their slices so no two benchmarks
+    share nodes. That keeps test_delete_nodes from removing nodes the attribute
+    benchmarks reported on, and keeps every round of a destructive benchmark doing the
+    same amount of work.
+
+    Note this is no longer needed to avoid a measurement artifact: "score" is populated
+    on every node at build time (see _graph_builders.make_tracks), so a node
+    that has never been written costs the same to update as one that has.
+    """
+    nodes = [source for source, _ in solution_edges]
+    batches = [nodes[r * per_round : (r + 1) * per_round] for r in range(n_rounds)]
+    assert all(len(batch) == per_round for batch in batches), (
+        f"need {n_rounds * per_round} source nodes to give each round a distinct batch, "
+        f"only have {len(nodes)}"
+    )
+    return batches
+
+
+def test_update_node_attrs_single(benchmark, tracks, solution_edges):
+    """N single-node attribute updates, each its own history entry."""
+    batches = iter(_node_batches(solution_edges))
+
+    def run():
+        for i, node in enumerate(next(batches)):
+            UserUpdateNodeAttrs(tracks, node, {"score": float(i)})
+
+    benchmark.pedantic(run, rounds=ROUNDS, iterations=1)
+
+
+def test_update_node_attrs_bulk(benchmark, tracks, solution_edges):
+    """The same N updates as one batched action.
+
+    UserUpdateNodesAttrs exists to collapse history and refresh into a single entry for
+    the whole batch; this pairs with test_update_node_attrs_single to show what that
+    batching is worth.
+    """
+    # Offset past the batches test_update_node_attrs_single consumed, so these updates
+    # also land on nodes not yet touched in this session.
+    batches = iter(_node_batches(solution_edges[ROUNDS * N_OPS :]))
+
+    def run():
+        nodes = next(batches)
+        UserUpdateNodesAttrs(
+            tracks, nodes, {"score": [float(i) for i in range(len(nodes))]}
+        )
+
+    benchmark.pedantic(run, rounds=ROUNDS, iterations=1)
+
+
+def test_add_delete_edges(benchmark, tracks, solution_edges):
+    """Alternating edge delete and re-add on one edge.
+
+    Both directions trigger UpdateTrackIDs: deleting orphans the downstream segment onto
+    a fresh track/lineage id, and re-adding merges it back. Self-inverting, so the graph
+    is unchanged afterwards.
+    """
+    edge = solution_edges[len(solution_edges) // 2]
+
+    def run():
+        for _ in range(N_OPS // 2):
+            UserDeleteEdge(tracks, edge)
+            UserAddEdge(tracks, edge)
+
+    benchmark.pedantic(run, rounds=ROUNDS, iterations=1)
+
+
+def test_update_segmentation(benchmark, tracks, solution_edges):
+    """Alternating paint-out and paint-back of one sub-mask patch.
+
+    Simulates the paint stroke path: UpdateNodeSeg plus the regionprops recompute
+    (position, area) and the IoU refresh on adjacent edges. Removing then restoring the
+    same patch keeps every operation on the UpdateNodeSeg branch and leaves the mask
+    exactly as it started -- the patch is a strict subset of the mask so the node is
+    never fully erased, and the paint-back targets an existing node id so it never falls
+    through to UserAddNode.
+    """
+    node = solution_edges[len(solution_edges) // 2][0]
+    mask_pixels = td_mask_to_pixels(
+        tracks.get_mask(node), tracks.get_time(node), ndim=tracks.ndim
+    )
+    n_patch = len(mask_pixels[0]) // PATCH_FRACTION
+    patch = tuple(dim_pixels[:n_patch] for dim_pixels in mask_pixels)
+
+    def run():
+        for i in range(N_OPS):
+            if i % 2 == 0:
+                # Erase the patch: old value is the node, new value is background.
+                UserUpdateSegmentation(
+                    tracks,
+                    new_value=0,
+                    updated_pixels=[(patch, node)],
+                    current_track_id=1,
+                )
+            else:
+                # Paint it back: old value is background, new value is the node.
+                UserUpdateSegmentation(
+                    tracks,
+                    new_value=node,
+                    updated_pixels=[(patch, 0)],
+                    current_track_id=1,
+                )
+
+    benchmark.pedantic(run, rounds=ROUNDS, iterations=1)
+
+
+def test_delete_nodes(benchmark, tracks, solution_edges):
+    """N node deletions, each cascading to its adjacent edges and track ids.
+
+    Destructive and not self-inverting, so this runs last in the module and takes a
+    fresh slice of nodes per round. Deleting a mid-track node makes UserDeleteNode
+    delete both adjacent edges and reconnect the neighbours with a skip edge, which is
+    the expensive case.
+    """
+    # Start past the nodes the attribute benchmarks consumed so deletions do not remove
+    # nodes those benchmarks already reported on.
+    batches = iter(_node_batches(solution_edges[2 * ROUNDS * N_OPS :]))
+
+    def run():
+        for node in next(batches):
+            UserDeleteNode(tracks, node)
+
+    benchmark.pedantic(run, rounds=ROUNDS, iterations=1)