Add time-based ordering shortcuts for DSSMLTaskSettings

Author: samronsin

dataikuapi/dss/ml.py

@@ -138,6 +138,42 @@ def get_split_params(self):
         """
         return PredictionSplitParamsHandler(self.mltask_settings)
 
+    def use_time_variable(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.remove_time_variable()
+        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"
+        self.mltask_settings['preprocessing']['per_feature'][feature_name]['numerical_handling'] = "REGULAR"
+        self.mltask_settings['preprocessing']['per_feature'][feature_name]['rescaling'] = "NONE"
+        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"
+
+    def remove_time_variable(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"
 
     def get_feature_preprocessing(self, feature_name):
         """