diff --git a/plotly/figure_factory/_quiver.py b/plotly/figure_factory/_quiver.py index 8d0de352baf..0dd74b9d8f4 100644 --- a/plotly/figure_factory/_quiver.py +++ b/plotly/figure_factory/_quiver.py @@ -8,7 +8,7 @@ def create_quiver(x, y, u, v, scale=.1, arrow_scale=.3, - angle=math.pi / 9, **kwargs): + angle=math.pi / 9, scaleratio = None, **kwargs): """ Returns data for a quiver plot. @@ -21,6 +21,9 @@ def create_quiver(x, y, u, v, scale=.1, arrow_scale=.3, :param (float in [0,1]) arrow_scale: value multiplied to length of barb to get length of arrowhead. Default = .3 :param (angle in radians) angle: angle of arrowhead. Default = pi/9 + :param (positive float) angle: the ratio between the scale of the y-axis + and the scale of the x-axis (scale_y / scale_x). Default = None, the + scale ratio is not fixed. :param kwargs: kwargs passed through plotly.graph_objs.Scatter for more information on valid kwargs call help(plotly.graph_objs.Scatter) @@ -80,6 +83,29 @@ def create_quiver(x, y, u, v, scale=.1, arrow_scale=.3, # Add title to layout fig['layout'].update(title='Quiver Plot') + # Plot + py.plot(fig, filename='quiver') + ``` + + Example 4: Forcing a fix scale ratio to maintain the arrow length + ``` + import plotly.plotly as py + from plotly.figure_factory import create_quiver + + import numpy as np + + # Add data + x,y = np.meshgrid(np.arange(0.5, 3.5, .5), np.arange(0.5, 4.5, .5)) + u = x + v = y + angle = np.arctan(v / u) + norm = 0.25 + u = norm * np.cos(angle) + v = norm * np.sin(angle) + + # Create quiver with a fix scale ratio + fig = create_quiver(x, y, u, v, scale = 1, scaleratio = 0.5) + # Plot py.plot(fig, filename='quiver') ``` @@ -87,26 +113,39 @@ def create_quiver(x, y, u, v, scale=.1, arrow_scale=.3, utils.validate_equal_length(x, y, u, v) utils.validate_positive_scalars(arrow_scale=arrow_scale, scale=scale) - barb_x, barb_y = _Quiver(x, y, u, v, scale, - arrow_scale, angle).get_barbs() - arrow_x, arrow_y = _Quiver(x, y, u, v, scale, - arrow_scale, angle).get_quiver_arrows() - quiver = graph_objs.Scatter(x=barb_x + arrow_x, + if scaleratio is None: + quiver_obj = _Quiver(x, y, u, v, scale, arrow_scale, angle) + else: + quiver_obj = _Quiver(x, y, u, v, scale, arrow_scale, angle, scaleratio) + + barb_x, barb_y = quiver_obj.get_barbs() + arrow_x, arrow_y = quiver_obj.get_quiver_arrows() + + quiver_plot = graph_objs.Scatter(x=barb_x + arrow_x, y=barb_y + arrow_y, mode='lines', **kwargs) - data = [quiver] - layout = graph_objs.Layout(hovermode='closest') + data = [quiver_plot] - return graph_objs.Figure(data=data, layout=layout) + if scaleratio is None: + layout = graph_objs.Layout(hovermode='closest') + else: + layout = graph_objs.Layout( + hovermode='closest', + yaxis=dict( + scaleratio = scaleratio, + scaleanchor = "x" + ) + ) + return graph_objs.Figure(data=data, layout=layout) class _Quiver(object): """ Refer to FigureFactory.create_quiver() for docstring """ def __init__(self, x, y, u, v, - scale, arrow_scale, angle, **kwargs): + scale, arrow_scale, angle, scaleratio = 1, **kwargs): try: x = utils.flatten(x) except exceptions.PlotlyError: @@ -132,6 +171,7 @@ def __init__(self, x, y, u, v, self.u = u self.v = v self.scale = scale + self.scaleratio = scaleratio self.arrow_scale = arrow_scale self.angle = angle self.end_x = [] @@ -148,7 +188,7 @@ def scale_uv(self): endpoints of the arrows so a smaller scale value will result in less overlap of arrows. """ - self.u = [i * self.scale for i in self.u] + self.u = [i * self.scale * self.scaleratio for i in self.u] self.v = [i * self.scale for i in self.v] def get_barbs(self): @@ -188,13 +228,13 @@ def get_quiver_arrows(self): point1, endpoint, point2 y_values separated by a None to create the barb of the arrow. """ - dif_x = [i - j for i, j in zip(self.end_x, self.x)] + dif_x = [i - j for i, j in zip(self.end_x, self.x)] dif_y = [i - j for i, j in zip(self.end_y, self.y)] # Get barb lengths(default arrow length = 30% barb length) barb_len = [None] * len(self.x) for index in range(len(barb_len)): - barb_len[index] = math.hypot(dif_x[index], dif_y[index]) + barb_len[index] = math.hypot(dif_x[index] / self.scaleratio, dif_y[index]) # Make arrow lengths arrow_len = [None] * len(self.x) @@ -203,7 +243,7 @@ def get_quiver_arrows(self): # Get barb angles barb_ang = [None] * len(self.x) for index in range(len(barb_ang)): - barb_ang[index] = math.atan2(dif_y[index], dif_x[index]) + barb_ang[index] = math.atan2(dif_y[index], dif_x[index] / self.scaleratio) # Set angles to create arrow ang1 = [i + self.angle for i in barb_ang] @@ -231,9 +271,9 @@ def get_quiver_arrows(self): # Set coordinates to create arrow for index in range(len(self.end_x)): - point1_x = [i - j for i, j in zip(self.end_x, seg1_x)] + point1_x = [i - j * self.scaleratio for i, j in zip(self.end_x, seg1_x)] point1_y = [i - j for i, j in zip(self.end_y, seg1_y)] - point2_x = [i - j for i, j in zip(self.end_x, seg2_x)] + point2_x = [i - j * self.scaleratio for i, j in zip(self.end_x, seg2_x)] point2_y = [i - j for i, j in zip(self.end_y, seg2_y)] # Combine lists to create arrow diff --git a/plotly/tests/test_optional/test_figure_factory/test_figure_factory.py b/plotly/tests/test_optional/test_figure_factory/test_figure_factory.py index 32c25113859..8a14be6242f 100644 --- a/plotly/tests/test_optional/test_figure_factory/test_figure_factory.py +++ b/plotly/tests/test_optional/test_figure_factory/test_figure_factory.py @@ -2848,3 +2848,43 @@ def test_full_choropleth(self): ] self.assertEqual(fig['data'][2]['x'][:50], exp_fig_head) + +class TestQuiver(TestCase): + + def test_scaleratio_param(self): + x,y = np.meshgrid(np.arange(0.5, 3.5, .5), np.arange(0.5, 4.5, .5)) + u = x + v = y + angle = np.arctan(v / u) + norm = 0.25 + u = norm * np.cos(angle) + v = norm * np.sin(angle) + fig = ff.create_quiver(x, y, u, v, scale = 1, scaleratio = 0.5) + + exp_fig_head = [( + 0.5, + 0.5883883476483185, + None, + 1.0, + 1.1118033988749896, + None, + 1.5, + 1.6185854122563141, + None, + 2.0), + (0.5, + 0.6767766952966369, + None, + 0.5, + 0.6118033988749895, + None, + 0.5, + 0.5790569415042095, + None, + 0.5)] + + fig_head = [fig['data'][0]['x'][:10], fig['data'][0]['y'][:10]] + + self.assertEqual(fig_head, exp_fig_head) + +