set and unset split time ordering from split_params instead of ML settings

Author: kgued

dataikuapi/dss/ml.py

@@ -119,6 +119,45 @@ def set_split_explicit(self, train_selection, test_selection, dataset_name=None,
             else:
                 test_split["filter"] = test_filter
 
+    def set_order_by(self, feature_name, ascending=True):
+        """
+        Uses a variable to sort the data for train/test split and hyperparameter optimization
+        :param str feature_name: Name of the variable to use
+        :param bool ascending: True iff the test set is expected to have larger time values than the train set
+        """
+        self.unset_order_by()
+        if not feature_name in self.mltask_settings["preprocessing"]["per_feature"]:
+            raise ValueError("Feature %s doesn't exist in this ML task, can't use as time" % feature_name)
+        self.mltask_settings['time']['enabled'] = True
+        self.mltask_settings['time']['timeVariable'] = feature_name
+        self.mltask_settings['time']['ascending'] = ascending
+        self.mltask_settings['preprocessing']['per_feature'][feature_name]['missing_handling'] = "DROP_ROW"
+        if self.mltask_settings['splitParams']['ttPolicy'] == "SPLIT_SINGLE_DATASET":
+            self.mltask_settings['splitParams']['ssdSplitMode'] = "SORTED"
+            self.mltask_settings['splitParams']['ssdColumn'] = feature_name
+        if self.mltask_settings['modeling']['gridSearchParams']['mode'] == "KFOLD":
+            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "TIME_SERIES_KFOLD"
+        elif self.mltask_settings['modeling']['gridSearchParams']['mode'] == "SHUFFLE":
+            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "TIME_SERIES_SINGLE_SPLIT"
+
+        return self
+
+    def unset_order_by(self):
+        """
+        Remove time-based ordering.
+        """
+        self.mltask_settings['time']['enabled'] = False
+        self.mltask_settings['time']['timeVariable'] = None
+        if self.mltask_settings['splitParams']['ttPolicy'] == "SPLIT_SINGLE_DATASET":
+            self.mltask_settings['splitParams']['ssdSplitMode'] = "RANDOM"
+            self.mltask_settings['splitParams']['ssdColumn'] = None
+        if self.mltask_settings['modeling']['gridSearchParams']['mode'] == "TIME_SERIES_KFOLD":
+            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "KFOLD"
+        elif self.mltask_settings['modeling']['gridSearchParams']['mode'] == "TIME_SERIES_SINGLE_SPLIT":
+            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "SHUFFLE"
+
+        return self
+
 
 class DSSMLTaskSettings(object):
     """
@@ -368,37 +407,20 @@ def split_ordered_by(self, feature_name, ascending=True):
         Uses a variable to sort the data for train/test split and hyperparameter optimization
         :param str feature_name: Name of the variable to use
         :param bool ascending: True iff the test set is expected to have larger time values than the train set
+
+        :rtype: self
         """
-        self.remove_ordered_split()
-        if not feature_name in self.mltask_settings["preprocessing"]["per_feature"]:
-            raise ValueError("Feature %s doesn't exist in this ML task, can't use as time" % feature_name)
-        self.mltask_settings['time']['enabled'] = True
-        self.mltask_settings['time']['timeVariable'] = feature_name
-        self.mltask_settings['time']['ascending'] = ascending
-        self.mltask_settings['preprocessing']['per_feature'][feature_name]['missing_handling'] = "DROP_ROW"
-        if self.mltask_settings['splitParams']['ttPolicy'] == "SPLIT_SINGLE_DATASET":
-            self.mltask_settings['splitParams']['ssdSplitMode'] = "SORTED"
-            self.mltask_settings['splitParams']['ssdColumn'] = feature_name
-        if self.mltask_settings['modeling']['gridSearchParams']['mode'] == "KFOLD":
-            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "TIME_SERIES_KFOLD"
-        elif self.mltask_settings['modeling']['gridSearchParams']['mode'] == "SHUFFLE":
-            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "TIME_SERIES_SINGLE_SPLIT"
+        self.split_params.set_order_by(feature_name, ascending=True)
 
         return self
 
     def remove_ordered_split(self):
         """
         Remove time-based ordering.
+
+        :rtype: self
         """
-        self.mltask_settings['time']['enabled'] = False
-        self.mltask_settings['time']['timeVariable'] = None
-        if self.mltask_settings['splitParams']['ttPolicy'] == "SPLIT_SINGLE_DATASET":
-            self.mltask_settings['splitParams']['ssdSplitMode'] = "RANDOM"
-            self.mltask_settings['splitParams']['ssdColumn'] = None
-        if self.mltask_settings['modeling']['gridSearchParams']['mode'] == "TIME_SERIES_KFOLD":
-            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "KFOLD"
-        elif self.mltask_settings['modeling']['gridSearchParams']['mode'] == "TIME_SERIES_SINGLE_SPLIT":
-            self.mltask_settings['modeling']['gridSearchParams']['mode'] = "SHUFFLE"
+        self.split_params.unset_order_by()
 
         return self