merging latest changes

Author: govarsha

ads/opctl/operator/lowcode/forecast/const.py

@@ -81,3 +81,4 @@ class ForecastOutputColumns(str, metaclass=ExtendedEnumMeta):
 DEFAULT_TRIALS = 10
 SUMMARY_METRICS_HORIZON_LIMIT = 10
 PROPHET_INTERNAL_DATE_COL = "ds"
+RENDER_LIMIT = 5000

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

@@ -7,6 +7,7 @@
 import pandas as pd
 import numpy as np
 import pmdarima as pm
+from joblib import Parallel, delayed
 
 from ads.opctl import logger
 
@@ -15,7 +16,7 @@
 from ..operator_config import ForecastOperatorConfig
 import traceback
 from .forecast_datasets import ForecastDatasets, ForecastOutput
-from ..const import ForecastOutputColumns, SupportedModels
+from ..const import ForecastOutputColumns
 
 
 class ArimaOperatorModel(ForecastOperatorBaseModel):
@@ -29,32 +30,31 @@ def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
         self.formatted_global_explanation = None
         self.formatted_local_explanation = None
 
-    def _build_model(self) -> pd.DataFrame:
-        full_data_dict = self.datasets.full_data_dict
-
-        # Extract the Confidence Interval Width and convert to arima's equivalent - alpha
-        if self.spec.confidence_interval_width is None:
-            self.spec.confidence_interval_width = 1 - self.spec.model_kwargs.get(
-                "alpha", 0.05
-            )
-        model_kwargs = self.spec.model_kwargs
-        model_kwargs["alpha"] = 1 - self.spec.confidence_interval_width
-        if "error_action" not in model_kwargs.keys():
-            model_kwargs["error_action"] = "ignore"
+    def _train_model(self, i, target, df):
+        """Trains the ARIMA model for a given target.
 
-        models = []
-        self.datasets.datetime_col = self.spec.datetime_column.name
-        self.forecast_output = ForecastOutput(
-            confidence_interval_width=self.spec.confidence_interval_width
-        )
+        Parameters
+        ----------
+        i: int
+            The index of the target
+        target: str
+            The name of the target
+        df: pd.DataFrame
+            The dataframe containing the target data
+        """
+        try:
+            # Extract the Confidence Interval Width and convert to arima's equivalent - alpha
+            if self.spec.confidence_interval_width is None:
+                self.spec.confidence_interval_width = 1 - self.spec.model_kwargs.get(
+                    "alpha", 0.05
+                )
+            model_kwargs = self.spec.model_kwargs
+            model_kwargs["alpha"] = 1 - self.spec.confidence_interval_width
+            if "error_action" not in model_kwargs.keys():
+                model_kwargs["error_action"] = "ignore"
 
-        outputs = dict()
-        outputs_legacy = []
-        fitted_values = dict()
-        actual_values = dict()
-        dt_columns = dict()
+            # models = []
 
-        for i, (target, df) in enumerate(full_data_dict.items()):
             # format the dataframe for this target. Dropping NA on target[df] will remove all future data
             le, df_encoded = utils._label_encode_dataframe(
                 df, no_encode={self.spec.datetime_column.name, target}
@@ -72,34 +72,28 @@ def _build_model(self) -> pd.DataFrame:
                 target,
                 self.spec.datetime_column.name,
             }
-            logger.debug(
-                f"Additional Regressors Detected {list(additional_regressors)}"
-            )
+            logger.debug(f"Additional Regressors Detected {list(additional_regressors)}")
 
             # Split data into X and y for arima tune method
             y = data_i[target]
             X_in = None
             if len(additional_regressors):
                 X_in = data_i.drop(target, axis=1)
 
-            model = self.loaded_models[i] if self.loaded_models is not None else None
-            if model is None:
-                # Build and fit model
-                model = pm.auto_arima(y=y, X=X_in, **self.spec.model_kwargs)
+            # Build and fit model
+            model = pm.auto_arima(y=y, X=X_in, **self.spec.model_kwargs)
 
-            fitted_values[target] = model.predict_in_sample(X=X_in)
-            actual_values[target] = y
-            actual_values[target].index = pd.to_datetime(y.index)
+            self.fitted_values[target] = model.predict_in_sample(X=X_in)
+            self.actual_values[target] = y
+            self.actual_values[target].index = pd.to_datetime(y.index)
 
             # Build future dataframe
             start_date = y.index.values[-1]
             n_periods = self.spec.horizon
             if len(additional_regressors):
                 X = df_clean[df_clean[target].isnull()].drop(target, axis=1)
             else:
-                X = pd.date_range(
-                    start=start_date, periods=n_periods, freq=self.spec.freq
-                )
+                X = pd.date_range(start=start_date, periods=n_periods, freq=self.spec.freq)
 
             # Predict and format forecast
             yhat, conf_int = model.predict(
@@ -110,7 +104,7 @@ def _build_model(self) -> pd.DataFrame:
             )
             yhat_clean = pd.DataFrame(yhat, index=yhat.index, columns=["yhat"])
 
-            dt_columns[target] = df_encoded[self.spec.datetime_column.name]
+            self.dt_columns[target] = df_encoded[self.spec.datetime_column.name]
             conf_int_clean = pd.DataFrame(
                 conf_int, index=yhat.index, columns=["yhat_lower", "yhat_upper"]
             )
@@ -119,25 +113,42 @@ def _build_model(self) -> pd.DataFrame:
             logger.debug(forecast[["yhat", "yhat_lower", "yhat_upper"]].tail())
 
             # Collect all outputs
-            if self.loaded_models is None:
-                models.append(model)
-            outputs_legacy.append(
+            # models.append(model)
+            self.outputs_legacy.append(
                 forecast.reset_index().rename(columns={"index": "ds"})
             )
-            outputs[target] = forecast
-
-            params = vars(model).copy()
-            for param in ['arima_res', 'endog_index_']:
-                if param in params:
-                    params.pop(param)
-            self.model_parameters[target] = {
-                "framework": SupportedModels.Arima,
-                **params,
-            }
+            self.outputs[target] = forecast
+
+            self.models_dict[target] = model
 
-        self.models = self.loaded_models if self.loaded_models is not None else models
+            logger.debug("===========Done===========")
+        except Exception as e:
+            self.errors_dict[target] = {"model_name": self.spec.model, "error": str(e)}
+
+    def _build_model(self) -> pd.DataFrame:
+        full_data_dict = self.datasets.full_data_dict
+
+        self.datasets.datetime_col = self.spec.datetime_column.name
+        self.forecast_output = ForecastOutput(
+            confidence_interval_width=self.spec.confidence_interval_width
+        )
+
+        self.outputs = dict()
+        self.outputs_legacy = []
+        self.fitted_values = dict()
+        self.actual_values = dict()
+        self.dt_columns = dict()
+        self.models_dict = dict()
+        self.errors_dict = dict()
+
+        Parallel(n_jobs=-1, require="sharedmem")(
+            delayed(ArimaOperatorModel._train_model)(self, i, target, df)
+            for self, (i, (target, df)) in zip(
+                [self] * len(full_data_dict), enumerate(full_data_dict.items())
+            )
+        )
 
-        logger.debug("===========Done===========")
+        self.models = [self.models_dict[target] for target in self.target_columns]
 
         # Merge the outputs from each model into 1 df with all outputs by target and category
         col = self.original_target_column
@@ -146,15 +157,15 @@ def _build_model(self) -> pd.DataFrame:
         yhat_lower_name = ForecastOutputColumns.LOWER_BOUND
         for cat in self.categories:
             output_i = pd.DataFrame()
-            output_i["Date"] = dt_columns[f"{col}_{cat}"]
+            output_i["Date"] = self.dt_columns[f"{col}_{cat}"]
             output_i["Series"] = cat
             output_i = output_i.set_index("Date")
 
-            output_i["input_value"] = actual_values[f"{col}_{cat}"]
-            output_i["fitted_value"] = fitted_values[f"{col}_{cat}"]
-            output_i["forecast_value"] = outputs[f"{col}_{cat}"]["yhat"]
-            output_i[yhat_upper_name] = outputs[f"{col}_{cat}"]["yhat_upper"]
-            output_i[yhat_lower_name] = outputs[f"{col}_{cat}"]["yhat_lower"]
+            output_i["input_value"] = self.actual_values[f"{col}_{cat}"]
+            output_i["fitted_value"] = self.fitted_values[f"{col}_{cat}"]
+            output_i["forecast_value"] = self.outputs[f"{col}_{cat}"]["yhat"]
+            output_i[yhat_upper_name] = self.outputs[f"{col}_{cat}"]["yhat_upper"]
+            output_i[yhat_lower_name] = self.outputs[f"{col}_{cat}"]["yhat_lower"]
 
             output_i = output_i.reset_index(drop=False)
             output_col = pd.concat([output_col, output_i])
@@ -196,7 +207,7 @@ def _generate_report(self):
                 global_explanation_df = pd.DataFrame(self.global_explanation)
 
                 self.formatted_global_explanation = (
-                        global_explanation_df / global_explanation_df.sum(axis=0) * 100
+                    global_explanation_df / global_explanation_df.sum(axis=0) * 100
                 )
 
                 # Create a markdown section for the global explainability
@@ -264,7 +275,7 @@ def _custom_predict_arima(self, data):
 
         """
         date_col = self.spec.datetime_column.name
-        data[date_col] = pd.to_datetime(data[date_col], unit='s')
+        data[date_col] = pd.to_datetime(data[date_col], unit="s")
         data = data.set_index(date_col)
         # Get the index of the current series id
         series_index = self.target_columns.index(self.series_id)