Merge pull request matplotlib#5148 from efiring/scatter2

FIX: scatter accepts 2-D x, y, c; 

closes matplotlib#5141

Author: tacaswell

lib/matplotlib/axes/_axes.py

@@ -3788,6 +3788,12 @@ def scatter(self, x, y, s=20, c=None, marker='o', cmap=None, norm=None,
         which case all masks will be combined and only unmasked points
         will be plotted.
 
+        Fundamentally, scatter works with 1-D arrays; `x`, `y`, `s`,
+        and `c` may be input as 2-D arrays, but within scatter
+        they will be flattened. The exception is `c`, which
+        will be flattened only if its size matches the size of `x`
+        and `y`.
+
         Examples
         --------
         .. plot:: mpl_examples/shapes_and_collections/scatter_demo.py
@@ -3839,13 +3845,13 @@ def scatter(self, x, y, s=20, c=None, marker='o', cmap=None, norm=None,
         # After this block, c_array will be None unless
         # c is an array for mapping.  The potential ambiguity
         # with a sequence of 3 or 4 numbers is resolved in
-        # favor mapping, not rgb or rgba.
+        # favor of mapping, not rgb or rgba.
         try:
             c_array = np.asanyarray(c, dtype=float)
-            if c_array.shape == x.shape:
+            if c_array.size == x.size:
                 c = np.ma.ravel(c_array)
             else:
-                # Wrong shape; it must not be intended for mapping.
+                # Wrong size; it must not be intended for mapping.
                 c_array = None
         except ValueError:
             # Failed to make a floating-point array; c must be color specs.

lib/matplotlib/tests/baseline_images/test_axes/scatter_2D.png

@@ -1241,6 +1241,17 @@ def test_scatter_marker():
                 marker=mmarkers.MarkerStyle('o', fillstyle='top'))
 
 
+@image_comparison(baseline_images=['scatter_2D'], remove_text=True,
+                  extensions=['png'])
+def test_scatter_2D():
+    x = np.arange(3)
+    y = np.arange(2)
+    x, y = np.meshgrid(x, y)
+    z = x + y
+    fig, ax = plt.subplots()
+    ax.scatter(x, y, c=z, s=200, edgecolors='face')
+
+
 @cleanup
 def test_as_mpl_axes_api():
     # tests the _as_mpl_axes api
@@ -2150,17 +2161,17 @@ def test_errorbar():
 def test_errorbar_shape():
     fig = plt.figure()
     ax = fig.gca()
-    
+
     x = np.arange(0.1, 4, 0.5)
     y = np.exp(-x)
     yerr1 = 0.1 + 0.2*np.sqrt(x)
     yerr = np.vstack((yerr1, 2*yerr1)).T
     xerr = 0.1 + yerr
-    
+
     assert_raises(ValueError, ax.errorbar, x, y, yerr=yerr, fmt='o')
     assert_raises(ValueError, ax.errorbar, x, y, xerr=xerr, fmt='o')
     assert_raises(ValueError, ax.errorbar, x, y, yerr=yerr, xerr=xerr, fmt='o')
-    
+
 
 @image_comparison(baseline_images=['errorbar_limits'])
 def test_errorbar_limits():