Repository: KIT-IAI/pyWATTS

Issue:
Storage of custom keras loss function
**Describe the bug**
If pipeline.to_folder is called on a pipeline which contains a KerasWrapper with a custom loss function, then pyWATTS fails. 

**To Reproduce**
Steps to reproduce the behavior:
1. Build a NN with Keras with a custom loss functoin
2. Add the NN and the custom loss function to the wrapper
3. Add the wrapper to the pipline
4. Call to_folder on the pipeline

**Expected behavior**
The pipeline with the keraswrapper and the custom loss function is stored and can be reloaded correctly.


Discussion:

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c9218b..373fece 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,7 @@
 * Enable Multiple RefitConditions ([#192](https://github.com/KIT-IAI/pyWATTS/issues/192))
 * Periodic Refit Condition ([#192](https://github.com/KIT-IAI/pyWATTS/issues/192))
 * Add AnomalyGeneration modules ([#195](https://github.com/KIT-IAI/pyWATTS/issues/195))
+* Add get_min_data method in modules to enable them to specify how much data they need. ([#147](https://github.com/KIT-IAI/pyWATTS/issues/147))
 
 ### Changed
 * Retraining is triggered after all steps are transformed ([#117](https://github.com/KIT-IAI/pyWATTS/issues/117))
@@ -26,7 +27,7 @@
 * Added missing summaries to pywatts.summaries.__init__.py ([#183](https://github.com/KIT-IAI/pyWATTS/issues/183))
 * Fixed wrong time dimension name in ComputationMode.Transform ([#199](https://github.com/KIT-IAI/pyWATTS/issues/199))
 * Fixed online execution of subpipeline ([#203](https://github.com/KIT-IAI/pyWATTS/issues/203))
-
+* Enable to store a Keras Model with custom objects as custom loss functions ([#17](https://github.com/KIT-IAI/pyWATTS/issues/17))
 
 ### Deprecated
 
diff --git a/examples/example_keras.py b/examples/example_keras.py
index b7eecc8..c91f031 100644
--- a/examples/example_keras.py
+++ b/examples/example_keras.py
@@ -18,7 +18,7 @@ from pywatts.core.pipeline import Pipeline
 # Import the pyWATTS pipeline and the required modules
 from pywatts.modules import ClockShift, LinearInterpolater, SKLearnWrapper, KerasWrapper
 from pywatts.summaries import RMSE
-
+from tensorflow.keras import backend as K
 
 def get_keras_model():
     # write the model with the Functional API, Sequential does not support multiple input tensors
@@ -59,8 +59,11 @@ if __name__ == "__main__":
     shift_power_statistics2 = ClockShift(lag=2, name="ClockShift_Lag2")(x=scale_power_statistics)
 
     keras_wrapper = KerasWrapper(keras_model,
+                                 custom_objects={"<lambda>": lambda x, y: K.sqrt(K.mean(K.square(x - y)))},
                                  fit_kwargs={"batch_size": 8, "epochs": 1},
-                                 compile_kwargs={"loss": "mse", "optimizer": "Adam", "metrics": ["mse"]}) \
+                                 compile_kwargs={"loss": lambda x, y: K.sqrt(K.mean(K.square(x - y))),
+                                                 "optimizer": "Adam",
+                                                 "metrics": ["mse"]}) \
         (ClockShift_Lag1=shift_power_statistics,
          ClockShift_Lag2=shift_power_statistics2,
          target=scale_power_statistics)
@@ -82,3 +85,6 @@ if __name__ == "__main__":
 
     pipeline.train(data)
     pipeline.to_folder("../results/pipe_keras")
+
+    pipeline = Pipeline.from_folder("../results/pipe_keras")
+    pipeline.train(data)
diff --git a/pywatts/core/base.py b/pywatts/core/base.py
index f4c8757..4a7da00 100644
--- a/pywatts/core/base.py
+++ b/pywatts/core/base.py
@@ -137,6 +137,9 @@ class Base(ABC):
         """
         return self.fit(**kwargs)
 
+    def get_min_data(self):
+        return 0
+
     def __call__(self,
                  use_inverse_transform: bool = False,
                  use_prob_transform: bool = False,
diff --git a/pywatts/core/base_step.py b/pywatts/core/base_step.py
index 072bc99..cc064d4 100644
--- a/pywatts/core/base_step.py
+++ b/pywatts/core/base_step.py
@@ -36,7 +36,6 @@ class BaseStep(ABC):
         self.input_steps: Dict[str, "BaseStep"] = dict() if input_steps is None else input_steps
         self.targets: Dict[str, "BaseStep"] = dict() if targets is None else targets
         self.condition = condition
-        self.cached_result = {"cached": None, "start": None, "end": None}
 
         self.name = name
 
@@ -49,7 +48,7 @@ class BaseStep(ABC):
         self.transform_time = SummaryObjectList(self.name + " Transform Time", category=SummaryCategory.TransformTime)
 
     def get_result(self, start: pd.Timestamp, end: Optional[pd.Timestamp], buffer_element: str = None,
-                   return_all=False):
+                   return_all=False, minimum_data=(0, pd.Timedelta(0))):
         """
         This method is responsible for providing the result of this step.
         Therefore,
@@ -74,9 +73,7 @@ class BaseStep(ABC):
         # Only execute the module if the step is not finished and the results are not yet calculated
         if not self.finished and not (end is not None and self._current_end is not None and end <= self._current_end):
             if not self.buffer or not self._current_end or end > self._current_end:
-                self.cached_result["cached"] = self._compute(start, end)
-                self.cached_result["start"] = start
-                self.cached_result["end"] = end
+                self._compute(start, end, minimum_data)
                 self._current_end = end
             if not end:
                 self.finished = True
@@ -87,15 +84,9 @@ class BaseStep(ABC):
             if self.finished:
                 self._callbacks()
 
-        # Check if the cached results fits to the request, if yes return it.
-        if self.cached_result["cached"] is not None and self.cached_result["start"] == start and self.cached_result[
-            "end"] == end:
-            return copy.deepcopy(self.cached_result["cached"]) if return_all else copy.deepcopy(self.cached_result["cached"][
-                    buffer_element]) if buffer_element is not None else copy.deepcopy(list(self.cached_result["cached"].values())[
-                    0])
-        return self._pack_data(start, end, buffer_element, return_all=return_all)
+        return self._pack_data(start, end, buffer_element, return_all=return_all, minimum_data=minimum_data)
 
-    def _compute(self, start, end) -> Dict[str, xr.DataArray]:
+    def _compute(self, start, end, minimum_data) -> Dict[str, xr.DataArray]:
         pass
 
     def further_elements(self, counter: pd.Timestamp) -> bool:
@@ -118,12 +109,16 @@ class BaseStep(ABC):
                 return False
         return True
 
-    def _pack_data(self, start, end, buffer_element=None, return_all=False):
+    def _pack_data(self, start, end, buffer_element=None, return_all=False, minimum_data=(0, pd.Timedelta(0))):
         # Provide requested data
         time_index = _get_time_indexes(self.buffer)
         if start:
             index = list(self.buffer.values())[0].indexes[time_index[0]]
-            start = start.to_numpy()
+            if len(index) > 1:
+                freq = index[1] - index[0]
+            else:
+                freq = 0
+            start = start.to_numpy() - pd.Timedelta(minimum_data[0] * freq) - minimum_data[1]
             # If end is not set, all values should be considered. Thus we add a small timedelta to the last index entry.
             end = end.to_numpy() if end is not None else (index[-1] + pd.Timedelta(nanoseconds=1)).to_numpy()
             # After sel copy is not needed, since it returns a new array.
@@ -164,7 +159,8 @@ class BaseStep(ABC):
             # Time dimension is mandatory, consequently there dim has to exist
             dim = _get_time_indexes(result)[0]
             for key in self.buffer.keys():
-                self.buffer[key] = xr.concat([self.buffer[key], result[key]], dim=dim)
+                last = self.buffer[key][dim].values[-1]
+                self.buffer[key] = xr.concat([self.buffer[key], result[key][result[key][dim] > last]], dim=dim)
         return result
 
     def get_json(self, fm: FileManager) -> Dict:
@@ -201,10 +197,10 @@ class BaseStep(ABC):
         :return: The restored step.
         """
 
-    def _get_input(self, start, batch):
+    def _get_input(self, start, batch, minimum_data=(0, pd.Timedelta(0))):
         return None
 
-    def _get_target(self, start, batch):
+    def _get_target(self, start, batch, minimum_data=(0, pd.Timedelta(0))):
         return None
 
     def _should_stop(self, start, end) -> bool:
diff --git a/pywatts/core/either_or_step.py b/pywatts/core/either_or_step.py
index c7e0c6e..8708dec 100644
--- a/pywatts/core/either_or_step.py
+++ b/pywatts/core/either_or_step.py
@@ -1,3 +1,4 @@
+import pandas as pd
 from pywatts.core.base_step import BaseStep
 
 
@@ -14,14 +15,14 @@ class EitherOrStep(BaseStep):
         super().__init__(input_steps)
         self.name = "EitherOr"
 
-    def _compute(self, start, end):
-        input_data = self._get_input(start, end)
+    def _compute(self, start, end, minimum_data):
+        input_data = self._get_input(start, end, minimum_data)
         return self._transform(input_data)
 
-    def _get_input(self, start, batch):
+    def _get_input(self, start, batch, minimum_data=(0, pd.Timedelta(0))):
         inputs = []
         for step in self.input_steps.values():
-            inp = step.get_result(start, batch)
+            inp = step.get_result(start, batch, minimum_data=minimum_data)
             inputs.append(inp)
         return inputs
 
diff --git a/pywatts/core/result_step.py b/pywatts/core/result_step.py
index 84c5a18..58576ca 100644
--- a/pywatts/core/result_step.py
+++ b/pywatts/core/result_step.py
@@ -15,14 +15,16 @@ class ResultStep(BaseStep):
         super().__init__(input_steps=input_steps)
         self.buffer_element = buffer_element
 
-    def get_result(self, start: pd.Timestamp, end: Optional[pd.Timestamp], *args, return_all=False, **kwargs):
+    def get_result(self, start: pd.Timestamp, end: Optional[pd.Timestamp], buffer_element: str = None,
+                   return_all=False, minimum_data=(0, pd.Timedelta(0))):
         """
         Returns the specified result of the previous step.
         """
+
         if not return_all:
-            return list(self.input_steps.values())[0].get_result(start, end, self.buffer_element)
+            return list(self.input_steps.values())[0].get_result(start, end, self.buffer_element, minimum_data=minimum_data)
         else:
-            return {self.buffer_element: list(self.input_steps.values())[0].get_result(start, end, self.buffer_element)}
+            return {self.buffer_element: list(self.input_steps.values())[0].get_result(start, end, self.buffer_element, minimum_data=minimum_data)}
 
     def get_json(self, fm: FileManager) -> Dict:
         """
diff --git a/pywatts/core/step.py b/pywatts/core/step.py
index cd89a8a..7d32b19 100644
--- a/pywatts/core/step.py
+++ b/pywatts/core/step.py
@@ -4,6 +4,7 @@ import warnings
 from typing import Optional, Dict, Union, Callable, List
 
 import cloudpickle
+import numpy as np
 import pandas as pd
 import xarray as xr
 
@@ -144,15 +145,15 @@ class Step(BaseStep):
 
         return step
 
-    def _compute(self, start, end):
-        input_data = self._get_input(start, end)
-        target = self._get_target(start, end)
+    def _compute(self, start, end, minimum_data):
+        input_data = self._get_input(start, end, minimum_data)
+        target = self._get_target(start, end, minimum_data)
         if self.current_run_setting.computation_mode in [ComputationMode.Default, ComputationMode.FitTransform,
                                                          ComputationMode.Train]:
             # Fetch input_data and target data
             if self.batch_size:
-                input_batch = self._get_input(end - self.batch_size, end)
-                target_batch = self._get_target(end - self.batch_size, end)
+                input_batch = self._get_input(end - self.batch_size, end, minimum_data)
+                target_batch = self._get_target(end - self.batch_size, end, minimum_data)
                 start_time = time.time()
                 self._fit(input_batch, target_batch)
                 self.training_time.set_kv("", time.time() - start_time)
@@ -173,14 +174,26 @@ class Step(BaseStep):
             result_dict[key] = res.sel(**{_get_time_indexes(res)[0]: index[(index >= start)]})
         return result_dict
 
-    def _get_target(self, start, batch):
+    def _get_target(self, start, batch, minimum_data=(0, pd.Timedelta(0))):
+        min_data_module = self.module.get_min_data()
+        if isinstance(min_data_module, (int, np.integer)):
+            minimum_data = minimum_data[0] + min_data_module, minimum_data[1]
+        else:
+            minimum_data = minimum_data[0], minimum_data[1] + min_data_module
         return {
-            key: target.get_result(start, batch) for key, target in self.targets.items()
+            key: target.get_result(start, batch, minimum_data=minimum_data)
+            for key, target in self.targets.items()
         }
 
-    def _get_input(self, start, batch):
+    def _get_input(self, start, batch, minimum_data=(0, pd.Timedelta(0))):
+        min_data_module = self.module.get_min_data()
+        if isinstance(min_data_module, (int, np.integer)):
+            minimum_data = minimum_data[0] + min_data_module, minimum_data[1]
+        else:
+            minimum_data = minimum_data[0], minimum_data[1] + min_data_module
         return {
-            key: input_step.get_result(start, batch) for key, input_step in self.input_steps.items()
+            key: input_step.get_result(start, batch, minimum_data=minimum_data) for
+            key, input_step in self.input_steps.items()
         }
 
     def get_json(self, fm: FileManager):
diff --git a/pywatts/modules/feature_extraction/rolling_base.py b/pywatts/modules/feature_extraction/rolling_base.py
index 56e841f..54573fb 100644
--- a/pywatts/modules/feature_extraction/rolling_base.py
+++ b/pywatts/modules/feature_extraction/rolling_base.py
@@ -2,6 +2,7 @@ from abc import ABC, abstractmethod
 from enum import IntEnum
 from typing import Dict, Optional
 
+import pandas as pd
 import xarray as xr
 
 from pywatts.core.base import BaseTransformer
@@ -129,3 +130,6 @@ class RollingBase(BaseTransformer, ABC):
     @abstractmethod
     def _get_rolling(self, df):
         pass
+
+    def get_min_data(self):
+        return pd.Timedelta(f"{self.window_size}{self.window_size_unit}")
diff --git a/pywatts/modules/feature_extraction/trend_extraction.py b/pywatts/modules/feature_extraction/trend_extraction.py
index 78b9f95..d4f9d3b 100644
--- a/pywatts/modules/feature_extraction/trend_extraction.py
+++ b/pywatts/modules/feature_extraction/trend_extraction.py
@@ -83,3 +83,6 @@ class TrendExtraction(BaseTransformer):
                   range(1, self.length + 1)]
         trend = xr.DataArray(np.stack(trends, axis=-1), dims=(*x.dims, "length"), coords=x.coords)
         return trend.transpose(_get_time_indexes(x)[0], "length", ...)
+
+    def get_min_data(self):
+        return self.period * (self.length + 1)
diff --git a/pywatts/modules/preprocessing/clock_shift.py b/pywatts/modules/preprocessing/clock_shift.py
index 4ed6f7b..667059b 100644
--- a/pywatts/modules/preprocessing/clock_shift.py
+++ b/pywatts/modules/preprocessing/clock_shift.py
@@ -1,5 +1,7 @@
 from typing import Dict, List
 
+import numpy as np
+import pandas as pd
 import xarray as xr
 
 from pywatts.core.base import BaseTransformer
@@ -75,3 +77,6 @@ class ClockShift(BaseTransformer):
                 f"Not all indexes ({indexes}) are in the indexes of x ({list(x.indexes.keys())}).",
                 "Perhaps you set the wrong indexes with set_params or during the initialization of the ClockShift.",
                 module=self.name) from exc
+
+    def get_min_data(self):
+        return np.abs(self.lag)
diff --git a/pywatts/modules/preprocessing/resampler.py b/pywatts/modules/preprocessing/resampler.py
index 780496c..bceeea1 100644
--- a/pywatts/modules/preprocessing/resampler.py
+++ b/pywatts/modules/preprocessing/resampler.py
@@ -1,5 +1,6 @@
 from typing import Dict, Optional, Any
 
+import pandas as pd
 import xarray as xr
 
 from pywatts.core.base import BaseTransformer
@@ -108,3 +109,6 @@ class Resampler(BaseTransformer):
         return getattr(
             x.resample(**{self.time_index: self.target_time}), self.method
         )(**args)
+
+    def get_min_data(self):
+        return pd.Timedelta(self.target_time)
diff --git a/pywatts/modules/preprocessing/sampler.py b/pywatts/modules/preprocessing/sampler.py
index 8bad874..572f72e 100644
--- a/pywatts/modules/preprocessing/sampler.py
+++ b/pywatts/modules/preprocessing/sampler.py
@@ -1,5 +1,6 @@
 from typing import Dict, List
 
+import pandas as pd
 import xarray as xr
 
 from pywatts.core.base import BaseTransformer
@@ -33,6 +34,9 @@ class Sampler(BaseTransformer):
         self.sample_size = sample_size
         self.indexes = indexes
 
+    def get_min_data(self):
+        return self.sample_size
+
     def get_params(self) -> Dict[str, object]:
         """
         Returns a list of all hyperparameters/ user defined parameters
diff --git a/pywatts/modules/wrappers/keras_wrapper.py b/pywatts/modules/wrappers/keras_wrapper.py
index 471994e..663315c 100644
--- a/pywatts/modules/wrappers/keras_wrapper.py
+++ b/pywatts/modules/wrappers/keras_wrapper.py
@@ -1,6 +1,7 @@
 import logging
 from typing import Tuple, Union, Dict
 
+import cloudpickle
 import tensorflow as tf
 import xarray as xr
 
@@ -21,10 +22,14 @@ class KerasWrapper(DlWrapper):
     :type fit_kwargs: dict
     :param compile_kwargs: The compile keyword arguments necessary for compiling the model.
     :type compile_kwargs: dict
+    :param custom_objects: This dict contains all custom objects needed by the keras model. Note,
+                           users that uses such customs objects (e.g. Custom Loss) need to specify this to enable
+                           the loading of the stored Keras model.
+    :type custom_objects: dict
     """
 
     def __init__(self, model: Union[tf.keras.Model, Tuple[tf.keras.Model, Dict[str, tf.keras.Model]]],
-                 name: str = "KerasWrapper", fit_kwargs=None, compile_kwargs=None):
+                 name: str = "KerasWrapper", fit_kwargs=None, compile_kwargs=None, custom_objects=None):
         self.aux_models = {}
         self.targets = []
         if isinstance(model, tuple):
@@ -35,6 +40,10 @@ class KerasWrapper(DlWrapper):
             self.compile_kwargs = {}
         else:
             self.compile_kwargs = compile_kwargs
+        if custom_objects is None:
+            self.custom_objects = {}
+        else:
+            self.custom_objects = custom_objects
 
     def fit(self, **kwargs: xr.DataArray):
         """
@@ -74,7 +83,22 @@ class KerasWrapper(DlWrapper):
         :param fm: The Filemanager, which contains the path where the model should be stored
         :return: The path where the model is stored.
         """
-        json = super().save(fm)
+        json = {"name": self.name,
+                "class": self.__class__.__name__,
+                "module": self.__module__}
+
+        params = self.get_params()
+        params_path = fm.get_path(f"{self.name}_params.pickle")
+        with open(params_path, "wb") as outfile:
+            cloudpickle.dump(params, outfile)
+        json["params"] = params_path
+        json["is_fitted"] = self.is_fitted
+
+        custom_path = fm.get_path(f"{self.name}_custom.pickle")
+        with open(custom_path, "wb") as outfile:
+            cloudpickle.dump(self.custom_objects, outfile)
+        json["custom_objects"] = custom_path
+
         json["targets"] = self.targets
         model_path = fm.get_path(f"{self.name}.h5")
         self.model.save(filepath=model_path)
@@ -98,10 +122,17 @@ class KerasWrapper(DlWrapper):
         (Note: This models should be taken from the pipeline json file)
         :return: A wrapped keras model.
         """
-        params = load_information["params"]
         name = load_information["name"]
+        params_path = load_information["params"]
+        with open(params_path, "rb") as infile:
+            params = cloudpickle.load(infile)
+
+        custom_path = load_information["custom_objects"]
+        with open(custom_path, "rb") as infile:
+            custom_objects = cloudpickle.load(infile)
+
         try:
-            model = tf.keras.models.load_model(filepath=load_information["model"])
+            model = tf.keras.models.load_model(filepath=load_information["model"], custom_objects=custom_objects)
         except Exception as exception:
             logging.error("No model found in %s.", load_information['model'])
             raise exception
@@ -109,7 +140,7 @@ class KerasWrapper(DlWrapper):
         if "aux_models" in load_information.keys():
             for aux_name, path in load_information["aux_models"]:
                 try:
-                    aux_models[aux_name] = tf.keras.models.load_model(filepath=path)
+                    aux_models[aux_name] = tf.keras.models.load_model(filepath=path, custom_objects=custom_objects)
                 except Exception as exception:
                     logging.error("No model found in path %s", path)
                     raise exception
@@ -128,16 +159,22 @@ class KerasWrapper(DlWrapper):
         """
         return {
             "fit_kwargs": self.fit_kwargs,
-            "compile_kwargs": self.compile_kwargs
+            "compile_kwargs": self.compile_kwargs,
+            "custom_objects": self.custom_objects
         }
 
-    def set_params(self, fit_kwargs=None, compile_kwargs=None):
+    def set_params(self, fit_kwargs=None, compile_kwargs=None, custom_objects=None):
         """
         Set the parameters of the deep learning wrappers
         :param fit_kwargs: keyword arguments for the fit method.
         :param compile_kwargs: keyword arguments for the compile methods.
+        :param custom_objects: This dict contains all custom objects needed by the keras model. Note,
+                               users that uses such customs objects (e.g. Custom Loss) need to specify this to enable
+                               the loading of the stored Keras model.
         """
         if fit_kwargs:
             self.fit_kwargs = fit_kwargs
         if compile_kwargs:
             self.compile_kwargs = compile_kwargs
+        if custom_objects:
+            self.custom_objects = custom_objects


Test patch:
diff --git a/tests/unit/core/test_step.py b/tests/unit/core/test_step.py
index 8453a97..d9de8f1 100644
--- a/tests/unit/core/test_step.py
+++ b/tests/unit/core/test_step.py
@@ -112,6 +112,7 @@ class TestStep(unittest.TestCase):
 
         self.module_mock.transform.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
                                                                coords={'time': time2})
+        self.module_mock.get_min_data.return_value = 2
         step = Step(self.module_mock, {"x": input_step}, file_manager=MagicMock())
         da = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})
         step.buffer = {"test": da}
@@ -122,8 +123,10 @@ class TestStep(unittest.TestCase):
 
         # Two calls, once in should_stop and once in _transform
         input_step.get_result.assert_has_calls(
-            [call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00')),
-             call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00'))])
+            [call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00'),
+                  minimum_data=(2, pd.Timedelta(0))),
+             call(pd.Timestamp('2000-01-07 00:00:00'), pd.Timestamp('2020-01-14 00:00:00'),
+                  minimum_data=(2, pd.Timedelta(0)))])
         xr_mock.concat.assert_called_once()
 
         xr.testing.assert_equal(da, list(xr_mock.concat.call_args_list[0])[0][0][0])
@@ -147,8 +150,10 @@ class TestStep(unittest.TestCase):
 
         # Two calls, once in should_stop and once in _transform
         input_step.get_result.assert_has_calls(
-            [call(pd.Timestamp('2000-01-01 '), pd.Timestamp('2020-12-12 ')),
-             call(pd.Timestamp('2000-01-01 '), pd.Timestamp('2020-12-12 '))])
+            [call(pd.Timestamp("2000-01-01"), pd.Timestamp('2020-12-12 '),
+                  minimum_data=(0, self.module_mock.get_min_data().__radd__())),
+             call(pd.Timestamp('2000-01-01 '), pd.Timestamp('2020-12-12 '),
+                  minimum_data=(0, self.module_mock.get_min_data().__radd__()))])
 
         self.module_mock.transform.assert_called_once_with(x=input_step_result_mock)
 
@@ -271,27 +276,33 @@ class TestStep(unittest.TestCase):
         assert step._should_stop(None, None) == False
         assert step.finished == False
 
+
+    @patch('pywatts.core.step.Step._get_target', return_value={"target": 1})
+    @patch('pywatts.core.step.Step._get_input', return_value={"x": 2})
     @patch('pywatts.core.step.isinstance', side_effect=[True, False, True])
-    def test_refit_refit_conditions_false(self, isinstance_mock):
+    def test_refit_refit_conditions_false(self, isinstance_mock, get_input_mock, get_target_mock):
         step = Step(self.module_mock, {"x": self.step_mock}, file_manager=None, refit_conditions=[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.Step._get_target', return_value={"target": 1})
+    @patch('pywatts.core.step.Step._get_input', return_value={"x": 2})
     @patch('pywatts.core.step.isinstance', side_effect=[True, False, True])
-    def test_refit_refit_conditions_true(self, isinstance_mock):
+    def test_refit_refit_conditions_true(self, isinstance_mock, get_input_mock, get_target_mock):
         step = Step(self.module_mock, {"x": self.step_mock},
                     targets={"target": self.step_mock}, file_manager=None, refit_conditions=[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())
+        self.module_mock.refit.assert_called_once_with(x=2, target=1)
 
+    @patch('pywatts.core.step.Step._get_target', return_value={"target": 1})
+    @patch('pywatts.core.step.Step._get_input', return_value={"x": 2})
     @patch('pywatts.core.step.isinstance', side_effect=[True, False, True, False, True])
-    def test_multiple_refit_conditions(self, isinstance_mock):
+    def test_multiple_refit_conditions(self, isinstance_mock, get_input_mock, get_target_mock):
         step = Step(self.module_mock, {"x": self.step_mock},
                     targets={"target": self.step_mock}, file_manager=None, refit_conditions=[lambda x, y: False, 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())
+        self.module_mock.refit.assert_called_once_with(x=2, target=1)
+
diff --git a/tests/unit/core/test_summary_step.py b/tests/unit/core/test_summary_step.py
index 644324a..0bd495c 100644
--- a/tests/unit/core/test_summary_step.py
+++ b/tests/unit/core/test_summary_step.py
@@ -30,9 +30,10 @@ class TestSummaryStep(unittest.TestCase):
         self.fm_mock = None
 
     def test_get_summary(self):
+        self.module_mock.get_min_data.return_value = pd.Timedelta(hours=1)
         result = self.summary.get_summary(None, None)
 
-        self.step_mock.get_result.assert_called_once_with(None, None)
+        self.step_mock.get_result.assert_called_once_with(None, None, minimum_data=(0, pd.Timedelta(hours=1)))
         self.module_mock.transform.assert_called_once_with(file_manager=self.fm_mock, input=self.result_mock)
         self.assertEqual(result, "#I AM MARKDOWN")
 
diff --git a/tests/unit/modules/test_clock_shift.py b/tests/unit/modules/test_clock_shift.py
index b256bfa..bb815a8 100644
--- a/tests/unit/modules/test_clock_shift.py
+++ b/tests/unit/modules/test_clock_shift.py
@@ -60,3 +60,7 @@ class TestClockShift(unittest.TestCase):
         self.assertEqual(context.exception.message,
                          "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 ClockShift.")
+
+    def test_get_min_data(self):
+        clock_shift = ClockShift(lag=24)
+        self.assertEqual(clock_shift.get_min_data(), 24)
diff --git a/tests/unit/modules/test_resampler.py b/tests/unit/modules/test_resampler.py
index fd8a935..f88fb75 100644
--- a/tests/unit/modules/test_resampler.py
+++ b/tests/unit/modules/test_resampler.py
@@ -88,3 +88,7 @@ class TestResampler(unittest.TestCase):
         expected_data = np.arange(0, self.times - 0.5, step=0.5)
         expected_time = pd.date_range(self.start_time, freq="1800s", periods=len(expected_data))
         self._test_base("linear", "1800s", "interpolate", expected_time, expected_data, method_args={"kind": "linear"})
+
+    def test_get_min_data(self):
+        resampler = Resampler(target_time="12h")
+        self.assertEqual(resampler.get_min_data(), pd.Timedelta(hours=12))
diff --git a/tests/unit/modules/test_rolling_mean.py b/tests/unit/modules/test_rolling_mean.py
index 05cd5a6..36ebc1c 100644
--- a/tests/unit/modules/test_rolling_mean.py
+++ b/tests/unit/modules/test_rolling_mean.py
@@ -83,3 +83,7 @@ class TestRollingMean(unittest.TestCase):
         expected_result = xr.DataArray([0., 0., 0., 0., 3, 3, 3.5, 3.5, 2, 2., 6.,6,4.5, 4.5], dims=["time"], coords={'time': time})
 
         xr.testing.assert_equal(result, expected_result)
+
+    def test_get_min_data(self):
+        r_mean = RollingMean(window_size=12, window_size_unit="h")
+        self.assertEqual(r_mean.get_min_data(), pd.Timedelta(hours=12))
diff --git a/tests/unit/modules/test_sampler.py b/tests/unit/modules/test_sampler.py
index fdac408..65f11ff 100644
--- a/tests/unit/modules/test_sampler.py
+++ b/tests/unit/modules/test_sampler.py
@@ -67,3 +67,7 @@ class TestSampler(unittest.TestCase):
         self.assertEqual(context.exception.message,
                          "Sample size cannot be less than or equal to zero. "
                          "Please define a sample size greater than zero.")
+
+    def test_get_min_data(self):
+        sampler = Sampler(sample_size=24)
+        self.assertEqual(sampler.get_min_data(), 24)
diff --git a/tests/unit/modules/test_trend_extraction.py b/tests/unit/modules/test_trend_extraction.py
index 4bcb85c..113f74b 100644
--- a/tests/unit/modules/test_trend_extraction.py
+++ b/tests/unit/modules/test_trend_extraction.py
@@ -49,3 +49,7 @@ class TestTrendExtraction(unittest.TestCase):
              [9, 7, 5]], dims=["time", "length"], coords={'time': time})
 
         xr.testing.assert_equal(result, expected_result)
+
+    def test_get_min_data(self):
+        trend = TrendExtraction(period=12, length=2)
+        self.assertEqual(trend.get_min_data(), 36)
diff --git a/tests/unit/wrapper/test_keras_wrapper.py b/tests/unit/wrapper/test_keras_wrapper.py
index d287788..6f44ed6 100644
--- a/tests/unit/wrapper/test_keras_wrapper.py
+++ b/tests/unit/wrapper/test_keras_wrapper.py
@@ -25,19 +25,8 @@ stored_model = {
     "module": "pywatts.wrappers.keras_wrapper",
     "name": "SimpleAE",
     'is_fitted': False,
-    "params": {
-        "compile_kwargs": {
-            "loss": "mse",
-            "metrics": [
-                "mse"
-            ],
-            "optimizer": "Adam"
-        },
-        "fit_kwargs": {
-            "batch_size": 512,
-            "epochs": 1
-        }
-    },
+    "params": "params_path",
+    "custom_objects" : "custom_path",
     "targets": []
 }
 
@@ -132,29 +121,41 @@ class TestKerasWrapper(unittest.TestCase):
     def test_get_params(self):
         self.assertEqual(self.keras_wrapper.get_params(),
                          {'compile_kwargs': {'test': 'arg1'},
-                          'fit_kwargs': {'42': 24}})
+                          'fit_kwargs': {'42': 24},
+                          'custom_objects' : {}
+                          })
 
     def test_set_params(self):
         self.assertEqual(self.keras_wrapper.get_params(),
                          {'compile_kwargs': {'test': 'arg1'},
-                          'fit_kwargs': {'42': 24}})
+                          'fit_kwargs': {'42': 24},
+                          'custom_objects': {}})
         self.keras_wrapper.set_params(fit_kwargs={"loss": "mse"},
                                       compile_kwargs={"optimizer": "Adam"})
         self.assertEqual(self.keras_wrapper.get_params(),
-                         {
-                             "fit_kwargs": {
-                                 "loss": "mse"
-                             }, "compile_kwargs": {
-                             "optimizer": "Adam"
-                         },
-                         })
+                         {"fit_kwargs": {"loss": "mse"},
+                          "compile_kwargs": {"optimizer": "Adam"},
+                          "custom_objects": {}})
 
-    def test_save(self):
+    @patch("pywatts.modules.wrappers.keras_wrapper.cloudpickle")
+    @patch("builtins.open")
+    def test_save(self, open_mock, pickle_mock):
         fm_mock = MagicMock()
-        fm_mock.get_path.return_value = os.path.join("new_path", "to_somewhere", "KerasWrapper.h5")
+        fm_mock.get_path.side_effect = [
+            "params_path",
+            "custom_path",
+            os.path.join("new_path", "to_somewhere", "KerasWrapper.h5")]
+        file_mock = MagicMock()
+        open_mock().__enter__.return_value = file_mock
+
         json = self.keras_wrapper.save(fm_mock)
         self.keras_mock.save.assert_called_once_with(
             filepath=os.path.join("new_path", "to_somewhere", "KerasWrapper.h5"))
+
+        open_mock.assert_has_calls([call("params_path", "wb"), call("custom_path", "wb")], any_order=True)
+
+        pickle_mock.dump.assert_has_calls([call(self.keras_wrapper.get_params(), file_mock),
+                                      call(self.keras_wrapper.custom_objects, file_mock)])
         fm_mock.get_path.has_calls(call(os.path.join("to_somewhere", "KerasWrapper.h5")),
                                    any_order=True)
         self.assertEqual(json, {'aux_models': [],
@@ -163,21 +164,28 @@ class TestKerasWrapper(unittest.TestCase):
                                 'model': os.path.join("new_path", "to_somewhere", "KerasWrapper.h5"),
                                 'module': 'pywatts.modules.wrappers.keras_wrapper',
                                 'name': 'KerasWrapper',
-                                'params': {'compile_kwargs': {"test": "arg1"}, 'fit_kwargs': {"42": 24}},
+                                'params': "params_path",
+                                "custom_objects" : "custom_path",
                                 "targets":[]
                                 })
 
+    @patch("pywatts.modules.wrappers.keras_wrapper.cloudpickle")
+    @patch("builtins.open")
     @patch('pywatts.modules.wrappers.keras_wrapper.tf.keras.models.load_model')
-    def test_load(self, load_model_mock):
+    def test_load(self, load_model_mock, open_mock, pickle_mock):
         new_keras_mock = MagicMock()
         load_model_mock.return_value = new_keras_mock
+        pickle_mock.load.side_effect = [{"compile_kwargs": {"loss": "mse", "metrics": ["mse"], "optimizer": "Adam"},
+                                         "fit_kwargs": {"batch_size": 512, "epochs": 1},
+                                         "custom_objects": {}, },
+                                        "custom_object"]
         new_keras_wrapper = KerasWrapper.load(stored_model)
-        calls_open = [call(filepath=os.path.join("pipe1", "SimpleAE_4decoder.h5")),
-                      call(filepath=os.path.join("pipe1", "SimpleAE_4encoder.h5")),
-                      call(filepath=os.path.join("pipe1", "SimpleAE_4.h5")),
+        calls_open = [call(filepath=os.path.join("pipe1", "SimpleAE_4decoder.h5"), custom_objects="custom_object"),
+                      call(filepath=os.path.join("pipe1", "SimpleAE_4encoder.h5"), custom_objects="custom_object"),
+                      call(filepath=os.path.join("pipe1", "SimpleAE_4.h5"), custom_objects="custom_object"),
                       ]
-
         load_model_mock.assert_has_calls(calls_open, any_order=True)
+        open_mock.assert_has_calls([call("params_path", "rb"), call("custom_path", "rb")], any_order=True)
         self.assertEqual(load_model_mock.call_count, 3)
         self.assertEqual(new_keras_mock, new_keras_wrapper.model)
         self.assertEqual(new_keras_wrapper.get_params(),
@@ -192,5 +200,6 @@ class TestKerasWrapper(unittest.TestCase):
                              "fit_kwargs": {
                                  "batch_size": 512,
                                  "epochs": 1
-                             }
+                             },
+                             "custom_objects": {}
                          })