def test_set_get_bitset(values_to_insert, expected_bitset): n_32bits_ints = 3 bitset = np.zeros(n_32bits_ints, dtype=np.uint32) for value in values_to_insert: set_bitset_memoryview(bitset, value) assert_allclose(expected_bitset, bitset) for value in range(32 * n_32bits_ints): if value in values_to_insert: assert in_bitset_memoryview(bitset, value) else: assert not in_bitset_memoryview(bitset, value)
def test_raw_bitset_from_binned_bitset(raw_categories, binned_cat_to_insert, expected_raw_bitset): binned_bitset = np.zeros(2, dtype=np.uint32) raw_bitset = np.zeros(2, dtype=np.uint32) raw_categories = np.asarray(raw_categories, dtype=X_DTYPE) for val in binned_cat_to_insert: set_bitset_memoryview(binned_bitset, val) set_raw_bitset_from_binned_bitset(raw_bitset, binned_bitset, raw_categories) assert_allclose(expected_raw_bitset, raw_bitset) for binned_cat_val, raw_cat_val in enumerate(raw_categories): if binned_cat_val in binned_cat_to_insert: assert in_bitset_memoryview(raw_bitset, raw_cat_val) else: assert not in_bitset_memoryview(raw_bitset, raw_cat_val)
def test_categorical_predictor(bins_go_left, expected_predictions): # Test predictor outputs are correct with categorical features X_binned = np.array([[0, 1, 2, 3, 4, 5]], dtype=X_BINNED_DTYPE).T categories = np.array([2, 5, 6, 8, 10, 15], dtype=X_DTYPE) bins_go_left = np.array(bins_go_left, dtype=X_BINNED_DTYPE) # We just construct a simple tree with 1 root and 2 children # parent node nodes = np.zeros(3, dtype=PREDICTOR_RECORD_DTYPE) nodes[0]['left'] = 1 nodes[0]['right'] = 2 nodes[0]['feature_idx'] = 0 nodes[0]['is_categorical'] = True nodes[0]['missing_go_to_left'] = True # left child nodes[1]['is_leaf'] = True nodes[1]['value'] = 1 # right child nodes[2]['is_leaf'] = True nodes[2]['value'] = 0 binned_cat_bitsets = np.zeros((1, 8), dtype=X_BITSET_INNER_DTYPE) raw_categorical_bitsets = np.zeros((1, 8), dtype=X_BITSET_INNER_DTYPE) for go_left in bins_go_left: set_bitset_memoryview(binned_cat_bitsets[0], go_left) set_raw_bitset_from_binned_bitset(raw_categorical_bitsets[0], binned_cat_bitsets[0], categories) predictor = TreePredictor(nodes, binned_cat_bitsets, raw_categorical_bitsets) # Check binned data gives correct predictions prediction_binned = predictor.predict_binned(X_binned, missing_values_bin_idx=6) assert_allclose(prediction_binned, expected_predictions) # manually construct bitset known_cat_bitsets = np.zeros((1, 8), dtype=np.uint32) known_cat_bitsets[0, 0] = np.sum(2**categories, dtype=np.uint32) f_idx_map = np.array([0], dtype=np.uint32) # Check with un-binned data predictions = predictor.predict(categories.reshape(-1, 1), known_cat_bitsets, f_idx_map) assert_allclose(predictions, expected_predictions) # Check missing goes left because missing_values_bin_idx=6 X_binned_missing = np.array([[6]], dtype=X_BINNED_DTYPE).T predictions = predictor.predict_binned(X_binned_missing, missing_values_bin_idx=6) assert_allclose(predictions, [1]) # missing and unknown go left predictions = predictor.predict( np.array([[np.nan, 17]], dtype=X_DTYPE).T, known_cat_bitsets, f_idx_map) assert_allclose(predictions, [1, 1])