surrogate.py

from joblib import load
import pandas as pd
import random


class CrabNetSurrogateModel(object):
    def __init__(self):
        self.models = load("surrogate_models.pkl")

    def prepare_params_for_eval(self, raw_params):
        raw_params["bias"] = int(raw_params["bias"])
        raw_params["use_RobustL1"] = raw_params["criterion"] == "RobustL1"
        raw_params.pop("criterion")

        raw_params.pop("losscurve")
        raw_params.pop("learningcurve")

        # raw_params["train_frac"] = random.uniform(0.01, 1)

        elem_prop = raw_params["elem_prop"]
        raw_params["elem_prop_magpie"] = 0
        raw_params["elem_prop_mat2vec"] = 0
        raw_params["elem_prop_onehot"] = 0
        raw_params[f"elem_prop_{elem_prop}"] = 1
        raw_params.pop("elem_prop")

        return raw_params

    def surrogate_evaluate(self, params):

        parameters = self.prepare_params_for_eval(params)
        parameters = pd.DataFrame([parameters])

        percentile = random.uniform(0, 1)  # generate random percentile

        # TODO: should percentile be different for each objective? (I guess depends on what is meant to be correlated vs. not)
        mae = self.models["mae"].predict(parameters.assign(mae_rank=[percentile]))
        rmse = self.models["rmse"].predict(parameters.assign(rmse_rank=[percentile]))
        runtime = self.models["runtime"].predict(
            parameters.assign(runtime_rank=[percentile])
        )
        model_size = self.models["model_size"].predict(parameters)

        return mae, rmse, runtime, model_size