Skip to content

Commit

Permalink
Fix angles
Browse files Browse the repository at this point in the history
  • Loading branch information
@ahuang11
ahuang11 committed Jun 27, 2023
1 parent 2246b4c commit bc5efea
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 50 deletions.
24 changes: 24 additions & 0 deletions geoviews/element/geo.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -351,6 +351,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):
"""
Expand Down
38 changes: 13 additions & 25 deletions geoviews/plotting/mpl/chart.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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

Expand Down
73 changes: 48 additions & 25 deletions geoviews/tests/plotting/mpl/test_chart.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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"
Expand All @@ -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"
Expand All @@ -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"
Expand All @@ -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,
Expand Down

0 comments on commit bc5efea

Please sign in to comment.