openstef 3.4.29__py3-none-any.whl → 3.4.44__py3-none-any.whl

openstef/model/model_creator.py

@@ -6,7 +6,7 @@ from typing import Union
 
 import structlog
 
-from openstef.enums import MLModelType
+from openstef.enums import ModelType
 from openstef.model.regressors.arima import ARIMAOpenstfRegressor
 from openstef.model.regressors.custom_regressor import is_custom_type, load_custom_model
 from openstef.model.regressors.lgbm import LGBMOpenstfRegressor
@@ -29,7 +29,7 @@ structlog.configure(
 logger = structlog.get_logger(__name__)
 
 valid_model_kwargs = {
-    MLModelType.XGB: [
+    ModelType.XGB: [
         "n_estimators",
         "objective",
         "max_depth",
@@ -60,7 +60,7 @@ valid_model_kwargs = {
         "validate_parameters",
         "early_stopping_rounds",
     ],
-    MLModelType.LGB: [
+    ModelType.LGB: [
         "boosting_type",
         "objective",
         "num_leaves",
@@ -82,7 +82,7 @@ valid_model_kwargs = {
         "importance_type",
         "early_stopping_rounds",
     ],
-    MLModelType.XGB_QUANTILE: [
+    ModelType.XGB_QUANTILE: [
         "quantiles",
         "gamma",
         "colsample_bytree",
@@ -91,7 +91,7 @@ valid_model_kwargs = {
         "max_depth",
         "early_stopping_rounds",
     ],
-    MLModelType.XGB_MULTIOUTPUT_QUANTILE: [
+    ModelType.XGB_MULTIOUTPUT_QUANTILE: [
         "quantiles",
         "gamma",
         "colsample_bytree",
@@ -101,23 +101,27 @@ valid_model_kwargs = {
         "early_stopping_rounds",
         "arctan_smoothing",
     ],
-    MLModelType.LINEAR: [
+    ModelType.LINEAR: [
         "missing_values",
         "imputation_strategy",
         "fill_value",
     ],
-    MLModelType.FLATLINER: [
+    ModelType.FLATLINER: [
         "quantiles",
     ],
-    MLModelType.LINEAR_QUANTILE: [
+    ModelType.LINEAR_QUANTILE: [
         "alpha",
         "quantiles",
         "solver",
         "missing_values",
         "imputation_strategy",
         "fill_value",
+        "weight_scale_percentile",
+        "weight_exponent",
+        "weight_floor",
+        "no_fill_future_values_features",
     ],
-    MLModelType.ARIMA: [
+    ModelType.ARIMA: [
         "backtest_max_horizon",
         "order",
         "seasonal_order",
@@ -131,18 +135,18 @@ class ModelCreator:
 
     # Set object mapping
     MODEL_CONSTRUCTORS = {
-        MLModelType.XGB: XGBOpenstfRegressor,
-        MLModelType.LGB: LGBMOpenstfRegressor,
-        MLModelType.XGB_QUANTILE: XGBQuantileOpenstfRegressor,
-        MLModelType.XGB_MULTIOUTPUT_QUANTILE: XGBMultiOutputQuantileOpenstfRegressor,
-        MLModelType.LINEAR: LinearOpenstfRegressor,
-        MLModelType.LINEAR_QUANTILE: LinearQuantileOpenstfRegressor,
-        MLModelType.ARIMA: ARIMAOpenstfRegressor,
-        MLModelType.FLATLINER: FlatlinerRegressor,
+        ModelType.XGB: XGBOpenstfRegressor,
+        ModelType.LGB: LGBMOpenstfRegressor,
+        ModelType.XGB_QUANTILE: XGBQuantileOpenstfRegressor,
+        ModelType.XGB_MULTIOUTPUT_QUANTILE: XGBMultiOutputQuantileOpenstfRegressor,
+        ModelType.LINEAR: LinearOpenstfRegressor,
+        ModelType.LINEAR_QUANTILE: LinearQuantileOpenstfRegressor,
+        ModelType.ARIMA: ARIMAOpenstfRegressor,
+        ModelType.FLATLINER: FlatlinerRegressor,
     }
 
     @staticmethod
-    def create_model(model_type: Union[MLModelType, str], **kwargs) -> OpenstfRegressor:
+    def create_model(model_type: Union[ModelType, str], **kwargs) -> OpenstfRegressor:
         """Create a machine learning model based on model type.
 
         Args:
@@ -163,7 +167,7 @@ class ModelCreator:
                 model_class = load_custom_model(model_type)
                 valid_kwargs = model_class.valid_kwargs()
             else:
-                model_type = MLModelType(model_type)
+                model_type = ModelType(model_type)
                 model_class = ModelCreator.MODEL_CONSTRUCTORS[model_type]
                 valid_kwargs = valid_model_kwargs[model_type]
                 # Check if model as imported
@@ -174,7 +178,7 @@ class ModelCreator:
                         "Please refer to the ReadMe for instructions"
                     )
         except ValueError as e:
-            valid_types = [t.value for t in MLModelType]
+            valid_types = [t.value for t in ModelType]
             raise NotImplementedError(
                 f"No constructor for '{model_type}', "
                 f"valid model_types are: {valid_types} "

openstef/model/objective.py

@@ -8,7 +8,7 @@ from typing import Any, Callable, Optional
 import optuna
 import pandas as pd
 
-from openstef.enums import MLModelType
+from openstef.enums import ModelType
 from openstef.metrics import metrics
 from openstef.metrics.reporter import Report, Reporter
 from openstef.model.regressors.regressor import OpenstfRegressor
@@ -245,7 +245,7 @@ class RegressorObjective:
 class XGBRegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.XGB
+        self.model_type = ModelType.XGB
 
     # extend the parameters with the model specific ones per implementation
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
@@ -282,7 +282,7 @@ class XGBRegressorObjective(RegressorObjective):
 class LGBRegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.LGB
+        self.model_type = ModelType.LGB
 
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
         """Get parameters for LGB Regressor Objective with objective specific parameters.
@@ -323,7 +323,7 @@ class LGBRegressorObjective(RegressorObjective):
 class XGBQuantileRegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.XGB_QUANTILE
+        self.model_type = ModelType.XGB_QUANTILE
 
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
         """Get parameters for XGBQuantile Regressor Objective with objective specific parameters.
@@ -352,7 +352,7 @@ class XGBQuantileRegressorObjective(RegressorObjective):
 class XGBMultioutputQuantileRegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.XGB_QUANTILE
+        self.model_type = ModelType.XGB_QUANTILE
 
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
         """Get parameters for XGB Multioutput Quantile Regressor Objective with objective specific parameters.
@@ -382,7 +382,7 @@ class XGBMultioutputQuantileRegressorObjective(RegressorObjective):
 class LinearRegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.LINEAR
+        self.model_type = ModelType.LINEAR
 
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
         """Get parameters for Linear Regressor Objective with objective specific parameters.
@@ -405,7 +405,7 @@ class LinearRegressorObjective(RegressorObjective):
 class ARIMARegressorObjective(RegressorObjective):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.model_type = MLModelType.ARIMA
+        self.model_type = ModelType.ARIMA
 
     def get_params(self, trial: optuna.trial.FrozenTrial) -> dict:
         """Get parameters for ARIMA Regressor Objective with objective specific parameters.

openstef/model/objective_creator.py

@@ -4,7 +4,7 @@
 
 from typing import Union
 
-from openstef.enums import MLModelType
+from openstef.enums import ModelType
 from openstef.model.objective import (
     ARIMARegressorObjective,
     LGBRegressorObjective,
@@ -22,17 +22,17 @@ from openstef.model.regressors.custom_regressor import (
 
 class ObjectiveCreator:
     OBJECTIVES = {
-        MLModelType.XGB: XGBRegressorObjective,
-        MLModelType.LGB: LGBRegressorObjective,
-        MLModelType.XGB_QUANTILE: XGBQuantileRegressorObjective,
-        MLModelType.XGB_MULTIOUTPUT_QUANTILE: XGBMultioutputQuantileRegressorObjective,
-        MLModelType.LINEAR: LinearRegressorObjective,
-        MLModelType.LINEAR_QUANTILE: LinearRegressorObjective,
-        MLModelType.ARIMA: ARIMARegressorObjective,
+        ModelType.XGB: XGBRegressorObjective,
+        ModelType.LGB: LGBRegressorObjective,
+        ModelType.XGB_QUANTILE: XGBQuantileRegressorObjective,
+        ModelType.XGB_MULTIOUTPUT_QUANTILE: XGBMultioutputQuantileRegressorObjective,
+        ModelType.LINEAR: LinearRegressorObjective,
+        ModelType.LINEAR_QUANTILE: LinearRegressorObjective,
+        ModelType.ARIMA: ARIMARegressorObjective,
     }
 
     @staticmethod
-    def create_objective(model_type: Union[MLModelType, str]) -> RegressorObjective:
+    def create_objective(model_type: Union[ModelType, str]) -> RegressorObjective:
         """Create an objective function based on model type.
 
         Args:
@@ -51,10 +51,10 @@ class ObjectiveCreator:
             if is_custom_type(model_type):
                 objective = create_custom_objective(model_type)
             else:
-                model_type = MLModelType(model_type)
+                model_type = ModelType(model_type)
                 objective = ObjectiveCreator.OBJECTIVES[model_type]
         except ValueError as e:
-            valid_types = [t.value for t in MLModelType]
+            valid_types = [t.value for t in ModelType]
             raise NotImplementedError(
                 f"No objective for '{model_type}', "
                 f"valid model_types are: {valid_types}"

openstef/model/regressors/flatliner.py

@@ -2,18 +2,13 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 import re
-from typing import Dict, Union, Set, Optional, List
+from typing import List
 
 import numpy as np
 import pandas as pd
 from sklearn.base import RegressorMixin
-from sklearn.linear_model import QuantileRegressor
-from sklearn.preprocessing import MinMaxScaler
 from sklearn.utils.validation import check_is_fitted
 
-from openstef.feature_engineering.missing_values_transformer import (
-    MissingValuesTransformer,
-)
 from openstef.model.regressors.regressor import OpenstfRegressor
 
 
@@ -23,9 +18,9 @@ class FlatlinerRegressor(OpenstfRegressor, RegressorMixin):
     def __init__(self, quantiles=None):
         """Initialize FlatlinerRegressor.
 
-        The model always predicts 0.0, regardless of the input features. The model is
-        meant to be used for flatliner locations that still expect a prediction while
-        preserving the prediction interface.
+        The model always predicts 0.0, regardless of the input features. The model is meant to be used for flatliner
+        locations that still expect a prediction while preserving the prediction interface.
+
         """
         super().__init__()
         self.quantiles = quantiles

openstef/model/regressors/linear_quantile.py

@@ -2,13 +2,13 @@
 #
 # SPDX-License-Identifier: MPL-2.0
 import re
-from typing import Dict, Union, Set, Optional
+from typing import Dict, Union, Set, Optional, List
 
 import numpy as np
 import pandas as pd
 from sklearn.base import RegressorMixin
 from sklearn.linear_model import QuantileRegressor
-from sklearn.preprocessing import MinMaxScaler
+from sklearn.preprocessing import StandardScaler
 from sklearn.utils.validation import check_is_fitted
 
 from openstef.feature_engineering.missing_values_transformer import (
@@ -25,8 +25,8 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
     solver: str
 
     imputer_: MissingValuesTransformer
-    x_scaler_: MinMaxScaler
-    y_scaler_: MinMaxScaler
+    x_scaler_: StandardScaler
+    y_scaler_: StandardScaler
     models_: Dict[float, QuantileRegressor]
 
     is_fitted_: bool = False
@@ -47,6 +47,10 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
         missing_values: Union[int, float, str, None] = np.nan,
         imputation_strategy: Optional[str] = "mean",
         fill_value: Union[str, int, float] = None,
+        weight_scale_percentile: int = 95,
+        weight_exponent: float = 1,
+        weight_floor: float = 0.1,
+        no_fill_future_values_features: List[str] = None,
     ):
         """Initialize LinearQuantileOpenstfRegressor.
 
@@ -69,6 +73,12 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
             missing_values: Value to be considered as missing value
             imputation_strategy: Imputation strategy
             fill_value: Fill value
+            weight_scale_percentile: Percentile used in scaling of the samples
+            weight_exponent: Exponent used in sample weighing
+            weight_floor: Minimum weight for samples
+            no_fill_future_values_features: The features for which it does not make sense
+                to fill future values. Rows that contain trailing null values for these
+                features will be removed from the data.
 
         """
         super().__init__()
@@ -82,13 +92,17 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
         self.quantiles = quantiles
         self.alpha = alpha
         self.solver = solver
+        self.weight_scale_percentile = weight_scale_percentile
+        self.weight_exponent = weight_exponent
+        self.weight_floor = weight_floor
         self.imputer_ = MissingValuesTransformer(
             missing_values=missing_values,
             imputation_strategy=imputation_strategy,
             fill_value=fill_value,
+            no_fill_future_values_features=no_fill_future_values_features,
         )
-        self.x_scaler_ = MinMaxScaler(feature_range=(-1, 1))
-        self.y_scaler_ = MinMaxScaler(feature_range=(-1, 1))
+        self.x_scaler_ = StandardScaler()
+        self.y_scaler_ = StandardScaler()
         self.models_ = {
             quantile: QuantileRegressor(alpha=alpha, quantile=quantile, solver=solver)
             for quantile in quantiles
@@ -165,7 +179,7 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
         x = self._remove_ignored_features(x)
 
         # Fix nan columns
-        x = self.imputer_.fit_transform(x)
+        x, y = self.imputer_.fit_transform(x, y)
         if x.isna().any().any():
             raise ValueError(
                 "There are nan values in the input data. Set "
@@ -177,7 +191,7 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
         y_scaled = self.y_scaler_.fit_transform(y.to_frame())[:, 0]
 
         # Add more focus on extreme / peak values
-        sample_weight = np.abs(y_scaled)
+        sample_weight = self._calculate_sample_weights(y.values.squeeze())
 
         # Fit quantile regressors
         for quantile in self.quantiles:
@@ -191,6 +205,33 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
 
         return self
 
+    def _calculate_sample_weights(self, y: np.array):
+        """Calculate sample weights based on the y values of arbitrary scale.
+
+        The resulting weights are in the range [0,1] and are used to put more emphasis
+        on certain samples. The sample weighting function does:
+
+        * Rescale data to a [-1, 1] range using quantile scaling. 90% of the data will
+          be within this range. Rest is outside.
+        * Calculate the weight by taking the exponent of scaled data.
+          * exponent=0: Results in uniform weights for all samples.
+          * exponent=1: Results in linearly increasing weights for samples that are
+            closer to the extremes.
+          * exponent>1: Results in exponentially increasing weights for samples that are
+            closer to the extremes.
+        * Clip the data to [0, 1] range with weight_floor as the minimum weight.
+          * Weight floor is used to make sure that all the samples are considered.
+
+        """
+        return np.clip(
+            _weight_exp(
+                _scale_percentile(y, percentile=self.weight_scale_percentile),
+                exponent=self.weight_exponent,
+            ),
+            a_min=self.weight_floor,
+            a_max=1,
+        )
+
     def predict(self, x: pd.DataFrame, quantile: float = 0.5, **kwargs) -> np.array:
         """Makes a prediction for a desired quantile.
 
@@ -231,7 +272,7 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
         return np.array(
             [
                 reg_feature_importances_dict.get(c, 0)
-                for c in self.imputer_.in_feature_names
+                for c in self.imputer_.non_null_feature_names
             ]
         )
 
@@ -245,3 +286,11 @@ class LinearQuantileOpenstfRegressor(OpenstfRegressor, RegressorMixin):
 
     def __sklearn_is_fitted__(self) -> bool:
         return self.is_fitted_
+
+
+def _scale_percentile(x: np.ndarray, percentile: int = 95):
+    return np.abs(x / np.percentile(np.abs(x), percentile))
+
+
+def _weight_exp(x: np.ndarray, exponent: float = 1):
+    return np.abs(x) ** exponent

openstef/model/regressors/xgb.py

@@ -1,6 +1,10 @@
 # SPDX-FileCopyrightText: 2017-2023 Contributors to the OpenSTEF project <korte.termijn.prognoses@alliander.com> # noqa E501>
 #
 # SPDX-License-Identifier: MPL-2.0
+from typing import Optional
+
+import numpy as np
+from sklearn.base import RegressorMixin
 
 from xgboost import XGBRegressor
 
@@ -27,3 +31,22 @@ class XGBOpenstfRegressor(XGBRegressor, OpenstfRegressor):
             "gain_importance_name": "total_gain",
             "weight_importance_name": "weight",
         }
+
+    def fit(
+        self,
+        x: np.array,
+        y: np.array,
+        *,
+        early_stopping_rounds: Optional[int] = None,
+        callbacks: Optional[list] = None,
+        eval_metric: Optional[str] = None,
+        **kwargs
+    ):
+        if early_stopping_rounds is not None:
+            self.set_params(early_stopping_rounds=early_stopping_rounds)
+        if callbacks is not None:
+            self.set_params(callbacks=callbacks)
+        if eval_metric is not None:
+            self.set_params(eval_metric=eval_metric)
+
+        super().fit(x, y, **kwargs)

openstef/model_selection/model_selection.py

@@ -106,7 +106,7 @@ def split_data_train_validation_test(
     validation_fraction: float = 0.15,
     back_test: bool = False,
     stratification_min_max: bool = True,
-) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
     """Split input data into train, test and validation set.
 
     Function for splitting data with features in a train, test and

openstef/pipeline/create_component_forecast.py

@@ -108,7 +108,7 @@ def create_components_forecast_pipeline(
 
     # Make component forecasts
     try:
-        input_data = create_input(pj, input_data, weather_data)
+        dazls_input_data = create_input(pj, input_data, weather_data)
 
         # Save and load the model as .sav file (or as .z file)
         # For the code contact: korte.termijn.prognoses@alliander.com
@@ -119,13 +119,13 @@ def create_components_forecast_pipeline(
 
         # Use the predict function of Dazls model
         # As input data we use the input_data function which takes into consideration what we want as an input for the forecast and what Dazls can accept as an input
-        forecasts = dazls_model.predict(x=input_data)
+        forecasts = dazls_model.predict(x=dazls_input_data)
 
         # Set the columns for the output forecast dataframe
         forecasts = pd.DataFrame(
             forecasts,
             columns=["forecast_wind_on_shore", "forecast_solar"],
-            index=input_data.index,
+            index=dazls_input_data.index,
         )
 
         # Make post-processed forecasts for solar and wind power
@@ -140,18 +140,25 @@ def create_components_forecast_pipeline(
 
         # Make forecast for the component: "forecast_other"
         forecasts["forecast_other"] = (
-            input_data["total_load"]
+            dazls_input_data["total_load"]
             - forecasts["forecast_solar"]
             - forecasts["forecast_wind_on_shore"]
         )
+
+        # Make sure the forecasts have the same form as the input data. Pad with 0 if necessary
+        forecasts = forecasts.reindex(index=input_data.index, fill_value=0)
     except Exception as e:
-        # In case something goes wrong we fall back on aan empty dataframe
+        # In case something goes wrong we fall back on an a zero-filled dataframe
         logger.warning(
             f"Could not make component forecasts: {e}, falling back on series of"
             " zeros!",
             exc_info=e,
         )
-        forecasts = pd.DataFrame()
+        forecasts = pd.DataFrame(
+            data=0,
+            index=input_data.index,
+            columns=["forecast_wind_on_shore", "forecast_solar", "forecast_other"],
+        )
 
     # Prepare for output
     # Add more prediction properties to the forecast ("pid","customer","description","type","algtype)

openstef/pipeline/train_model.py

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: MPL-2.0
 import logging
 import os
-from typing import Optional, Union
+from typing import Optional, Union, Tuple
 
 import pandas as pd
 import structlog
@@ -155,7 +155,7 @@ def train_model_pipeline_core(
     input_data: pd.DataFrame,
     old_model: OpenstfRegressor = None,
     horizons: list[float] = DEFAULT_TRAIN_HORIZONS_HOURS,
-) -> Union[
+) -> Tuple[
     OpenstfRegressor,
     Report,
     ModelSpecificationDataClass,
@@ -246,7 +246,9 @@ def train_pipeline_common(
     test_fraction: float = 0.0,
     backtest: bool = False,
     test_data_predefined: pd.DataFrame = pd.DataFrame(),
-) -> tuple[OpenstfRegressor, Report, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+) -> tuple[
+    OpenstfRegressor, Report, pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame
+]:
     """Common pipeline shared with operational training and backtest training.
 
     Args:
@@ -314,7 +316,8 @@ def train_pipeline_common(
 
 def train_pipeline_step_load_model(
     pj: PredictionJobDataClass, serializer: MLflowSerializer
-) -> tuple[OpenstfRegressor, ModelSpecificationDataClass, Union[int, float]]:
+) -> Tuple[OpenstfRegressor, ModelSpecificationDataClass, Union[int, float]]:
+    old_model: Optional[OpenstfRegressor]
     try:
         old_model, model_specs = serializer.load_model(experiment_name=str(pj.id))
         old_model_age = old_model.age  # Age attribute is openstef specific
@@ -509,7 +512,7 @@ def train_pipeline_step_split_data(
     test_fraction: float,
     backtest: bool = False,
     test_data_predefined: pd.DataFrame = pd.DataFrame(),
-) -> Union[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
+) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
     """The default way to perform train, val, test split.
 
     Args:

openstef/tasks/calculate_kpi.py

@@ -29,7 +29,7 @@ import pandas as pd
 import structlog
 
 from openstef.data_classes.prediction_job import PredictionJobDataClass
-from openstef.enums import MLModelType
+from openstef.enums import ModelType
 from openstef.exceptions import NoPredictedLoadError, NoRealisedLoadError
 from openstef.metrics import metrics
 from openstef.settings import Settings
@@ -42,7 +42,7 @@ THRESHOLD_RETRAINING = 0.25
 THRESHOLD_OPTIMIZING = 0.50
 
 
-def main(model_type: MLModelType = None, config=None, database=None) -> None:
+def main(model_type: ModelType = None, config=None, database=None) -> None:
     taskname = Path(__file__).name.replace(".py", "")
 
     if database is None or config is None:
@@ -52,7 +52,7 @@ def main(model_type: MLModelType = None, config=None, database=None) -> None:
         )
 
     if model_type is None:
-        model_type = [ml.value for ml in MLModelType]
+        model_type = [ml.value for ml in ModelType]
 
     with TaskContext(taskname, config, database) as context:
         # Set start and end time

openstef/tasks/create_basecase_forecast.py

@@ -97,7 +97,7 @@ def create_basecase_forecast_task(
     context.database.write_forecast(basecase_forecast, t_ahead_series=True)
 
 
-def main(config: object = None, database: object = None):
+def main(config: object = None, database: object = None, **kwargs):
     taskname = Path(__file__).name.replace(".py", "")
 
     if database is None or config is None:
@@ -110,7 +110,7 @@ def main(config: object = None, database: object = None):
         model_type = ["xgb", "xgb_quantile", "lgb"]
 
         PredictionJobLoop(context, model_type=model_type).map(
-            create_basecase_forecast_task, context
+            create_basecase_forecast_task, context, **kwargs
         )
 
 

openstef/tasks/create_components_forecast.py

@@ -29,7 +29,7 @@ import pandas as pd
 import structlog
 
 from openstef.data_classes.prediction_job import PredictionJobDataClass
-from openstef.enums import MLModelType
+from openstef.enums import ModelType
 from openstef.exceptions import ComponentForecastTooShortHorizonError
 from openstef.pipeline.create_component_forecast import (
     create_components_forecast_pipeline,
@@ -140,7 +140,7 @@ def create_components_forecast_task(
         )
 
 
-def main(config: object = None, database: object = None):
+def main(config: object = None, database: object = None, **kwargs):
     taskname = Path(__file__).name.replace(".py", "")
 
     if database is None or config is None:
@@ -150,12 +150,12 @@ def main(config: object = None, database: object = None):
         )
 
     with TaskContext(taskname, config, database) as context:
-        model_type = [ml.value for ml in MLModelType]
+        model_type = [ml.value for ml in ModelType]
 
         PredictionJobLoop(
             context,
             model_type=model_type,
-        ).map(create_components_forecast_task, context)
+        ).map(create_components_forecast_task, context, **kwargs)
 
 
 if __name__ == "__main__":

openstef/tasks/create_forecast.py

@@ -24,7 +24,7 @@ from datetime import datetime, timedelta
 from pathlib import Path
 
 from openstef.data_classes.prediction_job import PredictionJobDataClass
-from openstef.enums import MLModelType, PipelineType
+from openstef.enums import ModelType, PipelineType
 from openstef.exceptions import InputDataOngoingZeroFlatlinerError
 from openstef.pipeline.create_forecast import create_forecast_pipeline
 from openstef.tasks.utils.predictionjobloop import PredictionJobLoop
@@ -118,7 +118,7 @@ def create_forecast_task(
     context.database.write_forecast(forecast, t_ahead_series=True)
 
 
-def main(model_type=None, config=None, database=None):
+def main(model_type=None, config=None, database=None, **kwargs):
     taskname = Path(__file__).name.replace(".py", "")
 
     if database is None or config is None:
@@ -129,10 +129,10 @@ def main(model_type=None, config=None, database=None):
 
     with TaskContext(taskname, config, database) as context:
         if model_type is None:
-            model_type = [ml.value for ml in MLModelType]
+            model_type = [ml.value for ml in ModelType]
 
         PredictionJobLoop(context, model_type=model_type).map(
-            create_forecast_task, context
+            create_forecast_task, context, **kwargs
         )