Merge pull request #2199 from ajdawson/2d-datetime-coords

efiring · efiring · commit e96d35371c23 · 2013-07-23T11:30:28.000-07:00
Fix support for datetime axes in 2d plots.
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,12 @@
+2013-07-12 Added support for datetime axes to 2d plots. Axis values are passed
+           through Axes.convert_xunits/Axes.convert_yunits before being used by
+           contour/contourf, pcolormesh and pcolor.
+
+2013-07-12 Allowed matplotlib.dates.date2num, matplotlib.dates.num2date,
+           and matplotlib.dates.datestr2num to accept n-d inputs. Also
+           factored in support for n-d arrays to matplotlib.dates.DateConverter
+           and matplotlib.units.Registry.
+
 2013-06-26 Refactored the axes module: the axes module is now a folder,
 	   containing the following submodule:
 	      - _subplots.py, containing all the subplots helper methods
diff --git a/doc/users/whats_new.rst b/doc/users/whats_new.rst
@@ -18,6 +18,32 @@ revision, see the :ref:`github-stats`.
 .. contents:: Table of Contents
    :depth: 3
 
+.. _whats-new-1-4:
+
+new in matplotlib-1.4
+=====================
+
+New plotting features
+---------------------
+
+Support for datetime axes in 2d plots
+`````````````````````````````````````
+Andrew Dawson added support for datetime axes to
+:func:`~matplotlib.pyplot.contour`, :func:`~matplotlib.pyplot.contourf`, 
+:func:`~matplotlib.pyplot.pcolormesh` and :func:`~matplotlib.pyplot.pcolor`. 
+
+
+Date handling
+-------------
+
+n-d array support for date conversion
+``````````````````````````````````````
+Andrew Dawson added support for n-d array handling to
+:func:`matplotlib.dates.num2date`, :func:`matplotlib.dates.date2num`
+and :func:`matplotlib.dates.datestr2num`. Support is also added to the unit
+conversion interfaces :class:`matplotlib.dates.DateConverter` and
+:class:`matplotlib.units.Registry`.
+
 .. _whats-new-1-3:
 
 new in matplotlib-1.3
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -4509,10 +4509,16 @@ def pcolor(self, *args, **kwargs):
         X, Y, C = self._pcolorargs('pcolor', *args, allmatch=False)
         Ny, Nx = X.shape
 
+        # unit conversion allows e.g. datetime objects as axis values
+        self._process_unit_info(xdata=X, ydata=Y, kwargs=kwargs)
+        X = self.convert_xunits(X)
+        Y = self.convert_yunits(Y)
+
         # convert to MA, if necessary.
         C = ma.asarray(C)
         X = ma.asarray(X)
         Y = ma.asarray(Y)
+
         mask = ma.getmaskarray(X) + ma.getmaskarray(Y)
         xymask = (mask[0:-1, 0:-1] + mask[1:, 1:] +
                   mask[0:-1, 1:] + mask[1:, 0:-1])
@@ -4705,6 +4711,11 @@ def pcolormesh(self, *args, **kwargs):
         X = X.ravel()
         Y = Y.ravel()
 
+        # unit conversion allows e.g. datetime objects as axis values
+        self._process_unit_info(xdata=X, ydata=Y, kwargs=kwargs)
+        X = self.convert_xunits(X)
+        Y = self.convert_yunits(Y)
+
         coords = np.zeros(((Nx * Ny), 2), dtype=float)
         coords[:, 0] = X
         coords[:, 1] = Y
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -203,6 +203,10 @@ def _to_ordinalf(dt):
     return base
 
 
+# a version of _to_ordinalf that can operate on numpy arrays
+_to_ordinalf_np_vectorized = np.vectorize(_to_ordinalf)
+
+
 def _from_ordinalf(x, tz=None):
     """
     Convert Gregorian float of the date, preserving hours, minutes,
@@ -229,6 +233,10 @@ def _from_ordinalf(x, tz=None):
     return dt
 
 
+# a version of _from_ordinalf that can operate on numpy arrays
+_from_ordinalf_np_vectorized = np.vectorize(_from_ordinalf)
+
+
 class strpdate2num:
     """
     Use this class to parse date strings to matplotlib datenums when
@@ -246,6 +254,10 @@ def __call__(self, s):
         return date2num(datetime.datetime(*time.strptime(s, self.fmt)[:6]))
 
 
+# a version of dateutil.parser.parse that can operate on nump0y arrays
+_dateutil_parser_parse_np_vectorized = np.vectorize(dateutil.parser.parse)
+
+
 def datestr2num(d):
     """
     Convert a date string to a datenum using
@@ -256,7 +268,10 @@ def datestr2num(d):
         dt = dateutil.parser.parse(d)
         return date2num(dt)
     else:
-        return date2num([dateutil.parser.parse(s) for s in d])
+        d = np.asarray(d)
+        if not d.size:
+            return d
+        return date2num(_dateutil_parser_parse_np_vectorized(d))
 
 
 def date2num(d):
@@ -273,7 +288,10 @@ def date2num(d):
     if not cbook.iterable(d):
         return _to_ordinalf(d)
     else:
-        return np.asarray([_to_ordinalf(val) for val in d])
+        d = np.asarray(d)
+        if not d.size:
+            return d
+        return _to_ordinalf_np_vectorized(d)
 
 
 def julian2num(j):
@@ -310,7 +328,10 @@ def num2date(x, tz=None):
     if not cbook.iterable(x):
         return _from_ordinalf(x, tz)
     else:
-        return [_from_ordinalf(val, tz) for val in x]
+        x = np.asarray(x)
+        if not x.size:
+            return x
+        return _from_ordinalf_np_vectorized(x, tz).tolist()
 
 
 def drange(dstart, dend, delta):
@@ -1292,6 +1313,9 @@ def convert(value, unit, axis):
     @staticmethod
     def default_units(x, axis):
         'Return the tzinfo instance of *x* or of its first element, or None'
+        if isinstance(x, np.ndarray):
+            x = x.ravel()
+
         try:
             x = x[0]
         except (TypeError, IndexError):
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/pcolor_datetime_axis.png b/lib/matplotlib/tests/baseline_images/test_axes/pcolor_datetime_axis.png
diff --git a/lib/matplotlib/tests/baseline_images/test_axes/pcolormesh_datetime_axis.png b/lib/matplotlib/tests/baseline_images/test_axes/pcolormesh_datetime_axis.png
diff --git a/lib/matplotlib/tests/baseline_images/test_contour/contour_datetime_axis.png b/lib/matplotlib/tests/baseline_images/test_contour/contour_datetime_axis.png
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1,5 +1,6 @@
 from nose.tools import assert_equal
 from nose.tools import assert_raises
+import datetime
 
 import numpy as np
 from numpy import ma
@@ -663,6 +664,59 @@ def test_pcolormesh():
     ax = fig.add_subplot(133)
     ax.pcolormesh(Qx,Qz,Z, shading="gouraud")
 
+
+@image_comparison(baseline_images=['pcolormesh_datetime_axis'],
+                  extensions=['png'], remove_text=False)
+def test_pcolormesh_datetime_axis():
+    fig = plt.figure()
+    fig.subplots_adjust(hspace=0.4, top=0.98, bottom=.15)
+    base = datetime.datetime(2013, 1, 1)
+    x = np.array([base + datetime.timedelta(days=d) for d in range(21)])
+    y = np.arange(21)
+    z1, z2 = np.meshgrid(np.arange(20), np.arange(20))
+    z = z1 * z2
+    plt.subplot(221)
+    plt.pcolormesh(x[:-1], y[:-1], z)
+    plt.subplot(222)
+    plt.pcolormesh(x, y, z)
+    x = np.repeat(x[np.newaxis], 21, axis=0)
+    y = np.repeat(y[:, np.newaxis], 21, axis=1)
+    plt.subplot(223)
+    plt.pcolormesh(x[:-1, :-1], y[:-1, :-1], z)
+    plt.subplot(224)
+    plt.pcolormesh(x, y, z)
+    for ax in fig.get_axes():
+        for label in ax.get_xticklabels():
+            label.set_ha('right')
+            label.set_rotation(30)
+
+
+@image_comparison(baseline_images=['pcolor_datetime_axis'],
+                  extensions=['png'], remove_text=False)
+def test_pcolor_datetime_axis():
+    fig = plt.figure()
+    fig.subplots_adjust(hspace=0.4, top=0.98, bottom=.15)
+    base = datetime.datetime(2013, 1, 1)
+    x = np.array([base + datetime.timedelta(days=d) for d in range(21)])
+    y = np.arange(21)
+    z1, z2 = np.meshgrid(np.arange(20), np.arange(20))
+    z = z1 * z2
+    plt.subplot(221)
+    plt.pcolor(x[:-1], y[:-1], z)
+    plt.subplot(222)
+    plt.pcolor(x, y, z)
+    x = np.repeat(x[np.newaxis], 21, axis=0)
+    y = np.repeat(y[:, np.newaxis], 21, axis=1)
+    plt.subplot(223)
+    plt.pcolor(x[:-1, :-1], y[:-1, :-1], z)
+    plt.subplot(224)
+    plt.pcolor(x, y, z)
+    for ax in fig.get_axes():
+        for label in ax.get_xticklabels():
+            label.set_ha('right')
+            label.set_rotation(30)
+
+
 def test_pcolorargs():
     n = 12
     x = np.linspace(-1.5, 1.5, n)
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -1,3 +1,5 @@
+import datetime
+
 import numpy as np
 
 from matplotlib.testing.decorators import cleanup, image_comparison
@@ -177,6 +179,32 @@ def test_given_colors_levels_and_extends():
         plt.colorbar()
 
 
+@image_comparison(baseline_images=['contour_datetime_axis'],
+                  extensions=['png'], remove_text=False)
+def test_contour_datetime_axis():
+    fig = plt.figure()
+    fig.subplots_adjust(hspace=0.4, top=0.98, bottom=.15)
+    base = datetime.datetime(2013, 1, 1)
+    x = np.array([base + datetime.timedelta(days=d) for d in range(20)])
+    y = np.arange(20)
+    z1, z2 = np.meshgrid(np.arange(20), np.arange(20))
+    z = z1 * z2
+    plt.subplot(221)
+    plt.contour(x, y, z)
+    plt.subplot(222)
+    plt.contourf(x, y, z)
+    x = np.repeat(x[np.newaxis], 20, axis=0)
+    y = np.repeat(y[:, np.newaxis], 20, axis=1)
+    plt.subplot(223)
+    plt.contour(x, y, z)
+    plt.subplot(224)
+    plt.contourf(x, y, z)
+    for ax in fig.get_axes():
+        for label in ax.get_xticklabels():
+            label.set_ha('right')
+            label.set_rotation(30)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule(argv=['-s', '--with-doctest'], exit=False)
diff --git a/lib/matplotlib/units.py b/lib/matplotlib/units.py
@@ -44,6 +44,7 @@ def default_units(x, axis):
 """
 from __future__ import print_function
 from matplotlib.cbook import iterable, is_numlike
+import numpy as np
 
 
 class AxisInfo:
@@ -132,6 +133,10 @@ def get_converter(self, x):
         if classx is not None:
             converter = self.get(classx)
 
+        if isinstance(x, np.ndarray) and x.size:
+            converter = self.get_converter(x.ravel()[0])
+            return converter
+
         if converter is None and iterable(x):
             for thisx in x:
                 # Make sure that recursing might actually lead to a solution,
