Ok, here's something of a weird patch, because I don't know how to use
subversion properly. It has changes to axes.py which update quiver so
that it accepts arbitrary X,Y data; it doesn't demand the data be on a
grid anymore.
The other changes are to collections.py; I updated LineCollection so it
inherits from ScalarMappable. I did this just by copying what looked
like the relevant code from PatchCollection and then I tested it with my
LineCollection based version of quiver. I think I got it right, but I
don't really know what I'm doing here so someone should check it over.
Jordan
Index: lib/matplotlib/axes.py
===================================================================
--- lib/matplotlib/axes.py (revision 2421)
+++ lib/matplotlib/axes.py (working copy)
@@ -3206,16 +3206,16 @@
QUIVER( U, V )
QUIVER( X, Y, U, V, S)
QUIVER( U, V, S )
- QUIVER( ..., color=None, width=1.0, cmap=None,norm=None )
+ QUIVER( ..., color=None, width=1.0, cmap=None, norm=None )
Make a vector plot (U, V) with arrows on a grid (X, Y)
- The optional arguments color and width are used to specify the color
and width
- of the arrow. color can be an array of colors in which case the arrows
can be
- colored according to another dataset.
+ If X and Y are not specified, U and V must be 2D arrays. Equally
spaced
+ X and Y grids are then generated using the meshgrid command.
- If cmap is specified and color is 'length', the colormap is
- used to give a color according to the vector's length.
+ color can be a color value or an array of colors, so that the arrows
can be
+ colored according to another dataset. If cmap is specified and color
is 'length',
+ the colormap is used to give a color according to the vector's length.
If color is a scalar field, the colormap is used to map the scalar to
a color
If a colormap is specified and color is an array of color triplets,
then the
@@ -3263,10 +3263,15 @@
assert X.shape == Y.shape
assert U.shape == X.shape
+ U = ravel(U)
+ V = ravel(V)
+ X = ravel(X)
+ Y = ravel(Y)
+
arrows = []
N = sqrt( U**2+V**2 )
if do_scale:
- Nmax = maximum.reduce(maximum.reduce(N)) or 1 # account for div by
zero
+ Nmax = maximum.reduce(N) or 1 # account for div by zero
U = U*(S/Nmax)
V = V*(S/Nmax)
N = N*Nmax
@@ -3281,7 +3286,6 @@
shading = kwargs.get('shading', 'faceted')
C = None
- I,J = U.shape
if color == 'length' or color is True:
if color is True:
warnings.warn('''Use "color='length'",
@@ -3290,12 +3294,15 @@
elif color is None:
color = (0,0,0,1)
else:
- clr = asarray(color)
+ clr = ravel(asarray(color))
if clr.shape == U.shape:
C = clr
- arrows = [ Arrow(X[i,j],Y[i,j],U[i,j],V[i,j],0.1*S ).get_verts()
- for i in xrange(I) for j in xrange(J) ]
+ I = U.shape[0]
+ arrows = []
+ for i in xrange(I):
+ arrows.append( FancyArrow(X[i],Y[i],U[i],V[i],0.1*S ).get_verts() )
+
collection = PolyCollection(
arrows,
edgecolors = 'None',
@@ -3311,6 +3318,7 @@
else:
collection.set_facecolor(color)
self.add_collection( collection )
+
lims = asarray(arrows)
_max = maximum.reduce( maximum.reduce( lims ))
_min = minimum.reduce( minimum.reduce( lims ))
Index: lib/matplotlib/collections.py
===================================================================
--- lib/matplotlib/collections.py (revision 2421)
+++ lib/matplotlib/collections.py (working copy)
@@ -367,7 +367,7 @@
raise NotImplementedError('Vertices in data coordinates are
calculated\n'
+ 'only with offsets and only if _transOffset == dataTrans.')
-class LineCollection(Collection):
+class LineCollection(Collection, ScalarMappable):
"""
All parameters must be sequences. The property of the ith line
segment is the prop[i % len(props)], ie the properties cycle if
@@ -381,6 +381,8 @@
linestyle = 'solid',
offsets = None,
transOffset = None,#identity_transform(),
+ norm = None, # optional for ScalarMappable
+ cmap = None, # ditto
):
"""
segments is a sequence of ( line0, line1, line2), where
@@ -412,6 +414,7 @@
"""
Collection.__init__(self)
+ ScalarMappable.__init__(self, norm, cmap)
if linewidths is None :
linewidths = (rcParams['lines.linewidth'], )
@@ -419,7 +422,7 @@
if colors is None :
colors = (rcParams['lines.color'],)
if antialiaseds is None :
- antialiaseds = (rcParams['lines.antialiased'], )
+ antialiaseds = (rcParams['lines.antialiased'], )
self._segments = list(segments)
self._colors = colorConverter.to_rgba_list(colors)
@@ -453,6 +456,7 @@
renderer.open_group('linecollection')
self._transform.freeze()
if self._transOffset is not None: self._transOffset.freeze()
+ self.update_scalarmappable()
renderer.draw_line_collection(
self._segments, self._transform, self.clipbox,
self._colors, self._lw, self._ls, self._aa, self._offsets,
@@ -461,6 +465,7 @@
if self._transOffset is not None: self._transOffset.thaw()
renderer.close_group('linecollection')
+
def set_linewidth(self, lw):
"""
Set the linewidth(s) for the collection. lw can be a scalar or a
@@ -546,10 +551,22 @@
verts = []
if self._offsets is None:
for seg in self._segments:
- verts.extend(seg)
+ verts.append(seg)
return [tuple(xy) for xy in verts]
if self._transOffset == dataTrans:
return [tuple(xy) for xy in self._offsets]
raise NotImplementedError('Vertices in data coordinates are
calculated\n'
+ 'with offsets only if _transOffset == dataTrans.')
+ def update_scalarmappable(self):
+ """
+ If the scalar mappable array is not none, update facecolors
+ from scalar data
+ """
+ if self._A is None: return
+ if len(self._A.shape)>1:
+ raise ValueError('LineCollections can only map rank 1 arrays')
+ self._colors = self.to_rgba(self._A, self._alpha)
+ #print self._A.shape, type(R), R.shape
+ #self._facecolors = [(r,g,b,a) for r,g,b,a in R]
+
