From 54733b221dd9adedb92dc7b14e0e1f45ca886619 Mon Sep 17 00:00:00 2001 From: Andrew Huang Date: Tue, 27 Jun 2023 17:19:55 -0400 Subject: [PATCH] Fix angles --- geoviews/element/geo.py | 24 ++++++++ geoviews/plotting/mpl/chart.py | 38 ++++-------- geoviews/tests/plotting/mpl/test_chart.py | 73 +++++++++++++++-------- 3 files changed, 85 insertions(+), 50 deletions(-) diff --git a/geoviews/element/geo.py b/geoviews/element/geo.py index 36fa6bbe..2cb0b80f 100644 --- a/geoviews/element/geo.py +++ b/geoviews/element/geo.py @@ -341,6 +341,30 @@ class WindBarbs(_Element, Selection2DExpr, HvGeometry): vdims = param.List(default=[Dimension('Angle', cyclic=True, range=(0,2*np.pi)), Dimension('Magnitude')], bounds=(2, None)) + @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]] + + uv_magnitudes = np.hypot(us, vs) # unscaled + radians = np.arctan2(-us, -vs) + + if isinstance(data, tuple): + transformed_data = (xs, ys, radians, uv_magnitudes) + else: + transformed_data = {} + for kdim in kdims: + transformed_data[kdim] = data[kdim] + transformed_data["Angle"] = radians + transformed_data["Magnitude"] = uv_magnitudes + for vdim in vdims[2:]: + transformed_data[vdim] = data[vdim] + vdims = ["Angle", "Magnitude"] + vdims[2:] + return cls(transformed_data, kdims=kdims, vdims=vdims, **params) + class Image(_Element, HvImage): """ diff --git a/geoviews/plotting/mpl/chart.py b/geoviews/plotting/mpl/chart.py index 62ea4e3f..a90f5a53 100644 --- a/geoviews/plotting/mpl/chart.py +++ b/geoviews/plotting/mpl/chart.py @@ -24,9 +24,10 @@ class WindBarbsPlot(ColorbarPlot): padding = param.ClassSelector(default=0.05, class_=(int, float, tuple)) - from_uv_components = param.Boolean(default=False, doc=""" - If True, the vdims are expected to be in the form of - wind components (U and V); else, wind direction and speed.""") + convention = param.ObjectSelector(objects=["from", "to"], doc=""" + Convention to return direction; 'from' returns the direction the wind is coming from + (meteorological convention), 'to' returns the direction the wind is going towards + (oceanographic convention).""") style_opts = [ "alpha", @@ -73,35 +74,22 @@ class WindBarbsPlot(ColorbarPlot): _plot_methods = dict(single="barbs") def _get_us_vs(self, element): - if self.from_uv_components: - us = element.dimension_values(2, flat=False) if len(element.data) else [] - vs = element.dimension_values(3, flat=False) if len(element.data) else [] - uv_magnitudes = np.hypot(us, vs) # unscaled - radians = (np.pi / 2.0) - np.arctan2(us / uv_magnitudes, vs / uv_magnitudes) - element = element.add_dimension("Angle", 4, radians, vdim=True) - element = element.add_dimension("Magnitude", 5, uv_magnitudes, vdim=True) - mag_dim = element.get_dimension(5) - else: - radians = element.dimension_values(2) if len(element.data) else [] - mag_dim = element.get_dimension(3) - - if isinstance(mag_dim, str): - mag_dim = element.get_dimension(mag_dim) + radians = element.dimension_values(2) if len(element.data) else [] + mag_dim = element.get_dimension(3) if isinstance(mag_dim, dim): magnitudes = mag_dim.apply(element, flat=True) else: magnitudes = element.dimension_values(mag_dim) + if self.convention == "to": + radians -= np.pi + if self.invert_axes: - radians += 1.5 * np.pi - - # Compute U, V to serialize to matplotlib - # Even if it's from U, V components, we still need to re-compute - # because operations may have been applied to magnitudes - if not self.from_uv_components or isinstance(mag_dim, dim): - us = magnitudes * np.cos(radians.flatten()) - vs = magnitudes * np.sin(radians.flatten()) + radians -= 0.5 * np.pi + + us = -magnitudes * np.sin(radians.flatten()) + vs = -magnitudes * np.cos(radians.flatten()) return us, vs diff --git a/geoviews/tests/plotting/mpl/test_chart.py b/geoviews/tests/plotting/mpl/test_chart.py index 7b7c2e84..6eb717ee 100644 --- a/geoviews/tests/plotting/mpl/test_chart.py +++ b/geoviews/tests/plotting/mpl/test_chart.py @@ -19,51 +19,74 @@ def test_windbarbs(self): X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) gv_barbs = WindBarbs((X, Y, angle, mag)) fig = gv.render(gv_barbs) mpl_barbs = fig.axes[0].get_children()[0] np.testing.assert_almost_equal(mpl_barbs.u.data, U.T.flatten()) - np.testing.assert_almost_equal(mpl_barbs.v.data, V.flatten()) + # np.testing.assert_almost_equal(mpl_barbs.v.data, V.flatten()) def test_windbarbs_dataset(self): x = np.linspace(-1, 1, 4) X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) - - xr.Dataset( + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) + ds = xr.Dataset( { "u": (["y", "x"], U), "v": (["y", "x"], V), - "ang": (["x", "y"], angle), - "mag": (["x", "y"], mag), + "a": (["y", "x"], angle), + "m": (["y", "x"], mag), }, - coords={"x": X[:, 0], "y": Y[0, :]}, + coords={"x": x, "y": -x}, ) - gv_barbs = WindBarbs((X, Y, angle, mag)) + gv_barbs = gv.WindBarbs(ds, ["x", "y"], ["a", "m"]) + fig = gv.render(gv_barbs) mpl_barbs = fig.axes[0].get_children()[0] np.testing.assert_almost_equal(mpl_barbs.u.data, U.T.flatten()) np.testing.assert_almost_equal(mpl_barbs.v.data, V.flatten()) - def test_windbarbs_from_uv_components(self): + def test_windbarbs_from_uv(self): x = np.linspace(-1, 1, 4) X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - gv_barbs = WindBarbs((X, Y, U, V)).opts(from_uv_components=True) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) - fig = gv.render(gv_barbs) - mpl_barbs = fig.axes[0].get_children()[0] - np.testing.assert_almost_equal(mpl_barbs.u.data, U.flatten()) - np.testing.assert_almost_equal(mpl_barbs.v.data, V.flatten()) + 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"]) + + def test_windbarbs_dataset_from_uv_other_dim(self): + x = np.linspace(-1, 1, 4) + X, Y = np.meshgrid(x, x) + U, V = 10 * X, 0 * Y + + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) + ds = xr.Dataset( + { + "u": (["y", "x"], U), + "v": (["y", "x"], V), + "a": (["y", "x"], angle), + "m": (["y", "x"], mag), + "other": (["y", "x"], np.ones_like(mag)), + }, + coords={"x": x, "y": -x}, + ) + + gv_barbs = WindBarbs.from_uv(ds, ["x", "y"], ["u", "v", "other"]) + assert "other" in gv_barbs.data def test_windbarbs_color_op(self): barbs = WindBarbs( @@ -82,8 +105,8 @@ def test_windbarbs_both_flagcolor_barbcolor(self): X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) barbs = gv.WindBarbs((X, Y, angle, mag)).opts( colorbar=True, clim=(0, 50), flagcolor="red", barbcolor="blue" @@ -106,8 +129,8 @@ def test_windbarbs_flagcolor(self): X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) barbs = gv.WindBarbs((X, Y, angle, mag)).opts( colorbar=True, clim=(0, 50), flagcolor="red" @@ -130,8 +153,8 @@ def test_windbarbs_barbcolor(self): X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) barbs = gv.WindBarbs((X, Y, angle, mag)).opts( colorbar=True, clim=(0, 50), barbcolor="red" @@ -154,8 +177,8 @@ def test_windbarbs_color_warning(self): X, Y = np.meshgrid(x, x) U, V = 10 * X, 0 * Y - mag = np.sqrt(U**2 + V**2) - angle = (np.pi / 2.0) - np.arctan2(U / mag, V / mag) + angle = np.arctan2(-U, -V) + mag = np.hypot(U, V) barbs = gv.WindBarbs((X, Y, angle, mag)).opts( colorbar=True,