From a143ad304f382e61320161ec9d17b8217e195b8d Mon Sep 17 00:00:00 2001 From: Stephan Hoyer Date: Tue, 14 Apr 2015 16:48:09 -0400 Subject: [PATCH] Fix for scalar timedelta conversion --- xray/conventions.py | 21 +++++++++++++++------ xray/test/test_conventions.py | 35 +++++++++++++++++++++++++++-------- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/xray/conventions.py b/xray/conventions.py index b029864ddbd..31871296408 100644 --- a/xray/conventions.py +++ b/xray/conventions.py @@ -139,12 +139,25 @@ def decode_cf_datetime(num_dates, units, calendar=None): return dates.reshape(num_dates.shape) +def _asarray_or_scalar(x): + x = np.asarray(x) + if x.ndim > 0: + return x + else: + return x.item() + + def decode_cf_timedelta(num_timedeltas, units): """Given an array of numeric timedeltas in netCDF format, convert it into a numpy timedelta64[ns] array. """ + num_timedeltas = _asarray_or_scalar(num_timedeltas) units = _netcdf_to_numpy_timeunit(units) - return pd.to_timedelta(np.asarray(num_timedeltas), unit=units, box=False) + result = pd.to_timedelta(num_timedeltas, unit=units, box=False) + # NaT is returned unboxed with wrong units; this should be fixed in pandas + if result.dtype != 'timedelta64[ns]': + result = result.astype('timedelta64[ns]') + return result TIME_UNITS = set(['days', 'hours', 'minutes', 'seconds']) @@ -266,11 +279,7 @@ def encode_cf_timedelta(timedeltas, units=None): np_unit = _netcdf_to_numpy_timeunit(units) num = 1.0 * timedeltas / np.timedelta64(1, np_unit) - - missing = pd.isnull(timedeltas) - if np.any(missing): - num = np.asarray(num, dtype=float) - num[missing] = np.nan + num = np.where(pd.isnull(timedeltas), np.nan, num) int_num = np.asarray(num, dtype=np.int64) if (num == int_num).all(): diff --git a/xray/test/test_conventions.py b/xray/test/test_conventions.py index 39ce5ad3e60..ee184672f60 100644 --- a/xray/test/test_conventions.py +++ b/xray/test/test_conventions.py @@ -119,6 +119,7 @@ def test_cf_datetime(self): ([10], 'daYs since 2000-01-01'), ([[10]], 'days since 2000-01-01'), ([10, 10], 'days since 2000-01-01'), + (np.array(10), 'days since 2000-01-01'), (0, 'days since 1000-01-01'), ([0], 'days since 1000-01-01'), ([[0]], 'days since 1000-01-01'), @@ -352,18 +353,36 @@ def test_infer_datetime_units(self): self.assertEqual(expected, conventions.infer_datetime_units(dates)) def test_cf_timedelta(self): - for timedeltas, units, expected in [ - (['NaT'], 'days', [np.nan]), - (['1D', '2D', '3D'], 'days', np.array([1, 2, 3], 'int64')), - (['NaT', '0s', '1s'], None, [np.nan, 0, 1]), - (['30m', '60m'], 'hours', [0.5, 1.0]), - ]: - timedeltas = pd.to_timedelta(timedeltas) - expected = np.array(expected) + examples = [ + ('1D', 'days', np.int64(1)), + (['1D', '2D', '3D'], 'days', np.array([1, 2, 3], 'int64')), + (['NaT', '0s', '1s'], None, [np.nan, 0, 1]), + (['30m', '60m'], 'hours', [0.5, 1.0]), + ] + if pd.__version__ >= '0.16': + # not quite sure why, but these examples don't work on older pandas + examples.extend([(np.timedelta64('NaT', 'ns'), 'days', np.nan), + (['NaT', 'NaT'], 'days', [np.nan, np.nan])]) + + for timedeltas, units, numbers in examples: + timedeltas = pd.to_timedelta(timedeltas, box=False) + numbers = np.array(numbers) + + expected = numbers actual, _ = conventions.encode_cf_timedelta(timedeltas, units) self.assertArrayEqual(expected, actual) self.assertEqual(expected.dtype, actual.dtype) + if units is not None: + expected = timedeltas + actual = conventions.decode_cf_timedelta(numbers, units) + self.assertArrayEqual(expected, actual) + self.assertEqual(expected.dtype, actual.dtype) + + expected = np.timedelta64('NaT', 'ns') + actual = conventions.decode_cf_timedelta(np.array(np.nan), 'days') + self.assertArrayEqual(expected, actual) + def test_infer_timedelta_units(self): for deltas, expected in [ (pd.to_timedelta(['1 day', '2 days']), 'days'),