Merge pull request #172 from tobac-project/Hotfix

Merging `hotfix` into `main` for v1.3.3

CHANGELOG.md

@@ -1,5 +1,11 @@
 ### Tobac Changelog
 
+_**Version 1.3.3:**_
+
+**Bug fixes**
+
+- Added a workaround to a bug in trackpy that fixes predictive tracking [#170](https://github.com/tobac-project/tobac/pull/170)
+
 _**Version 1.3.2:**_
 
 **Bug fixes**

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name="tobac",
-    version="1.3.2",
+    version="1.3.3",
     description="Tracking and object-based analysis of clouds",
     url="http://github.com/tobac-project/tobac",
     author=[

tobac/tests/test_tracking.py

@@ -5,6 +5,7 @@
 import tobac.testing
 import tobac.tracking
 import copy
+import pandas as pd
 from pandas.testing import assert_frame_equal
 import numpy as np
 
@@ -46,3 +47,60 @@ def test_linking_trackpy():
     assert_frame_equal(
         expected_out_feature.sort_index(axis=1), actual_out_feature.sort_index(axis=1)
     )
+
+
+def test_trackpy_predict():
+    """Function to test if linking_trackpy() with method='predict' correctly links two
+    features at constant speeds crossing each other.
+    """
+
+    cell_1 = tobac.testing.generate_single_feature(
+        1,
+        1,
+        min_h1=0,
+        max_h1=100,
+        min_h2=0,
+        max_h2=100,
+        frame_start=0,
+        num_frames=5,
+        spd_h1=20,
+        spd_h2=20,
+    )
+
+    cell_1_expected = copy.deepcopy(cell_1)
+    cell_1_expected["cell"] = 1
+
+    cell_2 = tobac.testing.generate_single_feature(
+        1,
+        100,
+        min_h1=0,
+        max_h1=100,
+        min_h2=0,
+        max_h2=100,
+        frame_start=0,
+        num_frames=5,
+        spd_h1=20,
+        spd_h2=-20,
+    )
+
+    cell_2_expected = copy.deepcopy(cell_2)
+    cell_2_expected["cell"] = 2
+
+    features = pd.concat([cell_1, cell_2])
+    expected_output = pd.concat([cell_1_expected, cell_2_expected])
+
+    output = tobac.linking_trackpy(
+        features, None, 1, 1, d_max=100, method_linking="predict"
+    )
+
+    output_random = tobac.linking_trackpy(
+        features, None, 1, 1, d_max=100, method_linking="random"
+    )
+
+    # check that the two methods of linking produce different results for this case
+    assert not output_random.equals(output)
+
+    # sorting and dropping indices for comparison with the expected output
+    output = output[["hdim_1", "hdim_2", "frame", "time", "feature", "cell"]]
+
+    assert_frame_equal(expected_output.sort_index(), output.sort_index())

tobac/tracking.py

@@ -113,12 +113,20 @@ def linking_trackpy(
         )
     elif method_linking == "predict":
 
+        # avoid setting pos_columns by renaimng to default values to avoid trackpy bug
+        features.rename(columns={"hdim_1": "y", "hdim_2": "x"}, inplace=True)
+
+        # generate list of features as input for df_link_iter to avoid bug in df_link
+        features_linking_list = [
+            frame for i, frame in features.groupby("frame", sort=True)
+        ]
+
         pred = tp.predict.NearestVelocityPredict(span=1)
-        trajectories_unfiltered = pred.link_df(
-            features_linking,
+        trajectories_unfiltered = pred.link_df_iter(
+            features_linking_list,
             search_range=search_range,
             memory=memory,
-            pos_columns=["hdim_1", "hdim_2"],
+            # pos_columns=["hdim_1", "hdim_2"], # not working atm
             t_column="frame",
             neighbor_strategy="KDTree",
             link_strategy="auto",
@@ -128,6 +136,14 @@ def linking_trackpy(
             #                                 hash_size=None, box_size=None, verify_integrity=True,
             #                                 retain_index=False
         )
+        # recreate a single dataframe from the list
+        trajectories_unfiltered = pd.concat(trajectories_unfiltered)
+
+        # change to column names back
+        trajectories_unfiltered.rename(
+            columns={"y": "hdim_1", "x": "hdim_2"}, inplace=True
+        )
+        features.rename(columns={"y": "hdim_1", "x": "hdim_2"}, inplace=True)
     else:
         raise ValueError("method_linking unknown")