[ENH] enhance FeatureAugmentReuser for classification; add corresponding test; add explanations.

learnware/reuse/feature_augment_reuser.py

@@ -1,38 +1,95 @@
-from typing import List
 import numpy as np
 
-from sklearn.linear_model import RidgeCV
-
+from sklearn.linear_model import RidgeCV, LogisticRegressionCV
 from .base import BaseReuser
 from learnware.learnware import Learnware
 
-
 class FeatureAugmentReuser(BaseReuser):
-    def __init__(self, learnware: Learnware = None, task_type: str = None):
-        self.learnware=learnware
-        assert task_type in ["classification", "regression"]
-        self.task_type=task_type
-
-    def predict(self, x_test: np.ndarray) -> np.ndarray:
-        x_test=self._fill_data(x_test)
-        y_pred=self.learnware.predict(x_test)
-        x_test_aug=np.concatenate((x_test, y_pred.reshape(-1, 1)), axis=1)
-        y_pred_aug=self.output_aligner.predict(x_test_aug)
+    """
+    FeatureAugmentReuser is a class for augmenting features using predictions of a given learnware model and applying regression or classification on the augmented dataset.
+
+    This class supports two modes:
+    - "regression": Uses RidgeCV for regression tasks.
+    - "classification": Uses LogisticRegressionCV for classification tasks.
+    """
+
+    def __init__(self, learnware: Learnware = None, mode: str = None):
+        """
+        Initializes the FeatureAugmentReuser with a learnware model and a mode.
+
+        Parameters
+        ----------
+        learnware : Learnware
+            A learnware model used for initial predictions.
+        mode : str
+            The mode of operation, either "regression" or "classification".
+        """
+        self.learnware = learnware
+        assert mode in ["classification", "regression"], "Mode must be either 'classification' or 'regression'"
+        self.mode = mode
+
+    def predict(self, user_data: np.ndarray) -> np.ndarray:
+        """
+        Predicts the output for user data using the trained output aligner model.
+
+        Parameters
+        ----------
+        user_data : np.ndarray
+            Input data for making predictions.
+
+        Returns
+        -------
+        np.ndarray
+            Predicted output from the output aligner model.
+        """
+        user_data = self._fill_data(user_data)
+        y_pred = self.learnware.predict(user_data)
+        user_data_aug = np.concatenate((user_data, y_pred.reshape(-1, 1)), axis=1)
+        y_pred_aug = self.output_aligner.predict(user_data_aug)
         return y_pred_aug
 
-    def fit(self, x_train, y_train):
-        x_train=self._fill_data(x_train)
-        y_pred=self.learnware.predict(x_train)
-        x_train_aug=np.concatenate((x_train, y_pred.reshape(-1, 1)), axis=1)
-        if self.task_type=="regression":
+    def fit(self, x_train: np.ndarray, y_train: np.ndarray):
+        """
+        Trains the output aligner model using the training data augmented with predictions from the learnware model.
+
+        Parameters
+        ----------
+        x_train : np.ndarray
+            Training data features.
+        y_train : np.ndarray
+            Training data labels.
+        """
+        x_train = self._fill_data(x_train)
+        y_pred = self.learnware.predict(x_train)
+        x_train_aug = np.concatenate((x_train, y_pred.reshape(-1, 1)), axis=1)
+        if self.mode == "regression":
             alpha_list = [0.01, 0.1, 1.0, 10, 100]
             ridge_cv = RidgeCV(alphas=alpha_list, store_cv_values=True)
             ridge_cv.fit(x_train_aug, y_train)
-            self.output_aligner=ridge_cv
-        elif self.task_type=="classification":
-            raise NotImplementedError("Not implemented yet!")
+            self.output_aligner = ridge_cv
+        elif self.mode == "classification":
+            self.output_aligner = LogisticRegressionCV()
+            self.output_aligner.fit(x_train_aug, y_train)
 
     def _fill_data(self, X: np.ndarray):
+        """
+        Fills missing data (NaN, Inf) in the input array with the mean of the column.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Input data array that may contain missing values.
+
+        Returns
+        -------
+        np.ndarray
+            Data array with missing values filled.
+
+        Raises
+        ------
+        ValueError
+            If a column in X contains only exceptional values (NaN, Inf).
+        """
         X[np.isinf(X) | np.isneginf(X) | np.isposinf(X) | np.isneginf(X)] = np.nan
         if np.any(np.isnan(X)):
             for col in range(X.shape[1]):
@@ -40,7 +97,6 @@ class FeatureAugmentReuser(BaseReuser):
                 if np.any(is_nan):
                     if np.all(is_nan):
                         raise ValueError(f"All values in column {col} are exceptional, e.g., NaN and Inf.")
-                    # Fill np.nan with np.nanmean
                     col_mean = np.nanmean(X[:, col])
                     X[:, col] = np.where(is_nan, col_mean, X[:, col])
-        return X
+        return X

learnware/reuse/hetero_reuser/__init__.py

@@ -6,20 +6,20 @@ from ..feature_augment_reuser import FeatureAugmentReuser
 
 class HeteroMapTableReuser(BaseReuser):
 
-    def __init__(self, learnware: Learnware = None, task_type: str = None, cuda_idx=0, **align_arguments):
+    def __init__(self, learnware: Learnware = None, mode: str = None, cuda_idx=0, **align_arguments):
         self.learnware=learnware
-        assert task_type in ["classification", "regression"]
-        self.task_type=task_type
+        assert mode in ["classification", "regression"]
+        self.mode=mode
         self.cuda_idx=cuda_idx
         self.align_arguments=align_arguments
 
     def fit(self, user_rkme):
-        self.feature_aligner=FeatureAligner(learnware=self.learnware, task_type=self.task_type, cuda_idx=self.cuda_idx, **self.align_arguments)
+        self.feature_aligner=FeatureAligner(learnware=self.learnware, mode=self.mode, cuda_idx=self.cuda_idx, **self.align_arguments)
         self.feature_aligner.fit(user_rkme)
         self.reuser=self.feature_aligner
 
     def finetune(self, x_train,y_train):
-        self.reuser=FeatureAugmentReuser(learnware=self.feature_aligner, task_type=self.task_type)
+        self.reuser=FeatureAugmentReuser(learnware=self.feature_aligner, mode=self.mode)
         self.reuser.fit(x_train, y_train)
 
     def predict(self, user_data):

learnware/reuse/hetero_reuser/feature_alignment.py

@@ -17,10 +17,10 @@ from ..base import BaseReuser
 
 class FeatureAligner(BaseReuser):
 
-    def __init__(self, learnware: Learnware = None, task_type: str = None, cuda_idx=0, **align_arguments):
+    def __init__(self, learnware: Learnware = None, mode: str = None, cuda_idx=0, **align_arguments):
         self.learnware=learnware
-        assert task_type in ["classification", "regression"]
-        self.task_type=task_type
+        assert mode in ["classification", "regression"]
+        self.mode=mode
         self.align_arguments=align_arguments
         self.cuda_idx=cuda_idx
         self.device = choose_device(cuda_idx=cuda_idx)

tests/test_market/test_hetero_market/test_hetero.py

@@ -365,15 +365,8 @@ class TestMarket(unittest.TestCase):
             print(f"score: {score}, learnware_id: {learnware.id}")
 
         # model reuse
-        reuser=HeteroMapTableReuser(single_learnware_list[0], task_type='regression')
+        reuser=HeteroMapTableReuser(single_learnware_list[0], mode='regression')
         reuser.fit(user_spec)
-        y_pred=reuser.predict(X)
-        
-        # calculate rmse
-        rmse=mean_squared_error(y, y_pred, squared=False)
-        print(f"rmse not finetune: {rmse}")
-
-        # finetune
         reuser.finetune(X[:100], y[:100])
         y_pred=reuser.predict(X)
         rmse=mean_squared_error(y, y_pred, squared=False)
@@ -388,7 +381,7 @@ def suite():
     # _suite.addTest(TestMarket("test_train_market_model"))
     # _suite.addTest(TestMarket("test_search_semantics"))
     _suite.addTest(TestMarket("test_stat_search"))
-    # _suite.addTest(TestMarket("test_model_reuse"))
+    _suite.addTest(TestMarket("test_model_reuse"))
     return _suite
 
 

tests/test_workflow/test_workflow.py

@@ -13,7 +13,7 @@ from shutil import copyfile, rmtree
 import learnware
 from learnware.market import instantiate_learnware_market, BaseUserInfo
 from learnware.specification import RKMETableSpecification, generate_rkme_spec
-from learnware.reuse import JobSelectorReuser, AveragingReuser, EnsemblePruningReuser
+from learnware.reuse import JobSelectorReuser, AveragingReuser, EnsemblePruningReuser, FeatureAugmentReuser
 
 curr_root = os.path.dirname(os.path.abspath(__file__))
 
@@ -219,17 +219,23 @@ class TestWorkflow(unittest.TestCase):
         reuse_ensemble.fit(train_X[-200:], train_y[-200:])
         ensemble_pruning_predict_y = reuse_ensemble.predict(user_data=data_X)
 
+        # Use feature augment reuser to reuse the searched learnwares to make prediction
+        reuse_feature_augment = FeatureAugmentReuser(learnware=reuse_ensemble, mode="classification")
+        reuse_feature_augment.fit(train_X[-200:], train_y[-200:])
+        feature_augment_predict_y = reuse_feature_augment.predict(user_data=data_X)
+
         print("Job Selector Acc:", np.sum(np.argmax(job_selector_predict_y, axis=1) == data_y) / len(data_y))
         print("Averaging Reuser Acc:", np.sum(np.argmax(ensemble_predict_y, axis=1) == data_y) / len(data_y))
         print("Ensemble Pruning Reuser Acc:", np.sum(ensemble_pruning_predict_y == data_y) / len(data_y))
+        print("Feature Augment Reuser Acc:", np.sum(feature_augment_predict_y == data_y) / len(data_y))
 
 
 def suite():
     _suite = unittest.TestSuite()
-    _suite.addTest(TestWorkflow("test_prepare_learnware_randomly"))
-    _suite.addTest(TestWorkflow("test_upload_delete_learnware"))
-    _suite.addTest(TestWorkflow("test_search_semantics"))
-    _suite.addTest(TestWorkflow("test_stat_search"))
+    # _suite.addTest(TestWorkflow("test_prepare_learnware_randomly"))
+    # _suite.addTest(TestWorkflow("test_upload_delete_learnware"))
+    # _suite.addTest(TestWorkflow("test_search_semantics"))
+    # _suite.addTest(TestWorkflow("test_stat_search"))
     _suite.addTest(TestWorkflow("test_learnware_reuse"))
     return _suite