ENH: support np.datenum64 in dates.py

jklymak · jklymak · commit 71262959eb33 · 2017-11-14T17:18:15.000-08:00
diff --git a/lib/matplotlib/dates.py b/lib/matplotlib/dates.py
@@ -243,6 +243,27 @@ def _to_ordinalf(dt):
 _to_ordinalf_np_vectorized = np.vectorize(_to_ordinalf)
 
 
+def _dt64_to_ordinalf(d):
+    """
+    Convert `numpy.datetime64` or an ndarray of those types to Gregorian
+    date as UTC float.  Roundoff is via float64 precision.  Practically:
+    microseconds for dates between 290301 BC, 294241 AD, milliseconds for
+    larger dates (see `numpy.datetime64`).  Nanoseconds aren't possible
+    because we do times compared to ``0001-01-01T00:00:00`` (plus one day).
+    """
+
+    # the "extra" ensures that we at least allow the dynamic range out to
+    # seconds.  That should get out to +/-2e11 years.
+    extra = d - d.astype('datetime64[s]')
+    extra = extra.astype('timedelta64[ns]')
+    t0 = np.datetime64('0001-01-01T00:00:00').astype('datetime64[s]')
+    dt = (d.astype('datetime64[s]') - t0).astype(np.float64)
+    dt += extra.astype(np.float64) / 1.0e9
+    dt = dt / SEC_PER_DAY + 1.0
+
+    return dt
+
+
 def _from_ordinalf(x, tz=None):
     """
     Convert Gregorian float of the date, preserving hours, minutes,
@@ -354,12 +375,13 @@ def date2num(d):
 
     Parameters
     ----------
-    d : :class:`datetime` or sequence of :class:`datetime`
+    d : :class:`datetime` or :class:`numpy.datetime64`, or sequences of
+        these classes.
 
     Returns
     -------
     float or sequence of floats
-        Number of days (fraction part represents hours, minutes, seconds)
+        Number of days (fraction part represents hours, minutes, seconds, ms)
         since 0001-01-01 00:00:00 UTC, plus one.
 
     Notes
@@ -368,6 +390,10 @@ def date2num(d):
     Gregorian calendar is assumed; this is not universal practice.
     For details see the module docstring.
     """
+
+    if ((isinstance(d, np.ndarray) and np.issubdtype(d.dtype, np.datetime64))
+            or isinstance(d, np.datetime64)):
+        return _dt64_to_ordinalf(d)
     if not cbook.iterable(d):
         return _to_ordinalf(d)
     else:
@@ -488,8 +514,8 @@ def drange(dstart, dend, delta):
     *dend* are :class:`datetime` instances.  *delta* is a
     :class:`datetime.timedelta` instance.
     """
-    f1 = _to_ordinalf(dstart)
-    f2 = _to_ordinalf(dend)
+    f1 = date2num(dstart)
+    f2 = date2num(dend)
     step = delta.total_seconds() / SEC_PER_DAY
 
     # calculate the difference between dend and dstart in times of delta
@@ -504,7 +530,7 @@ def drange(dstart, dend, delta):
         dinterval_end -= delta
         num -= 1
 
-    f2 = _to_ordinalf(dinterval_end)  # new float-endpoint
+    f2 = date2num(dinterval_end)  # new float-endpoint
     return np.linspace(f1, f2, num + 1)
 
 ### date tickers and formatters ###
@@ -1630,5 +1656,6 @@ def default_units(x, axis):
         return None
 
 
+units.registry[np.datetime64] = DateConverter()
 units.registry[datetime.date] = DateConverter()
 units.registry[datetime.datetime] = DateConverter()
diff --git a/lib/matplotlib/tests/test_dates.py b/lib/matplotlib/tests/test_dates.py
@@ -3,13 +3,14 @@
 
 from six.moves import map
 
-import datetime
-import warnings
-import tempfile
-import pytest
 
+import datetime
 import dateutil
+import numpy as np
+import pytest
 import pytz
+import tempfile
+import warnings
 
 try:
     # mock in python 3.3+
@@ -22,6 +23,52 @@
 import matplotlib.dates as mdates
 
 
+def test_date_numpyx():
+    # test that numpy dates work properly...
+    base = datetime.datetime(2017, 1, 1)
+    time = [base + datetime.timedelta(days=x) for x in range(0, 3)]
+    timenp = np.array(time, dtype='datetime64[ns]')
+    data = np.array([0., 2., 1.])
+    fig = plt.figure(figsize=(10, 2))
+    ax = fig.add_subplot(1, 1, 1)
+    h, = ax.plot(time, data)
+    hnp, = ax.plot(timenp, data)
+    assert np.array_equal(h.get_xdata(orig=False), hnp.get_xdata(orig=False))
+    fig = plt.figure(figsize=(10, 2))
+    ax = fig.add_subplot(1, 1, 1)
+    h, = ax.plot(data, time)
+    hnp, = ax.plot(data, timenp)
+    assert np.array_equal(h.get_ydata(orig=False), hnp.get_ydata(orig=False))
+
+
+def test_date_date2num_numpy():
+    t0 = datetime.datetime(2017, 1, 1, 0, 1, 1)
+    time = mdates.date2num(t0)
+    tnp = np.array(t0, dtype='datetime64[s]')
+    nptime = mdates.date2num(tnp)
+    assert time == nptime
+    tnp = np.array(t0, dtype='datetime64[us]')
+    nptime = mdates.date2num(tnp)
+    assert time == nptime
+    tnp = np.array(t0, dtype='datetime64[ns]')
+    nptime = mdates.date2num(tnp)
+    assert time == nptime
+    t0 = [datetime.datetime(2017, 1, 1, 0, 1, 1),
+          datetime.datetime(2017, 1, 1, 1, 1, 1)]
+    tnp = np.array(t0, dtype='datetime64[ms]')
+    time = mdates.date2num(t0)
+    nptime = mdates.date2num(tnp)
+    assert np.array_equal(time, nptime)
+    t0 = [[datetime.datetime(2017, 1, 1, 0, 1, 1),
+          datetime.datetime(2017, 1, 1, 1, 1, 1)],
+          [datetime.datetime(2017, 1, 1, 2, 1, 1),
+           datetime.datetime(2017, 1, 1, 3, 1, 1)]]
+    tnp = np.array(t0, dtype='datetime64[ms]')
+    time = mdates.date2num(t0)
+    nptime = mdates.date2num(tnp)
+    assert np.array_equal(time, nptime)
+
+
 @image_comparison(baseline_images=['date_empty'], extensions=['png'])
 def test_date_empty():
     # make sure mpl does the right thing when told to plot dates even
