Adding offset to assign_track_ids and new tracklet_graph for napari (#116)

* adding optional track_id_offset to assign_track_id

* adding rustworkx tracklet_graph method

* moving to tracklet_graph to BaseGraph

* adding sql graph tracklet_graph

* fixing edge filtering

* adding digraph to napari dict helper function

* fixing tests

* Update src/tracksdata/graph/_base_graph.py

Co-authored-by: Yohsuke T. Fukai <[email protected]>

---------

Co-authored-by: Yohsuke T. Fukai <[email protected]>

Author: JoOkuma

src/tracksdata/functional/__init__.py

@@ -1,5 +1,5 @@
 """Functional utilities for graph operations."""
 
-from tracksdata.functional._napari import to_napari_format
+from tracksdata.functional._napari import rx_digraph_to_napari_dict, to_napari_format
 
-__all__ = ["to_napari_format"]
+__all__ = ["rx_digraph_to_napari_dict", "to_napari_format"]

src/tracksdata/functional/_napari.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING, overload
 
 import polars as pl
+import rustworkx as rx
 
 from tracksdata.attrs import EdgeAttr, NodeAttr
 from tracksdata.constants import DEFAULT_ATTR_KEYS
@@ -119,3 +120,27 @@ def to_napari_format(
         return tracks_data, dict_graph, array_view
 
     return tracks_data, dict_graph
+
+
+def rx_digraph_to_napari_dict(
+    tracklet_graph: rx.PyDiGraph,
+) -> dict[int, list[int]]:
+    """
+    Convert a tracklet graph to a napari-ready dictionary.
+    The input is a (child -> parent) graph (forward in time) and it is converted
+    to a (parent -> child) dictionary (backward in time).
+
+    Parameters
+    ----------
+    tracklet_graph : rx.PyDiGraph
+        The tracklet graph to convert.
+
+    Returns
+    -------
+    dict[int, list[int]]
+        A dictionary of parent -> child relationships.
+    """
+    dict_graph = {}
+    for parent, child in tracklet_graph.edges():
+        dict_graph.setdefault(child, []).append(parent)
+    return dict_graph

src/tracksdata/functional/_rx.py

@@ -45,6 +45,7 @@ def _fast_path_transverse(
 def _fast_dag_transverse(
     starts: np.ndarray,
     dag: dict[int, int],
+    track_id_offset: int,
 ) -> tuple[list[np.ndarray], np.ndarray, np.ndarray, dict[int, int], dict[int, int]]:
     """
     Traverse the tracks DAG creating a distinct id to each linear path.
@@ -61,6 +62,8 @@ def _fast_dag_transverse(
     dag : dict[int, int]
         Directed acyclic graph mapping parent → child for linear paths only.
         Dividing edges are excluded and handled separately.
+    track_id_offset : int
+        The starting track id, useful when assigning track ids to a subgraph.
 
     Returns
     -------
@@ -77,7 +80,7 @@ def _fast_dag_transverse(
     last_to_track_id = {}
     first_to_track_id = {}
 
-    track_id = 1
+    track_id = track_id_offset
 
     for start in starts:
         path = _fast_path_transverse(start, dag)
@@ -247,6 +250,7 @@ def _track_id_edges_from_long_edges(
 
 def _assign_track_ids(
     graph: rx.PyDiGraph,
+    track_id_offset: int,
 ) -> tuple[np.ndarray, np.ndarray, rx.PyDiGraph]:
     """
     Assigns an unique `track_id` to each simple path in the graph and
@@ -256,6 +260,8 @@ def _assign_track_ids(
     ----------
     graph : rx.PyDiGraph
         Directed acyclic graph of tracks.
+    track_id_offset : int
+        The starting track id, useful when assigning track ids to a subgraph.
 
     Returns
     -------
@@ -272,7 +278,9 @@ def _assign_track_ids(
     # was it better (faster) when using a numpy array for the digraph as in ultrack?
     linear_dag, starts, long_edges_df = _rx_graph_to_dict_dag(graph)
 
-    paths, track_ids, lengths, last_to_track_id, first_to_track_id = _fast_dag_transverse(starts, linear_dag)
+    paths, track_ids, lengths, last_to_track_id, first_to_track_id = _fast_dag_transverse(
+        starts, linear_dag, track_id_offset
+    )
 
     n_tracks = len(track_ids)
 

src/tracksdata/functional/_test/test_rx.py

@@ -10,7 +10,7 @@ def test_empty_graph() -> None:
     """Test that empty graph raises ValueError."""
     graph = rx.PyDiGraph()
     with pytest.raises(ValueError, match="Graph is empty"):
-        _assign_track_ids(graph)
+        _assign_track_ids(graph, track_id_offset=1)
 
 
 def test_single_path() -> None:
@@ -22,7 +22,7 @@ def test_single_path() -> None:
     graph.add_edge(nodes[0], nodes[1], None)
     graph.add_edge(nodes[1], nodes[2], None)
 
-    node_ids, track_ids, tracks_graph = _assign_track_ids(graph)
+    node_ids, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
 
     assert np.array_equal(node_ids, [0, 1, 2])
     assert np.array_equal(track_ids, [1, 1, 1])
@@ -46,7 +46,7 @@ def test_symmetric_branching_path() -> None:
     graph.add_edge(nodes[0], nodes[1], None)
     graph.add_edge(nodes[0], nodes[2], None)
 
-    node_ids, track_ids, tracks_graph = _assign_track_ids(graph)
+    node_ids, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
 
     # Should create 2 tracks: one for each branch
     assert len(node_ids) == 3
@@ -76,7 +76,7 @@ def test_asymmetric_branching_path() -> None:
     graph.add_edge(nodes[1], nodes[2], None)
     graph.add_edge(nodes[0], nodes[3], None)
 
-    node_ids, track_ids, tracks_graph = _assign_track_ids(graph)
+    node_ids, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
 
     # Should create 2 tracks: one for each branch
     assert len(node_ids) == 4
@@ -103,7 +103,7 @@ def test_invalid_multiple_parents() -> None:
     graph.add_edge(nodes[1], nodes[2], None)
 
     with pytest.raises(RuntimeError, match="Invalid graph structure"):
-        _assign_track_ids(graph)
+        _assign_track_ids(graph, track_id_offset=1)
 
 
 def test_complex_valid_branching() -> None:
@@ -132,7 +132,7 @@ def test_complex_valid_branching() -> None:
     graph.add_edge(nodes[3], nodes[4], None)
     graph.add_edge(nodes[2], nodes[5], None)
 
-    node_ids, track_ids, tracks_graph = _assign_track_ids(graph)
+    node_ids, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
 
     # this order is an implementation detail, it could change
     # then the track ids should change accordingly
@@ -168,7 +168,7 @@ def test_three_children() -> None:
     graph.add_edge(nodes[0], nodes[2], None)
     graph.add_edge(nodes[0], nodes[3], None)
 
-    _, track_ids, tracks_graph = _assign_track_ids(graph)
+    _, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
     assert set(tracks_graph.successor_indices(track_ids[0])) == set(track_ids[1:])
 
 
@@ -182,7 +182,7 @@ def test_multiple_roots() -> None:
     graph.add_edge(nodes[0], nodes[1], None)
     graph.add_edge(nodes[2], nodes[3], None)
 
-    node_ids, track_ids, tracks_graph = _assign_track_ids(graph)
+    node_ids, track_ids, tracks_graph = _assign_track_ids(graph, track_id_offset=1)
 
     assert len(node_ids) == 4
     assert len(track_ids) == 4

src/tracksdata/graph/_base_graph.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 import polars as pl
+import rustworkx as rx
 from numpy.typing import ArrayLike
 
 from tracksdata.attrs import AttrComparison, NodeAttr
@@ -948,3 +949,79 @@ def spatial_filter(self, attrs_keys: list[str] | None = None) -> "SpatialFilter"
         from tracksdata.graph.filters._spatial_filter import SpatialFilter
 
         return SpatialFilter(self, attrs_keys=attrs_keys)
+
+    def tracklet_graph(
+        self,
+        track_id_key: str = DEFAULT_ATTR_KEYS.TRACK_ID,
+        ignore_track_id: int | None = None,
+    ) -> rx.PyDiGraph:
+        """
+        Create a compressed tracklet graph where each node is a tracklet
+        and each edge is a transition between tracklets.
+
+        IMPORTANT:
+        rx.PyDiGraph does not allow arbitrary indices, so we use the tracklet ids as node values.
+        And edge values are the tuple of source and target tracklet ids.
+
+        Parameters
+        ----------
+        track_id_key : str
+            The key of the track id attribute.
+        ignore_track_id : int | None
+            The track id to ignore. If None, all track ids are used.
+
+        Returns
+        -------
+        rx.PyDiGraph
+            A compressed tracklet graph.
+
+        See Also
+        --------
+        [rx_digraph_to_napari_dict][tracksdata.functional.rx_digraph_to_napari_dict]:
+            Convert a tracklet graph to a napari-ready dictionary.
+        """
+
+        if track_id_key not in self.node_attr_keys:
+            raise ValueError(f"Track id key '{track_id_key}' not found in graph. Expected '{self.node_attr_keys}'")
+
+        nodes_df = self.node_attrs(attr_keys=[DEFAULT_ATTR_KEYS.NODE_ID, track_id_key])
+        edges_df = self.edge_attrs(attr_keys=[])
+
+        if ignore_track_id is not None:
+            nodes_df = nodes_df.filter(pl.col(track_id_key) != ignore_track_id)
+
+        nodes_df = nodes_df.unique(subset=[track_id_key])
+
+        graph = rx.PyDiGraph()
+        nodes_df = nodes_df.with_columns(
+            pl.Series(
+                np.asarray(graph.add_nodes_from(nodes_df[track_id_key].to_list()), dtype=int),
+            ).alias("rx_id"),
+        )
+
+        edges_df = (
+            edges_df.join(
+                nodes_df.rename({track_id_key: "source_track_id", "rx_id": "source_rx_id"}),
+                left_on=DEFAULT_ATTR_KEYS.EDGE_SOURCE,
+                right_on=DEFAULT_ATTR_KEYS.NODE_ID,
+                how="right",
+            )
+            .join(
+                nodes_df.rename({track_id_key: "target_track_id", "rx_id": "target_rx_id"}),
+                left_on=DEFAULT_ATTR_KEYS.EDGE_TARGET,
+                right_on=DEFAULT_ATTR_KEYS.NODE_ID,
+                how="right",
+            )
+            .filter(~pl.col(DEFAULT_ATTR_KEYS.EDGE_ID).is_null())
+        )
+
+        graph.add_edges_from(
+            zip(
+                edges_df["source_rx_id"].to_list(),
+                edges_df["target_rx_id"].to_list(),
+                zip(edges_df["source_track_id"].to_list(), edges_df["target_track_id"].to_list(), strict=False),
+                strict=True,
+            )
+        )
+
+        return graph

src/tracksdata/graph/_graph_view.py

@@ -447,6 +447,7 @@ def assign_track_ids(
         self,
         output_key: str = DEFAULT_ATTR_KEYS.TRACK_ID,
         reset: bool = True,
+        track_id_offset: int = 1,
     ) -> rx.PyDiGraph:
         """
         Compute and assign track ids to nodes.
@@ -457,14 +458,16 @@ def assign_track_ids(
             The key of the output track id attribute.
         reset : bool
             Whether to reset all track ids before assigning new ones.
+        track_id_offset : int
+            The starting track id, useful when assigning track ids to a subgraph.
 
         Returns
         -------
         rx.PyDiGraph
             A compressed graph (parent -> child) with track ids lineage relationships.
         """
         try:
-            node_ids, track_ids, tracks_graph = _assign_track_ids(self.rx_graph)
+            node_ids, track_ids, tracks_graph = _assign_track_ids(self.rx_graph, track_id_offset)
         except RuntimeError as e:
             raise RuntimeError(
                 "Are you sure this graph is a valid lineage graph?\n"

src/tracksdata/graph/_sql_graph.py

@@ -1341,3 +1341,92 @@ def __setstate__(self, state: dict) -> None:
         # recreate deleted objects
         self._engine = sa.create_engine(self._url, **self._engine_kwargs)
         self._define_schema(overwrite=False)
+
+    def tracklet_graph(
+        self,
+        track_id_key: str = DEFAULT_ATTR_KEYS.TRACK_ID,
+        ignore_track_id: int | None = None,
+    ) -> rx.PyDiGraph:
+        """
+        Create a compressed tracklet graph where each node is a tracklet
+        and each edge is a transition between tracklets.
+
+        IMPORTANT:
+        rx.PyDiGraph does not allow arbitrary indices, so we use the tracklet ids as node values.
+        And edge values are the tuple of source and target tracklet ids.
+
+        Parameters
+        ----------
+        track_id_key : str
+            The key of the track id attribute.
+        ignore_track_id : int | None
+            The track id to ignore. If None, all track ids are used.
+
+        Returns
+        -------
+        rx.PyDiGraph
+            A compressed tracklet graph.
+        """
+
+        if track_id_key not in self.node_attr_keys:
+            raise ValueError(f"Track id key '{track_id_key}' not found in graph. Expected '{self.node_attr_keys}'")
+
+        with Session(self._engine) as session:
+            node_query = sa.select(getattr(self.Node, track_id_key)).distinct()
+
+            SourceNode = aliased(self.Node)
+            TargetNode = aliased(self.Node)
+
+            edge_query = (
+                sa.select(
+                    getattr(self.Edge, DEFAULT_ATTR_KEYS.EDGE_SOURCE),
+                    getattr(self.Edge, DEFAULT_ATTR_KEYS.EDGE_TARGET),
+                )
+                .join(
+                    SourceNode,
+                    SourceNode.node_id == self.Edge.source_id,
+                )
+                .join(
+                    TargetNode,
+                    TargetNode.node_id == self.Edge.target_id,
+                )
+                .filter(
+                    getattr(SourceNode, track_id_key) != getattr(TargetNode, track_id_key),
+                )
+            )
+
+            if ignore_track_id is not None:
+                node_query = node_query.filter(getattr(self.Node, track_id_key) != ignore_track_id)
+                edge_query = edge_query.filter(
+                    getattr(SourceNode, track_id_key) != ignore_track_id,
+                    getattr(TargetNode, track_id_key) != ignore_track_id,
+                )
+
+            edge_query = edge_query.with_only_columns(
+                getattr(SourceNode, track_id_key).label("source_track_id"),
+                getattr(TargetNode, track_id_key).label("target_track_id"),
+            )
+
+            nodes_df = pl.read_database(
+                self._raw_query(node_query),
+                connection=session.connection(),
+            )
+
+            edges_df = pl.read_database(
+                self._raw_query(edge_query),
+                connection=session.connection(),
+            )
+
+        graph = rx.PyDiGraph()
+        tracklet_ids = nodes_df[track_id_key].to_list()
+        tracklet_id_to_rx = dict(zip(tracklet_ids, graph.add_nodes_from(tracklet_ids), strict=False))
+        graph.add_edges_from(
+            zip(
+                edges_df["source_track_id"].map_elements(tracklet_id_to_rx.__getitem__, return_dtype=int).to_list(),
+                edges_df["target_track_id"].map_elements(tracklet_id_to_rx.__getitem__, return_dtype=int).to_list(),
+                zip(edges_df["source_track_id"].to_list(), edges_df["target_track_id"].to_list(), strict=False),
+                strict=True,
+            )
+        )
+
+        return graph