def test_grouping_data(self): """should correctly group data without a filter""" expect_two = dict(counts=[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]], ranks=[0, 1, 2, 3, 4], labels=['0', '1', '2', '3', '4']) coll = RegionCollection(**expect_two) grouped = coll.getGrouped(2) expect_counts = [[[0, 1], [2, 3]], [[4, 5], [6, 7]]] expect_ranks = [[0, 1], [2, 3]] expect_labels = [['0', '1'], ['2', '3']] self.assertEqual(grouped[0].tolist(), expect_counts) self.assertEqual(grouped[1].tolist(), expect_ranks) self.assertEqual(grouped[2].tolist(), expect_labels)
def test_only_counts(self): """correctly handle only being provided with counts data""" input_two = dict(counts=[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]]) # check write/read self.check_write_load(input_two) # check filter works expect_counts = [[[0, 1], [4, 5]], [[6, 7], [8, 9]]] func = lambda x: x.min() != 2 and x.max() != 3 coll = RegionCollection(**input_two) coll = coll.filtered(func) got_counts, got_ranks, got_labels = coll.getGrouped(2) self.assertEqual(got_counts.tolist(), expect_counts) self.assertEqual(got_ranks, None) self.assertEqual(got_labels, None)