def test_coargsort_categorical(self): string = ak.array(['a', 'b', 'a', 'b', 'c']) cat = ak.Categorical(string) cat_from_codes = ak.Categorical.from_codes( codes=ak.array([0, 1, 0, 1, 2]), categories=ak.array(['a', 'b', 'c'])) for algo in ak.SortingAlgorithm: str_perm = ak.coargsort([string], algo) str_sorted = string[str_perm].to_ndarray() # coargsort on categorical cat_perm = ak.coargsort([cat], algo) cat_sorted = cat[cat_perm].to_ndarray() self.assertTrue((str_sorted == cat_sorted).all()) # coargsort on categorical.from_codes # coargsort sorts using codes, the order isn't guaranteed, only grouping from_codes_perm = ak.coargsort([cat_from_codes], algo) from_codes_sorted = cat_from_codes[from_codes_perm].to_ndarray() self.assertTrue((['a', 'a', 'b', 'b', 'c'] == from_codes_sorted).all()) # coargsort on 2 categoricals (one from_codes) cat_perm = ak.coargsort([cat, cat_from_codes], algo) cat_sorted = cat[cat_perm].to_ndarray() self.assertTrue((str_sorted == cat_sorted).all()) # coargsort on mixed strings and categoricals mixed_perm = ak.coargsort([cat, string, cat_from_codes], algo) mixed_sorted = cat_from_codes[mixed_perm].to_ndarray() self.assertTrue((str_sorted == mixed_sorted).all())
def check_int_float(N): f = ak.randint(0, 2**63, N, dtype=ak.float64) i = ak.randint(0, 2**63, N, dtype=ak.int64) perm = ak.coargsort([f, i]) assert ak.is_sorted(f[perm]) perm = ak.coargsort([i, f]) assert ak.is_sorted(i[perm])
def test_error_handling(self): ones = ak.ones(100) short_ones = ak.ones(10) with self.assertRaises(ValueError): ak.coargsort([ones, short_ones]) with self.assertRaises(ValueError): ak.coargsort([list(range(0,10)), [0]])
def test_error_handling(self): ones = ak.ones(100) short_ones = ak.ones(10) for algo in ak.SortingAlgorithm: with self.assertRaises(ValueError): ak.coargsort([ones, short_ones], algo) for algo in ak.SortingAlgorithm: with self.assertRaises(TypeError): ak.coargsort([list(range(0, 10)), [0]], algo)
def check_int_uint_float(N, algo): f = ak.randint(0, 2**63, N, dtype=ak.float64) u = ak.randint(0, 2**63, N, dtype=ak.uint64) i = ak.randint(0, 2**63, N, dtype=ak.int64) perm = ak.coargsort([f, u, i], algo) assert ak.is_sorted(f[perm]) perm = ak.coargsort([u, i, f], algo) assert ak.is_sorted(u[perm]) perm = ak.coargsort([i, f, u], algo) assert ak.is_sorted(i[perm])
def check_float(N): a = ak.randint(0, 1, N, dtype=ak.float64) n = ak.randint(-1, 1, N, dtype=ak.float64) z = ak.zeros(N, dtype=ak.float64) perm = ak.coargsort([a]) assert ak.is_sorted(a[perm]) perm = ak.coargsort([a, n]) assert ak.is_sorted(a[perm]) perm = ak.coargsort([n, a]) assert ak.is_sorted(n[perm]) perm = ak.coargsort([z, a]) assert ak.is_sorted(a[perm]) perm = ak.coargsort([z, n]) assert ak.is_sorted(n[perm])
def check_coargsort(N): # create np version a = np.arange(N) a = a[::-1] iv = np.lexsort([a, a]) a = a[iv] # create ak version b = ak.arange(N) b = b[::-1] iv = ak.coargsort([b, b]) b = b[iv] # print(a,b) c = a == b.to_ndarray() # print(type(c),c) return pass_fail(c.all())
def check_int(N): z = ak.zeros(N, dtype=ak.int64) a2 = ak.randint(0, 2**16, N) b2 = ak.randint(0, 2**16, N) c2 = ak.randint(0, 2**16, N) d2 = ak.randint(0, 2**16, N) n2 = ak.randint(-(2**15), 2**15, N) perm = ak.coargsort([a2]) assert ak.is_sorted(a2[perm]) perm = ak.coargsort([n2]) assert ak.is_sorted(n2[perm]) perm = ak.coargsort([a2, b2, c2, d2]) assert ak.is_sorted(a2[perm]) perm = ak.coargsort([z, b2, c2, d2]) assert ak.is_sorted(b2[perm]) perm = ak.coargsort([z, z, c2, d2]) assert ak.is_sorted(c2[perm]) perm = ak.coargsort([z, z, z, d2]) assert ak.is_sorted(d2[perm]) a4 = ak.randint(0, 2**32, N) b4 = ak.randint(0, 2**32, N) n4 = ak.randint(-(2**31), 2**31, N) perm = ak.coargsort([a4]) assert ak.is_sorted(a4[perm]) perm = ak.coargsort([n4]) assert ak.is_sorted(n4[perm]) perm = ak.coargsort([a4, b4]) assert ak.is_sorted(a4[perm]) perm = ak.coargsort([b4, a4]) assert ak.is_sorted(b4[perm]) a8 = ak.randint(0, 2**64, N) b8 = ak.randint(0, 2**64, N) n8 = ak.randint(-(2**63), 2**64, N) perm = ak.coargsort([a8]) assert ak.is_sorted(a8[perm]) perm = ak.coargsort([n8]) assert ak.is_sorted(n8[perm]) perm = ak.coargsort([b8, a8]) assert ak.is_sorted(b8[perm]) from itertools import permutations all_perm = permutations([a2, a4, a8]) for p in all_perm: perm = ak.coargsort(p) assert ak.is_sorted(p[0][perm])
def check_large(N): l = [ak.randint(0, 2**63, N) for _ in range(10)] perm = ak.coargsort(l) assert ak.is_sorted(l[0][perm])