Repository: KIT-IAI/pyWATTS

Issue:
Rework train_if functionality in online learning
The train_if argument is evaluated when the model is fitted, which happens before the transformation. However, refit should be applied after a transformation. Moreover, instead of using a function for train_if we should use a condition object, which gets the output of steps as input and no complicated workarounds has to be done by the user like storing the result of a FunctionModule in global variable which is then evaluated by the train_if function. See #116 

Discussion:

Patch:
diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index 8e87b65..0f83ee1 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -42,5 +42,6 @@ jobs:
         python example.py
         python example_keras.py
         python example_day_and_night.py
+        python example_batch_execution.py
         python example_pytorch.py
         python example_transform_module.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1130d3e..30be6e8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,13 +10,16 @@
 
 
 ### Changed
+* Retraining is triggered after all steps are transformed ([#117](https://github.com/KIT-IAI/pyWATTS/issues/117))
 
 ### Fixed
 * Try to handle non matching shapes in metrics. If not possible raise an exception.  ([#154](https://github.com/KIT-IAI/pyWATTS/issues/154))
+* Fixed Sampler-module, samples now in forwarding direction ([#174](https://github.com/KIT-IAI/pyWATTS/issues/174))
 * Fixed wrong check of if summary should be returned in _run method of pipeline ([#157](https://github.com/KIT-IAI/pyWATTS/issues/157))
 * Fixed Pipeline.to_folder() for modules that store parameters in numpy format by adding a special JSON-encoder for types that are not json-serializable by default ([#175](https://github.com/KIT-IAI/pyWATTS/issues/175))
 * Added missing summaries to pywatts.summaries.__init__.py ([#183](https://github.com/KIT-IAI/pyWATTS/issues/183))
 
+
 ### Deprecated
 
 ## 0.2.0 - 2021-30-09
diff --git a/examples/example_batch_execution.py b/examples/example_batch_execution.py
new file mode 100644
index 0000000..e439c4b
--- /dev/null
+++ b/examples/example_batch_execution.py
@@ -0,0 +1,108 @@
+# -----------------------------------------------------------
+# This example presents the code for the batch/online execution
+# of the pipeline.
+# Additionally, the regressor in this pipeline is retrained
+# if the datapoint corresponds to midnight.
+# -----------------------------------------------------------
+
+import pandas as pd
+from sklearn.preprocessing import StandardScaler
+from sklearn.svm import SVR
+
+# Import the pyWATTS pipeline and the required modules
+from pywatts.callbacks import CSVCallback, LinePlotCallback
+from pywatts.core.computation_mode import ComputationMode
+from pywatts.core.pipeline import Pipeline
+from pywatts.modules import ClockShift, LinearInterpolater, RollingRMSE, SKLearnWrapper, Sampler, FunctionModule
+
+
+# This function creates and returns the preprocessing pipeline
+from pywatts.utils._xarray_time_series_utils import numpy_to_xarray
+
+
+def create_preprocessing_pipeline(power_scaler):
+    pipeline = Pipeline(path="../results/preprocessing")
+
+    # Deal with missing values through linear interpolation
+    imputer_power_statistics = LinearInterpolater(method="nearest", dim="time",
+                                                  name="imputed_power")(x=pipeline["scaler_power"])
+    # Scale the data using a standard SKLearn scaler
+    scale_power_statistics = power_scaler(x=imputer_power_statistics)
+
+    # Create lagged time series to later be used in the regression
+    historical_data = ClockShift(lag=24)(x=scale_power_statistics)
+    Sampler(24, name="sampled_data")(x=historical_data)
+    return pipeline
+
+
+# This function creates the pipeline which we use for testing.
+# The test pipeline works on batches with one hour
+def create_test_pipeline(regressor_svr):
+    # Create test pipeline which works on a batch size of one hour.
+    pipeline = Pipeline("../results/test_pipeline", batch=pd.Timedelta("1h"))
+
+    # Add the svr regressor to the pipeline. This regressor should be called if it is not daytime
+    regressor_svr_power_statistics = regressor_svr(historical_input=pipeline["historical_input"],
+                                                   target=pipeline["load_power_statistics"],
+                                                   computation_mode=ComputationMode.Refit,
+                                                   callbacks=[LinePlotCallback('SVR')],
+                                                   lag=pd.Timedelta(hours=24),
+                                                   train_if=lambda x, _: len(x["historical_input"].indexes["time"]) > 0 and
+                                                                         x["historical_input"].indexes["time"][0].hour == 0)
+    RollingRMSE(window_size=1, window_size_unit="d")(
+        y_hat=regressor_svr_power_statistics, y=pipeline["load_power_statistics"],
+        callbacks=[LinePlotCallback('RMSE'), CSVCallback('RMSE')])
+    return pipeline
+
+
+if __name__ == "__main__":
+    # Read the data via pandas.
+    data = pd.read_csv("../data/getting_started_data.csv", parse_dates=["time"], infer_datetime_format=True,
+                       index_col="time")
+
+    # Split the data into train and test data.
+    train = data[:6000]
+    test = data[8700:]
+
+    # Create all modules which are used multiple times.
+    regressor_svr = SKLearnWrapper(module=SVR(), name="regression")
+    power_scaler = SKLearnWrapper(module=StandardScaler(), name="scaled_power")
+
+    # Build a train pipeline. In this pipeline, each step processes all data at once.
+    train_pipeline = Pipeline(path="../results/train")
+
+    # Create preprocessing pipeline for the preprocessing steps
+    preprocessing_pipeline = create_preprocessing_pipeline(power_scaler)
+    preprocessing_pipeline = preprocessing_pipeline(scaler_power=train_pipeline["load_power_statistics"])
+
+    target = FunctionModule(lambda x: numpy_to_xarray(
+        x.values.reshape((-1,)), x, "target"
+    ))(x=train_pipeline["load_power_statistics"])
+
+    # Addd the regressors to the train pipeline
+    regressor_svr(hist_input=preprocessing_pipeline["sampled_data"],
+                  target=target,
+                  callbacks=[LinePlotCallback('SVR')])
+
+    print("Start training")
+    train_pipeline.train(data)
+    print("Training finished")
+
+    # Create a second pipeline. Necessary, since this pipeline has additional steps in contrast to the train pipeline.
+    pipeline = Pipeline(path="../results")
+
+    # Get preprocessing pipeline
+    preprocessing_pipeline = create_preprocessing_pipeline(power_scaler)
+    preprocessing_pipeline = preprocessing_pipeline(scaler_power=pipeline["load_power_statistics"])
+
+    # Get the test pipeline, the arguments are the modules, from the training pipeline, which should be reused
+    test_pipeline = create_test_pipeline(regressor_svr)
+    test_pipeline(historical_input=preprocessing_pipeline["sampled_data"],
+                  load_power_statistics=pipeline["load_power_statistics"],
+                  callbacks=[LinePlotCallback('Pipeline'), CSVCallback('Pipeline')])
+
+    # Now, the pipeline is complete so we can run it and explore the results
+    # Start the pipeline
+    print("Start testing")
+    result = pipeline.test(test)
+    print("Testing finished")
diff --git a/pywatts/core/base.py b/pywatts/core/base.py
index ab95362..d053545 100644
--- a/pywatts/core/base.py
+++ b/pywatts/core/base.py
@@ -130,6 +130,13 @@ class Base(ABC):
         name = load_information["name"]
         return cls(name=name, **params)
 
+    def refit(self, **kwargs):
+        """
+        This method refits the module. If not overwritten it is the same as fit.
+        :param kwargs: key word arguments as input. If the key word starts with target, then it is a target variable.
+        """
+        return self.fit(**kwargs)
+
     def __call__(self,
                  use_inverse_transform: bool = False,
                  use_prob_transform: bool = False,
@@ -138,6 +145,8 @@ class Base(ABC):
                  computation_mode: ComputationMode = ComputationMode.Default,
                  batch_size: Optional[pd.Timedelta] = None,
                  train_if: Optional[Union[Callable, bool]] = None,
+                 lag: Optional[int] = pd.Timedelta(hours=0),
+                 retrain_batch: Optional[int] = pd.Timedelta(hours=24),
                  **kwargs: Union[StepInformation, Tuple[StepInformation, ...]]
                  ) -> StepInformation:
         """
@@ -167,6 +176,12 @@ class Base(ABC):
         :param computation_mode: Determines the computation mode of the step. Could be ComputationMode.Train,
                                  ComputationMode.Transform, and Computation.FitTransform
         :type computation_mode: ComputationMode
+        :param lag: Needed for online learning. Determines what data can be used for retraining.
+                    E.g., when 24 hour forecasts are performed, a lag of 24 hours is needed, else the retraining would
+                    use future values as target values.
+        :type lag: pd.Timedelta
+        :param retrain_batch: Needed for online learning. Determines how much data should be used for retraining.
+        :type retrain_batch: pd.Timedelta
         :return: a step information.
         :rtype: StepInformation
         """
@@ -179,7 +194,9 @@ class Base(ABC):
                                          condition=condition,
                                          callbacks=callbacks,
                                          computation_mode=computation_mode, batch_size=batch_size,
-                                         train_if=train_if
+                                         train_if=train_if,
+                                         retrain_batch=retrain_batch,
+                                         lag=lag
                                          )
 
 
diff --git a/pywatts/core/base_step.py b/pywatts/core/base_step.py
index 4fd5b3d..8dbb77f 100644
--- a/pywatts/core/base_step.py
+++ b/pywatts/core/base_step.py
@@ -153,9 +153,7 @@ class BaseStep(ABC):
         pass
 
     def _post_transform(self, result):
-        if isinstance(result, dict) and len(result) <= 1:
-            result = {self.name: list(result.values())[0]}
-        elif not isinstance(result, dict):
+        if not isinstance(result, dict):
             result = {self.name: result}
 
         if not self.buffer:
diff --git a/pywatts/core/computation_mode.py b/pywatts/core/computation_mode.py
index 74984b6..17c9dd2 100644
--- a/pywatts/core/computation_mode.py
+++ b/pywatts/core/computation_mode.py
@@ -9,3 +9,4 @@ class ComputationMode(IntEnum):
     Train = 2
     FitTransform = 3
     Default = 4
+    Refit = 5
diff --git a/pywatts/core/pipeline.py b/pywatts/core/pipeline.py
index 9a75469..c7125b8 100644
--- a/pywatts/core/pipeline.py
+++ b/pywatts/core/pipeline.py
@@ -100,6 +100,7 @@ class Pipeline(BaseTransformer):
                     message = f"From {self.counter} until {self.counter + self.batch} no data are calculated"
                     warnings.warn(message)
                     logger.info(message)
+            self.refit(self.counter, self.counter + self.batch if self.batch is not None else self.counter)
             self.counter += self.batch
         return result
 
@@ -420,3 +421,15 @@ class Pipeline(BaseTransformer):
                 summaries.append(step.get_summary())
             summaries.extend([step.transform_time, step.training_time])
         return summary_formatter.create_summary(summaries, self.file_manager)
+
+    def refit(self, start, end):
+        """
+        Refits all steps inside of the pipeline.
+        :param start: The date of the first data used for retraining.
+        :param end: The date of the last data used for retraining.
+        """
+        for step in self.id_to_step.values():
+            # A lag is needed, since if we have a 24 hour forecast we can evaluate the forecast not until 24 hours
+            # are gone, since before not all target variables are available
+            if isinstance(step, Step):
+                step.refit(start - step.lag, end - step.lag)
diff --git a/pywatts/core/pipeline_step.py b/pywatts/core/pipeline_step.py
index 9dbcbf6..041c070 100644
--- a/pywatts/core/pipeline_step.py
+++ b/pywatts/core/pipeline_step.py
@@ -42,7 +42,7 @@ class PipelineStep(Step):
 
     def _post_transform(self, result):
         self.module._create_summary(self.current_run_setting.summary_formatter)
-        super()._post_transform(result)
+        return super()._post_transform(result)
 
     def reset(self):
         """
diff --git a/pywatts/core/step.py b/pywatts/core/step.py
index e1a8919..825a090 100644
--- a/pywatts/core/step.py
+++ b/pywatts/core/step.py
@@ -3,6 +3,7 @@ import time
 from typing import Optional, Dict, Union, Callable, List
 
 import cloudpickle
+import pandas as pd
 import xarray as xr
 
 from pywatts.callbacks import BaseCallback
@@ -13,6 +14,7 @@ from pywatts.core.computation_mode import ComputationMode
 from pywatts.core.exceptions.not_fitted_exception import NotFittedException
 from pywatts.core.filemanager import FileManager
 from pywatts.core.result_step import ResultStep
+from pywatts.utils._xarray_time_series_utils import _get_time_indexes
 
 logger = logging.getLogger(__name__)
 
@@ -38,6 +40,12 @@ class Step(BaseStep):
     :type condition: Callable[xr.DataArray, xr.DataArray, bool]
     :param train_if: A callable which checks if the train_if step should be executed or not.
     :type train_if: Callable[xr.DataArray, xr.DataArray, bool]
+    :param lag: Needed for online learning. Determines what data can be used for retraining.
+            E.g., when 24 hour forecasts are performed, a lag of 24 hours is needed, else the retraining would
+            use future values as target values.
+    :type lag: pd.Timedelta
+    :param retrain_batch: Needed for online learning. Determines how much data should be used for retraining.
+    :type retrain_batch: pd.Timedelta
     """
 
     def __init__(self, module: Base, input_steps: Dict[str, BaseStep], file_manager, *,
@@ -46,14 +54,17 @@ class Step(BaseStep):
                  callbacks: List[Union[BaseCallback, Callable[[Dict[str, xr.DataArray]], None]]] = [],
                  condition=None,
                  batch_size: Optional[None] = None,
-                 train_if=None):
-
+                 train_if=None,
+                 retrain_batch=pd.Timedelta(hours=24),
+                 lag=pd.Timedelta(hours=24)):
         super().__init__(input_steps, targets, condition=condition,
                          computation_mode=computation_mode, name=module.name)
         self.file_manager = file_manager
         self.module = module
+        self.retrain_batch = retrain_batch
         self.callbacks = callbacks
         self.batch_size = batch_size
+        self.lag = lag
         self.train_if = train_if
         self.result_steps: Dict[str, ResultStep] = {}
 
@@ -122,24 +133,30 @@ class Step(BaseStep):
         input_data = self._get_input(start, end)
         target = self._get_target(start, end)
         if self.current_run_setting.computation_mode in [ComputationMode.Default, ComputationMode.FitTransform,
-                                                         ComputationMode.Train] and (
-                not self.train_if or self.train_if(input_data, target)):
+                                                         ComputationMode.Train]:
             # Fetch input_data and target data
-            start_time = time.time()
             if self.batch_size:
                 input_batch = self._get_input(end - self.batch_size, end)
                 target_batch = self._get_target(end - self.batch_size, end)
+                start_time = time.time()
                 self._fit(input_batch, target_batch)
+                self.training_time.set_kv("", time.time() - start_time)
             else:
+                start_time = time.time()
                 self._fit(input_data, target)
-            self.training_time.set_kv("", time.time() - start_time)
+                self.training_time.set_kv("", time.time() - start_time)
         elif self.module is BaseEstimator:
             logger.info("%s not fitted in Step %s", self.module.name, self.name)
 
         start_time = time.time()
         result = self._transform(input_data)
         self.transform_time.set_kv("", time.time() - start_time)
-        return result
+        result_dict = {}
+        for key, res in result.items():
+            index = res.indexes[_get_time_indexes(result)[0]]
+            start = max(index[0], start.to_numpy())
+            result_dict[key] = res.sel(**{_get_time_indexes(res)[0]: index[(index >= start)]})
+        return result_dict
 
     def _get_target(self, start, batch):
         return {
@@ -175,6 +192,22 @@ class Step(BaseStep):
                      "batch_size": self.batch_size})
         return json
 
+    def refit(self, start: pd.Timestamp, end: pd.Timestamp):
+        """
+        Refits the module of the step.
+        :param start: The date of the first data used for retraining.
+        :param end: The date of the last data used for retraining.
+        """
+        if self.current_run_setting.computation_mode in [ComputationMode.Refit] and isinstance(self.module,
+                                                                                               BaseEstimator):
+            if self.train_if:
+                input_data = self._get_input(start, end)
+                target = self._get_target(start, end)
+                if self.train_if(input_data, target):
+                    refit_input = self._get_input(end - self.retrain_batch, end)
+                    refit_target = self._get_target(end - self.retrain_batch, end)
+                    self.module.refit(**refit_input, **refit_target)
+
     def get_result_step(self, item: str):
         if item not in self.result_steps:
             self.result_steps[item] = ResultStep(input_steps={"result": self}, buffer_element=item)
diff --git a/pywatts/core/step_factory.py b/pywatts/core/step_factory.py
index fee22fe..66c6b68 100644
--- a/pywatts/core/step_factory.py
+++ b/pywatts/core/step_factory.py
@@ -31,7 +31,9 @@ class StepFactory:
                     condition,
                     batch_size,
                     computation_mode,
-                    train_if):
+                    train_if,
+                    retrain_batch,
+                    lag):
         """
         Creates a appropriate step for the current situation.
 
@@ -45,6 +47,10 @@ class StepFactory:
         :param batch_size: The size of the past time range which should be used for relearning the module
         :param computation_mode: The computation mode of the step
         :param train_if: A method for determining if the step should be fitted at a specific timestamp.
+        :param retrain_batch: Determines how much past data should be used for relearning.
+        :param lag: Needed for online learning. Determines what data can be used for retraining.
+                    E.g., when 24 hour forecasts are performed, a lag of 24 hours is needed, else the retraining would
+                    use future values as target values.
         :return: StepInformation
         """
 
@@ -68,17 +74,19 @@ class StepFactory:
         if isinstance(module, Pipeline):
             step = PipelineStep(module, input_steps, pipeline.file_manager, targets=target_steps,
                                 callbacks=callbacks, computation_mode=computation_mode, condition=condition,
-                                batch_size=batch_size, train_if=train_if)
+                                batch_size=batch_size, train_if=train_if, retrain_batch=retrain_batch, lag=lag)
         elif use_inverse_transform:
             step = InverseStep(module, input_steps, pipeline.file_manager, targets=target_steps,
-                               callbacks=callbacks, computation_mode=computation_mode, condition=condition)
+                               callbacks=callbacks, computation_mode=computation_mode, condition=condition,
+                               retrain_batch=retrain_batch, lag=lag)
         elif use_predict_proba:
             step = ProbablisticStep(module, input_steps, pipeline.file_manager, targets=target_steps,
-                                    callbacks=callbacks, computation_mode=computation_mode, condition=condition)
+                                    callbacks=callbacks, computation_mode=computation_mode, condition=condition,
+                                    retrain_batch=retrain_batch, lag=lag)
         else:
             step = Step(module, input_steps, pipeline.file_manager, targets=target_steps,
                         callbacks=callbacks, computation_mode=computation_mode, condition=condition,
-                        batch_size=batch_size, train_if=train_if)
+                        batch_size=batch_size, train_if=train_if, retrain_batch=retrain_batch, lag=lag)
 
         step_id = pipeline.add(module=step,
                                input_ids=[step.id for step in input_steps.values()],
diff --git a/pywatts/modules/preprocessing/sampler.py b/pywatts/modules/preprocessing/sampler.py
index fe22e57..8bad874 100644
--- a/pywatts/modules/preprocessing/sampler.py
+++ b/pywatts/modules/preprocessing/sampler.py
@@ -7,6 +7,7 @@ from pywatts.core.exceptions.wrong_parameter_exception import WrongParameterExce
 from pywatts.utils._xarray_time_series_utils import _get_time_indexes
 import numpy as np
 
+
 class Sampler(BaseTransformer):
     """
     This module creates samples with a size specified by sample_size. I.e., if sample_size is 24h. It creates for each
@@ -24,6 +25,11 @@ class Sampler(BaseTransformer):
         super().__init__(name)
         if indexes is None:
             indexes = []
+        if sample_size <= 0:
+            raise WrongParameterException(
+                "Sample size cannot be less than or equal to zero.",
+                "Please define a sample size greater than zero.",
+                module=self.name)
         self.sample_size = sample_size
         self.indexes = indexes
 
@@ -51,6 +57,11 @@ class Sampler(BaseTransformer):
         """
         if sample_size:
             self.sample_size = sample_size
+            if sample_size <= 0:
+                raise WrongParameterException(
+                    "Sample size cannot be less than or equal to zero.",
+                    "Please define a sample size greater than zero.",
+                    module=self.name)
         if indexes is not None:
             # Do not use if indexes here, since this would be false if indexes is empty.
             self.indexes = indexes
@@ -68,7 +79,7 @@ class Sampler(BaseTransformer):
         if not indexes:
             indexes = _get_time_indexes(x)
         try:
-            r = [x.shift({index: i for index in indexes}, fill_value=0) for i in range(0, self.sample_size)]
+            r = [x.shift({index: i for index in indexes}, fill_value=0) for i in range(self.sample_size - 1, -1, -1)]
         except ValueError as exc:
             raise WrongParameterException(
                 f"Not all indexes ({indexes}) are in the indexes of x ({list(x.indexes.keys())}).",


Test patch:
diff --git a/tests/unit/core/test_inverse_step.py b/tests/unit/core/test_inverse_step.py
index ef4a67f..612554d 100644
--- a/tests/unit/core/test_inverse_step.py
+++ b/tests/unit/core/test_inverse_step.py
@@ -32,12 +32,12 @@ class TestInverseTransform(unittest.TestCase):
     def test_get_result(self):
         # This test checks if the get_result methods works corerctly, i.e. if it returns the correct result of the step and
         # calculate it if necessary.
-        self.inverse_step.get_result(None, None)
+        self.inverse_step.get_result(pd.Timestamp("2000.01.01"), None)
         self.inverse_module.inverse_transform.assert_called_once_with(input=self.input_step.get_result())
 
     def test_get_result_stop(self):
         self.input_step.get_result.return_value = None
-        self.inverse_step.get_result(None, None)
+        self.inverse_step.get_result(pd.Timestamp("2000.01.01"), None)
 
         self.inverse_module.inverse_transform.assert_not_called()
         self.assertTrue(self.inverse_step._should_stop(None, None))
diff --git a/tests/unit/core/test_pipeline.py b/tests/unit/core/test_pipeline.py
index 691322c..96063dc 100644
--- a/tests/unit/core/test_pipeline.py
+++ b/tests/unit/core/test_pipeline.py
@@ -571,3 +571,13 @@ class TestPipeline(unittest.TestCase):
         open_mock().__enter__.return_value.write.assert_called_once_with(summary)
 
         self.assertTrue("target" in result.keys())
+
+    @patch('pywatts.core.pipeline.isinstance', return_value=True)
+    def test_refit(self, isinstance_mock):
+        first_step = MagicMock()
+        first_step.lag = pd.Timedelta("1d")
+
+        self.pipeline.add(module=first_step)
+        self.pipeline.refit(pd.Timestamp("2000.01.02"), pd.Timestamp("2022.01.02"))
+
+        first_step.refit.assert_called_once_with(pd.Timestamp("2000.01.01"), pd.Timestamp("2022.01.01"))
diff --git a/tests/unit/core/test_step.py b/tests/unit/core/test_step.py
index 94b52d9..7c637aa 100644
--- a/tests/unit/core/test_step.py
+++ b/tests/unit/core/test_step.py
@@ -270,3 +270,19 @@ class TestStep(unittest.TestCase):
         assert step.current_run_setting.computation_mode == ComputationMode.Default
         assert step._should_stop(None, None) == False
         assert step.finished == False
+
+    @patch('pywatts.core.step.isinstance', return_value=True)
+    def test_refit_train_if_false(self, isinstance_mock):
+        step = Step(self.module_mock, {"x": self.step_mock}, file_manager=None, train_if=lambda x, y: False,
+                    computation_mode=ComputationMode.Refit)
+        step.refit(pd.Timestamp("2000.01.01"), pd.Timestamp("2020.01.01"))
+        self.module_mock.refit.assert_not_called()
+
+    @patch('pywatts.core.step.isinstance', return_value=True)
+    def test_refit_train_if_true(self, isinstance_m):
+        step = Step(self.module_mock, {"x": self.step_mock},
+                    targets={"target": self.step_mock}, file_manager=None, train_if=lambda x, y: True,
+                    computation_mode=ComputationMode.Refit)
+        step.refit(pd.Timestamp("2000.01.01"), pd.Timestamp("2020.01.01"))
+        self.module_mock.refit.assert_called_once_with(x=self.step_mock.get_result(),
+                                                       target=self.step_mock.get_result())
diff --git a/tests/unit/modules/test_sampler.py b/tests/unit/modules/test_sampler.py
index 22aff19..fdac408 100644
--- a/tests/unit/modules/test_sampler.py
+++ b/tests/unit/modules/test_sampler.py
@@ -8,7 +8,7 @@ from pywatts.modules import Sampler
 
 class TestSampler(unittest.TestCase):
     def setUp(self) -> None:
-        self.sampler = Sampler(sample_size=2)
+        self.sampler = Sampler(sample_size=3)
 
     def tearDown(self) -> None:
         self.sampler = None
@@ -18,14 +18,14 @@ class TestSampler(unittest.TestCase):
 
         self.assertEqual(params,
                          {
-                             "sample_size": 2,
+                             "sample_size": 3,
                              "indexes": [],
                          })
 
     def test_set_params(self):
         self.assertEqual(self.sampler.get_params(),
                          {
-                             "sample_size": 2,
+                             "sample_size": 3,
                              "indexes": [],
                          })
         self.sampler.set_params(indexes=["Foo"], sample_size=12)
@@ -38,14 +38,15 @@ class TestSampler(unittest.TestCase):
     def test_transform(self):
         time = pd.date_range('2000-01-01', freq='24H', periods=7)
 
-        da = xr.DataArray([2, 3, 4, 5, 6, 7, 8], dims=['time'], coords={"time": time})
+        da = xr.DataArray([1, 2, 3, 4, 5, 6, 7], dims=['time'], coords={"time": time})
 
+        self.sampler.set_params(sample_size=3)
         result = self.sampler.transform(da)
 
         time = pd.date_range('2000-01-01', freq='24H', periods=7)
 
-        expected_result = xr.DataArray([[2, 0], [3, 2], [4,3], [5,4], [6,5], [7,6], [8,7]], dims=["time", "horizon"],
-                                       coords={"time":time})
+        expected_result = xr.DataArray([[0, 0, 1], [0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5], [4, 5, 6], [5, 6, 7]],
+                                       dims=["time", "horizon"], coords={"time": time})
 
         xr.testing.assert_equal(result, expected_result)
 
@@ -60,3 +61,9 @@ class TestSampler(unittest.TestCase):
                          "Not all indexes (['FOO']) are in the indexes of x (['time']). "
                          "Perhaps you set the wrong indexes with set_params or during the initialization of the Sampler.")
 
+    def test_set_params_exception(self):
+        with self.assertRaises(WrongParameterException) as context:
+            self.sampler.set_params(sample_size=-3)
+        self.assertEqual(context.exception.message,
+                         "Sample size cannot be less than or equal to zero. "
+                         "Please define a sample size greater than zero.")