def test_2d_bool(self):
arr = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=bool)
result = common.take_2d(arr, [0, 2, 2, 1])
expected = arr.take([0, 2, 2, 1], axis=0)
self.assert_(np.array_equal(result, expected))
result = common.take_2d(arr, [0, 2, 2, 1], axis=1)
expected = arr.take([0, 2, 2, 1], axis=1)
self.assert_(np.array_equal(result, expected))
result = common.take_2d(arr, [0, 2, -1])
self.assert_(result.dtype == np.object_)
def test_2d_bool(self):
arr = np.array([[0, 1, 0], [1, 0, 1], [0, 1, 1]], dtype=bool)
result = com.take_2d(arr, [0, 2, 2, 1])
expected = arr.take([0, 2, 2, 1], axis=0)
self.assert_(np.array_equal(result, expected))
result = com.take_2d(arr, [0, 2, 2, 1], axis=1)
expected = arr.take([0, 2, 2, 1], axis=1)
self.assert_(np.array_equal(result, expected))
result = com.take_2d(arr, [0, 2, -1])
self.assert_(result.dtype == np.object_)
def test_2d_float32(self):
arr = np.random.randn(4, 3).astype(np.float32)
indexer = [0, 2, -1, 1, -1]
# axis=0
result = com.take_2d(arr, indexer)
result2 = np.empty_like(result)
com.take_2d(arr, indexer, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=0)
expected[[2, 4]] = np.nan
tm.assert_almost_equal(result, expected)
# test with float64 out buffer
out = np.empty((len(indexer), arr.shape[1]), dtype='f8')
com.take_2d(arr, indexer, out=out) # it works!
# axis=1
result = com.take_2d(arr, indexer, axis=1)
result2 = np.empty_like(result)
com.take_2d(arr, indexer, axis=1, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=1)
expected[:, [2, 4]] = np.nan
tm.assert_almost_equal(result, expected)
def test_2d_float32(self):
arr = np.random.randn(4, 3).astype(np.float32)
indexer = [0, 2, -1, 1, -1]
# axis=0
result = com.take_2d(arr, indexer)
result2 = np.empty_like(result)
com.take_2d(arr, indexer, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=0)
expected[[2, 4]] = np.nan
tm.assert_almost_equal(result, expected)
# test with float64 out buffer
out = np.empty((len(indexer), arr.shape[1]), dtype='f8')
com.take_2d(arr, indexer, out=out) # it works!
# axis=1
result = com.take_2d(arr, indexer, axis=1)
result2 = np.empty_like(result)
com.take_2d(arr, indexer, axis=1, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=1)
expected[:, [2, 4]] = np.nan
tm.assert_almost_equal(result, expected)
def reindex_items_from(self, new_ref_items, copy=True):
"""
Reindex to only those items contained in the input set of items
E.g. if you have ['a', 'b'], and the input items is ['b', 'c', 'd'],
then the resulting items will be ['b']
Returns
-------
reindexed : Block
"""
new_ref_items, indexer = self.items.reindex(new_ref_items)
if indexer is None:
new_items = new_ref_items
new_values = self.values.copy() if copy else self.values
else:
mask = indexer != -1
masked_idx = indexer[mask]
if self.values.ndim == 2:
new_values = com.take_2d(self.values, masked_idx, axis=0,
needs_masking=False)
else:
new_values = self.values.take(masked_idx, axis=0)
new_items = self.items.take(masked_idx)
return make_block(new_values, new_items, new_ref_items)
def test_2d_other_dtypes(self):
arr = np.random.randn(10, 5).astype(np.float32)
indexer = [1, 2, 3, -1]
# axis=0
result = com.take_2d(arr, indexer, axis=0)
expected = arr.take(indexer, axis=0)
expected[-1] = np.nan
tm.assert_almost_equal(result, expected)
# axis=1
result = com.take_2d(arr, indexer, axis=1)
expected = arr.take(indexer, axis=1)
expected[:, -1] = np.nan
tm.assert_almost_equal(result, expected)
def test_2d_other_dtypes(self):
arr = np.random.randn(10, 5).astype(np.float32)
indexer = [1, 2, 3, -1]
# axis=0
result = com.take_2d(arr, indexer, axis=0)
expected = arr.take(indexer, axis=0)
expected[-1] = np.nan
tm.assert_almost_equal(result, expected)
# axis=1
result = com.take_2d(arr, indexer, axis=1)
expected = arr.take(indexer, axis=1)
expected[:, -1] = np.nan
tm.assert_almost_equal(result, expected)
def reindex_items_from(self, new_ref_items, copy=True):
"""
Reindex to only those items contained in the input set of items
E.g. if you have ['a', 'b'], and the input items is ['b', 'c', 'd'],
then the resulting items will be ['b']
Returns
-------
reindexed : Block
"""
new_ref_items, indexer = self.items.reindex(new_ref_items)
if indexer is None:
new_items = new_ref_items
new_values = self.values.copy() if copy else self.values
else:
mask = indexer != -1
masked_idx = indexer[mask]
if self.values.ndim == 2:
new_values = com.take_2d(self.values,
masked_idx,
axis=0,
needs_masking=False)
else:
new_values = self.values.take(masked_idx, axis=0)
new_items = self.items.take(masked_idx)
return make_block(new_values, new_items, new_ref_items)
def get_result(self):
# TODO: find a better way than this masking business
values, value_mask = self.get_new_values()
columns = self.get_new_columns()
index = self.get_new_index()
# filter out missing levels
if values.shape[1] > 0:
col_inds, obs_ids = _compress_group_index(self.sorted_labels[-1])
# rare case, level values not observed
if len(obs_ids) < self.full_shape[1]:
inds = (value_mask.sum(0) > 0).nonzero()[0]
values = com.take_2d(values, inds, axis=1)
columns = columns[inds]
return DataFrame(values, index=index, columns=columns)
def _make_sorted_values_labels(self):
v = self.level
labs = self.index.labels
levs = self.index.levels
to_sort = labs[:v] + labs[v + 1:] + [labs[v]]
sizes = [len(x) for x in levs[:v] + levs[v + 1:] + [levs[v]]]
group_index = get_group_index(to_sort, sizes)
comp_index, obs_ids = _compress_group_index(group_index)
ngroups = len(obs_ids)
indexer = algos.groupsort_indexer(comp_index, ngroups)[0]
indexer = _ensure_platform_int(indexer)
self.sorted_values = com.take_2d(self.values, indexer, axis=0)
self.sorted_labels = [l.take(indexer) for l in to_sort]
def _make_sorted_values_labels(self):
v = self.level
labs = self.index.labels
levs = self.index.levels
to_sort = labs[:v] + labs[v + 1:] + [labs[v]]
sizes = [len(x) for x in levs[:v] + levs[v + 1:] + [levs[v]]]
group_index = get_group_index(to_sort, sizes)
comp_index, obs_ids = _compress_group_index(group_index)
ngroups = len(obs_ids)
indexer = lib.groupsort_indexer(comp_index, ngroups)[0]
indexer = _ensure_platform_int(indexer)
self.sorted_values = com.take_2d(self.values, indexer, axis=0)
self.sorted_labels = [l.take(indexer) for l in to_sort]
def get_result(self):
# TODO: find a better way than this masking business
values, value_mask = self.get_new_values()
columns = self.get_new_columns()
index = self.get_new_index()
# filter out missing levels
if values.shape[1] > 0:
col_inds, obs_ids = _compress_group_index(self.sorted_labels[-1])
# rare case, level values not observed
if len(obs_ids) < self.full_shape[1]:
inds = (value_mask.sum(0) > 0).nonzero()[0]
values = com.take_2d(values, inds, axis=1)
columns = columns[inds]
return DataFrame(values, index=index, columns=columns)
def test_2d_float32(self):
arr = np.random.randn(4, 3).astype(np.float32)
indexer = [0, 2, -1, 1, -1]
# axis=0
result = common.take_2d(arr, indexer)
result2 = np.empty_like(result)
common.take_2d(arr, indexer, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=0)
expected[[2, 4]] = np.nan
tm.assert_almost_equal(result, expected)
# axis=1
result = common.take_2d(arr, indexer, axis=1)
result2 = np.empty_like(result)
common.take_2d(arr, indexer, axis=1, out=result2)
tm.assert_almost_equal(result, result)
expected = arr.take(indexer, axis=1)
expected[:, [2, 4]] = np.nan
tm.assert_almost_equal(result, expected)