Add VectorField/WindBarbs project operation (#296)

Co-authored-by: Andrew Huang <ahuang@anaconda.com>

geoviews/operation/__init__.py

@@ -6,7 +6,8 @@
 from ..element import _Element
 from .projection import ( # noqa (API import)
     project_image, project_path, project_shape, project_points,
-    project_graph, project_quadmesh, project_geom, project)
+    project_graph, project_quadmesh, project_geom,
+    project_vectorfield, project_windbarbs, project)
 from .resample import resample_geometry # noqa (API import)
 
 geo_ops = [contours, bivariate_kde]

geoviews/operation/projection.py

@@ -15,7 +15,7 @@
 from ..data import GeoPandasInterface
 from ..element import (Image, Shape, Polygons, Path, Points, Contours,
                        RGB, Graph, Nodes, EdgePaths, QuadMesh, VectorField,
-                       HexTiles, Labels, Rectangles, Segments)
+                       HexTiles, Labels, Rectangles, Segments, WindBarbs)
 from ..util import (
     project_extents, path_to_geom_dicts, polygons_to_geom_dicts,
     geom_dict_to_array_dict
@@ -157,7 +157,7 @@ def _process_element(self, element):
 
 class project_points(_project_operation):
 
-    supported_types = [Points, Nodes, VectorField, HexTiles, Labels]
+    supported_types = [Points, Nodes, HexTiles, Labels]
 
     def _process_element(self, element):
         if not len(element):
@@ -194,8 +194,6 @@ class project_geom(_project_operation):
     supported_types = [Rectangles, Segments]
 
     def _process_element(self, element):
-        if not len(element):
-            return element.clone(crs=self.p.projection)
         x0d, y0d, x1d, y1d = element.kdims
         x0, y0, x1, y1 = (element.dimension_values(i) for i in range(4))
         p1 = self.p.projection.transform_points(element.crs, x0, y0)
@@ -222,6 +220,55 @@ def _process_element(self, element):
         return element.clone(tuple(new_data[d.name] for d in element.dimensions()),
                              crs=self.p.projection)
 
+class project_vectorfield(_project_operation):
+
+    supported_types = [VectorField]
+
+    def _calc_angles(self, ut, vt):
+        # mathematical convention; follows matplotlib
+        return np.arctan2(vt, ut)
+
+    def _process_element(self, element):
+        if not len(element):
+            return element.clone(crs=self.p.projection)
+
+        xdim, ydim, adim, mdim = element.dimensions()[:4]
+        xs, ys, ang, ms = (element.dimension_values(i) for i in range(4))
+        coordinates = self.p.projection.transform_points(element.crs, xs, ys)
+        mask = np.isfinite(coordinates[:, 0])
+        new_data = {k: v[mask] for k, v in element.columns().items()}
+        new_data[xdim.name] = coordinates[mask, 0]
+        new_data[ydim.name] = coordinates[mask, 1]
+        datatype = [element.interface.datatype]+element.datatype
+        us = np.sin(ang) * -ms
+        vs = np.cos(ang) * -ms
+        ut, vt = self.p.projection.transform_vectors(element.crs, xs, ys, us, vs)
+        with np.errstate(divide='ignore', invalid='ignore'):
+            angle = self._calc_angles(ut, vt)
+        mag = np.hypot(ut, vt)
+
+        new_data[adim.name] = angle[mask]
+        new_data[mdim.name] = mag[mask]
+
+        if len(new_data[xdim.name]) == 0:
+            self.warning('While projecting a {} element from a {} coordinate '
+                         'reference system (crs) to a {} projection none of '
+                         'the projected paths were contained within the bounds '
+                         'specified by the projection. Ensure you have specified '
+                         'the correct coordinate system for your data.'.format(type(element).__name__, type(element.crs).__name__,
+                          type(self.p.projection).__name__))
+
+        return element.clone(tuple(new_data[d.name] for d in element.dimensions()),
+                             crs=self.p.projection, datatype=datatype)
+
+class project_windbarbs(project_vectorfield):
+
+    supported_types = [WindBarbs]
+
+    def _calc_angles(self, ut, vt):
+        # meteorological convention; follows matplotlib
+        return np.pi / 2 - np.arctan2(-vt, -ut)
+
 
 class project_graph(_project_operation):
 
@@ -445,7 +492,7 @@ class project(Operation):
 
     _operations = [project_path, project_image, project_shape,
                    project_graph, project_quadmesh, project_points,
-                   project_geom]
+                   project_vectorfield, project_windbarbs, project_geom]
 
     def _process(self, element, key=None):
         for op in self._operations:

geoviews/plotting/bokeh/__init__.py

@@ -23,7 +23,7 @@
 )
 from ...operation import (
     project_image, project_points, project_path, project_graph,
-    project_quadmesh, project_geom
+    project_quadmesh, project_geom, project_vectorfield
 )
 from ...tile_sources import _ATTRIBUTIONS
 from ...util import poly_types, line_types
@@ -106,7 +106,7 @@ class GeoPointPlot(GeoPlot, PointPlot):
 
 class GeoVectorFieldPlot(GeoPlot, VectorFieldPlot):
 
-    _project_operation = project_points
+    _project_operation = project_vectorfield
 
 
 class GeoQuadMeshPlot(GeoPlot, QuadMeshPlot):

geoviews/plotting/mpl/__init__.py

@@ -37,7 +37,7 @@
 
 from ...operation import (
     project_points, project_path, project_graph, project_quadmesh,
-    project_geom
+    project_geom, project_vectorfield, project_windbarbs
 )
 from .chart import WindBarbsPlot
 
@@ -325,7 +325,7 @@ class GeoVectorFieldPlot(GeoPlot, VectorFieldPlot):
 
     apply_ranges = param.Boolean(default=True)
 
-    _project_operation = project_points
+    _project_operation = project_vectorfield
 
 
 class GeoWindBarbsPlot(GeoPlot, WindBarbsPlot):
@@ -335,7 +335,7 @@ class GeoWindBarbsPlot(GeoPlot, WindBarbsPlot):
 
     apply_ranges = param.Boolean(default=True)
 
-    _project_operation = project_points
+    _project_operation = project_windbarbs
 
 
 class GeometryPlot(GeoPlot):

geoviews/tests/test_projection.py

@@ -1,7 +1,7 @@
 import numpy as np
 import cartopy.crs as ccrs
 
-from geoviews.element import Image
+from geoviews.element import Image, VectorField, WindBarbs
 from geoviews.element.comparison import ComparisonTestCase
 from geoviews.operation import project
 
@@ -31,3 +31,45 @@ def test_image_project_latlon_to_mercator(self):
             [     0.,      0.,      0.,      0.,      0.],
             [ 12960.,  17280.,  21600.,   4320.,   8640.]
         ]))
+
+    def test_project_vectorfield(self):
+        xs = np.linspace(10, 50, 2)
+        X, Y = np.meshgrid(xs, xs)
+        U, V = 5 * X, 1 * Y
+        A = np.arctan2(V, U)
+        M = np.hypot(U, V)
+        crs = ccrs.PlateCarree()
+        vectorfield = VectorField((X, Y, A, M), crs=crs)
+        projection = ccrs.Orthographic()
+        projected = project(vectorfield, projection=projection)
+        assert projected.crs == projection
+
+        xs, ys, ang, ms = (vectorfield.dimension_values(i) for i in range(4))
+        us = np.sin(ang) * -ms
+        vs = np.cos(ang) * -ms
+        u, v = projection.transform_vectors(crs, xs, ys, us, vs)
+        a, m = np.arctan2(v, u).T, np.hypot(u, v).T
+
+        np.testing.assert_allclose(projected.dimension_values("Angle"), a.flatten())
+        np.testing.assert_allclose(projected.dimension_values("Magnitude"), m.flatten())
+
+    def test_project_windbarbs(self):
+        xs = np.linspace(10, 50, 2)
+        X, Y = np.meshgrid(xs, xs)
+        U, V = 5 * X, 1 * Y
+        A = np.arctan2(V, U)
+        M = np.hypot(U, V)
+        crs = ccrs.PlateCarree()
+        windbarbs = WindBarbs((X, Y, A, M), crs=crs)
+        projection = ccrs.Orthographic()
+        projected = project(windbarbs, projection=projection)
+        assert projected.crs == projection
+
+        xs, ys, ang, ms = (windbarbs.dimension_values(i) for i in range(4))
+        us = np.sin(ang) * -ms
+        vs = np.cos(ang) * -ms
+        u, v = projection.transform_vectors(crs, xs, ys, us, vs)
+        a, m = np.pi / 2 - np.arctan2(-v, -u).T, np.hypot(u, v).T
+
+        np.testing.assert_allclose(projected.dimension_values("Angle"), a.flatten())
+        np.testing.assert_allclose(projected.dimension_values("Magnitude"), m.flatten())