def testDigitize(self): x = tensor(np.array([0.2, 6.4, 3.0, 1.6]), chunks=2) bins = np.array([0.0, 1.0, 2.5, 4.0, 10.0]) inds = digitize(x, bins) self.assertEqual(inds.shape, (4, )) self.assertIsNotNone(inds.dtype) inds.tiles() self.assertEqual(len(inds.chunks), 2)
def testDigitizeExecution(self): data = np.array([0.2, 6.4, 3.0, 1.6]) x = tensor(data, chunk_size=2) bins = np.array([0.0, 1.0, 2.5, 4.0, 10.0]) inds = digitize(x, bins) res = self.executor.execute_tensor(inds, concat=True)[0] expected = np.digitize(data, bins) np.testing.assert_equal(res, expected) b = tensor(bins, chunk_size=2) inds = digitize(x, b) res = self.executor.execute_tensor(inds, concat=True)[0] expected = np.digitize(data, bins) np.testing.assert_equal(res, expected) data = np.array([1.2, 10.0, 12.4, 15.5, 20.]) x = tensor(data, chunk_size=2) bins = np.array([0, 5, 10, 15, 20]) inds = digitize(x, bins, right=True) res = self.executor.execute_tensor(inds, concat=True)[0] expected = np.digitize(data, bins, right=True) np.testing.assert_equal(res, expected) inds = digitize(x, bins, right=False) res = self.executor.execute_tensor(inds, concat=True)[0] expected = np.digitize(data, bins, right=False) np.testing.assert_equal(res, expected) data = sps.random(10, 1, density=.1) * 12 x = tensor(data, chunk_size=2) bins = np.array([1.0, 2.0, 2.5, 4.0, 10.0]) inds = digitize(x, bins) res = self.executor.execute_tensor(inds, concat=True)[0] expected = np.digitize(data.toarray(), bins, right=False) np.testing.assert_equal(res.toarray(), expected)