[WIP] tracksdata backend

.gitignore

@@ -93,3 +93,6 @@ pixi.lock
 
 # uv environments
 uv.lock
+
+# Claude
+CLAUDE.md

docs/features.md

@@ -150,7 +150,7 @@ classDiagram
     }
 
     class Tracks {
-        +graph: nx.DiGraph
+        +graph: td.graph.GraphView
         +segmentation: ndarray|None
         +features: FeatureDict
         +annotators: AnnotatorRegistry
@@ -204,6 +204,7 @@ These features are **automatically checked** during initialization:
 **Scenario 1: Loading tracks from CSV with pre-computed features**
 ```python
 # CSV has columns: id, time, y, x, area, track_id
+# TODO: load_graph_from_csv no longer exists!
 graph = load_graph_from_csv(df)  # Nodes already have area, track_id
 tracks = SolutionTracks(graph, segmentation=seg)
 # Auto-detection: pos, area, track_id exist → activate without recomputing
@@ -212,8 +213,9 @@ tracks = SolutionTracks(graph, segmentation=seg)
 **Scenario 2: Creating tracks from raw segmentation**
 ```python
 # Graph has no features yet
-graph = nx.DiGraph()
-graph.add_node(1, time=0)
+#TODO: test these examples with the new tracksdata api
+graph = create_empty_graphview_graph()
+graph.add_node(index=1, attrs={"t": 0, "solution": 1})
 tracks = Tracks(graph, segmentation=seg)
 # Auto-detection: pos, area don't exist → compute them
 ```
@@ -272,9 +274,9 @@ tracks.disable_features(["area"])
         def compute(self, feature_keys=None):
             # Compute feature values in bulk
             if "custom" in self.features:
-                for node in self.tracks.graph.nodes():
+                for node in self.tracks.graph.node_ids():
                     value = self._compute_custom(node)
-                    self.tracks.graph.nodes[node]["custom"] = value
+                    self.tracks.graph[node]["custom"] = value
 
         def update(self, action):
             # Incremental update when graph changes

pyproject.toml

@@ -39,6 +39,7 @@ dependencies =[
     "pandas>=2.3.3",
     "zarr>=2.18,<4",
     "numcodecs>=0.13,<0.16",
+    "tracksdata[spatial]@git+https://github.com/royerlab/tracksdata@24d93ef8d3c4c44cedc346dc3229c98c1d3eaa5e",
 ]
 
 [project.urls]
@@ -107,6 +108,7 @@ unfixable = [
 [tool.mypy]
 ignore_missing_imports = true
 python_version = "3.10"
+explicit_package_bases = true
 
 [tool.coverage.report]
 exclude_also = [

src/funtracks/actions/add_delete_edge.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from ._base import BasicAction
 
@@ -10,6 +10,8 @@
     from funtracks.data_model import Tracks
     from funtracks.data_model.tracks import Edge
 
+import tracksdata as td
+
 
 class AddEdge(BasicAction):
     """Action for adding a new edge. Endpoints must exist already."""
@@ -52,7 +54,24 @@ def _apply(self) -> None:
                     f"Cannot add edge {self.edge}: endpoint {node} not in graph yet"
                 )
 
-        self.tracks.graph.add_edge(self.edge[0], self.edge[1], **self.attributes)
+        if self.tracks.graph.has_edge(*self.edge):
+            raise ValueError(f"Edge {self.edge} already exists in the graph")
+
+        # Add required solution attribute
+        attrs = self.attributes
+        attrs[td.DEFAULT_ATTR_KEYS.SOLUTION] = 1
+
+        required_attrs = self.tracks.graph.edge_attr_keys()
+        for attr in required_attrs:
+            if attr not in attrs:
+                attrs[attr] = None
+
+        # Create edge attributes for this specific edge
+        self.tracks.graph.add_edge(
+            source_id=self.edge[0],
+            target_id=self.edge[1],
+            attrs=attrs,
+        )
 
         # Notify annotators to recompute features (will overwrite computed ones)
         self.tracks.notify_annotators(self)

src/funtracks/actions/add_delete_node.py

@@ -7,9 +7,18 @@
 if TYPE_CHECKING:
     from typing import Any
 
-    from funtracks.data_model import SolutionTracks
+    from funtracks.data_model.solution_tracks import SolutionTracks
     from funtracks.data_model.tracks import Node, SegMask
 
+import numpy as np
+import tracksdata as td
+
+from funtracks.utils.tracksdata_utils import (
+    compute_node_attrs_from_masks,
+    compute_node_attrs_from_pixels,
+    pixels_to_td_mask,
+)
+
 
 class AddNode(BasicAction):
     """Action for adding new nodes. If a segmentation should also be added, the
@@ -78,14 +87,112 @@ def inverse(self) -> BasicAction:
 
     def _apply(self) -> None:
         """Apply the action, and set segmentation if provided in self.pixels"""
+        attrs = self.attributes
+
+        final_pos: np.ndarray
+        if self.tracks.segmentation is not None:
+            if self.pixels is not None:
+                computed_attrs = compute_node_attrs_from_pixels(
+                    [self.pixels], self.tracks.ndim, self.tracks.scale
+                )
+                # Extract single values from lists (since we passed one pixel set)
+                computed_attrs = {key: value[0] for key, value in computed_attrs.items()}
+                # if masks are not given, calculate them from the pixels
+                if "mask" not in attrs:
+                    mask_obj, _ = pixels_to_td_mask(
+                        self.pixels, self.tracks.ndim, self.tracks.scale
+                    )
+                    attrs[td.DEFAULT_ATTR_KEYS.MASK] = mask_obj
+                    attrs[td.DEFAULT_ATTR_KEYS.BBOX] = mask_obj.bbox
+            elif "mask" in attrs:
+                computed_attrs = compute_node_attrs_from_masks(
+                    attrs["mask"], self.tracks.ndim, self.tracks.scale
+                )
+                # Extract single values from lists (since we passed one mask)
+                computed_attrs = {key: value[0] for key, value in computed_attrs.items()}
+            # Handle position_key safely using the same pattern as in tracks.py
+            if isinstance(self.tracks.features.position_key, list):
+                # Multi-axis position keys - check if any are missing from attrs
+                missing_keys = [
+                    k for k in self.tracks.features.position_key if k not in attrs
+                ]
+                if missing_keys:
+                    # Use computed position from segmentation
+                    final_pos = np.array(computed_attrs["pos"])
+                    # Set individual components in attrs
+                    for i, key in enumerate(self.tracks.features.position_key):
+                        attrs[key] = (
+                            final_pos[i] if final_pos.ndim == 1 else final_pos[:, i]
+                        )
+                else:
+                    # All position components provided, combine them
+                    final_pos = np.stack(
+                        [attrs[key] for key in self.tracks.features.position_key], axis=0
+                    )
+            else:
+                # Single position key
+                pos_key = self.tracks.features.position_key
+                if pos_key is not None and pos_key not in attrs:
+                    final_pos = np.array(computed_attrs["pos"])
+                    attrs[pos_key] = final_pos
+                elif pos_key is not None:
+                    final_pos = np.array(attrs[pos_key])
+                else:
+                    raise ValueError("Position key is None")
+            # Set area using string literal since FeatureDict doesn't have area_key
+            attrs["area"] = computed_attrs["area"]
+        else:
+            # No segmentation - handle position_key safely
+            if isinstance(self.tracks.features.position_key, list):
+                # Multi-axis position keys - check if any are missing
+                missing_keys = [
+                    k for k in self.tracks.features.position_key if k not in attrs
+                ]
+                if missing_keys:
+                    raise ValueError(
+                        f"Must provide positions {missing_keys} or segmentation"
+                    )
+                # All position components provided, combine them
+                final_pos = np.stack(
+                    [attrs[key] for key in self.tracks.features.position_key], axis=0
+                )
+            else:
+                # Single position key
+                if (
+                    self.tracks.features.position_key is None
+                    or self.tracks.features.position_key not in attrs
+                ):
+                    raise ValueError("Must provide positions or segmentation and ids")
+                final_pos = np.array(attrs[self.tracks.features.position_key])
+
+        # Position is already set in attrs above
+        # Add nodes to td graph
+        required_attrs = self.tracks.graph.node_attr_keys().copy()
+        if td.DEFAULT_ATTR_KEYS.NODE_ID in required_attrs:
+            required_attrs.remove(td.DEFAULT_ATTR_KEYS.NODE_ID)
+        if td.DEFAULT_ATTR_KEYS.SOLUTION not in attrs:
+            attrs[td.DEFAULT_ATTR_KEYS.SOLUTION] = 1
+        for attr in required_attrs:
+            if attr not in attrs:
+                attrs[attr] = None
+
+        node_dict = {
+            attr: np.array(values) if attr == "pos" else values
+            for attr, values in attrs.items()
+        }
+
+        self.tracks.graph.add_node(attrs=node_dict, index=self.node)
+
         if self.pixels is not None:
             self.tracks.set_pixels(self.pixels, self.node)
-        attrs = self.attributes
-        self.tracks.graph.add_node(self.node)
 
-        # set all user provided attributes including time and position
-        for attr, value in attrs.items():
-            self.tracks._set_node_attr(self.node, attr, value)
+        if type(self.tracks).__name__ == "SolutionTracks":
+            tracklet_key = self.tracks.features.tracklet_key
+            if tracklet_key is not None and tracklet_key in attrs:
+                track_id = attrs[tracklet_key]
+                if track_id not in self.tracks.track_id_to_node:
+                    self.tracks.track_id_to_node[track_id] = []
+                self.tracks.track_id_to_node[track_id].append(self.node)
 
         # Always notify annotators - they will check their own preconditions
         self.tracks.notify_annotators(self)
@@ -114,6 +221,17 @@ def __init__(
             if val is not None:
                 self.attributes[key] = val
 
+        if td.DEFAULT_ATTR_KEYS.MASK in self.tracks.graph.node_attr_keys():
+            self.attributes[td.DEFAULT_ATTR_KEYS.MASK] = self.tracks.get_nodes_attr(
+                [self.node], td.DEFAULT_ATTR_KEYS.MASK
+            )[0]
+            self.attributes[td.DEFAULT_ATTR_KEYS.BBOX] = self.tracks.get_nodes_attr(
+                [self.node], td.DEFAULT_ATTR_KEYS.BBOX
+            )[0]
+        self.attributes[td.DEFAULT_ATTR_KEYS.SOLUTION] = self.tracks.get_nodes_attr(
+            [self.node], td.DEFAULT_ATTR_KEYS.SOLUTION
+        )[0]
+
         self.pixels = self.tracks.get_pixels(node) if pixels is None else pixels
         self._apply()
 
@@ -130,8 +248,6 @@ def _apply(self) -> None:
             set pixels to 0 if self.pixels is provided
         - Remove nodes from graph
         """
-        if self.pixels is not None:
-            self.tracks.set_pixels(self.pixels, 0)
 
         self.tracks.graph.remove_node(self.node)
         self.tracks.notify_annotators(self)

src/funtracks/actions/update_segmentation.py

@@ -48,5 +48,5 @@ def inverse(self) -> BasicAction:
     def _apply(self) -> None:
         """Set new attributes"""
         value = self.node if self.added else 0
-        self.tracks.set_pixels(self.pixels, value)
+        self.tracks.set_pixels(self.pixels, value, self.node)
         self.tracks.notify_annotators(self)

src/funtracks/annotators/_compute_ious.py

@@ -31,5 +31,5 @@ def _compute_ious(frame1: np.ndarray, frame2: np.ndarray) -> list[tuple[int, int
         intersection = counts[index]
         id1, id2 = pair
         union = frame1_label_sizes[id1] + frame2_label_sizes[id2] - intersection
-        ious.append((id1, id2, intersection / union))
+        ious.append((int(id1), int(id2), intersection / union))
     return ious

src/funtracks/annotators/_edge_annotator.py

@@ -86,13 +86,16 @@ def compute(self, feature_keys: list[str] | None = None) -> None:
         # TODO: add skip edges
         if self.iou_key in keys_to_compute:
             nodes_by_frame = defaultdict(list)
-            for n in self.tracks.nodes():
+            for n in self.tracks.graph.node_ids():
                 nodes_by_frame[self.tracks.get_time(n)].append(n)
 
             for t in range(seg.shape[0] - 1):
                 nodes_in_t = nodes_by_frame[t]
-                edges = list(self.tracks.graph.out_edges(nodes_in_t))
-                self._iou_update(edges, seg[t], seg[t + 1])
+                edges = []
+                for node in nodes_in_t:
+                    for succ in self.tracks.graph.successors(node):
+                        edges.append((node, succ))
+                self._iou_update(edges, np.asarray(seg[t]), np.asarray(seg[t + 1]))
 
     def _iou_update(
         self,
@@ -119,7 +122,7 @@ def _iou_update(
 
         # anything left has IOU of 0
         for edge in edges:
-            self.tracks._set_edge_attr(edge, self.iou_key, 0)
+            self.tracks._set_edge_attr(edge, self.iou_key, 0.0)
 
     def update(self, action: BasicAction):
         """Update the edge features based on the action.
@@ -146,17 +149,23 @@ def update(self, action: BasicAction):
         else:  # UpdateNodeSeg
             # Get all incident edges to the modified node
             node = action.node
-            edges_to_update = list(self.tracks.graph.in_edges(node)) + list(
-                self.tracks.graph.out_edges(node)
-            )
+
+            edges_to_update = []
+            for node in self.tracks.graph.node_ids():
+                # Add edges from predecessors
+                for pred in self.tracks.graph.predecessors(node):
+                    edges_to_update.append((pred, node))
+                # Add edges from successors
+                for succ in self.tracks.graph.successors(node):
+                    edges_to_update.append((node, succ))
 
         # Update IoU for each edge
         for edge in edges_to_update:
             source, target = edge
             start_time = self.tracks.get_time(source)
             end_time = self.tracks.get_time(target)
-            start_seg = self.tracks.segmentation[start_time]
-            end_seg = self.tracks.segmentation[end_time]
+            start_seg = np.asarray(self.tracks.segmentation[start_time])
+            end_seg = np.asarray(self.tracks.segmentation[end_time])
             masked_start = np.where(start_seg == source, source, 0)
             masked_end = np.where(end_seg == target, target, 0)
             if np.max(masked_start) == 0 or np.max(masked_end) == 0:
@@ -165,10 +174,10 @@ def update(self, action: BasicAction):
                     f"in frame {end_time}: updating edge IOU value to 0",
                     stacklevel=2,
                 )
-                self.tracks._set_edge_attr(edge, self.iou_key, 0)
+                self.tracks._set_edge_attr(edge, self.iou_key, 0.0)
             else:
                 iou_list = _compute_ious(masked_start, masked_end)
-                iou = 0 if len(iou_list) == 0 else iou_list[0][2]
+                iou = 0.0 if len(iou_list) == 0 else iou_list[0][2]
                 self.tracks._set_edge_attr(edge, self.iou_key, iou)
 
     def change_key(self, old_key: str, new_key: str) -> None:

src/funtracks/annotators/_graph_annotator.py

@@ -66,6 +66,22 @@ def activate_features(self, keys: list[str]) -> None:
                 feat, _ = self.all_features[key]
                 self.all_features[key] = (feat, True)
 
+                # Ensure attribute key exists in graph schema
+                if (
+                    feat["feature_type"] == "node"
+                    and key not in self.tracks.graph.node_attr_keys()
+                ):
+                    self.tracks.graph.add_node_attr_key(
+                        key, default_value=feat["default_value"]
+                    )
+                elif (
+                    feat["feature_type"] == "edge"
+                    and key not in self.tracks.graph.edge_attr_keys()
+                ):
+                    self.tracks.graph.add_edge_attr_key(
+                        key, default_value=feat["default_value"]
+                    )
+
     def deactivate_features(self, keys: list[str]) -> None:
         """Deactivate computation of the given features in the annotation process.