From 8385e38a9a7caf962e1487c5760ea9773de42883 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Simon=20H=C3=B8xbro=20Hansen?= <simon.hansen@me.com>
Date: Wed, 20 Sep 2023 10:29:04 +0200
Subject: [PATCH] Correctly merge field_df and region_df

---
 holonote/annotate/annotator.py    | 29 +++++++++----------
 holonote/tests/test_indicators.py | 46 +++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 14 deletions(-)
 create mode 100644 holonote/tests/test_indicators.py

diff --git a/holonote/annotate/annotator.py b/holonote/annotate/annotator.py
index 8c4b624..3c4c7ac 100644
--- a/holonote/annotate/annotator.py
+++ b/holonote/annotate/annotator.py
@@ -71,24 +71,25 @@ def ranges_1d(cls, region_df, field_df, invert_axes=False, extra_params=None):
 
     @classmethod
     def _range_indicators(cls, region_df, field_df, dimensionality, invert_axes=False, extra_params=None):
-        rect_data = []
+        # TODO: Clean this up VSpans/HSpans/VLines/HLines
+        data = region_df.merge(field_df, left_on="_id", right_index=True)
+        index_col_name = 'id' if field_df.index.name is None else field_df.index.name
 
-        mdata_vals = ([None] * len(region_df['_id'])
-                      if len(field_df.columns)==0 else field_df.to_dict('records'))
-        for id_val, value, mdata in zip(region_df['_id'], region_df["value"], mdata_vals):
-            if dimensionality=='1d':
-                coords = (value[0], extra_params['rect_min'], value[1], extra_params['rect_max'])
-            else:
-                coords = (value[0], value[2], value[1], value[3]) # LBRT format
-
-            if None in coords: continue
+        values = pd.DataFrame.from_records(data["value"])
+        id_vals = data["_id"].rename({"_id": index_col_name})
+        mdata_vals = data[field_df.columns]
 
-            mdata_tuple =  () if len(field_df.columns)==0 else tuple(mdata.values())
-            rect_data.append(coords + mdata_tuple + (id_val,))
+        # TODO: Add check for none in values
 
-        index_col_name = ['id'] if field_df.index.name is None else [field_df.index.name]
-        return hv.Rectangles(rect_data, vdims=list(field_df.columns)+index_col_name) # kdims?
+        if dimensionality=='1d':
+            coords = values[[0, 0, 1, 1]].copy()
+            coords.iloc[:, 1] = extra_params["rect_min"]
+            coords.iloc[:, 3] = extra_params["rect_max"]
+        else:
+            coords = values[[0, 2, 1, 3]] # LBRT format
 
+        rect_data = list(pd.concat([coords, mdata_vals, id_vals], axis=1).itertuples(index=False))
+        return hv.Rectangles(rect_data, vdims=[*field_df.columns, index_col_name]) # kdims?
 
 
 class AnnotatorInterface(param.Parameterized):
diff --git a/holonote/tests/test_indicators.py b/holonote/tests/test_indicators.py
new file mode 100644
index 0000000..69ec9c2
--- /dev/null
+++ b/holonote/tests/test_indicators.py
@@ -0,0 +1,46 @@
+import numpy as np
+import pandas as pd
+
+from holonote.annotate.annotator import Indicator
+
+
+def test_range2d_id_matches() -> None:
+    value = np.arange(8).reshape(2, 4)
+    region_df = pd.DataFrame({"value": list(value), "_id": ["A", "B"]})
+    field_df = pd.DataFrame(["B", "A"], index=["B", "A"], columns=["description"])
+
+    # id and description should match
+    output = Indicator.ranges_2d(region_df, field_df).data
+    expected = pd.DataFrame(
+        {
+            "x0": {0: 0, 1: 4},
+            "y0": {0: 2, 1: 6},
+            "x1": {0: 1, 1: 5},
+            "y1": {0: 3, 1: 7},
+            "description": {0: "A", 1: "B"},
+            "id": {0: "A", 1: "B"},
+        }
+    )
+    pd.testing.assert_frame_equal(output, expected)
+
+
+def test_range1d_id_matches() -> None:
+    value = np.arange(4).reshape(2, 2)
+    region_df = pd.DataFrame({"value": list(value), "_id": ["A", "B"]})
+    field_df = pd.DataFrame(["B", "A"], index=["B", "A"], columns=["description"])
+
+    # id and description should match
+    output = Indicator.ranges_1d(
+        region_df, field_df, extra_params={"rect_min": -2, "rect_max": -2}
+    ).data
+    expected = pd.DataFrame(
+        {
+            "x0": {0: 0, 1: 2},
+            "y0": {0: -2, 1: -2},
+            "x1": {0: 1, 1: 3},
+            "y1": {0: -2, 1: -2},
+            "description": {0: "A", 1: "B"},
+            "id": {0: "A", 1: "B"},
+        }
+    )
+    pd.testing.assert_frame_equal(output, expected)