Support variety of data inputs for wind barbs (#662)

holoviz · Jul 19, 2023 · 7fa6484 · 7fa6484
1 parent f114de0
commit 7fa6484
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 21 deletions.
diff --git a/geoviews/element/geo.py b/geoviews/element/geo.py
@@ -353,29 +353,31 @@ class WindBarbs(_Element, Selection2DExpr, HvGeometry):
 
     @classmethod
     def from_uv(cls, data, kdims=None, vdims=None, **params):
-        if isinstance(data, tuple):
-            xs, ys, us, vs = data
-        else:
-            us = data[vdims[0]]
-            vs = data[vdims[1]]
+        if kdims is None:
+            kdims = ['x', 'y']
+        if vdims is None:
+            vdims = ['u', 'v']
+        dataset = Dataset(data, kdims=kdims, vdims=vdims, **params)
+        us, vs = (dataset.dimension_values(i) for i in (2, 3))
 
         uv_magnitudes = np.hypot(us, vs)  # unscaled
+        # using meteorological convention (direction FROM which wind blows)
         radians = np.pi / 2 - np.arctan2(-vs, -us)
 
-        if isinstance(data, tuple):
-            reorganized_data = (xs, ys, radians, uv_magnitudes)
-        else:
-            # calculations on this data could mutate the original data
-            # here we do not do any calculations; we only store the data
-            reorganized_data = {}
-            for kdim in kdims:
-                reorganized_data[kdim] = data[kdim]
-            reorganized_data["Angle"] = radians
-            reorganized_data["Magnitude"] = uv_magnitudes
-            for vdim in vdims[2:]:
-                reorganized_data[vdim] = data[vdim]
-            vdims = ["Angle", "Magnitude"] + vdims[2:]
-        return cls(reorganized_data, kdims=kdims, vdims=vdims, **params)
+        # calculations on this data could mutate the original data
+        # here we do not do any calculations; we only store the data
+        repackaged_dataset = {}
+        for kdim in kdims:
+            repackaged_dataset[kdim] = dataset[kdim]
+        repackaged_dataset["Angle"] = radians
+        repackaged_dataset["Magnitude"] = uv_magnitudes
+        for vdim in vdims[2:]:
+            repackaged_dataset[vdim] = dataset[vdim]
+        vdims = [
+            Dimension('Angle', cyclic=True, range=(0, 2 * np.pi)),
+            Dimension('Magnitude')
+        ] + vdims[2:]
+        return cls(repackaged_dataset, kdims=kdims, vdims=vdims, **params)
 
 
 class Image(_Element, HvImage):

diff --git a/geoviews/tests/plotting/mpl/test_chart.py b/geoviews/tests/plotting/mpl/test_chart.py
@@ -64,8 +64,8 @@ def test_windbarbs_from_uv(self):
         gv_barbs = WindBarbs((X, Y, angle, mag))
         gv_barbs_uv = WindBarbs.from_uv((X, Y, U, V))
 
-        np.testing.assert_almost_equal(gv_barbs.data["Angle"], gv_barbs_uv.data["Angle"])
-        np.testing.assert_almost_equal(gv_barbs.data["Magnitude"], gv_barbs_uv.data["Magnitude"])
+        np.testing.assert_almost_equal(gv_barbs.data["Angle"].T.flatten(), gv_barbs_uv.data["Angle"])
+        np.testing.assert_almost_equal(gv_barbs.data["Magnitude"].T.flatten(), gv_barbs_uv.data["Magnitude"])
 
     def test_windbarbs_dataset_from_uv_other_dim(self):
         x = np.linspace(-1, 1, 4)