简体   繁体   中英

matplotlib arrowheads and aspect ratio

If I run this script:

import matplotlib.pyplot as plt
import pylab as plab

plt.figure()
plt.plot([0,2], [2,0], color='c', lw=0.5)
plt.plot([1,2], [2,1], color='k', lw=0.5)
plt.arrow(1,1,0.5,0.5, head_width=0.1, width=0.01, head_length=0.1, color='r')
plt.arrow(1.25,0.75,0.5,0.5, head_width=0.1, width=0.01, head_length=0.1, color='g')

plab.axes().set_aspect(0.5)

plt.show()

I get this: 在此处输入图像描述 Notice that the backs of the arrowheads are flush with that black line that stretches from (1,2) to (2,1). That makes sense, but I want the backs of the arrowheads to be VISUALLY perpendicular to the tails WITHOUT changing the aspect ratio. How do I do this?

I've been annoyed with that problem for (almost) ever, and mostly end up using annotate() to draw arrows. For example (lacking a lot of tweaking to end up with the identical result as your plot...):

import matplotlib.pyplot as plt
import pylab as plab

plt.figure()
ax=plt.subplot(211)
plt.plot([0,2], [2,0], color='c', lw=0.5)
plt.plot([1,2], [2,1], color='k', lw=0.5)
plt.arrow(1,1,0.5,0.5, head_width=0.1, width=0.01, head_length=0.1, color='r')
ax.set_aspect(0.5)

ax=plt.subplot(212)
plt.plot([0,2], [2,0], color='c', lw=0.5)
plt.plot([1,2], [2,1], color='k', lw=0.5)
ax.annotate("",
            xy=(1.5, 1.5), xycoords='data',
            xytext=(1, 1), textcoords='data',
            arrowprops=dict(arrowstyle="-|>",
                            connectionstyle="arc3"),
            )
ax.set_aspect(0.5)

plt.show()

在此输入图像描述

It took me some time but I created a very basic class called WarpArrow to get arrows in matplotlib that behave similar to plt.arrow() and that take the aspectratio into consideration. Note that the aspectratio should be fixed. I didn't test it a ton so be careful with it and feel free to add more keyword arguments.

Basicly what it does it takes the corner points of some dummy arrow lying on the x-axis and rotates them by amount given as dx and dy. The rotation takes the aspectratio into consideration and there you have it. The rest is simply tweeking the corner points according to the input arguments, filling a polygon consisting of the cornerpoints and plotting a line to the arrowhead ('stem').

You can tell it to use 'x' or 'y' coordinates as a measurment unit for the head length and head width.

Tipp: draw() an arrow in x or y direction to find your desired length and width for the arrow head.

usage:

# creating arrow object (default: head_length='x' head_width_unit='y')
# set matplotlib axes and use all other standard arrow kwargs

my_arrow = WarpArrow(axes, x, y, dx, dy, head_length='x' head_width_unit='y')
my_arrow.draw()   # draw arrow onto axes

here is the full code to copy, mind the package names

import math
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

class WarpArrow:
    def __init__(self, axes, x, y, dx, dy, head_width_unit='y', head_length_unit='x', width_unit='y', **kwargs):
        # handle inputs
        self.axes = axes
        self.aspect = self.axes.get_aspect()
        if self.aspect == 'auto':
            self.aspect = 1

        self.x = x
        self.y = y
        self.dx = dx
        self.dy = dy
        self.len = math.sqrt(self.dx ** 2 + (self.dy * self.aspect) ** 2)

        self.kwargs = kwargs

        if self.dx == 0:  # computing arctan for rotation angle
            self.angle_rad = ((self.dy < 0) + .5) * math.pi
        else:
            _arctan = np.arctan(self.dy * self.aspect / self.dx)
            self.angle_rad = ((self.dx > 0) * (self.dy > 0) + (self.dx > 0) * (self.dy < 0)) * _arctan + ((self.dx < 0) * (self.dy > 0) + (self.dx < 0)*(self.dy <= 0)) * (math.pi + _arctan)
        self.angle_deg = 180 * self.angle_rad/math.pi

        # units for width and length of arrowhead
        self.head_width_unit = head_width_unit
        self.head_length_unit = head_length_unit
        self.width_unit = width_unit

        _err1, _err2 = '', ''                             # Error managment for units
        if self.head_width_unit not in ['x', 'y']:
            _err1 = 'head_width_unit'
        if self.head_length_unit not in ['x', 'y']:
            _err2 = 'head_length_unit'
        if _err1 + _err2 != '':
            raise ValueError(
                'choose \'x\' or \'y\' for \'' + _err1 + (_err1 != '') * (_err2 != '') * '\' and \'' + _err2 + '\'')

        if self.width_unit not in ['x', 'y']:
            raise ValueError('choose \'x\' or \'y\' for \'width_unit\'')

        # setting arrow kwargs arguments (and defaults)
        if 'width' in kwargs:
            self.width = kwargs['width']
        else:
            self.width = 0.001

        if 'head_width' in kwargs:
            self.head_width = kwargs['head_width']
        else:
            self.head_width = 3*self.width

        if 'head_length' in kwargs:
            self.head_length= kwargs['head_length']
        else:
            self.head_length = 3*self.head_width

        # convert length and width for computation
        if self.width_unit == 'x':
            self.norm_head_width = self.head_width * self.aspect
        if self.width_unit == 'y':
            self.norm_head_width = self.head_width

        if self.head_length_unit == 'x':
            self.norm_head_length = self.head_length
        if self.head_width_unit == 'y':
            self.norm_head_width = self.head_width

        if self.head_width_unit == 'x':
            self.norm_head_width = self.head_width / self.aspect
        if self.head_length_unit == 'y':
            self.norm_head_length = self.head_length * self.aspect

        if 'length_includes_head' in kwargs:
            self.length_includes_head = kwargs['length_includes_head']
        else:
            self.length_includes_head = False

        # further arrow kwargs and line plot kwargs
        if 'shape' in kwargs:                       # {'full', 'left', 'right'}
            self.shape = kwargs['shape']
        else:
            self.shape = 'full'

        if 'overhang' in kwargs:
            self.overhang = kwargs['overhang']
        else:
            self.overhang = 0

        if 'head_starts_at_zero' in kwargs:
            self.head_starts_at_zero = kwargs['head_starts_at_zero']
        else:
            self.head_starts_at_zero = False

        if 'c' in kwargs:
            self.color = kwargs['c']
        else:
            self.color = 'black'

        if 'color' in kwargs:
            self.color = kwargs['color']

        if 'linewidth' in kwargs:
            self.linewidth = kwargs['linewidth']

        if 'linestyle' in kwargs:
            self.linestyle = kwargs['linestyle']
        else:
            self.linestyle = '-'

        if 'zorder' in kwargs:
            self.zorder = kwargs['zorder']
        else:
            self.zorder = 2

        if 'lw' in kwargs:
            self.linewidth = kwargs['lw']
        else:
            self.linewidth = plt.rcParams['lines.linewidth']

        if 'linewidth' in kwargs:
            self.linewidth = kwargs['linewidth']

        # dummy arrow on x-axis
        self.arrow_head_origin = {}
        self.arrow_head_origin['tip'] = [self.len + self.norm_head_length * (self.length_includes_head is False), 0]
        self.arrow_head_origin['left'] = [self.arrow_head_origin['tip'][0] - self.norm_head_length, .5 * self.norm_head_width * (self.shape != 'right')]
        self.arrow_head_origin['left_stem'] = [self.arrow_head_origin['tip'][0] - (1 - self.overhang) * self.norm_head_length, .5 * self.width]
        self.arrow_head_origin['stem'] = [self.arrow_head_origin['tip'][0] - (1 - self.overhang) * self.norm_head_length, 0]
        self.arrow_head_origin['right_stem'] = [self.arrow_head_origin['tip'][0] - (1 - self.overhang) * self.norm_head_length, - .5 * self.width]
        self.arrow_head_origin['right'] = [self.arrow_head_origin['tip'][0] - self.norm_head_length, - .5 * self.norm_head_width * (self.shape != 'left')]

        # compute arrow head
        self.arrow_head = {}

        self.x_out, self.y_out = [], []
        for key in self.arrow_head_origin:
            new_x, new_y = self.rotate_2d(self.arrow_head_origin[key])
            self.arrow_head[key] = [new_x, new_y]
            self.x_out.append(new_x)
            self.y_out.append(new_y)
        # create polygon
        self.polygon = Polygon(np.array([[i, j] for i, j in zip(self.x_out, self.y_out)]), True, facecolor=self.color, zorder=self.zorder)
        # create plot coordinates for line plot for tail
        self.linepoints = [np.array([self.x, self.arrow_head['stem'][0]]), np.array([self.y, self.arrow_head['stem'][1]])]

    def draw(self):
        self.axes.plot(self.linepoints[0], self.linepoints[1], lw=self.linewidth, c=self.color, linestyle=self.linestyle, zorder=self.zorder)
        self.axes.add_patch(self.polygon)

    def rotate_2d(self,coordinates):
        '''
        :param coordinates:  [x,y]
        :type coordinates: list
        :param angle_deg: angle in degrees
        :type angle_deg: int, float, double

        '''
        x = coordinates[0]
        y = coordinates[1]

        newx = self.x + (x * np.cos(self.angle_deg * math.pi / 180) - y * self.aspect * np.sin(
            self.angle_deg * math.pi / 180))
        newy = self.y + (y * self.aspect * np.cos(self.angle_deg * math.pi / 180) + x * np.sin(
            self.angle_deg * math.pi / 180)) / self.aspect

        return newx, newy

example:

fig, ax = plt.subplots()
ax.set_xlim(0, 18)
ax.set_ylim(0.4, 1.6)
ax.set_aspect(6)
ax.grid()
 
arrow1 = WarpArrow(ax, 0, 1, 3, 0, head_width=0.04, head_length=1)
arrow1.draw()
 
arrow2 = WarpArrow(ax, 0, 1.2, 3, 0, head_width=0.04, head_length=1, length_includes_head=True)
arrow2.draw()

arrow3 = WarpArrow(ax, 0, 1.4, 3, 0, head_width=0.04, head_length=1, head_length_unit='x', head_width_unit='y', length_includes_head=True, )
arrow3.draw()

arrow4 = WarpArrow(ax, 4, 1, 3, .4, head_width=.5, head_length=0.2, head_length_unit='y', head_width_unit='x' ,length_includes_head=True, overhang=.5, color='lightskyblue')
arrow4.draw()


arrow5 = WarpArrow(ax, 5, 1, 3, .4, head_width=.5, head_length=0.2, head_length_unit='y', head_width_unit='x', length_includes_head=True, overhang=.5, color='rosybrown')
arrow5.draw()

arrow6 = WarpArrow(ax, 11, 1, -3, .4, head_width=.5/ax.get_aspect(), head_length=0.2*ax.get_aspect(), head_length_unit='x', head_width_unit='y', length_includes_head=True, overhang=.5, color='red', linestyle='--')
arrow6.draw()

plt.show()

output: Arrow Examples

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM