Matplotlib風羽自定義

時間 2019-12-11

標籤 matplotlib 自定義简体版

原文原文鏈接

【前言】對於氣象專業的小學生來講，風場是預報重要的參考數據，咱們所知的風羽有四種：短線表明風速2m/s，長線表明風速4m/s，空心三角表明風速20m/s，實心三角表明風速50m/s。而matplotlib的風羽只有短線、長線、三角三種，而這裏的三角不分空心實心，可是可經過改變風羽顏色爲白色使三角變爲空心形狀，雖然這三種能夠自定義各自表明的風速，可是仍與咱們的使用習慣不符，即便把三角設成20m/s，本來一個實心三角就能表示的50m/s的風在matplotlib中須要兩個三角外加兩條長線一條短線。爲了迎合預報員的需求，我在研究了matplotlib的風場函數barbs()的源代碼quiver.py文件後，對quiver.py作了適當的調整，使得matplotlib也有了空心三角和實心三角之分。python

1、函數`barbs`的使用

barb(X, Y, U, V,, **kw)

X：風場數據X座標
Y：風場數據Y座標
U：風的水平方向份量
V：風的垂直方向份量微信

'''
Demonstration of wind barb plots
'''
import matplotlib.pyplot as plt
import numpy as np
x = np.linspace(-5, 5, 5)
X, Y = np.meshgrid(x, x)
U, V = 12*X, 12*Y
data = [(-1.5, .5, -6, -6),(1, -1, -46, 46),(-3, -1, 11, -11),(1, 1.5, 80, 80),(0.5, 0.25, 25, 15),(-1.5, -0.5, -5, 40)]
data = np.array(data, dtype=[('x', np.float32), ('y', np.float32), ('u', np.float32), ('v', np.float32)])
# Default parameters, uniform grid
ax = plt.subplot(2, 2, 1)
ax.barbs(X, Y, U, V)
# Arbitrary set of vectors, make them longer and change the pivot point
#(point around which they're rotated) to be the middle
ax = plt.subplot(2, 2, 2)
ax.barbs(data['x'], data['y'], data['u'], data['v'], length=8, pivot='middle')
# Showing colormapping with uniform grid. Fill the circle for an empty barb,
# don't round the values, and change some of the size parameters
ax = plt.subplot(2, 2, 3)
ax.barbs(X, Y, U, V, np.sqrt(U*U + V*V), fill_empty=True, rounding=False,sizes=dict(emptybarb=0.25, spacing=0.2, height=0.3))
# Change colors as well as the increments for parts of the barbs
ax = plt.subplot(2, 2, 4)
ax.barbs(data['x'], data['y'], data['u'], data['v'], flagcolor='r',barbcolor=['b', 'g'], barb_increments=dict(half=10, full=20, flag=100),flip_barb=True)
plt.show()

2、源代碼解讀

1.`class Barbs()`

class Barbs(mcollections.PolyCollection):
   @docstring.interpd
   def __init__(self, ax, *args, **kw):    
        '...'
   def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
        '...'
   def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,pivot, sizes, fill_empty, flip):
        '...'
   def set_UVC(self, U, V, C=None):
        '...'
   def set_offsets(self, xy):
        '...'

經過讀源代碼可知類Barbs有五個方法分別爲__init__、_find_tails、_make_barbs、set_UVC、set_offsets。

2.`init`

@docstring.interpd
    def __init__(self, ax, *args, **kw):
        """
        The constructor takes one required argument, an Axes
        instance, followed by the args and kwargs described
        by the following pylab interface documentation:
        %(barbs_doc)s
        """
        self._pivot = kw.pop('pivot', 'tip')
        self._length = kw.pop('length', 7)
        barbcolor = kw.pop('barbcolor', None)
        flagcolor = kw.pop('flagcolor', None)
        self.sizes = kw.pop('sizes', dict())
        self.fill_empty = kw.pop('fill_empty', False)
        self.barb_increments = kw.pop('barb_increments', dict())
        self.rounding = kw.pop('rounding', True)
        self.flip = kw.pop('flip_barb', False)
        transform = kw.pop('transform', ax.transData)

        # Flagcolor and and barbcolor provide convenience parameters for
        # setting the facecolor and edgecolor, respectively, of the barb
        # polygon.  We also work here to make the flag the same color as the
        # rest of the barb by default

        if None in (barbcolor, flagcolor):
            kw['edgecolors'] = 'face'
            if flagcolor:
                kw['facecolors'] = flagcolor
            elif barbcolor:
                kw['facecolors'] = barbcolor
            else:
                # Set to facecolor passed in or default to black
                kw.setdefault('facecolors', 'k')
        else:
            kw['edgecolors'] = barbcolor
            kw['facecolors'] = flagcolor

        # Parse out the data arrays from the various configurations supported
        x, y, u, v, c = _parse_args(*args)
        self.x = x
        self.y = y
        xy = np.hstack((x[:, np.newaxis], y[:, np.newaxis]))

        # Make a collection
        barb_size = self._length ** 2 / 4  # Empirically determined
        mcollections.PolyCollection.__init__(self, [], (barb_size,),
                                             offsets=xy,
                                             transOffset=transform, **kw)
        self.set_transform(transforms.IdentityTransform())

        self.set_UVC(u, v, c)

__init__()方法爲初始化方法，此方法中flagcolor、barbcolor爲設置風羽顏色的關鍵字，中間的說明文字提示顏色設置是針對全部的風羽的，因此經過顏色設置達不到風羽中既有空心白色三角又有實心黑色三角。初始化方法中在對一些參數進行了初始化賦值後執行了set_UVC()方法，因此咱們順着這個set_UVC()方法往下繼續讀。

3.`set_UVC()`

def set_UVC(self, U, V, C=None):
        self.u = ma.masked_invalid(U, copy=False).ravel()
        self.v = ma.masked_invalid(V, copy=False).ravel()
        if C is not None:
            c = ma.masked_invalid(C, copy=False).ravel()
            x, y, u, v, c = delete_masked_points(self.x.ravel(),
                                                 self.y.ravel(),
                                                 self.u, self.v, c)
        else:
            x, y, u, v = delete_masked_points(self.x.ravel(), self.y.ravel(),
                                              self.u, self.v)

        magnitude = np.hypot(u, v)
        flags, emptyflags,barbs, halves, empty = self._find_tails(magnitude,
                                                       self.rounding,
                                                       **self.barb_increments)

        # Get the vertices for each of the barbs

        plot_barbs = self._make_barbs(u, v, flags, emptyflags,barbs, halves, empty,
                                      self._length, self._pivot, self.sizes,
                                      self.fill_empty, self.flip)
        self.set_verts(plot_barbs)

        # Set the color array
        if C is not None:
            self.set_array(c)

        # Update the offsets in case the masked data changed
        xy = np.hstack((x[:, np.newaxis], y[:, np.newaxis]))
        self._offsets = xy
        self.stale = True

在此方法中，首先進行了變量的命名賦值，而後依次執行了方法_find_tails和_make_barbs。_make_barbs的輸入爲_find_tails的輸出，_find_tails的輸入中有一個爲magnitude = np.hypot(u, v)，np.hypot()爲勾股定理方法，所以可知magnitude爲風速。

4.`_find_tails`

def _find_tails(self, mag, rounding=True, half=5, full=10, flag=50):
        '''
        Find how many of each of the tail pieces is necessary.  Flag
        specifies the increment for a flag, barb for a full barb, and half for
        half a barb. Mag should be the magnitude of a vector (i.e., >= 0).

        This returns a tuple of:

            (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)

        *half_flag* is a boolean whether half of a barb is needed,
        since there should only ever be one half on a given
        barb. *empty_flag* flag is an array of flags to easily tell if
        a barb is empty (too low to plot any barbs/flags.
        '''

        # If rounding, round to the nearest multiple of half, the smallest
        # increment
        if rounding:
            mag = half * (mag / half + 0.5).astype(np.int)

        num_flags = np.floor(mag / flag).astype(np.int)
        mag = np.mod(mag, flag)

        num_barb = np.floor(mag / full).astype(np.int)
        mag = np.mod(mag, full)

        half_flag = mag >= half
        empty_flag = ~(half_flag | (num_flags > 0) | (num_emptyflags > 0) |(num_barb > 0))

        return num_flags,num_barb, half_flag, empty_flag

經過讀此方法的說明文檔可知，此方法做用爲根據輸入的風速、設置的短線長線三角的數值計算並返回三角、長線、短線的個數以及有沒有無風的狀況。

5.`_make_barbs`

def _make_barbs(self, u, v, nflags, nbarbs, half_barb, empty_flag, length,
                    pivot, sizes, fill_empty, flip):
        '''
        This function actually creates the wind barbs.  *u* and *v*
        are components of the vector in the *x* and *y* directions,
        respectively.

        *nflags*, *nbarbs*, and *half_barb*, empty_flag* are,
        *respectively, the number of flags, number of barbs, flag for
        *half a barb, and flag for empty barb, ostensibly obtained
        *from :meth:`_find_tails`.

        *length* is the length of the barb staff in points.

        *pivot* specifies the point on the barb around which the
        entire barb should be rotated.  Right now, valid options are
        'head' and 'middle'.

        *sizes* is a dictionary of coefficients specifying the ratio
        of a given feature to the length of the barb. These features
        include:

            - *spacing*: space between features (flags, full/half
               barbs)

            - *height*: distance from shaft of top of a flag or full
               barb

            - *width* - width of a flag, twice the width of a full barb

            - *emptybarb* - radius of the circle used for low
               magnitudes

        *fill_empty* specifies whether the circle representing an
        empty barb should be filled or not (this changes the drawing
        of the polygon).

        *flip* is a flag indicating whether the features should be flipped to
        the other side of the barb (useful for winds in the southern
        hemisphere.

        This function returns list of arrays of vertices, defining a polygon
        for each of the wind barbs.  These polygons have been rotated to
        properly align with the vector direction.
        '''

        # These control the spacing and size of barb elements relative to the
        # length of the shaft
        spacing = length * sizes.get('spacing', 0.125)
        full_height = length * sizes.get('height', 0.4)
        full_width = length * sizes.get('width', 0.25)
        empty_rad = length * sizes.get('emptybarb', 0.15)

        # Controls y point where to pivot the barb.
        pivot_points = dict(tip=0.0, middle=-length / 2.)

        # Check for flip
        if flip:
            full_height = -full_height

        endx = 0.0
        endy = pivot_points[pivot.lower()]

        # Get the appropriate angle for the vector components.  The offset is
        # due to the way the barb is initially drawn, going down the y-axis.
        # This makes sense in a meteorological mode of thinking since there 0
        # degrees corresponds to north (the y-axis traditionally)
        angles = -(ma.arctan2(v, u) + np.pi / 2)

        # Used for low magnitude.  We just get the vertices, so if we make it
        # out here, it can be reused.  The center set here should put the
        # center of the circle at the location(offset), rather than at the
        # same point as the barb pivot; this seems more sensible.
        circ = CirclePolygon((0, 0), radius=empty_rad).get_verts()
        if fill_empty:
            empty_barb = circ
        else:
            # If we don't want the empty one filled, we make a degenerate
            # polygon that wraps back over itself
            empty_barb = np.concatenate((circ, circ[::-1]))

        barb_list = []
        for index, angle in np.ndenumerate(angles):
            # If the vector magnitude is too weak to draw anything, plot an
            # empty circle instead
            if empty_flag[index]:
                # We can skip the transform since the circle has no preferred
                # orientation
                barb_list.append(empty_barb)
                continue

            poly_verts = [(endx, endy)]
            offset = length

            # Add vertices for each flag
            for i in range(nflags[index]):
                # The spacing that works for the barbs is a little to much for
                # the flags, but this only occurs when we have more than 1
                # flag.
                if offset != length:
                    offset += spacing / 2.
                poly_verts.extend(
                    [[endx, endy + offset],
                     [endx + full_height, endy - full_width / 2 + offset],
                     [endx, endy - full_width + offset]])

                offset -= full_width + spacing

            # Add vertices for each barb.  These really are lines, but works
            # great adding 3 vertices that basically pull the polygon out and
            # back down the line
            for i in range(nbarbs[index]):
                poly_verts.extend(
                    [(endx, endy + offset),
                     (endx + full_height, endy + offset + full_width / 2),
                     (endx, endy + offset)])

                offset -= spacing

            # Add the vertices for half a barb, if needed
            if half_barb[index]:
                # If the half barb is the first on the staff, traditionally it
                # is offset from the end to make it easy to distinguish from a
                # barb with a full one
                if offset == length:
                    poly_verts.append((endx, endy + offset))
                    offset -= 1.5 * spacing
                poly_verts.extend(
                    [(endx, endy + offset),
                     (endx + full_height / 2, endy + offset + full_width / 4),
                     (endx, endy + offset)])

            # Rotate the barb according the angle. Making the barb first and
            # then rotating it made the math for drawing the barb really easy.
            # Also, the transform framework makes doing the rotation simple.
            poly_verts = transforms.Affine2D().rotate(-angle).transform(
                poly_verts)
            barb_list.append(poly_verts)

        return barb_list

經過讀此方法的說明文檔可知，此方法做用爲根據輸入的風數據以及短線長線三角的個數繪製風羽風向杆。
繪製過程爲：判斷地圖座標點是否是無風，若是無風就繪製一個空心圓圈表明。若是有風就開始按照三角、長線、短線的順序繪製。
繪製方法爲：

建立一個用於存儲關鍵點座標的列表poly_verts
計算關鍵點座標
經過transform方法將關鍵點座標列表中的各個關鍵點依次用黑線鏈接起來，最終將風羽風向杆繪製出來

此方法的幾個關鍵變量：

spacing：風羽上短線長線以及三角間的距離
full_height：三角的高度
full_width ：三角的寬度
endx ：風羽繪製的起始點x座標
endy：風羽繪製的起始點y座標
angles：風向杆角度
poly_verts ：繪製風羽風向杆的關鍵點列表
offset：繪製完一個三角或線後下一個三角或線的關鍵起始座標

poly_verts = [(endx, endy)]
            offset = length

            # Add vertices for each flag
            for i in range(nflags[index]):
                # The spacing that works for the barbs is a little to much for
                # the flags, but this only occurs when we have more than 1
                # flag.
                if offset != length:
                    offset += spacing / 2.
                poly_verts.extend(
                    [[endx, endy + offset],
                     [endx + full_height, endy - full_width / 2 + offset],
                     [endx, endy - full_width + offset]])

                offset -= full_width + spacing

這一段是繪製風羽的主要代碼，利用圖片的形式說明
app

3、繪製空心實心三角

在瞭解了風羽的繪製過程後，發現能夠經過增長關鍵點直接繪製實心三角，經過原繪製方法繪製空心三角。ide

1.實心三角繪製

實心三角繪製代碼

# Add vertices for each flag
            for i in range(nflags[index]):
                # The spacing that works for the barbs is a little to much for
                # the flags, but this only occurs when we have more than 1
                # flag.
                if offset != length:
                    offset += spacing / 2.
                poly_verts.extend(
                    [[endx, endy + offset],
                     [endx + full_height/4, endy - full_width / 8 + offset],
                     [endx, endy - full_width / 8 + offset],
                     [endx + full_height/4, endy - full_width / 8 + offset],
                     [endx + full_height/2, endy - full_width / 4 + offset],
                     [endx, endy - full_width / 4 + offset],
                     [endx + full_height/2, endy - full_width / 4 + offset],
                     [endx + 3*full_height/4, endy - 3*full_width / 8 + offset],
                     [endx, endy - 3*full_width / 8 + offset],
                     [endx + 3*full_height/4, endy - 3*full_width / 8 + offset],
                     [endx + full_height, endy - full_width / 2 + offset],
                     [endx,endy-full_width/2+offset],
                     [endx + full_height, endy - full_width / 2 + offset],
                     [endx + 3*full_height/4, endy - 5*full_width / 8 + offset],
                     [endx, endy - 5*full_width / 8 + offset],
                     [endx + 3*full_height/4, endy - 5*full_width / 8 + offset],
                     [endx + full_height/2, endy - 3*full_width / 4 + offset],
                     [endx, endy - 3*full_width / 4 + offset],
                     [endx + full_height/2, endy - 3*full_width / 4 + offset],
                     [endx + full_height/4, endy - 7*full_width / 8 + offset],
                     [endx, endy - 7*full_width / 8 + offset],
                     [endx + full_height/4, endy - 7*full_width / 8 + offset],
                     [endx, endy - full_width + offset]])

                offset -= full_width + spacing

實心三角繪製示意圖
函數
方法參數中加入nfullflags

def _make_barbs(self, u, v, nfullflags, nflags,nbarbs, half_barb, empty_flag, length,pivot, sizes, fill_empty, flip):
    '...'

2.實心三角個數計算

def _find_tails(self, mag, rounding=True, half=2, full=4, flag=20,fullflag=50):
        '''
        Find how many of each of the tail pieces is necessary.  Flag
        specifies the increment for a flag, barb for a full barb, and half for
        half a barb. Mag should be the magnitude of a vector (i.e., >= 0).

        This returns a tuple of:

            (*number of flags*, *number of barbs*, *half_flag*, *empty_flag*)

        *half_flag* is a boolean whether half of a barb is needed,
        since there should only ever be one half on a given
        barb. *empty_flag* flag is an array of flags to easily tell if
        a barb is empty (too low to plot any barbs/flags.
        '''

        # If rounding, round to the nearest multiple of half, the smallest
        # increment
        if rounding:
            mag = half * (mag / half + 0.5).astype(np.int)

        num_fullflags = np.floor(mag / fullflag).astype(np.int)
        mag = np.mod(mag, fullflag)

        num_flags = np.floor(mag / flag).astype(np.int)
        mag = np.mod(mag, flag)

        num_barb = np.floor(mag / full).astype(np.int)
        mag = np.mod(mag, full)

        half_flag = mag >= half
        empty_flag = ~(half_flag | (num_flags > 0) | (num_fullflags > 0) |(num_barb > 0))

        return num_fullflags,num_flags,num_barb, half_flag, empty_flag

3. 調整`set_UVC`中相關方法使用

fullflags, flags,barbs, halves, empty = self._find_tails(magnitude,
                                                       self.rounding,
                                                       **self.barb_increments)

        # Get the vertices for each of the barbs

        plot_barbs = self._make_barbs(u, v, fullflags, flags,barbs, halves, empty,
                                      self._length, self._pivot, self.sizes,
                                      self.fill_empty, self.flip)

4、測試

import matplotlib.pyplot as plt

fig=plt.figure()
ax=fig.add_subplot(111);
ax.axis([-1,1,-1,1])
ax.set_xticks([])
ax.set_yticks([])
ax.barbs(0,0,30*1.5,40*1.5,length=8,linewidth=0.5)

plt.show()