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_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_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))