8000 refactor conversion functions and add timedelta to ordinalf · matplotlib/matplotlib@385620b · GitHub
[go: up one dir, main page]

Skip to content

Commit 385620b

Browse files
committed
refactor conversion functions and add timedelta to ordinalf
1 parent 0dc6167 commit 385620b

File tree

2 files changed

+147
-19
lines changed

2 files changed

+147
-19
lines changed

lib/matplotlib/dates.py

Lines changed: 91 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -325,10 +325,43 @@ def _dt64_to_ordinalf(d):
325325
dt += extra.astype(np.float64) / 1.0e9
326326
dt = dt / SEC_PER_DAY
327327

328+
return _nat_to_nan(dt, d)
329+
330+
331+
def _timedelta64_to_ordinalf(t):
332+
"""
333+
Convert `numpy.timedelta64` or an ndarray of those types to a timedelta
334+
as float. Roundoff is float64 precision.
335+
TODO: precision ok, can be improved? be more concrete here
336+
"""
337+
td = t.astype('timedelta64[D]').astype(np.float64)
338+
339+
return _nat_to_nan(td, t)
340+
341+
342+
def _nat_to_nan(ordf, timeval):
343+
"""
344+
Replace all values in the converted array `ordf` that were 'NaT'
345+
originally in `timeval` with `np.nan`.
346+
347+
Parameters
348+
----------
349+
ordf: datetime or timedelta converted to float (ndarray or float)
350+
timeval: `numpy.datetime64` or `numpy.timedelta64` (ndarray or
351+
single value)
352+
353+
Returns
354+
-------
355+
ordf with all originally 'NaT' replaced by `np.nan`
356+
"""
328357
NaT_int = np.datetime64('NaT').astype(np.int64)
329-
d_int = d.astype(np.int64)
330-
dt[d_int == NaT_int] = np.nan
331-
return dt
358+
t_int = timeval.astype(np.int64)
359+
try:
360+
ordf[t_int == NaT_int] = np.nan
361+
except TypeError:
362+
if t_int == NaT_int:
363+
ordf = np.nan
364+
return ordf
332365

333366

334367
def _from_ordinalf(x, tz=None):
@@ -423,35 +456,74 @@ def date2num(d):
423456
The Gregorian calendar is assumed; this is not universal practice.
424457
For details see the module docstring.
425458
"""
459+
return _timevalue2num(d, np.datetime64)
460+
461+
462+
def timedelta2num(t):
463+
"""
464+
Convert datetime objects to Matplotlib dates.
465+
466+
Parameters
467+
----------
468+
t : `datetime.timedelta` or `numpy.timedelta64` or sequences of these
469+
470+
Returns
471+
-------
472+
float or sequence of floats
473+
Number of days. For example, 1 day 12 hours returns 1.5.
474+
"""
475+
return _timevalue2num(t, np.timedelta64)
476+
477+
478+
def _timevalue2num(v, v_cls):
479+
"""
480+
Convert datetime or timedelta to Matplotlibs representation as days
481+
(since the epoch for datetime) as float.
482+
483+
Parameters
484+
----------
485+
v: `datetime.datetime`, `numpy.datetime64`, `datetime.timedelta` or
486+
`numpy.timedelta64` or sequences of these
487+
v_cls: class `numpy.timedelta64` or `numpy.datetime64` depending on
488+
whether to convert datetime or timedelta values
489+
"""
490+
if v_cls is np.datetime64:
491+
dtype = 'datetime64[us]'
492+
else:
493+
dtype = 'timedelta64[us]'
494+
426495
# Unpack in case of e.g. Pandas or xarray object
427-
d = cbook._unpack_to_numpy(d)
496+
v = cbook._unpack_to_numpy(v)
428497

429498
# make an iterable, but save state to unpack later:
430-
iterable = np.iterable(d)
499+
iterable = np.iterable(v)
431500
if not iterable:
432-
d = [d]
501+
v = [v]
433502

434-
masked = np.ma.is_masked(d)
435-
mask = np.ma.getmask(d)
436-
d = np.asarray(d)
503+
masked = np.ma.is_masked(v)
504+
mask = np.ma.getmask(v)
505+
v = np.asarray(v)
437506

438507
# convert to datetime64 arrays, if not already:
439-
if not np.issubdtype(d.dtype, np.datetime64):
508+
if not np.issubdtype(v.dtype, v_cls):
440509
# datetime arrays
441-
if not d.size:
510+
if not v.size:
442511
# deals with an empty array...
443-
return d
444-
tzi = getattr(d[0], 'tzinfo', None)
512+
return v
513+
tzi = getattr(v[0], 'tzinfo', None)
445514
if tzi is not None:
446515
# make datetime naive:
447-
d = [dt.astimezone(UTC).replace(tzinfo=None) for dt in d]
448-
d = np.asarray(d)
449-
d = d.astype('datetime64[us]')
516+
v = [dt.astimezone(UTC).replace(tzinfo=None) for dt in v]
517+
v = np.asarray(v)
518+
v = v.astype(dtype)
450519

451-
d = np.ma.masked_array(d, mask=mask) if masked else d
452-
d = _dt64_to_ordinalf(d)
520+
v = np.ma.masked_array(v, mask=mask) if masked else v
521+
if v_cls is np.datetime64:
522+
v = _dt64_to_ordinalf(v)
523+
else:
524+
v = _timedelta64_to_ordinalf(v)
453525

454-
return d if iterable else d[0]
526+
return v if iterable else v[0]
455527

456528

457529
def num2date(x, tz=None):

lib/matplotlib/tests/test_dates.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,62 @@ def test_date2num_masked():
8888
(False, True, True, False, False, False,
8989
True))
9090

91+
@pytest.mark.xfail # TODO: un-fail once formatters exist
92+
def test_timedelta_numpy():
93+
# TODO: merge with datetime test?
94+
# test that numpy timedeltas work properly...
95+
time = [datetime.timedelta(days=x) for x in range(0, 3)]
96+
timenp = np.array(time, dtype='timedelta64[ns]')
97+
data = np.array([0., 2., 1.])
98+
fig = plt.figure(figsize=(10, 2))
99+
ax = fig.add_subplot(1, 1, 1)
100+
h, = ax.plot(time, data)
101+
hnp, = ax.plot(timenp, data)
102+
np.testing.assert_equal(h.get_xdata(orig=False), hnp.get_xdata(orig=False))
103+
fig = plt.figure(figsize=(10, 2))
104+
ax = fig.add_subplot(1, 1, 1)
105+
h, = ax.plot(data, time)
106+
hnp, = ax.plot(data, timenp)
107+
np.testing.assert_equal(h.get_ydata(orig=False), hnp.get_ydata(orig=False))
108+
109+
110+
@pytest.mark.parametrize('t0', [datetime.timedelta(100),
111+
112+
[datetime.timedelta(100, 1),
113+
datetime.timedelta(101, 1)],
114+
115+
[[datetime.timedelta(100, 0, 1),
116+
datetime.timedelta(101, 0, 1)],
117+
[datetime.timedelta(102, 0, 1),
118+
datetime.timedelta(103, 0, 1)]]])
119+
@pytest.mark.parametrize('dtype', ['timedelta64[s]',
120+
'timedelta64[us]',
121+
'timedelta64[ms]',
122+
'timedelta64[ns]'])
123+
def test_timedelta_timedelta2num_numpy(t0, dtype):
124+
time = mdates.timedelta2num(t0)
125+
tnp = np.array(t0, dtype=dtype)
126+
nptime = mdates.timedelta2num(tnp)
127+
np.testing.assert_equal(time, nptime)
128+
129+
130+
@pytest.mark.parametrize('dtype', ['timedelta64[s]',
131+
'timedelta64[us]',
132+
'timedelta64[ms]',
133+
'timedelta64[ns]'])
134+
def test_timedelta2num_NaT(dtype):
135+
t0 = datetime.timedelta(100, 0, 1)
136+
tmpl = [mdates.timedelta2num(t0), np.nan]
137+
tnp = np.array([t0, 'NaT'], dtype=dtype)
138+
nptime = mdates.timedelta2num(tnp)
139+
np.testing.assert_array_equal(tmpl, nptime)
140+
141+
142+
@pytest.mark.parametrize('units', ['s', 'ms', 'us', 'ns'])
143+
def test_timedelta2num_NaT_scalar(units):
144+
tmpl = mdates.timedelta2num(np.timedelta64('NaT', units))
145+
assert np.isnan(tmpl)
146+
91147

92148
def test_date_empty():
93149
# make sure we do the right thing when told to plot dates even

0 commit comments

Comments
 (0)
0