Merge branch 'development' into behaviour_prep_bake

rvandewater · Sep 23, 2022 · 4a2e62b · 4a2e62b
2 parents a86a381 + b969a2a
commit 4a2e62b
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 11 deletions.
diff --git a/icu_benchmarks/recipes/ingredients.py b/icu_benchmarks/recipes/ingredients.py
@@ -7,9 +7,9 @@
 class Ingredients(pd.DataFrame):
     _metadata = ["roles"]
 
-    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=None,):
-        super().__init__(data, index, columns, dtype, )
-        self.roles = {}
+    def __init__(self, data=None, index=None, columns=None, dtype=None, copy=None, roles = {}):
+        super().__init__(data, index, columns, dtype, copy, )
+        self.roles = roles
 
     @property
     def _constructor(self):

diff --git a/icu_benchmarks/recipes/main.py b/icu_benchmarks/recipes/main.py
@@ -16,6 +16,6 @@
     rec.add_step(StepImputeFill(method='ffill'))
     rec.add_step(StepImputeFill(value=0))
 
-    rec.prep()
-    rec.bake()
+    rec.prep(df.iloc[:-10000, :])
+    rec.bake(df.iloc[10000:, :])
 
diff --git a/icu_benchmarks/recipes/recipe.py b/icu_benchmarks/recipes/recipe.py
@@ -39,6 +39,8 @@ def add_step(self, step):
     def _check_data(self, data):
         if data is None:
             data = self.data
+        elif data.__class__ == pd.DataFrame:
+            data = Ingredients(data, roles=self.data.roles)
         if not data.columns.equals(self.data.columns):
             raise ValueError('Columns of data argument differs from recipe data.')
         return data
@@ -60,7 +62,8 @@ def prep(self, data=None, refit=False):
         data = copy(data)
         self._apply_fit_transform(self, data, refit)
         return self
-
+        return pd.DataFrame(data)
+
     def bake(self, data=None):
         """
         Transforms, or bakes, the data if it has been prepped.
@@ -69,10 +72,6 @@ def bake(self, data=None):
         """
         data = self._check_data(data)
         data = copy(data)
-        for step in self.steps:
-            if not step.trained:
-                raise RuntimeError(f'Step {step} not trained. Run prep first.')
-
         self._apply_fit_transform(self, data)
         return data
 
@@ -84,7 +83,7 @@ def _apply_fit_transform(self, data=None, refit=False):
                 data = step.fit_transform(data)
             else:
                 data = step.transform(data)
-        return data
+        return pd.DataFrame(data)
 
     def __repr__(self):
         repr = 'Recipe\n\n'