Hi,
I've got (what seems to me) a nice clean, self-contained implementation
of wind barbs plots. I'd like to see if I can get this into matplotlib,
as it would be very useful to the meteorology community. I've borrowed
heavily from Quiver for rounding out rough edges (like multiple calling
signatures) as for templates for documentation. The base
implementation, though, seems much simpler (thanks to Mike's transforms)
and is based more on scatter.
Right now it monkey-patches Axes so that it can be a stand-alone file.
Just running the file should give a good example of the expected output.
My only concern up front is if a new axes method is appropriate, since
this is somewhat domain specific. But I'd like to at least get the new
Collections class included, since otherwise, I have no idea how to get
this distributed to the community at large.
I welcome any comments/criticism to help improve this.
Ryan
--
Ryan May
Graduate Research Assistant
School of Meteorology
University of Oklahoma
'''
Support for plotting a field of (wind) barbs. This is like quiver in that
you have something that points along a vector field giving direction, but
the magnitude of the vector is given schematically by the presence of barbs
or flags on the barb. This differs from quiver, which uses the size of the
arrow to indicate vector magnitude.
'''
import numpy as np
from numpy import ma
from matplotlib import rcParams
import matplotlib.collections as mcoll
import matplotlib.transforms as mtransforms
import matplotlib.artist as martist
_barbs_doc = """
Plot a 2-D field of barbs.
call signatures::
barb(U, V, **kw)
barb(U, V, C, **kw)
barb(X, Y, U, V, **kw)
barb(X, Y, U, V, C, **kw)
Arguments:
*X*, *Y*:
The x and y coordinates of the barb locations
(default is head of barb; see *pivot* kwarg)
*U*, *V*:
give the *x* and *y* components of the barb shaft
*C*:
an optional array used to map colors to the barbs
All arguments may be 1-D or 2-D arrays or sequences. If *X* and *Y*
are absent, they will be generated as a uniform grid. If *U* and *V*
are 2-D arrays but *X* and *Y* are 1-D, and if len(*X*) and len(*Y*)
match the column and row dimensions of *U*, then *X* and *Y* will be
expanded with :func:`numpy.meshgrid`.
*U*, *V*, *C* may be masked arrays, but masked *X*, *Y* are not
supported at present.
Keyword arguments:
*scale*:
Length of the barb in points; the other parts of the barb
are scaled against this.
Default is 9
*pivot*: [ 'tip' | 'middle' ]
The part of the arrow that is at the grid point; the arrow
rotates about this point, hence the name *pivot*.
*barbcolor*: [ color | color sequence ]
Specifies the color all parts of the barb except any flags.
This parameter is analagous to the *edgecolor* parameter
for polygons, which can be used instead. However this parameter
will override facecolor.
*flagcolor*: [ color | color sequence ]
Specifies the color of any flags on the barb.
This parameter is analagous to the *facecolor* parameter
for polygons, which can be used instead. However this parameter
will override facecolor. If this is not set (and *C* has not either)
then *flagcolor* will be set to match *barbcolor* so that the barb
has a uniform color. If *C* has been set, *flagcolor* has no effect.
*pivot*: [ 'tip' | 'middle' ]
The part of the arrow that is at the grid point; the arrow
rotates about this point, hence the name *pivot*.
Barbs are traditionally used in meteorology as a way to plot the speed
and direction of wind observations, but can technically be used to plot
any two dimensional vector quantity. As opposed to arrows, which give
vector magnitude by the length of the arrow, the barbs give more quantitative
information about the vector magnitude by putting slanted lines or a triangle
for various increments in magnitude, as show schematically below:
/\ \
/ \ \
/ \ \ \
/ \ \ \
------------------------------
The largest increment is given by a triangle (or "flag"), which usually
represents inrements of 50. After those come full lines, which represent 10.
The smallest increment is a half line, which represents 5. There is only, of
course, ever at most 1 half line. If the magnitude is small and only needs a
single half-line and no full lines or triangles, the half-line is offset from
the end of the barb so that it can be easily distinguished from barbs with a
single full line. The magnitude for the barb shown above would nominally be
65.
linewidths and edgecolors can be used to customize the barb.
Additional :class:`~matplotlib.collections.PolyCollection`
keyword arguments:
%(PolyCollection)s
""" % martist.kwdocd
class Barbs(mcoll.PolyCollection):
'''
Specialized PolyCollection for barbs.
There are no API methods. Everything is performed in __init__().
There is one internal function _find_tails() which finds exactly
what should be put on the barb given the vector magnitude. From there
_make_barbs() is used to find the vertices of the polygon to represent the
barb based on this information.
'''
#This may be an abuse of polygons here to render what is essentially maybe
#1 triangle and a series of lines. It works fine as far as I can tell
#however.
def __init__(self, ax, *args, **kw):
pivot = kw.pop('pivot', 'tip')
length = kw.pop('length', 9)
barbcolor = kw.pop('barbcolor', None)
flagcolor = kw.pop('flagcolor', None)
#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 flagcolor:
kw['facecolor'] = flagcolor
elif barbcolor:
kw.setdefault('facecolor', barbcolor)
else:
kw['facecolor'] = rcParams['axes.edgecolor']
if barbcolor:
kw['edgecolor'] = barbcolor
#Parse out the data arrays from the various configurations supported
x, y, u, v, c = self._parse_args(*args)
self.xy = np.hstack((x[:,np.newaxis], y[:,np.newaxis]))
self.u = u
self.v = v
magnitude = np.sqrt(u*u + v*v)
flags, barbs, halves = self._find_tails(magnitude)
#Get the vertices for each of the barbs
plot_barbs = self._make_barbs(u, v, flags, barbs, halves, length, pivot)
#Make a collection
barb_size = length**2 / 4 #Empirically determined
mcoll.PolyCollection.__init__(self, plot_barbs, (barb_size,),
offsets=self.xy, transOffset=ax.transData, **kw)
self.set_array(c)
self.set_transform(mtransforms.IdentityTransform())
__init__.__doc__ = """
The constructor takes one required argument, an Axes
instance, followed by the args and kwargs described
by the following pylab interface documentation:
%s""" % _barbs_doc
def _find_tails(self, mag, half=5, barb=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 (ie. >= 0).
This returns a tuple of:
(number of flags, number of barbs, half_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.'''
num_flags = np.floor(mag / flag).astype(np.int)
mag = np.mod(mag, flag)
num_barb = np.floor(mag / barb).astype(np.int)
mag = np.mod(mag, barb)
half_flag = mag >= half
return num_flags, num_barb, half_flag
def _make_barbs(self, u, v, nflags, nbarbs, half_barb, length, pivot):
'''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 are the number of flags, number of barbs, and flag for half
a barb, ostensibly obtained from _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'.
This function returns list of arraya 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 / 10.
full_height = length / 2.5
full_width = length / 3.5
#Controls y point where to pivot the barb.
pivot_points = dict(tip=0.0, middle=-length/2.)
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 = -(np.arctan2(v, u) + np.pi/2)
barb_list = []
for index, angle in np.ndenumerate(angles):
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 = mtransforms.Affine2D().rotate(-angle).transform(poly_verts)
barb_list.append(poly_verts)
return barb_list
#Taken shamelessly from quiver.py
def _parse_args(self, *args):
X, Y, U, V, C = [None]*5
args = list(args)
if len(args) == 3 or len(args) == 5:
C = ma.asarray(args.pop(-1)).ravel()
V = ma.asarray(args.pop(-1))
U = ma.asarray(args.pop(-1))
nn = np.shape(U)
nc = nn[0]
nr = 1
if len(nn) > 1:
nr = nn[1]
if len(args) == 2: # remaining after removing U,V,C
X, Y = [np.array(a).ravel() for a in args]
if len(X) == nc and len(Y) == nr:
X, Y = [a.ravel() for a in np.meshgrid(X, Y)]
else:
indexgrid = np.meshgrid(np.arange(nc), np.arange(nr))
X, Y = [np.ravel(a) for a in indexgrid]
return X, Y, U, V, C
barbs_doc = _barbs_doc
#Monkey patch a barbs method onto axes for demo purposes
from matplotlib.axes import Axes
def barbs(self, *args, **kw):
if not self._hold: self.cla()
b = Barbs(self, *args, **kw)
self.add_collection(b)
self.update_datalim(b.xy)
self.autoscale_view()
return b
barbs.__doc__ = Barbs.barbs_doc
Axes.barbs = barbs
if __name__ == '__main__':
import matplotlib.pyplot as plt
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)]
#Default parameters for arbitrary set of vectors
ax = plt.subplot(2,2,1)
ax.barbs(*zip(*data))
#Default parameters, uniform grid
ax = plt.subplot(2,2,2)
ax.barbs(X, Y, U, V)
plt.show()
#Change parameters for arbitrary set of vectors
ax = plt.subplot(2,2,3)
ax.barbs(flagcolor='r', barbcolor=['b','g'], *zip(*data))
#Showing colormapping with uniform grid. Unfortunately, only the flags
#on barbs get colormapped this way.
ax = plt.subplot(2,2,4)
ax.barbs(X, Y, U, V, np.sqrt(U*U + V*V))
plt.show()
-------------------------------------------------------------------------
This SF.Net email is sponsored by the Moblin Your Move Developer's challenge
Build the coolest Linux based applications with Moblin SDK & win great prizes
Grand prize is a trip for two to an Open Source event anywhere in the world
http://moblin-contest.org/redirect.php?banner_id=100&url=/
_______________________________________________
Matplotlib-devel mailing list
Matplotlib-devel@lists.sourceforge.net
https://lists.sourceforge.net/lists/listinfo/matplotlib-devel
