def test_multiple_where_clauses(self):
n1 = np.arange(1, 10)
n2 = np.ones(9, dtype=np.int64)
a1 = ak.array(n1)
a2 = ak.array(n2)
cond = n1 > 2, n1 < 8
result = np.where(cond, n1, n2)
self.assertTrue((np.array([[1, 1, 3, 4, 5, 6, 7, 8, 9],
[1, 2, 3, 4, 5, 6, 7, 1,
1]]) == result).all())
# Arkouda does not support multiple where clauses
cond = a1 > 5, a1 < 8
with self.assertRaises(TypeError):
ak.where(cond, a1, a2)
def test_where_equivalence(self):
npA = {
'int64': np.random.randint(0, 10, SIZE),
'float64': np.random.randn(SIZE),
'bool': np.random.randint(0, 2, SIZE, dtype='bool')
}
akA = {k: ak.array(v) for k, v in npA.items()}
npB = {
'int64': np.random.randint(10, 20, SIZE),
'float64': np.random.randn(SIZE) + 10,
'bool': np.random.randint(0, 2, SIZE, dtype='bool')
}
akB = {k: ak.array(v) for k, v in npB.items()}
npCond = np.random.randint(0, 2, SIZE, dtype='bool')
akCond = ak.array(npCond)
scA = {'int64': 42, 'float64': 2.71828, 'bool': True}
scB = {'int64': -1, 'float64': 3.14159, 'bool': False}
dtypes = set(npA.keys())
failures = 0
tests = 0
for dtype in dtypes:
for (ak1, ak2), (np1, np2) in zip(product((akA, scA), (akB, scB)),
product((npA, scA), (npB, scB))):
tests += 1
akres = ak.where(akCond, ak1[dtype], ak2[dtype]).to_ndarray()
npres = np.where(npCond, np1[dtype], np2[dtype])
if not np.allclose(akres, npres, equal_nan=True):
self.assertWarning(
warnings.warn("{} !=\n{}".format(akres, npres)))
failures += 1
self.assertEqual(0, failures)
def test_greater_than_where_clause_with_scalars(self):
n1 = np.arange(1, 10)
a1 = ak.array(n1)
condN = n1 > 5
result = np.where(condN, n1, 1)
self.assertTrue((np.array([1, 1, 1, 1, 1, 6, 7, 8,
9]) == result).all())
condA = a1 > 5
result = ak.where(condA, a1, 1)
self.assertTrue((ak.array([1, 1, 1, 1, 1, 6, 7, 8,
9]) == result).all())
result = np.where(condN, 1, n1)
self.assertTrue((np.array([1, 2, 3, 4, 5, 1, 1, 1,
1]) == result).all())
result = ak.where(condA, 1, a1)
self.assertTrue((ak.array([1, 2, 3, 4, 5, 1, 1, 1,
1]) == result).all())
def test_greater_than_where_clause(self):
n1 = np.arange(1, 10)
n2 = np.ones(9, dtype=np.int64)
a1 = ak.array(n1)
a2 = ak.array(n2)
cond = n1 > 5
result = np.where(cond, n1, n2)
self.assertTrue((np.array([1, 1, 1, 1, 1, 6, 7, 8,
9]) == result).all())
cond = a1 > 5
result = ak.where(cond, a1, a2)
self.assertTrue((ak.array([1, 1, 1, 1, 1, 6, 7, 8,
9]) == result).all())
def test_less_than_where_clause(self):
n1 = np.arange(1, 10)
n2 = np.ones(9, dtype=np.int64)
a1 = ak.array(n1)
a2 = ak.array(n2)
cond = n1 < 5
result = np.where(cond, n1, n2)
self.assertTrue((np.array([1, 2, 3, 4, 1, 1, 1, 1,
1]) == result).all())
cond = a1 < 5
result = ak.where(cond, a1, a2)
self.assertTrue((ak.array([1, 2, 3, 4, 1, 1, 1, 1,
1]) == result).all())
def test_where_equivalence(self):
failures = 0
tests = 0
for dtype in self.dtypes:
for (self.ak1, self.ak2), (self.np1, self.np2) in zip(
product((self.akA, self.scA), (self.akB, self.scB)),
product((self.npA, self.scA), (self.npB, self.scB))):
tests += 1
akres = ak.where(self.akCond, self.ak1[dtype],
self.ak2[dtype]).to_ndarray()
npres = np.where(self.npCond, self.np1[dtype], self.np2[dtype])
if not np.allclose(akres, npres, equal_nan=True):
self.assertWarning(
warnings.warn("{} !=\n{}".format(akres, npres)))
failures += 1
self.assertEqual(0, failures)
def test_error_handling(self):
with self.assertRaises(TypeError) as cm:
ak.where([0], ak.linspace(1, 10, 10), ak.linspace(1, 10, 10))
self.assertEqual((
'type of argument "condition" must be arkouda.pdarrayclass.pdarray;'
+ ' got list instead'), cm.exception.args[0])
with self.assertRaises(TypeError) as cm:
ak.where(ak.linspace(1, 10, 10), [0], ak.linspace(1, 10, 10))
self.assertEqual((
'type of argument "A" must be one of (int, float, int64, pdarray); got list instead'
), cm.exception.args[0])
with self.assertRaises(TypeError) as cm:
ak.where(ak.linspace(1, 10, 10), ak.linspace(1, 10, 10), [0])
self.assertEqual(
'both A and B must be an int, np.int64, float, np.float64, or pdarray',
cm.exception.args[0])
