Support for Non-Numeric Columns in What-If Analysis (#1065)

Author: codeloop

ads/opctl/operator/lowcode/forecast/model/arima.py

@@ -116,7 +116,10 @@ def _train_model(self, i, s_id, df, model_kwargs):
                 lower_bound=self.get_horizon(forecast["yhat_lower"]).values,
             )
 
-            self.models[s_id] = model
+            self.models[s_id] = {}
+            self.models[s_id]["model"] = model
+            self.models[s_id]["le"] = self.le[s_id]
+            self.models[s_id]["predict_component_cols"] = X_pred.columns
 
             params = vars(model).copy()
             for param in ["arima_res_", "endog_index_"]:
@@ -163,7 +166,7 @@ def _generate_report(self):
             sec5_text = rc.Heading("ARIMA Model Parameters", level=2)
             blocks = [
                 rc.Html(
-                    m.summary().as_html(),
+                    m['model'].summary().as_html(),
                     label=s_id if self.target_cat_col else None,
                 )
                 for i, (s_id, m) in enumerate(self.models.items())
@@ -251,7 +254,7 @@ def _generate_report(self):
     def get_explain_predict_fn(self, series_id):
         def _custom_predict(
             data,
-            model=self.models[series_id],
+            model=self.models[series_id]["model"],
             dt_column_name=self.datasets._datetime_column_name,
             target_col=self.original_target_column,
         ):

ads/opctl/operator/lowcode/forecast/model/automlx.py

@@ -56,8 +56,8 @@ def set_kwargs(self):
         )
         return model_kwargs_cleaned, time_budget
 
-    def preprocess(self, data):  # TODO: re-use self.le for explanations
-        _, df_encoded = _label_encode_dataframe(
+    def preprocess(self, data, series_id):  # TODO: re-use self.le for explanations
+        self.le[series_id], df_encoded = _label_encode_dataframe(
             data,
             no_encode={self.spec.datetime_column.name, self.original_target_column},
         )
@@ -125,7 +125,7 @@ def _build_model(self) -> pd.DataFrame:
                 self.forecast_output.init_series_output(
                     series_id=s_id, data_at_series=df
                 )
-                data = self.preprocess(df)
+                data = self.preprocess(df, s_id)
                 data_i = self.drop_horizon(data)
                 X_pred = self.get_horizon(data).drop(target, axis=1)
 
@@ -157,7 +157,9 @@ def _build_model(self) -> pd.DataFrame:
                     target
                 ].values
 
-                self.models[s_id] = model
+                self.models[s_id] = {}
+                self.models[s_id]["model"] = model
+                self.models[s_id]["le"] = self.le[s_id]
 
                 # In case of Naive model, model.forecast function call does not return confidence intervals.
                 if f"{target}_ci_upper" not in summary_frame:
@@ -218,7 +220,8 @@ def _generate_report(self):
         other_sections = []
 
         if len(self.models) > 0:
-            for s_id, m in models.items():
+            for s_id, artifacts in models.items():
+                m = artifacts["model"]
                 selected_models[s_id] = {
                     "series_id": s_id,
                     "selected_model": m.selected_model_,
@@ -320,7 +323,7 @@ def _generate_report(self):
         )
 
     def get_explain_predict_fn(self, series_id):
-        selected_model = self.models[series_id]
+        selected_model = self.models[series_id]["model"]
 
         # If training date, use method below. If future date, use forecast!
         def _custom_predict_fn(
@@ -338,12 +341,12 @@ def _custom_predict_fn(
             data[dt_column_name] = seconds_to_datetime(
                 data[dt_column_name], dt_format=self.spec.datetime_column.format
             )
-            data = self.preprocess(data)
+            data = self.preprocess(data, series_id)
             horizon_data = horizon_data.drop(target_col, axis=1)
             horizon_data[dt_column_name] = seconds_to_datetime(
                 horizon_data[dt_column_name], dt_format=self.spec.datetime_column.format
             )
-            horizon_data = self.preprocess(horizon_data)
+            horizon_data = self.preprocess(horizon_data, series_id)
 
             rows = []
             for i in range(data.shape[0]):
@@ -421,7 +424,7 @@ def explain_model(self):
                 if self.spec.explanations_accuracy_mode == SpeedAccuracyMode.AUTOMLX:
                     # Use the MLExplainer class from AutoMLx to generate explanations
                     explainer = automlx.MLExplainer(
-                        self.models[s_id],
+                        self.models[s_id]["model"],
                         self.datasets.additional_data.get_data_for_series(series_id=s_id)
                         .drop(self.spec.datetime_column.name, axis=1)
                         .head(-self.spec.horizon)

ads/opctl/operator/lowcode/forecast/model/base_model.py

@@ -794,7 +794,7 @@ def local_explainer(self, kernel_explainer, series_id, datetime_col_name) -> Non
     def get_explain_predict_fn(self, series_id, fcst_col_name="yhat"):
         def _custom_predict(
             data,
-            model=self.models[series_id],
+            model=self.models[series_id]["model"],
             dt_column_name=self.datasets._datetime_column_name,
         ):
             """

ads/opctl/operator/lowcode/forecast/model/neuralprophet.py

@@ -172,8 +172,10 @@ def _train_model(self, i, s_id, df, model_kwargs):
                 ).values,
             )
 
-            self.models[s_id] = model
             self.trainers[s_id] = model.trainer
+            self.models[s_id] = {}
+            self.models[s_id]["model"] = model
+            self.models[s_id]["le"] = self.le[s_id]
 
             self.model_parameters[s_id] = {
                 "framework": SupportedModels.NeuralProphet,
@@ -355,7 +357,8 @@ def _generate_report(self):
 
             sec5_text = rc.Heading("Neural Prophet Model Parameters", level=2)
             model_states = []
-            for s_id, m in self.models.items():
+            for s_id, artifacts in self.models.items():
+                m = artifacts["model"]
                 model_states.append(
                     pd.Series(
                         m.state_dict(),

ads/opctl/operator/lowcode/forecast/model/prophet.py

@@ -108,7 +108,10 @@ def _train_model(self, i, series_id, df, model_kwargs):
                 upper_bound=self.get_horizon(forecast["yhat_upper"]).values,
                 lower_bound=self.get_horizon(forecast["yhat_lower"]).values,
             )
-            self.models[series_id] = model
+
+            self.models[series_id] = {}
+            self.models[series_id]["model"] = model
+            self.models[series_id]["le"] = self.le[series_id]
 
             params = vars(model).copy()
             for param in ["history", "history_dates", "stan_fit"]:
@@ -252,7 +255,7 @@ def _generate_report(self):
         all_sections = []
         if len(series_ids) > 0:
             sec1 = _select_plot_list(
-                lambda s_id: self.models[s_id].plot(
+                lambda s_id: self.models[s_id]["model"].plot(
                     self.outputs[s_id], include_legend=True
                 ),
                 series_ids=series_ids,
@@ -267,7 +270,7 @@ def _generate_report(self):
             )
 
             sec2 = _select_plot_list(
-                lambda s_id: self.models[s_id].plot_components(self.outputs[s_id]),
+                lambda s_id: self.models[s_id]["model"].plot_components(self.outputs[s_id]),
                 series_ids=series_ids,
                 target_category_column=self.target_cat_col
             )
@@ -276,11 +279,11 @@ def _generate_report(self):
             )
 
             sec3_figs = {
-                s_id: self.models[s_id].plot(self.outputs[s_id]) for s_id in series_ids
+                s_id: self.models[s_id]["model"].plot(self.outputs[s_id]) for s_id in series_ids
             }
             for s_id in series_ids:
                 add_changepoints_to_plot(
-                    sec3_figs[s_id].gca(), self.models[s_id], self.outputs[s_id]
+                    sec3_figs[s_id].gca(), self.models[s_id]["model"], self.outputs[s_id]
                 )
             sec3 = _select_plot_list(
                 lambda s_id: sec3_figs[s_id],
@@ -294,7 +297,7 @@ def _generate_report(self):
             sec5_text = rc.Heading("Prophet Model Seasonality Components", level=2)
             model_states = []
             for s_id in series_ids:
-                m = self.models[s_id]
+                m = self.models[s_id]["model"]
                 model_states.append(
                     pd.Series(
                         m.seasonalities,

ads/opctl/operator/lowcode/forecast/whatifserve/score.py

@@ -151,34 +151,42 @@ def get_forecast(future_df, model_name, series_id, model_object, date_col, targe
         pred_obj = model_object.predict(future_regressor=future_reg)
         return pred_obj.forecast[series_id].tolist()
     elif model_name == SupportedModels.Prophet and series_id in model_object:
-        model = model_object[series_id]
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
         processed = future_df.rename(columns={date_col_name: 'ds', target_column: 'y'})
-        forecast = model.predict(processed)
+        encoded_df = label_encoder.transform(processed)
+        forecast = model.predict(encoded_df)
         return forecast['yhat'].tolist()
     elif model_name == SupportedModels.NeuralProphet and series_id in model_object:
-        model = model_object[series_id]
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
         model.restore_trainer()
         accepted_regressors = list(model.config_regressors.regressors.keys())
         data = future_df.rename(columns={date_col_name: 'ds', target_column: 'y'})
-        future = data[accepted_regressors + ["ds"]].reset_index(drop=True)
+        encoded_df = label_encoder.transform(data)
+        future = encoded_df[accepted_regressors + ["ds"]].reset_index(drop=True)
         future["y"] = None
         forecast = model.predict(future)
         return forecast['yhat1'].tolist()
     elif model_name == SupportedModels.Arima and series_id in model_object:
-        model = model_object[series_id]
-        future_df = future_df.set_index(date_col_name)
-        x_pred = future_df.drop(target_cat_col, axis=1)
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
+        predict_cols = model_object[series_id]["predict_component_cols"]
+        encoded_df = label_encoder.transform(future_df)
+        x_pred = encoded_df.set_index(date_col_name)
+        x_pred = x_pred.drop(target_cat_col, axis=1)
         yhat, conf_int = model.predict(
             n_periods=horizon,
-            X=x_pred,
+            X=x_pred[predict_cols],
             return_conf_int=True
         )
         yhat_clean = pd.DataFrame(yhat, index=yhat.index, columns=["yhat"])
         return yhat_clean['yhat'].tolist()
     elif model_name == SupportedModels.AutoMLX and series_id in model_object:
-        # automlx model
-        model = model_object[series_id]
-        x_pred = future_df.drop(target_cat_col, axis=1)
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
+        encoded_df = label_encoder.transform(future_df)
+        x_pred = encoded_df.drop(target_cat_col, axis=1)
         x_pred = x_pred.set_index(date_col_name)
         forecast = model.forecast(
             X=x_pred,

tests/operators/forecast/test_errors.py

@@ -875,10 +875,7 @@ def test_what_if_analysis(operator_setup, model):
     historical_filtered = historical_data[historical_data['Date'] > "2013-03-01"]
     additional_data = pd.read_csv(additional_data_path, parse_dates=["Date"])
     add_filtered = additional_data[additional_data['Date'] > "2013-03-01"]
-    numeric_columns = add_filtered.select_dtypes(include=['number', 'object', 'datetime64'])
-    non_constant_columns = numeric_columns.columns[(numeric_columns != numeric_columns.iloc[0]).any()]
-    df_non_constant = numeric_columns[non_constant_columns.union(['Store'])]
-    df_non_constant.to_csv(f'{additional_test_path}', index=False)
+    add_filtered.to_csv(f'{additional_test_path}', index=False)
     historical_filtered.to_csv(f'{historical_test_path}', index=False)
 
     yaml_i, output_data_path = populate_yaml(