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

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

@@ -7,6 +7,7 @@
 import numpy as np
 import optuna
 import pandas as pd
+from joblib import Parallel, delayed
 from torch import Tensor
 from torchmetrics.regression import (
     MeanAbsoluteError,
@@ -22,8 +23,18 @@ from ads.common.decorator.runtime_dependency import (
 )
 from ads.opctl import logger
 
-from ..const import DEFAULT_TRIALS, ForecastOutputColumns
-from .. import utils
+from ..const import DEFAULT_TRIALS, ForecastOutputColumns, SupportedModels
+from ads.opctl.operator.lowcode.forecast.utils import (
+    load_pkl,
+    write_pkl,
+    _select_plot_list,
+    _label_encode_dataframe,
+)
+from ads.opctl.operator.lowcode.common.utils import (
+    disable_print,
+    enable_print,
+    seconds_to_datetime,
+)
 from .base_model import ForecastOperatorBaseModel
 from ..operator_config import ForecastOperatorConfig
 from .forecast_datasets import ForecastDatasets, ForecastOutput
@@ -52,7 +63,11 @@ def _get_np_metrics_dict(selected_metric):
     install_from=OptionalDependency.FORECAST,
 )
 def _fit_model(data, params, additional_regressors, select_metric):
-    from neuralprophet import NeuralProphet
+    from neuralprophet import NeuralProphet, set_log_level
+
+    if logger.level > 10:
+        set_log_level(logger.level)
+        disable_print()
 
     m = NeuralProphet(**params)
     m.metrics = _get_np_metrics_dict(select_metric)
@@ -60,6 +75,8 @@ def _fit_model(data, params, additional_regressors, select_metric):
         m = m.add_future_regressor(name=add_reg)
     m.fit(df=data)
     accepted_regressors_config = m.config_regressors or dict()
+
+    enable_print()
     return m, list(accepted_regressors_config.keys())
 
 
@@ -68,21 +85,21 @@ class NeuralProphetOperatorModel(ForecastOperatorBaseModel):
 
     def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
         super().__init__(config=config, datasets=datasets)
-        self.train_metrics = True
         self.forecast_col_name = "yhat1"
+        self.loaded_trainers = None
+        self.trainers = None
+
+    def _load_model(self):
+        try:
+            self.loaded_models = load_pkl(self.spec.previous_output_dir + "/model.pkl")
+            self.loaded_trainers = load_pkl(
+                self.spec.previous_output_dir + "/trainer.pkl"
+            )
+        except:
+            logger.debug("model.pkl/trainer.pkl is not present")
 
-    def _build_model(self) -> pd.DataFrame:
-        from neuralprophet import NeuralProphet
-
-        full_data_dict = self.datasets.full_data_dict
-        models = []
-        outputs = dict()
-        outputs_legacy = []
-
-        # Extract the Confidence Interval Width and
-        # convert to neural prophets equivalent - quantiles
-        model_kwargs = self.spec.model_kwargs
-
+    def set_kwargs(self):
+        # Extract the Confidence Interval Width and convert to prophet's equivalent - interval_width
         if self.spec.confidence_interval_width is None:
             quantiles = model_kwargs.get("quantiles", [0.05, 0.95])
             self.spec.confidence_interval_width = float(quantiles[1]) - float(
@@ -92,243 +109,262 @@ class NeuralProphetOperatorModel(ForecastOperatorBaseModel):
             boundaries = round((1 - self.spec.confidence_interval_width) / 2, 2)
             quantiles = [boundaries, self.spec.confidence_interval_width + boundaries]
 
+        model_kwargs = self.spec.model_kwargs
         model_kwargs["quantiles"] = quantiles
-        self.forecast_output = ForecastOutput(
-            confidence_interval_width=self.spec.confidence_interval_width
-        )
-
-        for i, (target, df) in enumerate(full_data_dict.items()):
-            le, df_encoded = utils._label_encode_dataframe(
-                df, no_encode={self.spec.datetime_column.name, target}
-            )
-            model_kwargs_i = model_kwargs.copy()
-
-            # format the dataframe for this target. Dropping NA on target[df] will remove all future data
-            df_clean = self._preprocess(
-                df_encoded,
-                self.spec.datetime_column.name,
-                self.spec.datetime_column.format,
-            )
-            data_i = df_clean[df_clean[target].notna()]
-            data_i.rename({target: "y"}, axis=1, inplace=True)
-
-            # Assume that all columns passed in should be used as additional data
-            additional_regressors = set(data_i.columns) - {"y", "ds"}
-            training_data = data_i[["y", "ds"] + list(additional_regressors)]
-
-            if self.perform_tuning:
-
-                def objective(trial):
-                    params = {
-                        # 'seasonality_mode': trial.suggest_categorical('seasonality_mode', ['additive', 'multiplicative']),
-                        # 'seasonality_reg': trial.suggest_float('seasonality_reg', 0.1, 500, log=True),
-                        # 'learning_rate': trial.suggest_float('learning_rate',  0.0001, 0.1, log=True),
-                        "newer_samples_start": trial.suggest_float(
-                            "newer_samples_start", 0.001, 0.999
-                        ),
-                        "newer_samples_weight": trial.suggest_float(
-                            "newer_samples_weight", 0, 100
-                        ),
-                        "changepoints_range": trial.suggest_float(
-                            "changepoints_range", 0.8, 0.95
-                        ),
-                    }
-                    # trend_reg, trend_reg_threshold, ar_reg, impute_rolling/impute_linear,
-                    params.update(model_kwargs_i)
-
-                    folds = NeuralProphet(**params).crossvalidation_split_df(
-                        data_i, k=3
-                    )
-                    test_metrics_total_i = []
-                    for df_train, df_test in folds:
-                        m, accepted_regressors = _fit_model(
-                            data=df_train,
-                            params=params,
-                            additional_regressors=additional_regressors,
-                            select_metric=self.spec.metric,
-                        )
-                        df_test = df_test[["y", "ds"] + accepted_regressors]
-
-                        test_forecast_i = m.predict(df=df_test)
-                        fold_metric_i = (
-                            m.metrics[self.spec.metric]
-                            .forward(
-                                Tensor(test_forecast_i["yhat1"]),
-                                Tensor(test_forecast_i["y"]),
-                            )
-                            .item()
-                        )
-                        test_metrics_total_i.append(fold_metric_i)
-                    logger.debug(
-                        f"----------------------{np.asarray(test_metrics_total_i).mean()}----------------------"
-                    )
-                    return np.asarray(test_metrics_total_i).mean()
-
-                study = optuna.create_study(direction="minimize")
-                m_params = NeuralProphet().parameters()
-                study.enqueue_trial(
-                    {
-                        # 'seasonality_mode': m_params['seasonality_mode'],
-                        # 'seasonality_reg': m_params['seasonality_reg'],
-                        # 'learning_rate': m_params['learning_rate'],
-                        "newer_samples_start": m_params["newer_samples_start"],
-                        "newer_samples_weight": m_params["newer_samples_weight"],
-                        "changepoints_range": m_params["changepoints_range"],
-                    }
-                )
-                study.optimize(
-                    objective,
-                    n_trials=self.spec.tuning.n_trials
-                    if self.spec.tuning
-                    else DEFAULT_TRIALS,
-                    n_jobs=-1,
+        return model_kwargs
+
+    def _train_model(self, i, s_id, df, model_kwargs):
+        try:
+            self.forecast_output.init_series_output(series_id=s_id, data_at_series=df)
+
+            data = self.preprocess(df, s_id)
+            data_i = self.drop_horizon(data)
+
+            if self.loaded_models is not None:
+                model = self.loaded_models[s_id]
+                accepted_regressors_config = model.config_regressors or dict()
+                self.accepted_regressors[s_id] = list(accepted_regressors_config.keys())
+                if self.loaded_trainers is not None:
+                    model.trainer = self.loaded_trainers[s_id]
+            else:
+                if self.perform_tuning:
+                    model_kwargs = self.run_tuning(data_i, model_kwargs)
+
+                # Build and fit model
+                model, self.accepted_regressors[s_id] = _fit_model(
+                    data=data_i,
+                    params=model_kwargs,
+                    additional_regressors=self.additional_regressors,
+                    select_metric=self.spec.metric,
                 )
 
-                selected_params = study.best_params
-                selected_params.update(model_kwargs_i)
-                model_kwargs_i = selected_params
-
-            # Build and fit model
-            model, accepted_regressors = _fit_model(
-                data=training_data,
-                params=model_kwargs_i,
-                additional_regressors=additional_regressors,
-                select_metric=self.spec.metric,
-            )
-            logger.debug(
-                f"Found the following additional data columns: {additional_regressors}"
-            )
             logger.debug(
-                f"While fitting the model, some additional data may have been "
-                f"discarded. Only using the columns: {accepted_regressors}"
+                f"Found the following additional data columns: {self.additional_regressors}"
             )
-
+            if set(self.additional_regressors) - set(self.accepted_regressors[s_id]):
+                logger.debug(
+                    f"While fitting the model, some additional data may have been "
+                    f"discarded. Only using the columns: {self.accepted_regressors[s_id]}"
+                )
             # Build future dataframe
-            future = df_clean.reset_index(drop=True)
+            future = data[self.accepted_regressors[s_id] + ["ds"]].reset_index(
+                drop=True
+            )
             future["y"] = None
-            future = future[["y", "ds"] + list(accepted_regressors)]
 
-            # Forecast model and collect outputs
             forecast = model.predict(future)
             logger.debug(f"-----------------Model {i}----------------------")
             logger.debug(forecast.tail())
-            models.append(model)
-            outputs[target] = forecast
-            outputs_legacy.append(forecast)
-
-        self.models = models
-        self.outputs = outputs_legacy
-
-        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"] = outputs[f"{col}_{cat}"]["ds"]
-            output_i["Series"] = cat
-            output_i[f"input_value"] = full_data_dict[f"{col}_{cat}"][f"{col}_{cat}"]
-
-            output_i[f"fitted_value"] = float("nan")
-            output_i[f"forecast_value"] = float("nan")
-            output_i[yhat_lower_name] = float("nan")
-            output_i[yhat_upper_name] = float("nan")
-
-            output_i.iloc[
-                : -self.spec.horizon, output_i.columns.get_loc(f"fitted_value")
-            ] = (outputs[f"{col}_{cat}"]["yhat1"].iloc[: -self.spec.horizon].values)
-            output_i.iloc[
-                -self.spec.horizon :,
-                output_i.columns.get_loc(f"forecast_value"),
-            ] = (
-                outputs[f"{col}_{cat}"]["yhat1"].iloc[-self.spec.horizon :].values
-            )
-            output_i.iloc[
-                -self.spec.horizon :,
-                output_i.columns.get_loc(yhat_upper_name),
-            ] = (
-                outputs[f"{col}_{cat}"][f"yhat1 {quantiles[1]*100}%"]
-                .iloc[-self.spec.horizon :]
-                .values
+
+            # TODO; could also extract trend and seasonality?
+            cols_to_read = filter(
+                lambda x: x.startswith("future_regressor"), forecast.columns
             )
-            output_i.iloc[
-                -self.spec.horizon :,
-                output_i.columns.get_loc(yhat_lower_name),
-            ] = (
-                outputs[f"{col}_{cat}"][f"yhat1 {quantiles[0]*100}%"]
-                .iloc[-self.spec.horizon :]
-                .values
+            self.explanations_info[s_id] = forecast[cols_to_read]
+            self.explanations_info[s_id]["Date"] = forecast["ds"]
+            self.explanations_info[s_id] = self.explanations_info[s_id].set_index(
+                "Date"
             )
-            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.outputs[s_id] = forecast
+            self.forecast_output.populate_series_output(
+                series_id=s_id,
+                fit_val=self.drop_horizon(forecast["yhat1"]).values,
+                forecast_val=self.get_horizon(forecast["yhat1"]).values,
+                upper_bound=self.get_horizon(
+                    forecast[f"yhat1 {model_kwargs['quantiles'][1]*100}%"]
+                ).values,
+                lower_bound=self.get_horizon(
+                    forecast[f"yhat1 {model_kwargs['quantiles'][0]*100}%"]
+                ).values,
             )
 
-        output_col = output_col.reset_index(drop=True)
+            self.models[s_id] = model
+            self.trainers[s_id] = model.trainer
+
+            self.model_parameters[s_id] = {
+                "framework": SupportedModels.NeuralProphet,
+                "config": model.config,
+                "config_trend": model.config_trend,
+                "config_train": model.config_train,
+                "config_seasonality": model.config_seasonality,
+                "config_regressors": model.config_regressors,
+                "config_ar": model.config_ar,
+                "config_events": model.config_events,
+                "config_country_holidays": model.config_country_holidays,
+                "config_lagged_regressors": model.config_lagged_regressors,
+                "config_normalization": model.config_normalization,
+                "config_missing": model.config_missing,
+                "config_model": model.config_model,
+                "config_normalization": model.config_normalization,
+                "data_freq": model.data_freq,
+                "fitted": model.fitted,
+                "data_params": model.data_params,
+                "future_periods": model.future_periods,
+                "predict_steps": model.predict_steps,
+                "highlight_forecast_step_n": model.highlight_forecast_step_n,
+                "true_ar_weights": model.true_ar_weights,
+            }
+
+            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.trainers = dict()
+        self.outputs = dict()
+        self.errors_dict = dict()
+        self.explanations_info = dict()
+        self.accepted_regressors = 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,
+        )
+
+        for i, (s_id, df) in enumerate(full_data_dict.items()):
+            self._train_model(i, s_id, df, model_kwargs=model_kwargs.copy())
 
-        return output_col
+        # Parallel(n_jobs=-1, require="sharedmem")(
+        #     delayed(NeuralProphetOperatorModel._train_model)(self, i, s_id, df, model_kwargs=model_kwargs.copy())
+        #     for self, (i, (s_id, df)) in zip(
+        #         [self] * len(full_data_dict), enumerate(full_data_dict.items())
+        #     )
+        # )
 
-    def _generate_report(self):
-        import datapane as dp
+        return self.forecast_output.get_forecast_long()
 
-        sec1_text = dp.Text(
-            "## Forecast Overview \nThese plots show your "
-            "forecast in the context of historical data."
+    def run_tuning(self, data, model_kwargs):
+        from neuralprophet import NeuralProphet
+
+        def objective(trial):
+            params = {
+                # 'seasonality_mode': trial.suggest_categorical('seasonality_mode', ['additive', 'multiplicative']),
+                # 'seasonality_reg': trial.suggest_float('seasonality_reg', 0.1, 500, log=True),
+                # 'learning_rate': trial.suggest_float('learning_rate',  0.0001, 0.1, log=True),
+                "newer_samples_start": trial.suggest_float(
+                    "newer_samples_start", 0.001, 0.999
+                ),
+                "newer_samples_weight": trial.suggest_float(
+                    "newer_samples_weight", 0, 100
+                ),
+                "changepoints_range": trial.suggest_float(
+                    "changepoints_range", 0.8, 0.95
+                ),
+            }
+            # trend_reg, trend_reg_threshold, ar_reg, impute_rolling/impute_linear,
+            params.update(model_kwargs)
+
+            folds = NeuralProphet(**params).crossvalidation_split_df(data, k=3)
+            test_metrics_total_i = []
+            for df_train, df_test in folds:
+                m, accepted_regressors = _fit_model(
+                    data=df_train,
+                    params=params,
+                    additional_regressors=self.additional_regressors,
+                    select_metric=self.spec.metric,
+                )
+                df_test = df_test[["y", "ds"] + accepted_regressors]
+
+                test_forecast_i = m.predict(df=df_test)
+                fold_metric_i = (
+                    m.metrics[self.spec.metric]
+                    .forward(
+                        Tensor(test_forecast_i["yhat1"]),
+                        Tensor(test_forecast_i["y"]),
+                    )
+                    .item()
+                )
+                test_metrics_total_i.append(fold_metric_i)
+            logger.debug(
+                f"----------------------{np.asarray(test_metrics_total_i).mean()}----------------------"
+            )
+            return np.asarray(test_metrics_total_i).mean()
+
+        study = optuna.create_study(direction="minimize")
+        m_params = NeuralProphet().parameters()
+        study.enqueue_trial(
+            {
+                # 'seasonality_mode': m_params['seasonality_mode'],
+                # 'seasonality_reg': m_params['seasonality_reg'],
+                # 'learning_rate': m_params['learning_rate'],
+                "newer_samples_start": m_params["newer_samples_start"],
+                "newer_samples_weight": m_params["newer_samples_weight"],
+                "changepoints_range": m_params["changepoints_range"],
+            }
         )
-        sec1 = utils._select_plot_list(
-            lambda idx, *args: self.models[idx].plot(self.outputs[idx]),
-            target_columns=self.target_columns,
+        study.optimize(
+            objective,
+            n_trials=self.spec.tuning.n_trials if self.spec.tuning else DEFAULT_TRIALS,
+            n_jobs=-1,
         )
 
-        sec2_text = dp.Text(f"## Forecast Broken Down by Trend Component")
-        sec2 = utils._select_plot_list(
-            lambda idx, *args: self.models[idx].plot_components(self.outputs[idx]),
-            target_columns=self.target_columns,
-        )
+        selected_params = study.best_params
+        selected_params.update(model_kwargs)
+        return selected_params
 
-        sec3_text = dp.Text(f"## Forecast Parameter Plots")
-        sec3 = utils._select_plot_list(
-            lambda idx, *args: self.models[idx].plot_parameters(),
-            target_columns=self.target_columns,
-        )
+    def _generate_report(self):
+        import datapane as dp
+
+        all_sections = []
+
+        try:
+            sec1_text = dp.Text(
+                "## Forecast Overview \nThese plots show your "
+                "forecast in the context of historical data."
+            )
+            sec1 = _select_plot_list(
+                lambda s_id: self.models[s_id].plot(self.outputs[s_id]),
+                series_ids=self.datasets.list_series_ids(),
+            )
+            all_sections = all_sections + [sec1_text, sec1]
+        except Exception as e:
+            logger.debug(f"Failed to plot with exception: {e.args}")
+
+        try:
+            sec2_text = dp.Text(f"## Forecast Broken Down by Trend Component")
+            sec2 = _select_plot_list(
+                lambda s_id: self.models[s_id].plot_components(self.outputs[s_id]),
+                series_ids=self.datasets.list_series_ids(),
+            )
+            all_sections = all_sections + [sec2_text, sec2]
+        except Exception as e:
+            logger.debug(f"Failed to plot with exception: {e.args}")
+
+        try:
+            sec3_text = dp.Text(f"## Forecast Parameter Plots")
+            sec3 = _select_plot_list(
+                lambda s_id: self.models[s_id].plot_parameters(),
+                series_ids=self.datasets.list_series_ids(),
+            )
+            all_sections = all_sections + [sec3_text, sec3]
+        except Exception as e:
+            logger.debug(f"Failed to plot with exception: {e.args}")
 
         sec5_text = dp.Text(f"## Neural Prophet Model Parameters")
         model_states = []
-        for i, m in enumerate(self.models):
+        for i, (s_id, m) in enumerate(self.models.items()):
             model_states.append(
                 pd.Series(
                     m.state_dict(),
                     index=m.state_dict().keys(),
-                    name=self.target_columns[i],
+                    name=s_id,
                 )
             )
         all_model_states = pd.concat(model_states, axis=1)
         sec5 = dp.DataTable(all_model_states)
 
-        # return [sec4_text, sec4]
-        all_sections = [
-            sec1_text,
-            sec1,
-            sec2_text,
-            sec2,
-            sec3_text,
-            sec3,
-            sec5_text,
-            sec5,
-        ]
+        all_sections = all_sections + [sec5_text, sec5]
 
         if self.spec.generate_explanations:
             try:
                 # If the key is present, call the "explain_model" method
-                self.explain_model(
-                    datetime_col_name="ds",
-                    explain_predict_fn=self._custom_predict_neuralprophet,
-                )
+                self.explain_model()
 
                 # Create a markdown text block for the global explanation section
                 global_explanation_text = dp.Text(
@@ -336,32 +372,16 @@ class NeuralProphetOperatorModel(ForecastOperatorBaseModel):
                     "The following tables provide the feature attribution for the global explainability."
                 )
 
-                # Convert the global explanation data to a DataFrame
-                global_explanation_df = pd.DataFrame(self.global_explanation)
-
-                self.formatted_global_explanation = (
-                    global_explanation_df / global_explanation_df.sum(axis=0) * 100
-                )
-
                 # Create a markdown section for the global explainability
                 global_explanation_section = dp.Blocks(
                     "### Global Explainability ",
                     dp.DataTable(self.formatted_global_explanation),
                 )
 
-                aggregate_local_explanations = pd.DataFrame()
-                for s_id, local_ex_df in self.local_explanation.items():
-                    local_ex_df_copy = local_ex_df.copy()
-                    local_ex_df_copy["Series"] = s_id
-                    aggregate_local_explanations = pd.concat(
-                        [aggregate_local_explanations, local_ex_df_copy], axis=0
-                    )
-                self.formatted_local_explanation = aggregate_local_explanations
-
                 local_explanation_text = dp.Text(f"## Local Explanation of Models \n ")
                 blocks = [
                     dp.DataTable(
-                        local_ex_df.div(local_ex_df.abs().sum(axis=1), axis=0) * 100,
+                        local_ex_df.drop("Series", axis=1),
                         label=s_id,
                     )
                     for s_id, local_ex_df in self.local_explanation.items()
@@ -395,9 +415,45 @@ class NeuralProphetOperatorModel(ForecastOperatorBaseModel):
             other_sections,
         )
 
-    def _custom_predict_neuralprophet(self, data):
-        raise NotImplementedError("NeuralProphet does not yet support explanations.")
-        # data_prepped = data.reset_index()
-        # data_prepped['y'] = None
-        # data_prepped['ds'] = pd.to_datetime(data_prepped['ds'])
-        # return self.models[self.target_columns.index(self.series_id)].predict(data_prepped)["yhat1"]
+    def _save_model(self, output_dir, storage_options):
+        write_pkl(
+            obj=self.models,
+            filename="model.pkl",
+            output_dir=output_dir,
+            storage_options=storage_options,
+        )
+        write_pkl(
+            obj=self.trainers,
+            filename="trainer.pkl",
+            output_dir=output_dir,
+            storage_options=storage_options,
+        )
+
+    def explain_model(self):
+        self.local_explanation = dict()
+        global_expl = []
+        rename_cols = {
+            f"future_regressor_{col}": col
+            for col in self.datasets.get_additional_data_column_names()
+        }
+
+        for s_id, expl_df in self.explanations_info.items():
+            expl_df = expl_df.rename(rename_cols, axis=1)
+            # Local Expl
+            self.local_explanation[s_id] = self.get_horizon(expl_df).drop(
+                ["future_regressors_additive"], axis=1
+            )
+            self.local_explanation[s_id]["Series"] = s_id
+            self.local_explanation[s_id].index.rename(self.dt_column_name, inplace=True)
+            # Global Expl
+            g_expl = self.drop_horizon(expl_df).mean()
+            g_expl.name = s_id
+            global_expl.append(g_expl)
+        self.global_explanation = pd.concat(global_expl, axis=1)
+        self.global_explanation = self.global_explanation.drop(
+            index=["future_regressors_additive"], axis=0
+        )
+        self.formatted_global_explanation = (
+            self.global_explanation / self.global_explanation.sum(axis=0) * 100
+        )
+        self.formatted_local_explanation = pd.concat(self.local_explanation.values())