views_stepshifter 0.1.1__tar.gz

views_stepshifter-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 views platform
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

views_stepshifter-0.1.1/PKG-INFO

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: views_stepshifter
+Version: 0.1.1
+Summary: 
+Author: Xiaolong Sun
+Author-email: xiaolong.sun@pcr.uu.se
+Requires-Python: >=3.11,<3.15
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: darts (>=0.30.0,<0.31.0)
+Requires-Dist: numpy (>=1.25.2,<2.0.0)
+Requires-Dist: pandas (>=1.5.3,<2.0.0)
+Requires-Dist: scikit-learn (>=1.2.2,<2.0.0)
+Requires-Dist: views_pipeline_core (>=0.1.0,<0.2.0)
+Description-Content-Type: text/markdown
+
+

views_stepshifter-0.1.1/pyproject.toml

@@ -0,0 +1,19 @@
+[tool.poetry]
+name = "views_stepshifter"
+version = "0.1.1"
+description = ""
+authors = ["Xiaolong Sun <xiaolong.sun@pcr.uu.se>"]
+readme = "README.md"
+
+[tool.poetry.dependencies]
+python = ">=3.11,<3.15"
+views_pipeline_core = "^0.1.0"
+scikit-learn = "^1.2.2"
+pandas = "^1.5.3"
+numpy = "^1.25.2"
+darts = "^0.30.0"
+
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"

views_stepshifter-0.1.1/views_stepshifter/manager/stepshifter_manager.py

@@ -0,0 +1,287 @@
+from views_pipeline_core.managers.model_manager import ModelManager
+from views_pipeline_core.managers.path_manager import ModelPath
+from views_pipeline_core.models.outputs import generate_output_dict
+from views_pipeline_core.files.utils import read_log_file, create_log_file
+from views_pipeline_core.wandb.utils import add_wandb_monthly_metrics, generate_wandb_log_dict, log_wandb_log_dict
+from views_pipeline_core.evaluation.metrics import generate_metric_dict
+from views_stepshifter.models.stepshifter import StepshifterModel
+from views_stepshifter.models.hurdle_model import HurdleModel
+from views_forecasts.extensions import *
+import logging
+import pandas as pd
+import numpy as np
+import time
+import wandb
+from datetime import datetime
+from sklearn.metrics import mean_squared_error
+
+logger = logging.getLogger(__name__)
+
+
+class StepshifterManager(ModelManager):
+
+    def __init__(self, model_path: ModelPath) -> None:
+        super().__init__(model_path)
+        self._is_hurdle = self._config_meta["algorithm"] == "HurdleModel"
+
+    def _update_sweep_config(self, args):
+        """
+        Updates the configuration object with config_hyperparameters, config_meta, config_deployment, and the command line arguments.
+
+        Args:
+            args: Command line arguments
+
+        Returns:
+            The updated configuration object.
+        """
+
+        config = self._config_sweep
+        config["parameters"]["run_type"] = {"value": args.run_type}
+        config["parameters"]["sweep"] = {"value": True}
+        config["parameters"]["name"] = {"value": self._config_meta["name"]}
+        config["parameters"]["depvar"] = {"value": self._config_meta["depvar"]}
+        config["parameters"]["algorithm"] = {"value": self._config_meta["algorithm"]}
+
+        if self._is_hurdle:
+            config["parameters"]["model_clf"] = {"value": self._config_meta["model_clf"]}
+            config["parameters"]["model_reg"] = {"value": self._config_meta["model_reg"]}
+    
+        return config
+    
+    def _split_hurdle_parameters(self):
+        """
+        Split the parameters dictionary into two separate dictionaries, one for the
+        classification model and one for the regression model.
+
+        Returns:
+            A dictionary containing original config, the split classification and regression parameters.
+        """
+
+        cls_dict = {}
+        reg_dict = {}
+        config = self.config
+
+        for key, value in config.items():
+            if key.startswith("cls_"):
+                cls_key = key.replace("cls_", "")
+                cls_dict[cls_key] = value
+            elif key.startswith("reg_"):
+                reg_key = key.replace("reg_", "")
+                reg_dict[reg_key] = value
+
+        config["parameters"]["clf"] = cls_dict
+        config["parameters"]["reg"] = reg_dict
+
+        return config
+
+    def _execute_model_tasks(self, config=None, train=None, eval=None, forecast=None, artifact_name=None):
+        """
+        Executes various model-related tasks including training, evaluation, and forecasting.
+
+        This function manages the execution of different tasks such as training the model,
+        evaluating an existing model, or performing forecasting.
+        It also initializes the WandB project.
+
+        Args:
+            config: Configuration object containing parameters and settings.
+            project: The WandB project name.
+            train: Flag to indicate if the model should be trained.
+            eval: Flag to indicate if the model should be evaluated.
+            forecast: Flag to indicate if forecasting should be performed.
+            artifact_name (optional): Specific name of the model artifact to load for evaluation or forecasting.
+        """
+        start_t = time.time()
+
+        # Initialize WandB
+        try:
+            with wandb.init(project=self._project, entity=self._entity, config=config):  # project and config ignored when running a sweep
+
+                # add the monthly metrics to WandB
+                add_wandb_monthly_metrics()
+
+                # Update config from WandB initialization above
+                self.config = wandb.config
+
+                # W&B does not directly support nested dictionaries for hyperparameters
+                if self.config["sweep"] and self._is_hurdle:
+                    self.config["parameters"] = self._split_hurdle_parameters(self.config)
+
+                if self.config["sweep"]:
+                    logger.info(f"Sweeping model {self.config['name']}...")
+                    model = self._train_model_artifact()
+                    logger.info(f"Evaluating model {self.config['name']}...")
+                    self._evaluate_sweep(model)
+
+                if train:
+                    logger.info(f"Training model {self.config['name']}...")
+                    self._train_model_artifact()
+
+                if eval:
+                    logger.info(f"Evaluating model {self.config['name']}...")
+                    self._evaluate_model_artifact(artifact_name)
+
+                if forecast:
+                    logger.info(f"Forecasting model {self.config['name']}...")
+                    self._forecast_model_artifact(artifact_name)
+            wandb.finish()
+        except Exception as e:
+            logger.error(f"Error during model tasks execution: {e}")
+        end_t = time.time()
+        minutes = (end_t - start_t) / 60
+        logger.info(f"Done. Runtime: {minutes:.3f} minutes.\n")
+
+    def _get_model(self, partitioner_dict: dict):
+        """
+        Get the model based on the algorithm specified in the config
+
+        Args:
+            partitioner_dict: The dictionary of partitioners.
+
+        Returns:
+            The model object based on the algorithm specified in the config
+        """
+
+        if self._is_hurdle:
+            model = HurdleModel(self.config, partitioner_dict)
+        else:
+            self.config["model_reg"] = self.config["algorithm"]
+            model = StepshifterModel(self.config, partitioner_dict)
+
+        return model
+    
+    def _get_standardized_df(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+        Standardize the DataFrame based on the run type
+
+        Args:
+            df: The DataFrame to standardize
+
+        Returns:
+            The standardized DataFrame
+        """
+
+        run_type = self.config["run_type"]
+        depvar = self.config["depvar"]
+
+        # choose the columns to keep based on the run type and replace negative values with 0
+        if run_type in ["calibration", "testing"]:
+            cols = [depvar] + df.forecasts.prediction_columns
+        elif run_type == "forecasting":
+            cols = ["step_pred_combined", depvar]
+        df = df.replace([np.inf, -np.inf], 0)[cols]
+        df = df.mask(df < 0, 0)
+        return df
+    
+    def _train_model_artifact(self):
+        # print(config)
+        path_raw = self._model_path.data_raw
+        path_generated = self._model_path.data_generated
+        path_artifacts = self._model_path.artifacts
+        run_type = self.config["run_type"]
+        df_viewser = pd.read_pickle(path_raw / f"{run_type}_viewser_df.pkl")
+
+        partitioner_dict = self._data_loader.partition_dict
+        stepshift_model = self._get_model(partitioner_dict)
+        stepshift_model.fit(df_viewser)
+
+        if not self.config["sweep"]:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            model_filename = ModelManager._generate_model_file_name(run_type, timestamp)
+            stepshift_model.save(path_artifacts / model_filename)
+            data_fetch_timestamp = read_log_file(path_raw / f"{run_type}_data_fetch_log.txt").get("Data Fetch Timestamp", None)
+            create_log_file(path_generated, self.config, timestamp, None, data_fetch_timestamp)
+        return stepshift_model
+
+    def _evaluate_model_artifact(self, artifact_name):
+        path_raw = self._model_path.data_raw
+        path_generated = self._model_path.data_generated
+        path_artifacts = self._model_path.artifacts
+        run_type = self.config["run_type"]
+
+        # if an artifact name is provided through the CLI, use it.
+        # Otherwise, get the latest model artifact based on the run type
+        if artifact_name:
+            logger.info(f"Using (non-default) artifact: {artifact_name}")
+
+            if not artifact_name.endswith(".pkl"):
+                artifact_name += ".pkl"
+            path_artifact = path_artifacts / artifact_name
+        else:
+            # use the latest model artifact based on the run type
+            logger.info(f"Using latest (default) run type ({run_type}) specific artifact")
+            path_artifact = self._get_latest_model_artifact(path_artifacts, run_type)
+
+        self.config["timestamp"] = path_artifact.stem[-15:]
+        df_viewser = pd.read_pickle(path_raw / f"{run_type}_viewser_df.pkl")
+
+        try:
+            stepshift_model = pd.read_pickle(path_artifact)
+        except FileNotFoundError:
+            logger.exception(f"Model artifact not found at {path_artifact}")
+
+        df = stepshift_model.predict(run_type, df_viewser)
+        df = self._get_standardized_df(df)
+        data_generation_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        data_fetch_timestamp = read_log_file(path_raw / f"{run_type}_data_fetch_log.txt").get("Data Fetch Timestamp", None)
+
+        _, df_output = generate_output_dict(df, self.config)
+        evaluation, df_evaluation = generate_metric_dict(df, self.config)
+        log_wandb_log_dict(self.config, evaluation)
+
+        self._save_model_outputs(df_evaluation, df_output, path_generated)
+        self._save_predictions(df, path_generated)
+        create_log_file(path_generated, self.config, self.config["timestamp"], data_generation_timestamp, data_fetch_timestamp)
+    
+    def _forecast_model_artifact(self, artifact_name):
+        path_raw = self._model_path.data_raw
+        path_generated = self._model_path.data_generated
+        path_artifacts = self._model_path.artifacts
+        run_type = self.config["run_type"]
+
+        # if an artifact name is provided through the CLI, use it.
+        # Otherwise, get the latest model artifact based on the run type
+        if artifact_name:
+            logger.info(f"Using (non-default) artifact: {artifact_name}")
+
+            if not artifact_name.endswith(".pkl"):
+                artifact_name += ".pkl"
+            path_artifact = path_artifacts / artifact_name
+        else:
+            # use the latest model artifact based on the run type
+            logger.info(f"Using latest (default) run type ({run_type}) specific artifact")
+            path_artifact = self._get_latest_model_artifact(path_artifacts, run_type)
+
+        self.config["timestamp"] = path_artifact.stem[-15:]
+        df_viewser = pd.read_pickle(path_raw / f"{run_type}_viewser_df.pkl")
+
+        try:
+            stepshift_model = pd.read_pickle(path_artifact)
+        except FileNotFoundError:
+            logger.exception(f"Model artifact not found at {path_artifact}")
+
+        df_predictions = stepshift_model.predict(run_type, df_viewser)
+        df_predictions = self._get_standardized_df(df_predictions)
+        data_generation_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        data_fetch_timestamp = read_log_file(path_raw / f"{run_type}_data_fetch_log.txt").get("Data Fetch Timestamp", None)
+
+        self._save_predictions(df_predictions, path_generated)
+        create_log_file(path_generated, self.config, self.config["timestamp"], data_generation_timestamp, data_fetch_timestamp)
+
+    def _evaluate_sweep(self, model):
+        path_raw = self._model_path.data_raw
+        run_type = self.config["run_type"]
+        steps = self.config["steps"]
+
+        df_viewser = pd.read_pickle(path_raw / f"{run_type}_viewser_df.pkl")
+        df = model.predict(run_type, df_viewser)
+        df = self._get_standardized_df(df)
+
+        # Temporarily keep this because the metric to minimize is MSE
+        pred_cols = [f"step_pred_{str(i)}" for i in steps]
+        df["mse"] = df.apply(lambda row: mean_squared_error([row[self.config["depvar"]]] * len(steps),
+                                                            [row[col] for col in pred_cols]), axis=1)
+
+        wandb.log({"MSE": df["mse"].mean()})
+
+        evaluation, _ = generate_metric_dict(df, self.config)
+        log_wandb_log_dict(self.config, evaluation)

views_stepshifter-0.1.1/views_stepshifter/models/__init__.py

@@ -0,0 +1,2 @@
+from .stepshifter import StepshifterModel
+from darts.models import LightGBMModel, XGBModel, RandomForest

views_stepshifter-0.1.1/views_stepshifter/models/hurdle_model.py

@@ -0,0 +1,94 @@
+import numpy as np
+from .stepshifter import StepshifterModel
+from .validation import views_validate
+from views_forecasts.extensions import *
+from sklearn.utils.validation import check_is_fitted
+
+
+class HurdleModel(StepshifterModel):
+    """
+    Hurdle model for time series forecasting. The model consists of two stages:
+    1. Binary stage: Predicts whether the target variable is 0 or > 0.
+    2. Positive stage: Predicts the value of the target variable when it is > 0.
+
+    Note:
+    This algorithm uses a two-step approach. 
+
+    **Step 1: Classification Stage**  
+    In the first step, a regression model is used with a binary target (0 or 1), 
+    indicating the absence or presence of violence. This stage functions similarly 
+    to a linear probability model, estimating the likelihood of a positive outcome. 
+    Since the model is a regression rather than a classification model, 
+    these estimates are not strictly bounded between 0 and 1, 
+    but this is acceptable for the purpose of this step.
+
+    To determine whether an observation is classified as "positive," we apply a threshold. 
+    The default threshold is 1, meaning that predictions above this value 
+    are considered positive outcomes. This threshold can be adjusted as 
+    a tunable hyperparameter to better suit specific requirements.
+
+    **Step 2: Regression Stage**  
+    In the second step, we use a regression model to predict a continuous or count value 
+    (e.g., the expected number of conflict fatalities) for the selected time series. 
+    We include the entire time series for countries or PRIO grids where the 
+    classification stage yielded at least one "positive" prediction, 
+    rather than limiting the regression to just the predicted positive values.
+    """
+    
+    def __init__(self, config: Dict, partitioner_dict: Dict[str, List[int]], threshold: float = 1.0):
+        super().__init__(config, partitioner_dict)
+        self._clf = self._resolve_estimator(config['model_clf'])
+        self._reg = self._resolve_estimator(config['model_reg'])
+        self._clf_params = self._get_parameters(config)['clf']
+        self._reg_params = self._get_parameters(config)['reg']
+        self._threshold = threshold
+
+    @views_validate
+    def fit(self, df: pd.DataFrame):
+        df = self._process_data(df)
+        self._prepare_time_series(df)
+
+        # Binary outcome (event/no-event)
+        # According to the DARTS doc, if timeseries uses a numeric type different from np.float32 or np.float64, not all functionalities may work properly.
+        # So use astype(float) instead of astype(int) (we should have binary outputs 0,1 though)
+        target_binary = [s.map(lambda x: (x > self._threshold).astype(float)) for s in self._target_train]
+
+        # Positive outcome (for cases where target > threshold)
+        target_pos, past_cov_pos = zip(*[(t, p) for t, p in zip(self._target_train, self._past_cov_train)
+                                         if (t.values() > self._threshold).any()])
+
+        for step in self._steps:
+            # Fit binary-like stage using a regression model, but the target is binary (0 or 1)
+            binary_model = self._clf(lags_past_covariates=[-step], **self._clf_params)
+            binary_model.fit(target_binary, past_covariates=self._past_cov_train)
+
+            # Fit positive stage using the regression model
+            positive_model = self._reg(lags_past_covariates=[-step], **self._reg_params)
+            positive_model.fit(target_pos, past_covariates=past_cov_pos)
+            self._models[step] = (binary_model, positive_model)
+        self.is_fitted_ = True
+
+    @views_validate
+    def predict(self, run_type: str, df: pd.DataFrame) -> pd.DataFrame:
+        df = self._process_data(df)
+        check_is_fitted(self, 'is_fitted_')
+
+        if run_type == 'forecasting':
+            pred_by_step_binary = [self._predict_by_step_combined(self._models[step][0], step, self._target_train) 
+                                   for step in self._steps]
+            pred_by_step_positive = [self._predict_by_step_combined(self._models[step][1], step, self._target_train) 
+                                     for step in self._steps]
+            final_pred = pd.concat(pred_by_step_binary, axis=0) * pd.concat(pred_by_step_positive, axis=0)
+            # Add the target variable to the predictions to make sure it is a VIEWS prediction
+            # If it is a forecasting run, the target variable is not available in the input data so we fill it with NaN
+            final_pred[self._depvar] = np.nan
+
+        else:
+            pred_by_step_binary = [self._predict_by_step(self._models[step][0], step, self._target_train)
+                                   for step in self._steps]
+            pred_by_step_positive = [self._predict_by_step(self._models[step][1], step, self._target_train)
+                                     for step in self._steps]
+            final_pred = pd.concat(pred_by_step_binary, axis=1) * pd.concat(pred_by_step_positive, axis=1)
+            final_pred = pd.merge(final_pred, df[self._depvar], left_index=True, right_index=True)
+
+        return final_pred

views_stepshifter-0.1.1/views_stepshifter/models/stepshifter.py

@@ -0,0 +1,199 @@
+import pickle
+import numpy as np
+import logging
+from darts import TimeSeries
+from sklearn.utils.validation import check_is_fitted
+from typing import List, Dict
+from views_forecasts.extensions import *
+from .validation import views_validate
+
+logger = logging.getLogger(__name__)
+
+
+class StepshifterModel:
+    def __init__(self, config: Dict, partitioner_dict: Dict[str, List[int]]):
+        self._steps = config['steps']
+        self._depvar = config['depvar']
+        self._reg = self._resolve_estimator(config['model_reg'])
+        self._params = self._get_parameters(config)
+        self._train_start, self._train_end = partitioner_dict['train']
+        self._test_start, self._test_end = partitioner_dict['predict']
+        self._models = {}
+    
+    @staticmethod
+    def _resolve_estimator(func_name: str):
+        """ Lookup table for supported estimators.
+        This is necessary because sklearn estimator default arguments
+        must pass equality test, and instantiated sub-estimators are not equal. """
+
+        match func_name:
+            case 'LinearRegressionModel':
+                from darts.models import LinearRegressionModel
+                return LinearRegressionModel
+            case 'RandomForestModel':
+                from darts.models import RandomForest
+                return RandomForest
+            case 'LightGBMModel':
+                from darts.models import LightGBMModel
+                return LightGBMModel
+            case 'XGBModel':
+                from darts.models import XGBModel
+                return XGBModel
+            case _:
+                raise ValueError(f"Model {func_name} is not a valid Darts forecasting model or is not supported now. "
+                                f"Change the model in the config file.")
+
+    @staticmethod
+    def _get_parameters(config: Dict):
+        '''
+        Get the parameters from the config file.
+        If not sweep, then get directly from the config file, otherwise have to remove some parameters.
+        '''
+
+        if config["sweep"]:
+            keys_to_remove = ["algorithm", "depvar", "steps", "sweep", "run_type", "model_cls", "model_reg", "name"]
+            parameters = {k: v for k, v in config.items() if k not in keys_to_remove}
+        else:
+            parameters = config["parameters"]
+
+        return parameters
+
+    def _process_data(self, df: pd.DataFrame):
+        '''
+        Countries appear and disappear, so we are predicting countries that exist in the last month of the training data.
+        If the country appeared earlier but don't have data previously, we will fill the missing data with 0.
+        '''
+
+        # set up
+        self._time = df.index.names[0]
+        self._level = df.index.names[1]
+        self._independent_variables = [c for c in df.columns if c != self._depvar]
+
+        last_month_id = df.index.get_level_values(self._time).max()
+        existing_country_ids = df.loc[last_month_id].index.unique()
+        df = df[df.index.get_level_values(self._level).isin(existing_country_ids)]
+
+        all_months = df.index.get_level_values(self._time).unique()
+        all_combinations = pd.MultiIndex.from_product([all_months, existing_country_ids], names=[self._time, self._level])
+        missing_combinations = all_combinations.difference(df.index)
+
+        missing_df = pd.DataFrame(0, index=missing_combinations, columns=df.columns)
+        df = pd.concat([df, missing_df]).sort_index()
+
+        return df
+
+    def _prepare_time_series(self, df: pd.DataFrame):
+        '''
+        Prepare time series for training and prediction
+        '''
+
+        df_reset = df.reset_index(level=[1])
+        self._series = TimeSeries.from_group_dataframe(df_reset, group_cols=self._level,
+                                                       value_cols=self._independent_variables + [self._depvar])
+
+        self._target_train = [series.slice(self._train_start, self._train_end + 1)[self._depvar]
+                              for series in self._series]  # ts.slice is different from df.slice
+        self._past_cov_train = [series.slice(self._train_start, self._train_end + 1)[self._independent_variables]
+                                for series in self._series]
+        self._past_cov = [series[self._independent_variables] for series in self._series]
+
+    def _predict_by_step(self, model, step, target):
+        ts_pred = model.predict(n=self._test_end - self._test_start + 1,
+                                series=target,
+                                # darts automatically locates the time period of past_covariates
+                                past_covariates=self._past_cov,
+                                show_warnings=False)
+        
+        # process the predictions
+        index_tuples, df_list = [], []
+        test_period = slice(self._test_start, self._test_end)
+        
+        for pred in ts_pred:
+            df_pred = pred.pd_dataframe().loc[test_period]
+            level = pred.static_covariates.iat[0, 0]
+            index_tuples.extend([(month, level) for month in df_pred.index])
+            df_list.append(df_pred.values)
+
+        df_preds = pd.DataFrame(
+            data=np.concatenate(df_list),
+            index=pd.MultiIndex.from_tuples(index_tuples, names=[self._time, self._level]),
+            columns=[f"step_pred_{step}"]
+        )
+
+        return df_preds.sort_index()
+
+    def _predict_by_step_combined(self, model, step, target):
+        '''
+        For forecasting only need to keep predictions with last-month-with-data, i.e., diagonal prediction
+        '''
+
+        ts_pred = model.predict(n=step,
+                                series=target,
+                                # darts automatically locates the time period of past_covariates
+                                past_covariates=self._past_cov,
+                                show_warnings=False)
+        
+        # process the predictions
+        index_tuples, df_list = [], []
+        for pred in ts_pred:
+            df_pred = pred.pd_dataframe().loc[[self._test_start + step - 1]]
+            level = pred.static_covariates.iat[0, 0]
+            index_tuples.extend([(month, level) for month in df_pred.index])
+            df_list.append(df_pred.values)
+
+        df_preds = pd.DataFrame(
+            data=np.concatenate(df_list),
+            index=pd.MultiIndex.from_tuples(index_tuples, names=[self._time, self._level]),
+            columns=["step_pred_combined"]
+        )
+
+        return df_preds.sort_index()
+    
+    @views_validate
+    def fit(self, df: pd.DataFrame):
+        df = self._process_data(df)
+        self._prepare_time_series(df)
+        for step in self._steps:
+            model = self._reg(lags_past_covariates=[-step], **self._params)
+            model.fit(self._target_train, past_covariates=self._past_cov_train)
+            self._models[step] = model
+        self.is_fitted_ = True
+
+    @views_validate
+    def predict(self, run_type: str, df: pd.DataFrame) -> pd.DataFrame:
+        df = self._process_data(df)
+        check_is_fitted(self, 'is_fitted_')
+
+        if run_type == 'forecasting':
+            pred_by_step = [self._predict_by_step_combined(self._models[step], step, self._target_train) for step in self._steps]
+            pred = pd.concat(pred_by_step, axis=0)
+            # Add the target variable to the predictions to make sure it is a VIEWS prediction
+            # If it is a forecasting run, the target variable is not available in the input data so we fill it with NaN
+            pred[self._depvar] = np.nan
+
+        else:
+            pred_by_step = [self._predict_by_step(self._models[step], step, self._target_train) for step in self._steps]
+            pred = pd.concat(pred_by_step, axis=1)
+            pred = pd.merge(pred, df[self._depvar], left_index=True, right_index=True)
+
+        return pred
+    
+    def save(self, path: str):
+        try:
+            with open(path, "wb") as file:
+                pickle.dump(self, file)
+            logger.info(f"Model successfully saved to {path}")
+        except Exception as e:
+            logger.exception(f"Failed to save model: {e}")
+
+    @property
+    def models(self):
+        return list(self._models.values())
+
+    @property
+    def steps(self):
+        return self._steps
+
+    @property
+    def depvar(self):
+        return self._depvar

views_stepshifter-0.1.1/views_stepshifter/models/validation.py

@@ -0,0 +1,45 @@
+import functools
+import numpy as np
+import pandas as pd
+import logging
+from views_forecasts.extensions import *
+
+logger = logging.getLogger(__name__)
+
+
+def dataframe_is_right_format(dataframe: pd.DataFrame):
+    try:
+        assert len(dataframe.index.levels) == 2
+        # print("The dataframe has a two-level index")
+    except AssertionError:
+        logger.exception("Dataframe must have a two-level index")
+        raise AssertionError("Dataframe must have a two-level index")
+        
+    try:
+        assert dataframe.index.names[0] == "month_id"
+        # print("The first level of the index is correct")
+    except AssertionError:
+        logger.exception("The first level of the index must be 'month_id'")
+        raise AssertionError("The first level of the index must be 'month_id'")
+    
+    try:
+        assert dataframe.index.names[1] in ["country_id", "priogrid_gid"]
+        # print("The second level of the index is correct")
+    except AssertionError:
+        logger.exception("The second level of the index must be 'country_id' or 'priogrid_gid'")
+        raise AssertionError("The second level of the index must be 'country_id' or 'priogrid_gid'")
+
+    try:
+        assert set(dataframe.dtypes) == {np.dtype(float)}
+        # print("The dataframe contains only np.float64 floats")
+    except AssertionError:
+        logger.exception("The dataframe must contain only np.float64 floats")
+        raise AssertionError("The dataframe must contain only np.float64 floats")
+    
+
+def views_validate(fn):
+    @functools.wraps(fn)
+    def inner(*args,**kwargs):
+        dataframe_is_right_format(args[-1])
+        return fn(*args, **kwargs)
+    return inner