def test_where(self):
# cases to test
# no x and y
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=None)
cond = a > 3
wh = ht.where(cond)
self.assertEqual(wh.gshape, (6, 2))
self.assertEqual(wh.dtype, ht.int64)
self.assertEqual(wh.split, None)
# split
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1)
cond = a > 3
wh = ht.where(cond)
self.assertEqual(wh.gshape, (6, 2))
self.assertEqual(wh.dtype, ht.int64)
self.assertEqual(wh.split, 0)
# not split cond
a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]],
split=None)
res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]],
split=None)
wh = ht.where(a < 4.0, a, -1)
self.assertTrue(
ht.equal(a[ht.nonzero(a < 4)],
ht.array([0.0, 1.0, 2.0, 0.0, 2.0, 0.0, 3.0])))
self.assertTrue(ht.equal(wh, res))
self.assertEqual(wh.gshape, (3, 3))
self.assertEqual(wh.dtype, ht.float64)
# split cond
a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]],
split=0)
res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]],
split=0)
wh = ht.where(a < 4.0, a, -1)
self.assertTrue(ht.all(wh[ht.nonzero(a >= 4)] == -1))
self.assertTrue(ht.equal(wh, res))
self.assertEqual(wh.gshape, (3, 3))
self.assertEqual(wh.dtype, ht.float64)
self.assertEqual(wh.split, 0)
a = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, 4.0], [0.0, 3.0, 6.0]],
split=1)
res = ht.array([[0.0, 1.0, 2.0], [0.0, 2.0, -1.0], [0.0, 3.0, -1.0]],
split=1)
wh = ht.where(a < 4.0, a, -1.0)
self.assertTrue(ht.equal(wh, res))
self.assertEqual(wh.gshape, (3, 3))
self.assertEqual(wh.dtype, ht.float)
self.assertEqual(wh.split, 1)
with self.assertRaises(TypeError):
ht.where(cond, a)
with self.assertRaises(NotImplementedError):
ht.where(cond, ht.ones((3, 3), split=0), ht.ones((3, 3), split=1))
def test_nonzero(self):
# cases to test:
# not split
a = ht.array([[1, 2, 3], [4, 5, 2], [7, 8, 9]], split=None)
cond = a > 3
nz = ht.nonzero(cond)
self.assertEqual(nz.gshape, (5, 2))
self.assertEqual(nz.dtype, ht.int64)
self.assertEqual(nz.split, None)
# split
a = ht.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], split=1)
cond = a > 3
nz = cond.nonzero()
self.assertEqual(nz.gshape, (6, 2))
self.assertEqual(nz.dtype, ht.int64)
self.assertEqual(nz.split, 0)
a[nz] = 10.0
self.assertEqual(ht.all(a[nz] == 10), 1)