def test_select_dropped_dimensions_restoration(self):
d = np.random.randn(3, 8)
da = xr.DataArray(d, name='stuff', dims=['chain', 'value'],
coords=dict(chain=range(d.shape[0]), value=range(d.shape[1])))
ds = Dataset(da)
t = ds.select(chain=0)
self.assertEqual(t.data.dims , dict(chain=1,value=8))
self.assertEqual(t.data.stuff.shape , (1,8))
def test_select_dropped_dimensions_restoration(self):
d = np.random.randn(3, 8)
da = xr.DataArray(d, name='stuff', dims=['chain', 'value'],
coords=dict(chain=range(d.shape[0]), value=range(d.shape[1])))
ds = Dataset(da)
t = ds.select(chain=0)
self.assertEqual(t.data.dims , dict(chain=1,value=8))
self.assertEqual(t.data.stuff.shape , (1,8))
def test_dataset_groupby_multiple_dims(self):
dataset = Dataset((range(8), range(8), range(8), range(8),
da.from_array(np.random.rand(8, 8, 8, 8), 4)),
kdims=['a', 'b', 'c', 'd'], vdims=['Value'])
grouped = dataset.groupby(['c', 'd'])
keys = list(product(range(8), range(8)))
self.assertEqual(list(grouped.keys()), keys)
for c, d in keys:
self.assertEqual(grouped[c, d], dataset.select(c=c, d=d).reindex(['a', 'b']))
def test_dataset_groupby_multiple_dims(self):
dataset = Dataset((range(8), range(8), range(8), range(8),
da.from_array(np.random.rand(8, 8, 8, 8), 4)),
kdims=['a', 'b', 'c', 'd'], vdims=['Value'])
grouped = dataset.groupby(['c', 'd'])
keys = list(product(range(8), range(8)))
self.assertEqual(list(grouped.keys()), keys)
for c, d in keys:
self.assertEqual(grouped[c, d], dataset.select(c=c, d=d).reindex(['a', 'b']))
def test_select_expression_lazy(self):
df = pd.DataFrame({
'a': [1, 2, 3, 4, 5],
'b': [10, 10, 11, 11, 10],
})
ddf = dd.from_pandas(df, npartitions=2)
ds = Dataset(ddf)
new_ds = ds.select(selection_expr=dim('b') == 10)
# Make sure that selecting by expression didn't cause evaluation
self.assertIsInstance(new_ds.data, dd.DataFrame)
self.assertEqual(new_ds.data.compute(), df[df.b == 10])
def test_select_multi_index(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude=0, latitude=0), 5)
def test_select_set(self):
cube = Dataset(self.cube)
self.assertEqual(
cube.select(longitude={0, 1}).data.data,
np.array([[1, 2], [5, 6], [9, 10]], dtype=np.int32))
def test_select_index(self):
cube = Dataset(self.cube)
self.assertEqual(
cube.select(longitude=0).data.data,
np.array([[1, 5, 9]], dtype=np.int32))
def test_select_lazy(self):
import dask.array as da
arr = da.from_array(np.arange(1, 12), 3)
ds = Dataset({'x': range(11), 'y': arr}, 'x', 'y')
self.assertIsInstance(ds.select(x=(0, 5)).data['y'], da.Array)
def test_select_multi_slice2(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude={0, 2},
latitude={0, 2}).data.data,
np.array([[5, 7]], dtype=np.int32))
def test_select_lazy(self):
import dask.array as da
arr = da.from_array(np.arange(1, 12), 3)
ds = Dataset({'x': range(11), 'y': arr}, 'x', 'y')
self.assertIsInstance(ds.select(x=(0, 5)).data['y'], da.Array)
def test_select_multi_index(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude=0, latitude=0), 5)
def test_select_set(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude={0, 1}).data.data,
np.array([[1, 2], [5, 6], [9, 10]], dtype=np.int32))
def test_select_index(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude=0).data.data,
np.array([[1, 5, 9]], dtype=np.int32))
def test_select_multi_slice1(self):
cube = Dataset(self.cube)
self.assertEqual(
cube.select(longitude=(0, 1.01), latitude=(0, 1.01)).data.data,
np.array([[5, 6], [9, 10]], dtype=np.int32))
def test_select_multi_slice2(self):
cube = Dataset(self.cube)
self.assertEqual(
cube.select(longitude={0, 2}, latitude={0, 2}).data.data,
np.array([[5, 7]], dtype=np.int32))
def test_select_multi_slice1(self):
cube = Dataset(self.cube)
self.assertEqual(cube.select(longitude=(0, 1+self.epsilon),
latitude=(0, 1+self.epsilon)).data.data,
np.array([[5, 6], [9, 10]], dtype=np.int32))
