generate_heuristic.py

generate_heuristic.py

cd cpp/scripts/heuristics/select_k
python3 generate_heuristic.py ../../../../select_k_dataset_float_times.json

import sys
from select_k_dataset import load_dataframe, get_dataset
import sklearn.tree
import numpy as np

if len(sys.argv) < 2:
    print("Error: Please provide path to select_k_dataset json file.")
    sys.exit(1)

# Get the filename from the first parameter
filename = sys.argv[1]
print(f"Processing file: {filename}")

df = load_dataframe(filename)

print(df)

X, y, weights = get_dataset(df)
train_test_sets = sklearn.model_selection.train_test_split(
    X, y, weights, test_size=0.15, random_state=1
)
X_train, X_test, y_train, y_test, weights_train, weights_test = train_test_sets
print(X_train.shape, X_test.shape)

model = sklearn.tree.DecisionTreeClassifier(max_depth=4, max_leaf_nodes=8)
model.fit(X_train, y_train)  # , weights_train)

print(model.score(X_train, y_train, weights_train))

print(model.score(X_test, y_test, weights_test))

def convert_model_to_code(model):
    classes = model.classes_
    tree = model.tree_
    feature_names = ["k", "rows", "cols", "use_memory_pool"]

    def _get_label(nodeid):
        """returns the most frequent class name for the node"""
        return classes[np.argsort(tree.value[nodeid, 0])[-1]]

    def _is_leaf_node(nodeid):
        """returns whether or not the node is a leaf node in the tree"""
        # negative values here indicate we're a leaf
        if tree.feature[nodeid] < 0:
            return True

        # some nodes have both branches with the same label, combine those
        left, right = tree.children_left[nodeid], tree.children_right[nodeid]
        if (
            _is_leaf_node(left)
            and _is_leaf_node(right)
            and _get_label(left) == _get_label(right)
        ):
            return True

        return False

    code = []

    def _convert_node(nodeid, indent):
        if _is_leaf_node(nodeid):
            # we're a leaf node, just output the label of the most frequent algorithm
            class_name = _get_label(nodeid)
            code.append(" " * indent + f"return SelectAlgo::{class_name};")
        else:
            feature = feature_names[tree.feature[nodeid]]
            threshold = int(np.floor(tree.threshold[nodeid]))
            code.append(" " * indent + f"if ({feature} > {threshold}) " + "{")
            _convert_node(tree.children_right[nodeid], indent + 2)
            code.append(" " * indent + "} else {")
            _convert_node(tree.children_left[nodeid], indent + 2)
            code.append(" " * indent + "}")

    code.append(
        "inline SelectAlgo choose_select_k_algorithm(size_t rows, size_t cols, int k)"
    )
    code.append("{")
    _convert_node(0, indent=2)
    code.append("}")
    return "\n".join(code)


code = convert_model_to_code(model)
print(code)