def test_getitem_coords(self):
orig = DataArray([[10], [20]],
{'x': [1, 2], 'y': [3], 'z': 4,
'x2': ('x', ['a', 'b']),
'y2': ('y', ['c']),
'xy': (['y', 'x'], [['d', 'e']])},
dims=['x', 'y'])
self.assertDataArrayIdentical(orig, orig[:])
self.assertDataArrayIdentical(orig, orig[:, :])
self.assertDataArrayIdentical(orig, orig[...])
self.assertDataArrayIdentical(orig, orig[:2, :1])
self.assertDataArrayIdentical(orig, orig[[0, 1], [0]])
actual = orig[0, 0]
expected = DataArray(
10, {'x': 1, 'y': 3, 'z': 4, 'x2': 'a', 'y2': 'c', 'xy': 'd'})
self.assertDataArrayIdentical(expected, actual)
actual = orig[0, :]
expected = DataArray(
[10], {'x': 1, 'y': [3], 'z': 4, 'x2': 'a', 'y2': ('y', ['c']),
'xy': ('y', ['d'])},
dims='y')
self.assertDataArrayIdentical(expected, actual)
actual = orig[:, 0]
expected = DataArray(
[10, 20], {'x': [1, 2], 'y': 3, 'z': 4, 'x2': ('x', ['a', 'b']),
'y2': 'c', 'xy': ('x', ['d', 'e'])},
dims='x')
self.assertDataArrayIdentical(expected, actual)
def test_coord_coords(self):
orig = DataArray([10, 20], {
'x': [1, 2],
'x2': ('x', ['a', 'b']),
'z': 4
},
dims='x')
actual = orig.coords['x']
expected = DataArray([1, 2], {
'z': 4,
'x2': ('x', ['a', 'b'])
},
dims='x',
name='x')
self.assertDataArrayIdentical(expected, actual)
del actual.coords['x2']
self.assertDataArrayIdentical(expected.reset_coords('x2', drop=True),
actual)
actual.coords['x3'] = ('x', ['a', 'b'])
expected = DataArray([1, 2], {
'z': 4,
'x3': ('x', ['a', 'b'])
},
dims='x',
name='x')
self.assertDataArrayIdentical(expected, actual)
def test_groupby_count(self):
array = DataArray([0, 0, np.nan, np.nan, 0, 0],
coords={'cat': ('x', ['a', 'b', 'b', 'c', 'c', 'c'])},
dims='x')
actual = array.groupby('cat').count()
expected = DataArray([1, 1, 2], coords=[('cat', ['a', 'b', 'c'])])
self.assertDataArrayIdentical(actual, expected)
def test_reduce(self):
coords = {
'x': [-1, -2],
'y': ['ab', 'cd', 'ef'],
'lat': (['x', 'y'], [[1, 2, 3], [-1, -2, -3]]),
'c': -999
}
orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=['x', 'y'])
actual = orig.mean()
expected = DataArray(0, {'c': -999})
self.assertDataArrayIdentical(expected, actual)
actual = orig.mean(['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
actual = orig.mean('x')
expected = DataArray([-2, 0, 2], {'y': coords['y'], 'c': -999}, 'y')
self.assertDataArrayIdentical(expected, actual)
actual = orig.mean(['x'])
self.assertDataArrayIdentical(expected, actual)
actual = orig.mean('y')
expected = DataArray([0, 0], {'x': coords['x'], 'c': -999}, 'x')
self.assertDataArrayIdentical(expected, actual)
self.assertVariableEqual(self.dv.reduce(np.mean, 'x'),
self.v.reduce(np.mean, 'x'))
def test_is_null(self):
x = np.random.RandomState(42).randn(5, 6)
x[x < 0] = np.nan
original = DataArray(x, [-np.arange(5), np.arange(6)], ['x', 'y'])
expected = DataArray(pd.isnull(x), [-np.arange(5), np.arange(6)],
['x', 'y'])
self.assertDataArrayIdentical(expected, original.isnull())
self.assertDataArrayIdentical(~expected, original.notnull())
def test_align_dtype(self):
# regression test for #264
x1 = np.arange(30)
x2 = np.arange(5, 35)
a = DataArray(np.random.random((30,)).astype('f32'), {'x': x1})
b = DataArray(np.random.random((30,)).astype('f32'), {'x': x2})
c, d = align(a, b, join='outer')
self.assertEqual(c.dtype, np.float32)
def test_where_dispatching(self):
a = np.arange(10)
b = a > 3
x = da.from_array(a, 5)
y = da.from_array(b, 5)
expected = DataArray(a).where(b)
self.assertLazyAndIdentical(expected, DataArray(a).where(y))
self.assertLazyAndIdentical(expected, DataArray(x).where(b))
self.assertLazyAndIdentical(expected, DataArray(x).where(y))
def setUp(self):
da = DataArray(easy_array((10, 15), start=-1), dims=['y', 'x'])
# add 2d coords
ds = da.to_dataset(name='testvar')
x, y = np.meshgrid(da.x.values, da.y.values)
ds['x2d'] = DataArray(x, dims=['y', 'x'])
ds['y2d'] = DataArray(y, dims=['y', 'x'])
ds.set_coords(['x2d', 'y2d'], inplace=True)
# set darray and plot method
self.darray = ds.testvar
self.plotmethod = getattr(self.darray.plot, self.plotfunc.__name__)
def test_reset_coords(self):
data = DataArray(np.zeros((3, 4)), {
'bar': ('x', ['a', 'b', 'c']),
'baz': ('y', range(4))
},
dims=['x', 'y'],
name='foo')
actual = data.reset_coords()
expected = Dataset({
'foo': (['x', 'y'], np.zeros((3, 4))),
'bar': ('x', ['a', 'b', 'c']),
'baz': ('y', range(4))
})
self.assertDatasetIdentical(actual, expected)
actual = data.reset_coords(['bar', 'baz'])
self.assertDatasetIdentical(actual, expected)
actual = data.reset_coords('bar')
expected = Dataset(
{
'foo': (['x', 'y'], np.zeros((3, 4))),
'bar': ('x', ['a', 'b', 'c'])
}, {'baz': ('y', range(4))})
self.assertDatasetIdentical(actual, expected)
actual = data.reset_coords(['bar'])
self.assertDatasetIdentical(actual, expected)
actual = data.reset_coords(drop=True)
expected = DataArray(np.zeros((3, 4)), dims=['x', 'y'], name='foo')
self.assertDataArrayIdentical(actual, expected)
actual = data.copy()
actual.reset_coords(drop=True, inplace=True)
self.assertDataArrayIdentical(actual, expected)
actual = data.reset_coords('bar', drop=True)
expected = DataArray(np.zeros((3, 4)), {'baz': ('y', range(4))},
dims=['x', 'y'],
name='foo')
self.assertDataArrayIdentical(actual, expected)
with self.assertRaisesRegexp(ValueError, 'cannot reset coord'):
data.reset_coords(inplace=True)
with self.assertRaises(KeyError):
data.reset_coords('foo', drop=True)
with self.assertRaisesRegexp(ValueError, 'cannot be found'):
data.reset_coords('not_found')
with self.assertRaisesRegexp(ValueError, 'cannot remove index'):
data.reset_coords('y')
def test_constructor_invalid(self):
data = np.random.randn(3, 2)
with self.assertRaisesRegexp(ValueError, 'coords is not dict-like'):
DataArray(data, [[0, 1, 2]], ['x', 'y'])
with self.assertRaisesRegexp(ValueError, 'not a subset of the .* dim'):
DataArray(data, {'x': [0, 1, 2]}, ['a', 'b'])
with self.assertRaisesRegexp(ValueError, 'not a subset of the .* dim'):
DataArray(data, {'x': [0, 1, 2]})
with self.assertRaisesRegexp(TypeError, 'is not a string'):
DataArray(data, dims=['x', None])
def test_name(self):
arr = self.dv
self.assertEqual(arr.name, 'foo')
copied = arr.copy()
arr.name = 'bar'
self.assertEqual(arr.name, 'bar')
self.assertDataArrayEqual(copied, arr)
actual = DataArray(Coordinate('x', [3]))
actual.name = 'y'
expected = DataArray(Coordinate('y', [3]))
self.assertDataArrayIdentical(actual, expected)
def test_index_math(self):
orig = DataArray(range(3), dims='x', name='x')
actual = orig + 1
expected = DataArray(1 + np.arange(3), coords=[('x', range(3))])
self.assertDataArrayIdentical(expected, actual)
# regression tests for #254
actual = orig[0] < orig
expected = DataArray([False, True, True], coords=[('x', range(3))])
self.assertDataArrayIdentical(expected, actual)
actual = orig > orig[0]
self.assertDataArrayIdentical(expected, actual)
def test_constructor_from_self_described(self):
data = [[-0.1, 21], [0, 2]]
expected = DataArray(data,
indexes={
'x': ['a', 'b'],
'y': [-1, -2]
},
dimensions=['x', 'y'],
name='foobar',
attributes={'bar': 2},
encoding={'foo': 3})
actual = DataArray(expected)
self.assertDataArrayIdentical(expected, actual)
frame = pd.DataFrame(data,
index=pd.Index(['a', 'b'], name='x'),
columns=pd.Index([-1, -2], name='y'))
actual = DataArray(frame)
self.assertDataArrayEqual(expected, actual)
series = pd.Series(data[0], index=pd.Index([-1, -2], name='y'))
actual = DataArray(series)
self.assertDataArrayEqual(expected[0], actual)
panel = pd.Panel({0: frame})
actual = DataArray(panel)
expected = DataArray([data], expected.coordinates, ['dim_0', 'x', 'y'])
self.assertDataArrayIdentical(expected, actual)
expected = DataArray(['a', 'b'], name='foo')
actual = DataArray(pd.Index(['a', 'b'], name='foo'))
self.assertDataArrayIdentical(expected, actual)
def test_dataset_math(self):
# more comprehensive tests with multiple dataset variables
obs = Dataset(
{
'tmin': ('x', np.arange(5)),
'tmax': ('x', 10 + np.arange(5))
}, {
'x': ('x', 0.5 * np.arange(5)),
'loc': ('x', range(-2, 3))
})
actual = 2 * obs['tmax']
expected = DataArray(2 * (10 + np.arange(5)), obs.coords, name='tmax')
self.assertDataArrayIdentical(actual, expected)
actual = obs['tmax'] - obs['tmin']
expected = DataArray(10 * np.ones(5), obs.coords)
self.assertDataArrayIdentical(actual, expected)
sim = Dataset({
'tmin': ('x', 1 + np.arange(5)),
'tmax': ('x', 11 + np.arange(5)),
# does *not* include 'loc' as a coordinate
'x': ('x', 0.5 * np.arange(5))
})
actual = sim['tmin'] - obs['tmin']
expected = DataArray(np.ones(5), obs.coords, name='tmin')
self.assertDataArrayIdentical(actual, expected)
actual = -obs['tmin'] + sim['tmin']
self.assertDataArrayIdentical(actual, expected)
actual = sim['tmin'].copy()
actual -= obs['tmin']
self.assertDataArrayIdentical(actual, expected)
actual = sim.copy()
actual['tmin'] = sim['tmin'] - obs['tmin']
expected = Dataset(
{
'tmin': ('x', np.ones(5)),
'tmax': ('x', sim['tmax'].values)
}, obs.coords)
self.assertDatasetIdentical(actual, expected)
actual = sim.copy()
actual['tmin'] -= obs['tmin']
self.assertDatasetIdentical(actual, expected)
def test_reindex(self):
foo = self.dv
bar = self.dv[:2, :2]
self.assertDataArrayIdentical(foo.reindex_like(bar), bar)
expected = foo.copy()
expected[:] = np.nan
expected[:2, :2] = bar
self.assertDataArrayIdentical(bar.reindex_like(foo), expected)
# regression test for #279
expected = DataArray(np.random.randn(5), dims=["time"])
time2 = DataArray(np.arange(5), dims="time2")
actual = expected.reindex(time=time2)
self.assertDataArrayIdentical(actual, expected)
def test_convenient_facetgrid(self):
a = easy_array((10, 15, 4))
d = DataArray(a, dims=['y', 'x', 'z'])
g = self.plotfunc(d, x='x', y='y', col='z', col_wrap=2)
self.assertArrayEqual(g.axes.shape, [2, 2])
for (y, x), ax in np.ndenumerate(g.axes):
self.assertTrue(ax.has_data())
if x == 0:
self.assertEqual('y', ax.get_ylabel())
else:
self.assertEqual('', ax.get_ylabel())
if y == 1:
self.assertEqual('x', ax.get_xlabel())
else:
self.assertEqual('', ax.get_xlabel())
# Infering labels
g = self.plotfunc(d, col='z', col_wrap=2)
self.assertArrayEqual(g.axes.shape, [2, 2])
for (y, x), ax in np.ndenumerate(g.axes):
self.assertTrue(ax.has_data())
if x == 0:
self.assertEqual('y', ax.get_ylabel())
else:
self.assertEqual('', ax.get_ylabel())
if y == 1:
self.assertEqual('x', ax.get_xlabel())
else:
self.assertEqual('', ax.get_xlabel())
def test_verbose_facetgrid(self):
a = easy_array((10, 15, 3))
d = DataArray(a, dims=['y', 'x', 'z'])
g = xplt.FacetGrid(d, col='z')
g.map_dataarray(self.plotfunc, 'x', 'y')
for ax in g.axes.flat:
self.assertTrue(ax.has_data())
def test_coords(self):
coords = [Coordinate('x', [-1, -2]), Coordinate('y', [0, 1, 2])]
da = DataArray(np.random.randn(2, 3), coords, name='foo')
self.assertEquals(2, len(da.coords))
self.assertEqual(['x', 'y'], list(da.coords))
self.assertTrue(coords[0].identical(da.coords['x']))
self.assertTrue(coords[1].identical(da.coords['y']))
self.assertIn('x', da.coords)
self.assertNotIn(0, da.coords)
self.assertNotIn('foo', da.coords)
with self.assertRaises(KeyError):
da.coords[0]
with self.assertRaises(KeyError):
da.coords['foo']
expected = dedent("""\
Coordinates:
* x (x) int64 -1 -2
* y (y) int64 0 1 2""")
actual = repr(da.coords)
self.assertEquals(expected, actual)
def test_to_and_from_cdms2(self):
try:
import cdms2
except ImportError:
raise unittest.SkipTest('cdms2 not installed')
original = DataArray(np.arange(6).reshape(2, 3),
[('distance', [-2, 2], {'units': 'meters'}),
('time', pd.date_range('2000-01-01', periods=3))],
name='foo', attrs={'baz': 123})
expected_coords = [Coordinate('distance', [-2, 2]),
Coordinate('time', [0, 1, 2])]
actual = original.to_cdms2()
self.assertArrayEqual(actual, original)
self.assertEqual(actual.id, original.name)
self.assertItemsEqual(actual.getAxisIds(), original.dims)
for axis, coord in zip(actual.getAxisList(), expected_coords):
self.assertEqual(axis.id, coord.name)
self.assertArrayEqual(axis, coord.values)
self.assertEqual(actual.baz, original.attrs['baz'])
component_times = actual.getAxis(1).asComponentTime()
self.assertEqual(len(component_times), 3)
self.assertEqual(str(component_times[0]), '2000-1-1 0:0:0.0')
roundtripped = DataArray.from_cdms2(actual)
self.assertDataArrayIdentical(original, roundtripped)
def test_groupby_math(self):
array = self.make_groupby_example_array()
for squeeze in [True, False]:
grouped = array.groupby('x', squeeze=squeeze)
expected = array + array.coords['x']
actual = grouped + array.coords['x']
self.assertDataArrayIdentical(expected, actual)
actual = array.coords['x'] + grouped
self.assertDataArrayIdentical(expected, actual)
ds = array.coords['x'].to_dataset()
expected = array + ds
actual = grouped + ds
self.assertDatasetIdentical(expected, actual)
actual = ds + grouped
self.assertDatasetIdentical(expected, actual)
grouped = array.groupby('abc')
expected_agg = (grouped.mean() - np.arange(3)).rename(None)
actual = grouped - DataArray(range(3), [('abc', ['a', 'b', 'c'])])
actual_agg = actual.groupby('abc').mean()
self.assertDataArrayAllClose(expected_agg, actual_agg)
with self.assertRaisesRegexp(TypeError, 'only support arithmetic'):
grouped + 1
with self.assertRaisesRegexp(TypeError, 'only support arithmetic'):
grouped + grouped
def test_concat(self):
self.ds['bar'] = Variable(['x', 'y'], np.random.randn(10, 20))
foo = self.ds['foo']
bar = self.ds['bar']
# from dataset array:
expected = DataArray(np.array([foo.values, bar.values]),
dims=['w', 'x', 'y'])
actual = concat([foo, bar], 'w')
self.assertDataArrayEqual(expected, actual)
# from iteration:
grouped = [g for _, g in foo.groupby('x')]
stacked = concat(grouped, self.ds['x'])
self.assertDataArrayIdentical(foo, stacked)
# with an index as the 'dim' argument
stacked = concat(grouped, self.ds.indexes['x'])
self.assertDataArrayIdentical(foo, stacked)
actual = concat([foo[0], foo[1]], pd.Index([0, 1])).reset_coords(drop=True)
expected = foo[:2].rename({'x': 'concat_dim'})
self.assertDataArrayIdentical(expected, actual)
actual = concat([foo[0], foo[1]], [0, 1]).reset_coords(drop=True)
expected = foo[:2].rename({'x': 'concat_dim'})
self.assertDataArrayIdentical(expected, actual)
with self.assertRaisesRegexp(ValueError, 'not identical'):
concat([foo, bar], compat='identical')
def test_groupby_sum(self):
array = self.make_groupby_example_array()
grouped = array.groupby('abc')
expected_sum_all = Dataset(
{'foo': Variable(['abc'], np.array([self.x[:, :9].sum(),
self.x[:, 10:].sum(),
self.x[:, 9:10].sum()]).T),
'abc': Variable(['abc'], np.array(['a', 'b', 'c']))})['foo']
self.assertDataArrayAllClose(expected_sum_all, grouped.reduce(np.sum))
self.assertDataArrayAllClose(expected_sum_all, grouped.sum())
expected = DataArray([array['y'].values[idx].sum() for idx
in [slice(9), slice(10, None), slice(9, 10)]],
[['a', 'b', 'c']], ['abc'])
actual = array['y'].groupby('abc').apply(np.sum)
self.assertDataArrayAllClose(expected, actual)
actual = array['y'].groupby('abc').sum()
self.assertDataArrayAllClose(expected, actual)
expected_sum_axis1 = Dataset(
{'foo': (['x', 'abc'], np.array([self.x[:, :9].sum(1),
self.x[:, 10:].sum(1),
self.x[:, 9:10].sum(1)]).T),
'x': self.ds['x'],
'abc': Variable(['abc'], np.array(['a', 'b', 'c']))})['foo']
self.assertDataArrayAllClose(expected_sum_axis1,
grouped.reduce(np.sum, 'y'))
self.assertDataArrayAllClose(expected_sum_axis1, grouped.sum('y'))
def setUp(self):
x = np.arange(start=0, stop=10, step=2)
y = np.arange(start=9, stop=-7, step=-3)
xy = np.dstack(np.meshgrid(x, y))
distance = np.linalg.norm(xy, axis=2)
self.darray = DataArray(distance, list(zip(('y', 'x'), (y, x))))
self.data_min = distance.min()
self.data_max = distance.max()
def test_subplot_kws(self):
a = easy_array((10, 15, 4))
d = DataArray(a, dims=['y', 'x', 'z'])
d.coords['z'] = list('abcd')
g = d.plot(x='x', y='y', col='z', col_wrap=2, cmap='cool',
subplot_kws=dict(axisbg='r'))
for ax in g.axes.flat:
self.assertEqual(ax.get_axis_bgcolor(), 'r')
def test_to_dataframe(self):
# regression test for #260
arr = DataArray(np.random.randn(3, 4),
[('B', [1, 2, 3]), ('A', list('cdef'))])
expected = arr.to_series()
actual = arr.to_dataframe()[None]
self.assertArrayEqual(expected.values, actual.values)
self.assertArrayEqual(expected.index.values, actual.index.values)
def test_convenient_facetgrid(self):
a = easy_array((10, 15, 4))
d = DataArray(a, dims=['y', 'x', 'z'])
g = self.plotfunc(d, x='x', y='y', col='z', col_wrap=2)
self.assertArrayEqual(g.axes.shape, [2, 2])
for ax in g.axes.flat:
self.assertTrue(ax.has_data())
def test_convenient_facetgrid_4d(self):
a = easy_array((10, 15, 2, 3))
d = DataArray(a, dims=['y', 'x', 'columns', 'rows'])
g = self.plotfunc(d, x='x', y='y', col='columns', row='rows')
self.assertArrayEqual(g.axes.shape, [3, 2])
for ax in g.axes.flat:
self.assertTrue(ax.has_data())
def test_facetgrid_cmap(self):
# Regression test for GH592
data = (np.random.random(size=(20, 25, 12)) + np.linspace(-3, 3, 12))
d = DataArray(data, dims=['x', 'y', 'time'])
fg = d.plot.pcolormesh(col='time')
# check that all color limits are the same
self.assertTrue(len(set(m.get_clim() for m in fg._mappables)) == 1)
# check that all colormaps are the same
self.assertTrue(len(set(m.get_cmap().name for m in fg._mappables)) == 1)
def test_datetime_dimension(self):
nrow = 3
ncol = 4
time = pd.date_range('2000-01-01', periods=nrow)
a = DataArray(easy_array((nrow, ncol)),
coords=[('time', time), ('y', range(ncol))])
a.plot()
ax = plt.gca()
self.assertTrue(ax.has_data())
def setUp(self):
a = easy_array((10, 15, 3, 2))
darray = DataArray(a, dims=['y', 'x', 'col', 'row'])
darray.coords['col'] = np.array(
['col' + str(x) for x in darray.coords['col'].values])
darray.coords['row'] = np.array(
['row' + str(x) for x in darray.coords['row'].values])
self.darray = darray
