updated AndEstimator syntax and implementation

Author: yuanchi2807

codeflare/pipelines/tests/test_Datamodel.py

@@ -7,32 +7,29 @@
 from sklearn.pipeline import Pipeline
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
 from sklearn.tree import DecisionTreeClassifier
+import sklearn.base as base
 import codeflare.pipelines.Datamodel as dm
 import codeflare.pipelines.Runtime as rt
 from codeflare.pipelines.Datamodel import Xy
 from codeflare.pipelines.Runtime import ExecutionType
 
-
 class FeatureUnion(dm.AndEstimator):
     def __init__(self):
         pass
-
-    def fit_transform(self, xy_list: list):
-        return self.transform(xy_list)
-
     def get_estimator_type(self):
         return 'transform'
-
+    def clone(self):
+        return base.clone(self)
+    def fit_transform(self, xy_list):
+        return self.transform(xy_list)
     def transform(self, xy_list):
         X_list = []
-        y_list = []
-
+        y_vec = None
         for xy in xy_list:
             X_list.append(xy.get_x())
-        X_concat = np.concatenate(X_list, axis=0)
-
-        return Xy(X_concat, None)
-
+            y_vec = xy.get_y()
+        X_concat = np.concatenate(X_list, axis=1)
+        return Xy(X_concat, y_vec)
 
 class MultibranchTestCase(unittest.TestCase):
 

codeflare/pipelines/tests/test_and.py

@@ -2,6 +2,7 @@
 import ray
 import pandas as pd
 import numpy as np
+import sklearn.base as base
 from sklearn.preprocessing import StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler
 import codeflare.pipelines.Datamodel as dm
 import codeflare.pipelines.Runtime as rt
@@ -12,19 +13,19 @@
 class FeatureUnion(dm.AndEstimator):
     def __init__(self):
         pass
-
     def get_estimator_type(self):
         return 'transform'
-
+    def clone(self):
+        return base.clone(self)
+    def fit_transform(self, xy_list):
+        return self.transform(xy_list)
     def transform(self, xy_list):
         X_list = []
         y_vec = None
-
         for xy in xy_list:
             X_list.append(xy.get_x())
             y_vec = xy.get_y()
         X_concat = np.concatenate(X_list, axis=1)
-
         return Xy(X_concat, y_vec)
 
 def test_two_tier_and():

codeflare/pipelines/tests/test_multibranch.py

@@ -8,6 +8,7 @@
 from sklearn.preprocessing import StandardScaler, MinMaxScaler
 from sklearn.tree import DecisionTreeClassifier
 from sklearn.linear_model import LogisticRegression
+import sklearn.base as base
 import codeflare.pipelines.Datamodel as dm
 import codeflare.pipelines.Runtime as rt
 from codeflare.pipelines.Datamodel import Xy
@@ -17,17 +18,20 @@
 class FeatureUnion(dm.AndEstimator):
     def __init__(self):
         pass
-
+    def get_estimator_type(self):
+        return 'transform'
+    def clone(self):
+        return base.clone(self)
+    def fit_transform(self, xy_list):
+        return self.transform(xy_list)
     def transform(self, xy_list):
         X_list = []
         y_vec = None
-
         for xy in xy_list:
             X_list.append(xy.get_x())
             y_vec = xy.get_y()
         X_concat = np.concatenate(X_list, axis=1)
-
-        return Xy(X_concat, y_vec.values.ravel())
+        return Xy(X_concat, y_vec)
 
 def test_multibranch_1():