Add CrabNetSurrogateModel class to surrogate.py

surrogate.py

@@ -0,0 +1,44 @@
+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