def test_groupby(self):
data = Dataset({'x': ('x', list('abc')),
'c': ('x', [0, 1, 0]),
'z': (['x', 'y'], np.random.randn(3, 5))})
groupby = data.groupby('x')
self.assertEqual(len(groupby), 3)
expected_groups = {'a': 0, 'b': 1, 'c': 2}
self.assertEqual(groupby.groups, expected_groups)
expected_items = [('a', data.indexed(x=0)),
('b', data.indexed(x=1)),
('c', data.indexed(x=2))]
self.assertEqual(list(groupby), expected_items)
identity = lambda x: x
for k in ['x', 'c', 'y']:
actual = data.groupby(k, squeeze=False).apply(identity)
self.assertEqual(data, actual)
data = create_test_data()
for n, (t, sub) in enumerate(list(data.groupby('dim1'))[:3]):
self.assertEqual(data['dim1'][n], t)
self.assertVariableEqual(data['var1'][n], sub['var1'])
self.assertVariableEqual(data['var2'][n], sub['var2'])
self.assertVariableEqual(data['var3'][:, n], sub['var3'])
# TODO: test the other edge cases
with self.assertRaisesRegexp(ValueError, 'must be 1 dimensional'):
data.groupby('var1')
with self.assertRaisesRegexp(ValueError, 'length does not match'):
data.groupby(data['dim1'][:3])
def test_groupby(self):
data = Dataset({
'x': ('x', list('abc')),
'c': ('x', [0, 1, 0]),
'z': (['x', 'y'], np.random.randn(3, 5))
})
groupby = data.groupby('x')
self.assertEqual(len(groupby), 3)
expected_groups = {'a': 0, 'b': 1, 'c': 2}
self.assertEqual(groupby.groups, expected_groups)
expected_items = [('a', data.indexed(x=0)), ('b', data.indexed(x=1)),
('c', data.indexed(x=2))]
self.assertEqual(list(groupby), expected_items)
identity = lambda x: x
for k in ['x', 'c', 'y']:
actual = data.groupby(k, squeeze=False).apply(identity)
self.assertEqual(data, actual)
data = create_test_data()
for n, (t, sub) in enumerate(list(data.groupby('dim1'))[:3]):
self.assertEqual(data['dim1'][n], t)
self.assertVariableEqual(data['var1'][n], sub['var1'])
self.assertVariableEqual(data['var2'][n], sub['var2'])
self.assertVariableEqual(data['var3'][:, n], sub['var3'])
# TODO: test the other edge cases
with self.assertRaisesRegexp(ValueError, 'must be 1 dimensional'):
data.groupby('var1')
with self.assertRaisesRegexp(ValueError, 'length does not match'):
data.groupby(data['dim1'][:3])
class TestDataArray(TestCase):
def setUp(self):
self.attrs = {'attr1': 'value1', 'attr2': 2929}
self.x = np.random.random((10, 20))
self.v = Variable(['x', 'y'], self.x)
self.va = Variable(['x', 'y'], self.x, self.attrs)
self.ds = Dataset({'foo': self.v})
self.dv = self.ds['foo']
def test_repr(self):
v = Variable(['time', 'x'], [[1, 2, 3], [4, 5, 6]], {'foo': 'bar'})
data_array = Dataset({
'my_variable': v,
'other': ([], 0)
})['my_variable']
expected = dedent("""
<xray.DataArray 'my_variable' (time: 2, x: 3)>
array([[1, 2, 3],
[4, 5, 6]])
Coordinates:
time: Int64Index([0, 1], dtype='int64')
x: Int64Index([0, 1, 2], dtype='int64')
Linked dataset variables:
other
Attributes:
foo: bar
""").strip()
self.assertEqual(expected, repr(data_array))
def test_properties(self):
self.assertDatasetIdentical(self.dv.dataset, self.ds)
self.assertEqual(self.dv.name, 'foo')
self.assertVariableEqual(self.dv.variable, self.v)
self.assertArrayEqual(self.dv.values, self.v.values)
for attr in ['dimensions', 'dtype', 'shape', 'size', 'ndim', 'attrs']:
self.assertEqual(getattr(self.dv, attr), getattr(self.v, attr))
self.assertEqual(len(self.dv), len(self.v))
self.assertVariableEqual(self.dv, self.v)
self.assertEqual(list(self.dv.coordinates), list(self.ds.coordinates))
for k, v in iteritems(self.dv.coordinates):
self.assertArrayEqual(v, self.ds.coordinates[k])
with self.assertRaises(AttributeError):
self.dv.name = 'bar'
with self.assertRaises(AttributeError):
self.dv.dataset = self.ds
self.assertIsInstance(self.ds['x'].as_index, pd.Index)
with self.assertRaisesRegexp(ValueError, 'must be 1-dimensional'):
self.ds['foo'].as_index
def test_constructor(self):
data = np.random.random((2, 3))
actual = DataArray(data)
expected = Dataset({None: (['dim_0', 'dim_1'], data)})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, [['a', 'b'], [-1, -2, -3]])
expected = Dataset({
None: (['dim_0', 'dim_1'], data),
'dim_0': ('dim_0', ['a', 'b']),
'dim_1': ('dim_1', [-1, -2, -3])
})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(
data,
[pd.Index(['a', 'b'], name='x'),
pd.Index([-1, -2, -3], name='y')])
expected = Dataset({
None: (['x', 'y'], data),
'x': ('x', ['a', 'b']),
'y': ('y', [-1, -2, -3])
})[None]
self.assertDataArrayIdentical(expected, actual)
indexes = [['a', 'b'], [-1, -2, -3]]
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = [
pd.Index(['a', 'b'], name='A'),
pd.Index([-1, -2, -3], name='B')
]
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = {'x': ['a', 'b'], 'y': [-1, -2, -3]}
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = OrderedDict([('x', ['a', 'b']), ('y', [-1, -2, -3])])
actual = DataArray(data, indexes)
self.assertDataArrayIdentical(expected, actual)
expected = Dataset({
None: (['x', 'y'], data),
'x': ('x', ['a', 'b'])
})[None]
actual = DataArray(data, {'x': ['a', 'b']}, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
with self.assertRaisesRegexp(ValueError, 'but data has ndim'):
DataArray(data, [[0, 1, 2]], ['x', 'y'])
with self.assertRaisesRegexp(ValueError, 'not array dimensions'):
DataArray(data, {'x': [0, 1, 2]}, ['a', 'b'])
with self.assertRaisesRegexp(ValueError, 'must have the same length'):
DataArray(data, {'x': [0, 1, 2]})
actual = DataArray(data, dimensions=['x', 'y'])
expected = Dataset({None: (['x', 'y'], data)})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], name='foo')
expected = Dataset({'foo': (['x', 'y'], data)})['foo']
self.assertDataArrayIdentical(expected, actual)
with self.assertRaisesRegexp(TypeError, 'is not a string'):
DataArray(data, dimensions=['x', None])
actual = DataArray(data, name='foo')
expected = Dataset({'foo': (['dim_0', 'dim_1'], data)})['foo']
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], attributes={'bar': 2})
expected = Dataset({None: (['x', 'y'], data, {'bar': 2})})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], encoding={'bar': 2})
expected = Dataset({None: (['x', 'y'], data, {}, {'bar': 2})})[None]
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_equals_and_identical(self):
da2 = self.dv.copy()
self.assertTrue(self.dv.equals(da2))
self.assertTrue(self.dv.identical(da2))
da3 = self.dv.rename('baz')
self.assertTrue(self.dv.equals(da3))
self.assertFalse(self.dv.identical(da3))
da4 = self.dv.rename({'x': 'xxx'})
self.assertFalse(self.dv.equals(da4))
self.assertFalse(self.dv.identical(da4))
da5 = self.dv.copy()
da5.attrs['foo'] = 'bar'
self.assertTrue(self.dv.equals(da5))
self.assertFalse(self.dv.identical(da5))
da6 = self.dv.copy()
da6['x'] = ('x', -np.arange(10))
self.assertFalse(self.dv.equals(da6))
self.assertFalse(self.dv.identical(da6))
da2[0, 0] = np.nan
self.dv[0, 0] = np.nan
self.assertTrue(self.dv.equals(da2))
self.assertTrue(self.dv.identical(da2))
da2[:] = np.nan
self.assertFalse(self.dv.equals(da2))
self.assertFalse(self.dv.identical(da2))
def test_items(self):
# strings pull out dataarrays
self.assertDataArrayIdentical(self.dv, self.ds['foo'])
x = self.dv['x']
y = self.dv['y']
self.assertDataArrayIdentical(self.ds['x'], x)
self.assertDataArrayIdentical(self.ds['y'], y)
# integer indexing
I = ReturnItem()
for i in [
I[:], I[...], I[x.values], I[x.variable], I[x], I[x, y],
I[x.values > -1], I[x.variable > -1], I[x > -1], I[x > -1,
y > -1]
]:
self.assertVariableEqual(self.dv, self.dv[i])
for i in [
I[0], I[:, 0], I[:3, :2], I[x.values[:3]], I[x.variable[:3]],
I[x[:3]], I[x[:3], y[:4]], I[x.values > 3], I[x.variable > 3],
I[x > 3], I[x > 3, y > 3]
]:
self.assertVariableEqual(self.v[i], self.dv[i])
# make sure we always keep the array around, even if it's a scalar
self.assertVariableEqual(self.dv[0, 0], self.dv.variable[0, 0])
for k in ['x', 'y', 'foo']:
self.assertIn(k, self.dv[0, 0].dataset)
def test_indexed(self):
self.assertEqual(self.dv[0].dataset, self.ds.indexed(x=0))
self.assertEqual(self.dv[:3, :5].dataset,
self.ds.indexed(x=slice(3), y=slice(5)))
self.assertDataArrayIdentical(self.dv, self.dv.indexed(x=slice(None)))
self.assertDataArrayIdentical(self.dv[:3], self.dv.indexed(x=slice(3)))
def test_labeled(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
da = self.ds['foo']
self.assertDataArrayIdentical(da, da.labeled(x=slice(None)))
self.assertDataArrayIdentical(da[1], da.labeled(x='b'))
self.assertDataArrayIdentical(da[:3], da.labeled(x=slice('c')))
def test_loc(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
da = self.ds['foo']
self.assertDataArrayIdentical(da[:3], da.loc[:'c'])
self.assertDataArrayIdentical(da[1], da.loc['b'])
self.assertDataArrayIdentical(da[:3], da.loc[['a', 'b', 'c']])
self.assertDataArrayIdentical(da[:3, :4], da.loc[['a', 'b', 'c'],
np.arange(4)])
da.loc['a':'j'] = 0
self.assertTrue(np.all(da.values == 0))
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)
def test_rename(self):
renamed = self.dv.rename('bar')
self.assertEqual(renamed.dataset, self.ds.rename({'foo': 'bar'}))
self.assertEqual(renamed.name, 'bar')
renamed = self.dv.rename({'foo': 'bar'})
self.assertEqual(renamed.dataset, self.ds.rename({'foo': 'bar'}))
self.assertEqual(renamed.name, 'bar')
def test_dataset_getitem(self):
dv = self.ds['foo']
self.assertDataArrayIdentical(dv, self.dv)
def test_array_interface(self):
self.assertArrayEqual(np.asarray(self.dv), self.x)
# test patched in methods
self.assertArrayEqual(self.dv.astype(float), self.v.astype(float))
self.assertVariableEqual(self.dv.argsort(), self.v.argsort())
self.assertVariableEqual(self.dv.clip(2, 3), self.v.clip(2, 3))
# test ufuncs
expected = deepcopy(self.ds)
expected['foo'][:] = np.sin(self.x)
self.assertDataArrayEqual(expected['foo'], np.sin(self.dv))
self.assertDataArrayEqual(self.dv, np.maximum(self.v, self.dv))
bar = Variable(['x', 'y'], np.zeros((10, 20)))
self.assertDataArrayEqual(self.dv, np.maximum(self.dv, bar))
def test_math(self):
x = self.x
v = self.v
a = self.dv
# variable math was already tested extensively, so let's just make sure
# that all types are properly converted here
self.assertDataArrayEqual(a, +a)
self.assertDataArrayEqual(a, a + 0)
self.assertDataArrayEqual(a, 0 + a)
self.assertDataArrayEqual(a, a + 0 * v)
self.assertDataArrayEqual(a, 0 * v + a)
self.assertDataArrayEqual(a, a + 0 * x)
self.assertDataArrayEqual(a, 0 * x + a)
self.assertDataArrayEqual(a, a + 0 * a)
self.assertDataArrayEqual(a, 0 * a + a)
# test different indices
ds2 = self.ds.update({'x': ('x', 3 + np.arange(10))}, inplace=False)
b = ds2['foo']
with self.assertRaisesRegexp(ValueError, 'not aligned'):
a + b
with self.assertRaisesRegexp(ValueError, 'not aligned'):
b + a
with self.assertRaisesRegexp(TypeError, 'datasets do not support'):
a + a.dataset
def test_dataset_math(self):
# verify that mathematical operators keep around the expected variables
# when doing math with dataset arrays from one or more aligned datasets
obs = Dataset({
'tmin': ('x', np.arange(5)),
'tmax': ('x', 10 + np.arange(5)),
'x': ('x', 0.5 * np.arange(5))
})
actual = 2 * obs['tmax']
expected = Dataset({
'tmax2': ('x', 2 * (10 + np.arange(5))),
'x': obs['x']
})['tmax2']
self.assertDataArrayEqual(actual, expected)
actual = obs['tmax'] - obs['tmin']
expected = Dataset({
'trange': ('x', 10 * np.ones(5)),
'x': obs['x']
})['trange']
self.assertDataArrayEqual(actual, expected)
sim = Dataset({
'tmin': ('x', 1 + np.arange(5)),
'tmax': ('x', 11 + np.arange(5)),
'x': ('x', 0.5 * np.arange(5))
})
actual = sim['tmin'] - obs['tmin']
expected = Dataset({
'error': ('x', np.ones(5)),
'x': obs['x']
})['error']
self.assertDataArrayEqual(actual, expected)
# in place math shouldn't remove or conflict with other variables
actual = deepcopy(sim['tmin'])
actual -= obs['tmin']
expected = Dataset({
'tmin': ('x', np.ones(5)),
'tmax': sim['tmax'],
'x': sim['x']
})['tmin']
self.assertDataArrayEqual(actual, expected)
def test_math_name(self):
# Verify that name is preserved only when it can be done unambiguously.
# The rule (copied from pandas.Series) is keep the current name only if
# the other object has no name attribute and this object isn't a
# coordinate; otherwise reset to None.
ds = self.ds
a = self.dv
self.assertEqual((+a).name, 'foo')
self.assertEqual((a + 0).name, 'foo')
self.assertIs((a + a.rename(None)).name, None)
self.assertIs((a + a).name, None)
self.assertIs((+ds['x']).name, None)
self.assertIs((ds['x'] + 0).name, None)
self.assertIs((a + ds['x']).name, None)
def test_coord_math(self):
ds = Dataset({'x': ('x', 1 + np.arange(3))})
expected = ds.copy()
expected['x2'] = ('x', np.arange(3))
actual = ds['x'] - 1
self.assertDataArrayEqual(expected['x2'], actual)
def test_item_math(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
self.assertVariableEqual(self.dv + self.dv[0, 0],
self.dv + self.dv[0, 0].values)
new_data = self.x[0][None, :] + self.x[:, 0][:, None]
self.assertVariableEqual(self.dv[:, 0] + self.dv[0],
Variable(['x', 'y'], new_data))
self.assertVariableEqual(self.dv[0] + self.dv[:, 0],
Variable(['y', 'x'], new_data.T))
def test_inplace_math(self):
x = self.x
v = self.v
a = self.dv
b = a
b += 1
self.assertIs(b, a)
self.assertIs(b.variable, v)
self.assertArrayEqual(b.values, x)
self.assertIs(source_ndarray(b.values), x)
self.assertDatasetIdentical(b.dataset, self.ds)
def test_transpose(self):
self.assertVariableEqual(self.dv.variable.transpose(),
self.dv.transpose())
def test_squeeze(self):
self.assertVariableEqual(self.dv.variable.squeeze(), self.dv.squeeze())
def test_reduce(self):
self.assertVariableEqual(self.dv.reduce(np.mean, 'x'),
self.v.reduce(np.mean, 'x'))
# needs more...
# should check which extra dimensions are dropped
def test_reduce_keep_attrs(self):
# Test dropped attrs
vm = self.va.mean()
self.assertEqual(len(vm.attrs), 0)
self.assertEqual(vm.attrs, OrderedDict())
# Test kept attrs
vm = self.va.mean(keep_attrs=True)
self.assertEqual(len(vm.attrs), len(self.attrs))
self.assertEqual(vm.attrs, self.attrs)
def test_unselect(self):
with self.assertRaisesRegexp(ValueError, 'cannot unselect the name'):
self.dv.unselect('foo')
with self.assertRaisesRegexp(ValueError, 'must be a variable in'):
self.dv.unselect('y')
def test_groupby_iter(self):
for ((act_x, act_dv), (exp_x, exp_ds)) in \
zip(self.dv.groupby('y'), self.ds.groupby('y')):
self.assertEqual(exp_x, act_x)
self.assertDataArrayIdentical(exp_ds['foo'], act_dv)
for ((_, exp_dv), act_dv) in zip(self.dv.groupby('x'), self.dv):
self.assertDataArrayIdentical(exp_dv, act_dv)
def test_groupby(self):
agg_var = Variable(['y'], np.array(['a'] * 9 + ['c'] + ['b'] * 10))
self.dv['abc'] = agg_var
self.dv['y'] = 20 + 100 * self.ds['y'].variable
identity = lambda x: x
for g in ['x', 'y', 'abc']:
for shortcut in [False, True]:
for squeeze in [False, True]:
expected = self.dv
grouped = self.dv.groupby(g, squeeze=squeeze)
actual = grouped.apply(identity, shortcut=shortcut)
self.assertDataArrayIdentical(expected, actual)
grouped = self.dv.groupby('abc', squeeze=True)
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())
self.assertDataArrayAllClose(expected_sum_all, grouped.sum())
expected_unique = Variable('abc', ['a', 'b', 'c'])
self.assertVariableEqual(expected_unique, grouped.unique_coord)
self.assertEqual(3, len(grouped))
grouped = self.dv.groupby('abc', squeeze=False)
self.assertDataArrayAllClose(expected_sum_all,
grouped.sum(dimension=None))
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.variables['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 center(x):
return x - np.mean(x)
expected_ds = self.dv.dataset.copy()
exp_data = np.hstack([
center(self.x[:, :9]),
center(self.x[:, 9:10]),
center(self.x[:, 10:])
])
expected_ds['foo'] = (['x', 'y'], exp_data)
expected_centered = expected_ds['foo']
self.assertDataArrayAllClose(expected_centered, grouped.apply(center))
def test_concat(self):
self.ds['bar'] = Variable(['x', 'y'], np.random.randn(10, 20))
foo = self.ds['foo'].select()
bar = self.ds['bar'].rename('foo').select()
# from dataset array:
self.assertVariableEqual(
Variable(['w', 'x', 'y'], np.array([foo.values, bar.values])),
DataArray.concat([foo, bar], 'w'))
# from iteration:
grouped = [g for _, g in foo.groupby('x')]
stacked = DataArray.concat(grouped, self.ds['x'])
self.assertDataArrayIdentical(foo.select(), stacked)
def test_align(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
with self.assertRaises(ValueError):
self.dv + self.dv[:5]
dv1, dv2 = align(self.dv, self.dv[:5], join='inner')
self.assertDataArrayIdentical(dv1, self.dv[:5])
self.assertDataArrayIdentical(dv2, self.dv[:5])
def test_to_and_from_series(self):
expected = self.dv.to_dataframe()['foo']
actual = self.dv.to_series()
self.assertArrayEqual(expected.values, actual.values)
self.assertArrayEqual(expected.index.values, actual.index.values)
self.assertEqual('foo', actual.name)
# test roundtrip
self.assertDataArrayIdentical(self.dv, DataArray.from_series(actual))
# test name is None
actual.name = None
expected_da = self.dv.rename(None)
self.assertDataArrayIdentical(expected_da,
DataArray.from_series(actual))
class TestDataArray(TestCase):
def setUp(self):
self.attrs = {'attr1': 'value1', 'attr2': 2929}
self.x = np.random.random((10, 20))
self.v = Variable(['x', 'y'], self.x)
self.va = Variable(['x', 'y'], self.x, self.attrs)
self.ds = Dataset({'foo': self.v})
self.dv = self.ds['foo']
def test_repr(self):
v = Variable(['time', 'x'], [[1, 2, 3], [4, 5, 6]], {'foo': 'bar'})
data_array = Dataset({'my_variable': v, 'other': ([], 0)}
)['my_variable']
expected = dedent("""
<xray.DataArray 'my_variable' (time: 2, x: 3)>
array([[1, 2, 3],
[4, 5, 6]])
Coordinates:
time: Int64Index([0, 1], dtype='int64')
x: Int64Index([0, 1, 2], dtype='int64')
Linked dataset variables:
other
Attributes:
foo: bar
""").strip()
self.assertEqual(expected, repr(data_array))
def test_properties(self):
self.assertDatasetIdentical(self.dv.dataset, self.ds)
self.assertEqual(self.dv.name, 'foo')
self.assertVariableEqual(self.dv.variable, self.v)
self.assertArrayEqual(self.dv.values, self.v.values)
for attr in ['dimensions', 'dtype', 'shape', 'size', 'ndim', 'attrs']:
self.assertEqual(getattr(self.dv, attr), getattr(self.v, attr))
self.assertEqual(len(self.dv), len(self.v))
self.assertVariableEqual(self.dv, self.v)
self.assertEqual(list(self.dv.coordinates), list(self.ds.coordinates))
for k, v in iteritems(self.dv.coordinates):
self.assertArrayEqual(v, self.ds.coordinates[k])
with self.assertRaises(AttributeError):
self.dv.name = 'bar'
with self.assertRaises(AttributeError):
self.dv.dataset = self.ds
self.assertIsInstance(self.ds['x'].as_index, pd.Index)
with self.assertRaisesRegexp(ValueError, 'must be 1-dimensional'):
self.ds['foo'].as_index
def test_constructor(self):
data = np.random.random((2, 3))
actual = DataArray(data)
expected = Dataset({None: (['dim_0', 'dim_1'], data)})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, [['a', 'b'], [-1, -2, -3]])
expected = Dataset({None: (['dim_0', 'dim_1'], data),
'dim_0': ('dim_0', ['a', 'b']),
'dim_1': ('dim_1', [-1, -2, -3])})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, [pd.Index(['a', 'b'], name='x'),
pd.Index([-1, -2, -3], name='y')])
expected = Dataset({None: (['x', 'y'], data),
'x': ('x', ['a', 'b']),
'y': ('y', [-1, -2, -3])})[None]
self.assertDataArrayIdentical(expected, actual)
indexes = [['a', 'b'], [-1, -2, -3]]
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = [pd.Index(['a', 'b'], name='A'),
pd.Index([-1, -2, -3], name='B')]
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = {'x': ['a', 'b'], 'y': [-1, -2, -3]}
actual = DataArray(data, indexes, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
indexes = OrderedDict([('x', ['a', 'b']), ('y', [-1, -2, -3])])
actual = DataArray(data, indexes)
self.assertDataArrayIdentical(expected, actual)
expected = Dataset({None: (['x', 'y'], data),
'x': ('x', ['a', 'b'])})[None]
actual = DataArray(data, {'x': ['a', 'b']}, ['x', 'y'])
self.assertDataArrayIdentical(expected, actual)
with self.assertRaisesRegexp(ValueError, 'but data has ndim'):
DataArray(data, [[0, 1, 2]], ['x', 'y'])
with self.assertRaisesRegexp(ValueError, 'not array dimensions'):
DataArray(data, {'x': [0, 1, 2]}, ['a', 'b'])
with self.assertRaisesRegexp(ValueError, 'must have the same length'):
DataArray(data, {'x': [0, 1, 2]})
actual = DataArray(data, dimensions=['x', 'y'])
expected = Dataset({None: (['x', 'y'], data)})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], name='foo')
expected = Dataset({'foo': (['x', 'y'], data)})['foo']
self.assertDataArrayIdentical(expected, actual)
with self.assertRaisesRegexp(TypeError, 'is not a string'):
DataArray(data, dimensions=['x', None])
actual = DataArray(data, name='foo')
expected = Dataset({'foo': (['dim_0', 'dim_1'], data)})['foo']
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], attributes={'bar': 2})
expected = Dataset({None: (['x', 'y'], data, {'bar': 2})})[None]
self.assertDataArrayIdentical(expected, actual)
actual = DataArray(data, dimensions=['x', 'y'], encoding={'bar': 2})
expected = Dataset({None: (['x', 'y'], data, {}, {'bar': 2})})[None]
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_equals_and_identical(self):
da2 = self.dv.copy()
self.assertTrue(self.dv.equals(da2))
self.assertTrue(self.dv.identical(da2))
da3 = self.dv.rename('baz')
self.assertTrue(self.dv.equals(da3))
self.assertFalse(self.dv.identical(da3))
da4 = self.dv.rename({'x': 'xxx'})
self.assertFalse(self.dv.equals(da4))
self.assertFalse(self.dv.identical(da4))
da5 = self.dv.copy()
da5.attrs['foo'] = 'bar'
self.assertTrue(self.dv.equals(da5))
self.assertFalse(self.dv.identical(da5))
da6 = self.dv.copy()
da6['x'] = ('x', -np.arange(10))
self.assertFalse(self.dv.equals(da6))
self.assertFalse(self.dv.identical(da6))
da2[0, 0] = np.nan
self.dv[0, 0] = np.nan
self.assertTrue(self.dv.equals(da2))
self.assertTrue(self.dv.identical(da2))
da2[:] = np.nan
self.assertFalse(self.dv.equals(da2))
self.assertFalse(self.dv.identical(da2))
def test_items(self):
# strings pull out dataarrays
self.assertDataArrayIdentical(self.dv, self.ds['foo'])
x = self.dv['x']
y = self.dv['y']
self.assertDataArrayIdentical(self.ds['x'], x)
self.assertDataArrayIdentical(self.ds['y'], y)
# integer indexing
I = ReturnItem()
for i in [I[:], I[...], I[x.values], I[x.variable], I[x], I[x, y],
I[x.values > -1], I[x.variable > -1], I[x > -1],
I[x > -1, y > -1]]:
self.assertVariableEqual(self.dv, self.dv[i])
for i in [I[0], I[:, 0], I[:3, :2],
I[x.values[:3]], I[x.variable[:3]], I[x[:3]], I[x[:3], y[:4]],
I[x.values > 3], I[x.variable > 3], I[x > 3], I[x > 3, y > 3]]:
self.assertVariableEqual(self.v[i], self.dv[i])
# make sure we always keep the array around, even if it's a scalar
self.assertVariableEqual(self.dv[0, 0], self.dv.variable[0, 0])
for k in ['x', 'y', 'foo']:
self.assertIn(k, self.dv[0, 0].dataset)
def test_indexed(self):
self.assertEqual(self.dv[0].dataset, self.ds.indexed(x=0))
self.assertEqual(self.dv[:3, :5].dataset,
self.ds.indexed(x=slice(3), y=slice(5)))
self.assertDataArrayIdentical(self.dv, self.dv.indexed(x=slice(None)))
self.assertDataArrayIdentical(self.dv[:3], self.dv.indexed(x=slice(3)))
def test_labeled(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
da = self.ds['foo']
self.assertDataArrayIdentical(da, da.labeled(x=slice(None)))
self.assertDataArrayIdentical(da[1], da.labeled(x='b'))
self.assertDataArrayIdentical(da[:3], da.labeled(x=slice('c')))
def test_loc(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
da = self.ds['foo']
self.assertDataArrayIdentical(da[:3], da.loc[:'c'])
self.assertDataArrayIdentical(da[1], da.loc['b'])
self.assertDataArrayIdentical(da[:3], da.loc[['a', 'b', 'c']])
self.assertDataArrayIdentical(da[:3, :4],
da.loc[['a', 'b', 'c'], np.arange(4)])
da.loc['a':'j'] = 0
self.assertTrue(np.all(da.values == 0))
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)
def test_rename(self):
renamed = self.dv.rename('bar')
self.assertEqual(renamed.dataset, self.ds.rename({'foo': 'bar'}))
self.assertEqual(renamed.name, 'bar')
renamed = self.dv.rename({'foo': 'bar'})
self.assertEqual(renamed.dataset, self.ds.rename({'foo': 'bar'}))
self.assertEqual(renamed.name, 'bar')
def test_dataset_getitem(self):
dv = self.ds['foo']
self.assertDataArrayIdentical(dv, self.dv)
def test_array_interface(self):
self.assertArrayEqual(np.asarray(self.dv), self.x)
# test patched in methods
self.assertArrayEqual(self.dv.astype(float), self.v.astype(float))
self.assertVariableEqual(self.dv.argsort(), self.v.argsort())
self.assertVariableEqual(self.dv.clip(2, 3), self.v.clip(2, 3))
# test ufuncs
expected = deepcopy(self.ds)
expected['foo'][:] = np.sin(self.x)
self.assertDataArrayEqual(expected['foo'], np.sin(self.dv))
self.assertDataArrayEqual(self.dv, np.maximum(self.v, self.dv))
bar = Variable(['x', 'y'], np.zeros((10, 20)))
self.assertDataArrayEqual(self.dv, np.maximum(self.dv, bar))
def test_math(self):
x = self.x
v = self.v
a = self.dv
# variable math was already tested extensively, so let's just make sure
# that all types are properly converted here
self.assertDataArrayEqual(a, +a)
self.assertDataArrayEqual(a, a + 0)
self.assertDataArrayEqual(a, 0 + a)
self.assertDataArrayEqual(a, a + 0 * v)
self.assertDataArrayEqual(a, 0 * v + a)
self.assertDataArrayEqual(a, a + 0 * x)
self.assertDataArrayEqual(a, 0 * x + a)
self.assertDataArrayEqual(a, a + 0 * a)
self.assertDataArrayEqual(a, 0 * a + a)
# test different indices
ds2 = self.ds.update({'x': ('x', 3 + np.arange(10))}, inplace=False)
b = ds2['foo']
with self.assertRaisesRegexp(ValueError, 'not aligned'):
a + b
with self.assertRaisesRegexp(ValueError, 'not aligned'):
b + a
with self.assertRaisesRegexp(TypeError, 'datasets do not support'):
a + a.dataset
def test_dataset_math(self):
# verify that mathematical operators keep around the expected variables
# when doing math with dataset arrays from one or more aligned datasets
obs = Dataset({'tmin': ('x', np.arange(5)),
'tmax': ('x', 10 + np.arange(5)),
'x': ('x', 0.5 * np.arange(5))})
actual = 2 * obs['tmax']
expected = Dataset({'tmax2': ('x', 2 * (10 + np.arange(5))),
'x': obs['x']})['tmax2']
self.assertDataArrayEqual(actual, expected)
actual = obs['tmax'] - obs['tmin']
expected = Dataset({'trange': ('x', 10 * np.ones(5)),
'x': obs['x']})['trange']
self.assertDataArrayEqual(actual, expected)
sim = Dataset({'tmin': ('x', 1 + np.arange(5)),
'tmax': ('x', 11 + np.arange(5)),
'x': ('x', 0.5 * np.arange(5))})
actual = sim['tmin'] - obs['tmin']
expected = Dataset({'error': ('x', np.ones(5)),
'x': obs['x']})['error']
self.assertDataArrayEqual(actual, expected)
# in place math shouldn't remove or conflict with other variables
actual = deepcopy(sim['tmin'])
actual -= obs['tmin']
expected = Dataset({'tmin': ('x', np.ones(5)),
'tmax': sim['tmax'],
'x': sim['x']})['tmin']
self.assertDataArrayEqual(actual, expected)
def test_math_name(self):
# Verify that name is preserved only when it can be done unambiguously.
# The rule (copied from pandas.Series) is keep the current name only if
# the other object has no name attribute and this object isn't a
# coordinate; otherwise reset to None.
ds = self.ds
a = self.dv
self.assertEqual((+a).name, 'foo')
self.assertEqual((a + 0).name, 'foo')
self.assertIs((a + a.rename(None)).name, None)
self.assertIs((a + a).name, None)
self.assertIs((+ds['x']).name, None)
self.assertIs((ds['x'] + 0).name, None)
self.assertIs((a + ds['x']).name, None)
def test_coord_math(self):
ds = Dataset({'x': ('x', 1 + np.arange(3))})
expected = ds.copy()
expected['x2'] = ('x', np.arange(3))
actual = ds['x'] - 1
self.assertDataArrayEqual(expected['x2'], actual)
def test_item_math(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
self.assertVariableEqual(self.dv + self.dv[0, 0],
self.dv + self.dv[0, 0].values)
new_data = self.x[0][None, :] + self.x[:, 0][:, None]
self.assertVariableEqual(self.dv[:, 0] + self.dv[0],
Variable(['x', 'y'], new_data))
self.assertVariableEqual(self.dv[0] + self.dv[:, 0],
Variable(['y', 'x'], new_data.T))
def test_inplace_math(self):
x = self.x
v = self.v
a = self.dv
b = a
b += 1
self.assertIs(b, a)
self.assertIs(b.variable, v)
self.assertArrayEqual(b.values, x)
self.assertIs(source_ndarray(b.values), x)
self.assertDatasetIdentical(b.dataset, self.ds)
def test_transpose(self):
self.assertVariableEqual(self.dv.variable.transpose(),
self.dv.transpose())
def test_squeeze(self):
self.assertVariableEqual(self.dv.variable.squeeze(), self.dv.squeeze())
def test_reduce(self):
self.assertVariableEqual(self.dv.reduce(np.mean, 'x'),
self.v.reduce(np.mean, 'x'))
# needs more...
# should check which extra dimensions are dropped
def test_reduce_keep_attrs(self):
# Test dropped attrs
vm = self.va.mean()
self.assertEqual(len(vm.attrs), 0)
self.assertEqual(vm.attrs, OrderedDict())
# Test kept attrs
vm = self.va.mean(keep_attrs=True)
self.assertEqual(len(vm.attrs), len(self.attrs))
self.assertEqual(vm.attrs, self.attrs)
def test_unselect(self):
with self.assertRaisesRegexp(ValueError, 'cannot unselect the name'):
self.dv.unselect('foo')
with self.assertRaisesRegexp(ValueError, 'must be a variable in'):
self.dv.unselect('y')
def test_groupby_iter(self):
for ((act_x, act_dv), (exp_x, exp_ds)) in \
zip(self.dv.groupby('y'), self.ds.groupby('y')):
self.assertEqual(exp_x, act_x)
self.assertDataArrayIdentical(exp_ds['foo'], act_dv)
for ((_, exp_dv), act_dv) in zip(self.dv.groupby('x'), self.dv):
self.assertDataArrayIdentical(exp_dv, act_dv)
def test_groupby(self):
agg_var = Variable(['y'], np.array(['a'] * 9 + ['c'] + ['b'] * 10))
self.dv['abc'] = agg_var
self.dv['y'] = 20 + 100 * self.ds['y'].variable
identity = lambda x: x
for g in ['x', 'y', 'abc']:
for shortcut in [False, True]:
for squeeze in [False, True]:
expected = self.dv
grouped = self.dv.groupby(g, squeeze=squeeze)
actual = grouped.apply(identity, shortcut=shortcut)
self.assertDataArrayIdentical(expected, actual)
grouped = self.dv.groupby('abc', squeeze=True)
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())
self.assertDataArrayAllClose(
expected_sum_all, grouped.sum())
expected_unique = Variable('abc', ['a', 'b', 'c'])
self.assertVariableEqual(expected_unique, grouped.unique_coord)
self.assertEqual(3, len(grouped))
grouped = self.dv.groupby('abc', squeeze=False)
self.assertDataArrayAllClose(
expected_sum_all, grouped.sum(dimension=None))
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.variables['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 center(x):
return x - np.mean(x)
expected_ds = self.dv.dataset.copy()
exp_data = np.hstack([center(self.x[:, :9]),
center(self.x[:, 9:10]),
center(self.x[:, 10:])])
expected_ds['foo'] = (['x', 'y'], exp_data)
expected_centered = expected_ds['foo']
self.assertDataArrayAllClose(expected_centered,
grouped.apply(center))
def test_concat(self):
self.ds['bar'] = Variable(['x', 'y'], np.random.randn(10, 20))
foo = self.ds['foo'].select()
bar = self.ds['bar'].rename('foo').select()
# from dataset array:
self.assertVariableEqual(Variable(['w', 'x', 'y'],
np.array([foo.values, bar.values])),
DataArray.concat([foo, bar], 'w'))
# from iteration:
grouped = [g for _, g in foo.groupby('x')]
stacked = DataArray.concat(grouped, self.ds['x'])
self.assertDataArrayIdentical(foo.select(), stacked)
def test_align(self):
self.ds['x'] = ('x', np.array(list('abcdefghij')))
with self.assertRaises(ValueError):
self.dv + self.dv[:5]
dv1, dv2 = align(self.dv, self.dv[:5], join='inner')
self.assertDataArrayIdentical(dv1, self.dv[:5])
self.assertDataArrayIdentical(dv2, self.dv[:5])
def test_to_and_from_series(self):
expected = self.dv.to_dataframe()['foo']
actual = self.dv.to_series()
self.assertArrayEqual(expected.values, actual.values)
self.assertArrayEqual(expected.index.values, actual.index.values)
self.assertEqual('foo', actual.name)
# test roundtrip
self.assertDataArrayIdentical(self.dv, DataArray.from_series(actual))
# test name is None
actual.name = None
expected_da = self.dv.rename(None)
self.assertDataArrayIdentical(expected_da,
DataArray.from_series(actual))
