Adapting fit, predict and score fn args. Still working on tests.

Author: Alessandro Lucantonio

bench/bench.py

@@ -1,7 +1,7 @@
 # import matplotlib.pyplot as plt
 from deap import gp
 
-from alpine.gp import gpsymbreg as gps
+from alpine.gp import regressor as gps
 from alpine.data import Dataset
 from alpine.gp import util
 import numpy as np

examples/simple_sr.py

@@ -1,6 +1,6 @@
 import os
 from deap import gp
-from alpine.gp.gpsymbreg import GPSymbolicRegressor
+from alpine.gp.regressor import GPSymbolicRegressor
 from alpine.data import Dataset
 import numpy as np
 import ray
@@ -40,44 +40,44 @@ def get_features_batch(
     return individ_length, nested_trigs, num_trigs
 
 
-def eval_MSE_sol(individual, true_data):
+def eval_MSE_sol(individual, X, y):
     warnings.filterwarnings("ignore")
 
-    y_pred = individual(true_data.X)
-    MSE = np.mean(np.square(y_pred - true_data.y))
+    y_pred = individual(X)
+    MSE = np.mean(np.square(y_pred - y))
     if np.isnan(MSE):
         MSE = 1e5
     return MSE, y_pred
 
 
 @ray.remote
-def predict(individuals_str, toolbox, true_data, penalty):
+def predict(individuals_str, toolbox, X_test, penalty):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     u = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        _, u[i] = eval_MSE_sol(ind, true_data)
+        _, u[i] = eval_MSE_sol(ind, X_test, None)
 
     return u
 
 
 @ray.remote
-def score(individuals_str, toolbox, true_data, penalty):
+def score(individuals_str, toolbox, X_test, y_test, penalty):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     MSE = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        MSE[i], _ = eval_MSE_sol(ind, true_data)
+        MSE[i], _ = eval_MSE_sol(ind, X_test, y_test)
 
     return MSE
 
 
 @ray.remote
-def fitness(individuals_str, toolbox, true_data, penalty):
+def fitness(individuals_str, toolbox, X_train, y_train, penalty):
     callables = compile_individuals(toolbox, individuals_str)
 
     individ_length, nested_trigs, num_trigs = get_features_batch(individuals_str)
@@ -87,7 +87,7 @@ def fitness(individuals_str, toolbox, true_data, penalty):
         if individ_length[i] >= 50:
             fitnesses[i] = (1e8,)
         else:
-            MSE, _ = eval_MSE_sol(ind, true_data)
+            MSE, _ = eval_MSE_sol(ind, X_train, y_train)
 
             fitnesses[i] = (
                 MSE
@@ -131,8 +131,7 @@ def main():
         **regressor_params
     )
 
-    train_data = Dataset("true_data", x, y)
-    gpsr.fit(train_data)
+    gpsr.fit(x, y)
 
     ray.shutdown()
 

examples/simple_sr_noyaml.py

@@ -1,5 +1,5 @@
 from deap import gp
-from alpine.gp.gpsymbreg import GPSymbolicRegressor
+from alpine.gp.regressor import GPSymbolicRegressor
 from alpine.data import Dataset
 import numpy as np
 import ray
@@ -39,44 +39,44 @@ def get_features_batch(
     return individ_length, nested_trigs, num_trigs
 
 
-def eval_MSE_sol(individual, true_data):
+def eval_MSE_sol(individual, X, y):
     warnings.filterwarnings("ignore")
 
-    y_pred = individual(true_data.X)
-    MSE = np.mean(np.square(y_pred - true_data.y))
+    y_pred = individual(X)
+    MSE = np.mean(np.square(y_pred - y))
     if np.isnan(MSE):
         MSE = 1e5
     return MSE, y_pred
 
 
 @ray.remote
-def predict(individuals_str, toolbox, true_data, penalty):
+def predict(individuals_str, toolbox, X_test, penalty):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     u = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        _, u[i] = eval_MSE_sol(ind, true_data)
+        _, u[i] = eval_MSE_sol(ind, X_test, None)
 
     return u
 
 
 @ray.remote
-def score(individuals_str, toolbox, true_data, penalty):
+def score(individuals_str, toolbox, X_test, y_test, penalty):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     MSE = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        MSE[i], _ = eval_MSE_sol(ind, true_data)
+        MSE[i], _ = eval_MSE_sol(ind, X_test, y_test)
 
     return MSE
 
 
 @ray.remote
-def fitness(individuals_str, toolbox, true_data, penalty):
+def fitness(individuals_str, toolbox, X_train, y_train, penalty):
     callables = compile_individuals(toolbox, individuals_str)
 
     individ_length, nested_trigs, num_trigs = get_features_batch(individuals_str)
@@ -86,7 +86,7 @@ def fitness(individuals_str, toolbox, true_data, penalty):
         if individ_length[i] >= 50:
             fitnesses[i] = (1e8,)
         else:
-            MSE, _ = eval_MSE_sol(ind, true_data)
+            MSE, _ = eval_MSE_sol(ind, X_train, y_train)
 
             fitnesses[i] = (
                 MSE
@@ -145,8 +145,7 @@ def main():
         batch_size=100,
     )
 
-    train_data = Dataset("true_data", x, y)
-    gpsr.fit(train_data)
+    gpsr.fit(x, y)
 
     ray.shutdown()
 

src/alpine/gp/gpsymbreg.py src/alpine/gp/regressor.py

@@ -11,6 +11,7 @@
 import ray
 import random
 from itertools import chain
+from sklearn.base import BaseEstimator, RegressorMixin
 
 # reducing the number of threads launched by fitness evaluations
 os.environ["MKL_NUM_THREADS"] = "1"
@@ -24,7 +25,7 @@
 )
 
 
-class GPSymbolicRegressor:
+class GPSymbolicRegressor(RegressorMixin, BaseEstimator):
     """Symbolic regression problem via Genetic Programming.
 
     Args:
@@ -130,7 +131,7 @@ def __init__(
 
         if common_data is not None:
             # FIXME: does everything work when the functions do not have common args?
-            self.store_fit_error_common_args(common_data)
+            self.__store_fit_error_common_args(common_data)
 
         self.NINDIVIDUALS = NINDIVIDUALS
         self.NGEN = NGEN
@@ -161,9 +162,6 @@ def __init__(
 
         # config individual creator and toolbox
         self.__creator_toolbox_config()
-        # self.createIndividual = individualCreator
-
-        # self.toolbox = toolbox
 
         self.seed = seed
 
@@ -253,7 +251,7 @@ def __creator_toolbox_config(self):
 
         self.createIndividual = createIndividual
 
-    def store_fit_error_common_args(self, data: Dict):
+    def __store_fit_error_common_args(self, data: Dict):
         """Store names and values of the arguments that are in common between
         the fitness and the error metric functions in the common object space.
 
@@ -262,7 +260,7 @@ def store_fit_error_common_args(self, data: Dict):
         """
         self.__store_shared_objects("common", data)
 
-    def store_datasets(self, datasets: Dict[str, Dataset]):
+    def __store_datasets(self, datasets: Dict[str, Dataset]):
         """Store datasets with the corresponding label ("train", "val" or "test")
         in the common object space. The datasets are passed as parameters to
         the fitness, and possibly to the error metric and the prediction functions.
@@ -272,12 +270,12 @@ def store_datasets(self, datasets: Dict[str, Dataset]):
                 the validation and the test datasets, respectively. The associated
                 values are `Dataset` objects.
         """
-        for dataset_label in datasets.keys():
-            dataset_name_data = {datasets[dataset_label].name: datasets[dataset_label]}
-            self.__store_shared_objects(dataset_label, dataset_name_data)
+        for dataset_label, dataset_data in datasets.items():
+            self.__store_shared_objects(dataset_label, dataset_data)
 
     def __store_shared_objects(self, label: str, data: Dict):
         for key, value in data.items():
+            # replace each item of the dataset with its obj ref
             data[key] = ray.put(value)
         self.data_store[label] = data
 
@@ -414,31 +412,35 @@ def mapper(f, individuals, toolbox_ref):
         toolbox_ref = ray.put(self.toolbox)
         self.toolbox.register("map", mapper, toolbox_ref=toolbox_ref)
 
-    def fit(self, train_data: Dataset, val_data: Dataset | None = None):
+    def fit(self, X_train, y_train=None, X_val=None, y_val=None):
         """Fits the training data using GP-based symbolic regression."""
-        if self.validate and val_data is not None:
+        train_data = {"X_train": X_train, "y_train": y_train}
+        if self.validate and X_val is not None:
+            val_data = {"X_val": X_val, "y_val": y_val}
             datasets = {"train": train_data, "val": val_data}
         else:
             datasets = {"train": train_data}
-        self.store_datasets(datasets)
+        self.__store_datasets(datasets)
         self.__register_fitness_func()
         if self.validate and self.error_metric is not None:
             self.__register_val_funcs()
         self.__run()
 
-    def predict(self, test_data: Dataset):
+    def predict(self, X_test):
+        test_data = {"X_test": X_test}
         datasets = {"test": test_data}
-        self.store_datasets(datasets)
+        self.__store_datasets(datasets)
         self.__register_predict_func()
         u_best = self.toolbox.map(self.toolbox.evaluate_test_sols, (self.best,))[0]
         return u_best
 
-    def score(self, test_data: Dataset):
+    def score(self, X_test, y_test):
         """Computes the error metric (passed to the `GPSymbolicRegressor` constructor)
         on a given dataset.
         """
+        test_data = {"X_test": X_test, "y_test": y_test}
         datasets = {"test": test_data}
-        self.store_datasets(datasets)
+        self.__store_datasets(datasets)
         self.__register_score_func()
         score = self.toolbox.map(self.toolbox.evaluate_test_score, (self.best,))[0]
         return score

tests/test_basic_sr.py

@@ -1,7 +1,7 @@
 import os
 from dctkit import config
 from deap import gp
-from alpine.gp.gpsymbreg import GPSymbolicRegressor
+from alpine.gp.regressor import GPSymbolicRegressor
 from alpine.data import Dataset
 from alpine.gp import util
 import jax.numpy as jnp
@@ -21,53 +21,56 @@ def compile_individuals(toolbox, individuals_str_batch):
 y = x**4 + x**3 + x**2 + x
 
 
-def eval_MSE_sol(individual, true_data):
+def eval_MSE_sol(individual, X, y):
     import os
 
     os.environ["JAX_PLATFORMS"] = "cpu"
     config()
     # Evaluate the mean squared error between the expression
     # and the real function : x**4 + x**3 + x**2 + x
-    y_pred = individual(true_data.X)
-    MSE = jnp.sum(jnp.square(y_pred - true_data.y)) / len(true_data.X)
-    if jnp.isnan(MSE):
-        MSE = 1e5
+    y_pred = individual(X)
+    MSE = None
+
+    if y is not None:
+        MSE = jnp.mean(jnp.sum(jnp.square(y_pred - y)))
+        MSE = jnp.nan_to_num(MSE, nan=1e5)
+
     return MSE, y_pred
 
 
 @ray.remote
-def predict(individuals_str, toolbox, true_data):
+def predict(individuals_str, toolbox, X_test):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     u = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        _, u[i] = eval_MSE_sol(ind, true_data)
+        _, u[i] = eval_MSE_sol(ind, X_test, None)
 
     return u
 
 
 @ray.remote
-def score(individuals_str, toolbox, true_data):
+def score(individuals_str, toolbox, X_test, y_test):
 
     callables = compile_individuals(toolbox, individuals_str)
 
     MSE = [None] * len(individuals_str)
 
     for i, ind in enumerate(callables):
-        MSE[i], _ = eval_MSE_sol(ind, true_data)
+        MSE[i], _ = eval_MSE_sol(ind, X_test, y_test)
 
     return MSE
 
 
 @ray.remote
-def fitness(individuals_str, toolbox, true_data):
+def fitness(individuals_str, toolbox, X_train, y_train):
     callables = compile_individuals(toolbox, individuals_str)
 
     fitnesses = [None] * len(individuals_str)
     for i, ind in enumerate(callables):
-        MSE, _ = eval_MSE_sol(ind, true_data)
+        MSE, _ = eval_MSE_sol(ind, X_train, y_train)
 
         fitnesses[i] = (MSE,)
 
@@ -110,10 +113,12 @@ def test_basic_sr(set_test_dir):
         **regressor_params
     )
 
-    train_data = Dataset("true_data", x, y)
-    gpsr.fit(train_data)
+    # train_data = Dataset("true_data", x, y)
+    gpsr.fit(x, y)
+
+    fit_score = gpsr.score(x, y)
 
-    fit_score = gpsr.score(train_data)
+    y_pred = gpsr.predict(x)
 
     ray.shutdown()