oracle-ads 2.10.0__py3-none-any.whl → 2.11.0__py3-none-any.whl

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

@@ -7,15 +7,17 @@
 import pandas as pd
 import numpy as np
 import pmdarima as pm
+from joblib import Parallel, delayed
 
 from ads.opctl import logger
 
-from .. import utils
+from ads.opctl.operator.lowcode.forecast.utils import _label_encode_dataframe
+from ads.opctl.operator.lowcode.common.utils import seconds_to_datetime
 from .base_model import ForecastOperatorBaseModel
 from ..operator_config import ForecastOperatorConfig
 import traceback
 from .forecast_datasets import ForecastDatasets, ForecastOutput
-from ..const import ForecastOutputColumns
+from ..const import ForecastOutputColumns, SupportedModels
 
 
 class ArimaOperatorModel(ForecastOperatorBaseModel):
@@ -25,13 +27,10 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
         super().__init__(config, datasets=datasets)
         self.global_explanation = {}
         self.local_explanation = {}
-        self.train_metrics = True
         self.formatted_global_explanation = None
         self.formatted_local_explanation = None
 
-    def _build_model(self) -> pd.DataFrame:
-        full_data_dict = self.datasets.full_data_dict
-
+    def set_kwargs(self):
         # 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(
@@ -41,74 +40,57 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
         model_kwargs["alpha"] = 1 - self.spec.confidence_interval_width
         if "error_action" not in model_kwargs.keys():
             model_kwargs["error_action"] = "ignore"
+        return model_kwargs
 
-        models = []
-        self.datasets.datetime_col = self.spec.datetime_column.name
-        self.forecast_output = ForecastOutput(
-            confidence_interval_width=self.spec.confidence_interval_width
+    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},
         )
+        return df_encoded.set_index(self.spec.datetime_column.name)
 
-        outputs = dict()
-        outputs_legacy = []
-        fitted_values = dict()
-        actual_values = dict()
-        dt_columns = dict()
-
-        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}
-            )
+    def _train_model(self, i, s_id, df, model_kwargs):
+        """Trains the ARIMA model for a given series of the dataset.
 
-            df_encoded[self.spec.datetime_column.name] = pd.to_datetime(
-                df_encoded[self.spec.datetime_column.name],
-                format=self.spec.datetime_column.format,
-            )
-            df_clean = df_encoded.set_index(self.spec.datetime_column.name)
-            data_i = df_clean[df_clean[target].notna()]
+        Parameters
+        ----------
+        i: int
+            The index of the series
+        s_id: str
+            The name of the series
+        df: pd.DataFrame
+            The dataframe containing the target data
+        """
+        try:
+            target = self.original_target_column
+            self.forecast_output.init_series_output(series_id=s_id, data_at_series=df)
 
-            # Assume that all columns passed in should be used as additional data
-            additional_regressors = set(data_i.columns) - {
-                target,
-                self.spec.datetime_column.name,
-            }
-            logger.debug(
-                f"Additional Regressors Detected {list(additional_regressors)}"
-            )
+            # format the dataframe for this target. Dropping NA on target[df] will remove all future data
+            data = self.preprocess(df, s_id)
+            data_i = self.drop_horizon(data)
 
             # 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)
-
-            # 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)
-
-            # 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)
+            X_in = data_i.drop(target, axis=1) if len(data_i.columns) > 1 else None
+            X_pred = self.get_horizon(data).drop(target, axis=1)
+
+            if self.loaded_models is not None:
+                model = self.loaded_models[s_id]
             else:
-                X = pd.date_range(
-                    start=start_date, periods=n_periods, freq=self.spec.freq
-                )
+                # Build and fit model
+                model = pm.auto_arima(y=y, X=X_in, **model_kwargs)
+
+            fitted_values = model.predict_in_sample(X=X_in).values
 
             # Predict and format forecast
             yhat, conf_int = model.predict(
-                n_periods=n_periods,
-                X=X,
+                n_periods=self.spec.horizon,
+                X=X_pred,
                 return_conf_int=True,
                 alpha=model_kwargs["alpha"],
             )
             yhat_clean = pd.DataFrame(yhat, index=yhat.index, columns=["yhat"])
 
-            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"]
             )
@@ -116,43 +98,47 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
             logger.debug(f"-----------------Model {i}----------------------")
             logger.debug(forecast[["yhat", "yhat_lower", "yhat_upper"]].tail())
 
-            # Collect all outputs
-            models.append(model)
-            outputs_legacy.append(
-                forecast.reset_index().rename(columns={"index": "ds"})
-            )
-            outputs[target] = forecast
-
-        self.models = models
-
-        logger.debug("===========Done===========")
-
-        # Merge the outputs from each model into 1 df with all outputs by target and category
-        col = self.original_target_column
-        output_col = pd.DataFrame()
-        yhat_upper_name = ForecastOutputColumns.UPPER_BOUND
-        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["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 = output_i.reset_index(drop=False)
-            output_col = pd.concat([output_col, output_i])
-            self.forecast_output.add_category(
-                category=cat, target_category_column=f"{col}_{cat}", forecast=output_i
+            self.forecast_output.populate_series_output(
+                series_id=s_id,
+                fit_val=fitted_values,
+                forecast_val=self.get_horizon(forecast["yhat"]).values,
+                upper_bound=self.get_horizon(forecast["yhat_upper"]).values,
+                lower_bound=self.get_horizon(forecast["yhat_lower"]).values,
             )
 
-        output_col = output_col.reset_index(drop=True)
+            self.models[s_id] = model
+
+            params = vars(model).copy()
+            for param in ["arima_res_", "endog_index_"]:
+                if param in params:
+                    params.pop(param)
+            self.model_parameters[s_id] = {
+                "framework": SupportedModels.Arima,
+                **params,
+            }
 
-        return output_col
+            logger.debug("===========Done===========")
+        except Exception as e:
+            self.errors_dict[s_id] = {"model_name": self.spec.model, "error": str(e)}
+
+    def _build_model(self) -> pd.DataFrame:
+        full_data_dict = self.datasets.get_data_by_series()
+        self.models = dict()
+        self.additional_regressors = self.datasets.get_additional_data_column_names()
+        model_kwargs = self.set_kwargs()
+        self.forecast_output = ForecastOutput(
+            confidence_interval_width=self.spec.confidence_interval_width,
+            horizon=self.spec.horizon,
+            target_column=self.original_target_column,
+            dt_column=self.spec.datetime_column.name,
+        )
+
+        Parallel(n_jobs=-1, require="sharedmem")(
+            delayed(self._train_model)(i, s_id, df, model_kwargs.copy())
+            for (i, (s_id, df)) in enumerate(full_data_dict.items())
+        )
+
+        return self.forecast_output.get_forecast_long()
 
     def _generate_report(self):
         """The method that needs to be implemented on the particular model level."""
@@ -160,8 +146,11 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
 
         sec5_text = dp.Text(f"## ARIMA Model Parameters")
         blocks = [
-            dp.HTML(m.summary().as_html(), label=self.target_columns[i])
-            for i, m in enumerate(self.models)
+            dp.HTML(
+                m.summary().as_html(),
+                label=s_id,
+            )
+            for i, (s_id, m) in enumerate(self.models.items())
         ]
         sec5 = dp.Select(blocks=blocks) if len(blocks) > 1 else blocks[0]
         all_sections = [sec5_text, sec5]
@@ -169,11 +158,7 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
         if self.spec.generate_explanations:
             try:
                 # If the key is present, call the "explain_model" method
-                self.explain_model(
-                    datetime_col_name=self.spec.datetime_column.name,
-                    explain_predict_fn=self._custom_predict_arima,
-                )
-
+                self.explain_model()
                 # Create a markdown text block for the global explanation section
                 global_explanation_text = dp.Text(
                     f"## Global Explanation of Models \n "
@@ -186,6 +171,12 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
                 self.formatted_global_explanation = (
                     global_explanation_df / global_explanation_df.sum(axis=0) * 100
                 )
+                self.formatted_global_explanation = (
+                    self.formatted_global_explanation.rename(
+                        {self.spec.datetime_column.name: ForecastOutputColumns.DATE},
+                        axis=1,
+                    )
+                )
 
                 # Create a markdown section for the global explainability
                 global_explanation_section = dp.Blocks(
@@ -238,23 +229,21 @@ class ArimaOperatorModel(ForecastOperatorBaseModel):
             other_sections,
         )
 
-    def _custom_predict_arima(self, data):
-        """
-        Custom prediction function for ARIMA models.
-
-        Parameters
-        ----------
-            data (array-like): The input data to be predicted.
-
-        Returns
-        -------
-            array-like: The predicted values.
-
-        """
-        # Get the index of the current series id
-        series_index = self.target_columns.index(self.series_id)
-
-        # Use the ARIMA model to predict the values
-        predictions = self.models[series_index].predict(X=data, n_periods=len(data))
+    def get_explain_predict_fn(self, series_id):
+        def _custom_predict(
+            data,
+            model=self.models[series_id],
+            dt_column_name=self.datasets._datetime_column_name,
+            target_col=self.original_target_column,
+        ):
+            """
+            data: ForecastDatasets.get_data_at_series(s_id)
+            """
+            data = data.drop([target_col], axis=1)
+            data[dt_column_name] = seconds_to_datetime(
+                data[dt_column_name], dt_format=self.spec.datetime_column.format
+            )
+            data = self.preprocess(data, series_id)
+            return model.predict(X=data, n_periods=len(data))
 
-        return predictions
+        return _custom_predict