Merge pull request #330 from yfukai/index_offset

adding index_offset

src/laptrack/_tracking.py

@@ -657,6 +657,7 @@ def predict_dataframe(
         validate_frame: bool = True,
         only_coordinate_cols: bool = True,
         connected_edges: Optional[List[Tuple[Int, Int]]] = None,
+        index_offset: Int = 0,
     ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
         """Shorthand for the tracking with the dataframe input / output.
 
@@ -676,6 +677,8 @@ def predict_dataframe(
             The edges that is known to be connected.
             Must be a list of the tuples of the row numbers (not index, but `iloc`).
             If None, no edges are assumed to be connected.
+        index_offset : Int
+            The offset to add to the track and tree index.
 
         Returns
         -------
@@ -721,6 +724,16 @@ def predict_dataframe(
             track_df, split_df, merge_df = convert_tree_to_dataframe(
                 tree, dataframe=df, frame_index=frame_index
             )
+
+        track_df["track_id"] = track_df["track_id"] + index_offset
+        track_df["tree_id"] = track_df["tree_id"] + index_offset
+        if not split_df.empty:
+            split_df["parent_track_id"] = split_df["parent_track_id"] + index_offset
+            split_df["child_track_id"] = split_df["child_track_id"] + index_offset
+        if not merge_df.empty:
+            merge_df["parent_track_id"] = merge_df["parent_track_id"] + index_offset
+            merge_df["child_track_id"] = merge_df["child_track_id"] + index_offset
+
         return track_df, split_df, merge_df
 
 

tests/test_tracking.py

@@ -137,6 +137,28 @@ def test_reproducing_trackmate(testdata, tracker_class, parallel_backend) -> Non
     assert all(split_df == split_df2)
     assert all(merge_df == merge_df2)
 
+    # check index offset
+    track_df3, split_df3, merge_df3 = lt.predict_dataframe(
+        df, ["x", "y"], index_offset=2
+    )
+    assert min(track_df3["track_id"]) == 2
+    assert min(track_df3["tree_id"]) == 2
+
+    track_df4 = track_df2.copy()
+    split_df4 = split_df2.copy()
+    merge_df4 = merge_df2.copy()
+    track_df4["track_id"] = track_df4["track_id"] + 2
+    track_df4["tree_id"] = track_df4["tree_id"] + 2
+    if not split_df4.empty:
+        split_df4["parent_track_id"] = split_df4["parent_track_id"] + 2
+        split_df4["child_track_id"] = split_df4["child_track_id"] + 2
+    if not merge_df4.empty:
+        merge_df4["parent_track_id"] = merge_df4["parent_track_id"] + 2
+        merge_df4["child_track_id"] = merge_df4["child_track_id"] + 2
+    assert all(track_df3 == track_df2)
+    assert all(split_df3 == split_df2)
+    assert all(merge_df3 == merge_df2)
+
     track_df, split_df, merge_df = lt.predict_dataframe(
         df, ["x", "y"], only_coordinate_cols=False
     )