def test_further_elements_already_buffered(self):
time = pd.date_range('2000-01-01', freq='24H', periods=7)
step = Step(self.module_mock, {"x": self.step_mock}, file_manager=MagicMock())
step.buffer = {"STEP": xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})}
result = step.further_elements(pd.Timestamp("2000-01-05"))
self.step_mock.further_elements.assert_not_called()
self.assertEqual(result, True)
def test_store_load_of_step_with_train_if(self, cloudpickle_mock, open_mock):
train_if_mock = MagicMock()
step = Step(self.module_mock, self.step_mock, None, train_if=train_if_mock)
fm_mock = MagicMock()
fm_mock.get_path.return_value = os.path.join("folder", "test_train_if.pickle")
json = step.get_json(fm_mock)
reloaded_step = Step.load(json, [self.step_mock], targets=None, module=self.module_mock,
file_manager=MagicMock())
# One call in load and one in save
open_mock.assert_has_calls(
[call(os.path.join("folder", "test_train_if.pickle"), "wb"),
call(os.path.join("folder", "test_train_if.pickle"), "rb")],
any_order=True)
self.assertEqual(json, {
"target_ids": {},
# Same as for test_load.
"input_ids": {},
"id": -1,
'batch_size': None,
'default_run_setting': {'computation_mode': 4},
"train_if": os.path.join("folder", "test_train_if.pickle"),
"module": "pywatts.core.step",
"class": "Step",
"name": "test",
'callbacks': [],
"last": True,
'condition': None}, json),
self.assertEqual(reloaded_step.module, self.module_mock)
self.assertEqual(reloaded_step.input_steps, [self.step_mock])
cloudpickle_mock.load.assert_called_once_with(open_mock().__enter__.return_value)
cloudpickle_mock.dump.assert_called_once_with(train_if_mock, open_mock().__enter__.return_value)
def test_transform_batch_with_existing_buffer(self, xr_mock, *args):
# Check that data in batch learning are concatenated
input_step = MagicMock()
input_step._should_stop.return_value = False
time = pd.date_range('2000-01-01', freq='1D', periods=7)
time2 = pd.date_range('2000-01-14', freq='1D', periods=7)
time3 = pd.date_range('2000-01-01', freq='1D', periods=14)
self.module_mock.transform.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time2})
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}
xr_mock.concat.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2, 2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time3})
step.get_result(pd.Timestamp("2000.01.07"), pd.Timestamp("2020.01.14"))
# 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'))])
xr_mock.concat.assert_called_once()
xr.testing.assert_equal(da, list(xr_mock.concat.call_args_list[0])[0][0][0])
xr.testing.assert_equal(xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time2}), list(xr_mock.concat.call_args_list[0])[0][0][1])
assert {'dim': 'time'} == list(xr_mock.concat.call_args_list[0])[1]
def test_set_run_setting(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.set_run_setting(RunSetting(ComputationMode.FitTransform))
assert step.current_run_setting.computation_mode == ComputationMode.FitTransform
step.set_run_setting(RunSetting(ComputationMode.Transform))
assert step.current_run_setting.computation_mode == ComputationMode.Transform
def test_fit_with_targets(self):
input_dict = {'input_data': None}
target_mock = MagicMock()
step = Step(self.module_mock, self.step_mock, MagicMock(), targets={"target": MagicMock()})
step._fit(input_dict, {"target": target_mock})
self.module_mock.fit.assert_called_once_with(**input_dict, target=target_mock)
def test_set_computation_mode(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.set_computation_mode(ComputationMode.FitTransform)
assert step.computation_mode == ComputationMode.FitTransform
step.set_computation_mode(ComputationMode.Transform)
assert step.computation_mode == ComputationMode.Transform
def test_further_elements_target_false(self):
target_step = MagicMock()
target_step.further_elements.return_value = False
time = pd.date_range('2000-01-01', freq='1H', periods=7)
step = Step(self.module_mock, {"x": self.step_mock}, targets={"target": target_step}, file_manager=MagicMock())
step.buffer = {"STEP": xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"], coords={'time': time})}
result = step.further_elements(pd.Timestamp("2000.12.12"))
target_step.further_elements.assert_called_once_with(pd.Timestamp("2000.12.12"))
self.assertFalse(result)
def test_reset(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.buffer = MagicMock()
step.current_run_setting = RunSetting(computation_mode=ComputationMode.Transform)
step.finished = True
step.reset()
self.assertIsNone(None)
assert step.current_run_setting.computation_mode == ComputationMode.Default
assert step._should_stop(None, None) == False
assert step.finished == False
def test_reset(self):
step = Step(MagicMock(), MagicMock(), MagicMock())
step.buffer = MagicMock()
step.computation_mode = ComputationMode.Transform
step.stop = True
step.finished = True
step.reset()
self.assertIsNone(None)
assert step.computation_mode == ComputationMode.Default
assert step.stop == False
assert step.finished == False
def test_get_json(self):
step = Step(self.module_mock, self.step_mock, None)
json = step.get_json("file")
self.assertEqual({'batch_size': None,
'callbacks': [],
'class': 'Step',
'default_run_setting': {'computation_mode': 4},
'condition': None,
'id': -1,
'input_ids': {},
'last': True,
'module': 'pywatts.core.step',
'name': 'test',
'target_ids': {},
'train_if': None}, json)
def test_get_result(self):
# Tests if the get_result method calls correctly the previous step and the module
input_step = MagicMock()
input_step_result_mock = MagicMock()
input_step.get_result.return_value = input_step_result_mock
input_step._should_stop.return_value = False
time = pd.date_range('2000-01-01', freq='1H', periods=7)
self.module_mock.transform.return_value = xr.DataArray([2, 3, 4, 3, 3, 1, 2], dims=["time"],
coords={'time': time})
step = Step(self.module_mock, {"x": input_step}, file_manager=MagicMock())
step.get_result(pd.Timestamp("2000.01.01"), pd.Timestamp("2020.12.12"))
# 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 '))])
self.module_mock.transform.assert_called_once_with(x=input_step_result_mock)
def create_step(self,
module: Base,
kwargs: Dict[str, Union[StepInformation, Tuple[StepInformation, ...]]],
use_inverse_transform: bool, use_predict_proba: bool,
callbacks: List[Union[BaseCallback, Callable[[Dict[str, xr.DataArray]], None]]],
condition,
batch_size,
computation_mode,
train_if):
"""
Creates a appropriate step for the current situation.
:param module: The module which should be added to the pipeline
:param kwargs: The input steps for the current step
:param targets: The target steps for the currrent step
:param use_inverse_transform: Should inverse_transform be called instead of transform
:param use_predict_proba: Should probabilistic_transform be called instead of transform
:param callbacks: Callbacks to use after results are processed.
:param condition: A function returning True or False which indicates if the step should be performed
: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.
:return: StepInformation
"""
arguments = inspect.signature(module.transform).parameters.keys()
if "kwargs" not in arguments and not isinstance(module, Pipeline):
for argument in arguments:
if argument not in kwargs.keys():
raise StepCreationException(
f"The module {module.name} miss {argument} as input. The module needs {arguments} as input. "
f"{kwargs} are given as input."
f"Add {argument}=<desired_input> when adding {module.name} to the pipeline.",
module
)
# TODO needs to check that inputs are unambigious -> I.e. check that each input has only one output
pipeline = self._check_ins(kwargs)
input_steps, target_steps = self._split_input_target_steps(kwargs, pipeline)
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)
elif use_inverse_transform:
step = InverseStep(module, input_steps, pipeline.file_manager, targets=target_steps,
callbacks=callbacks, computation_mode=computation_mode, condition=condition)
elif use_predict_proba:
step = ProbablisticStep(module, input_steps, pipeline.file_manager, targets=target_steps,
callbacks=callbacks, computation_mode=computation_mode, condition=condition)
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)
step_id = pipeline.add(module=step,
input_ids=[step.id for step in input_steps.values()],
target_ids=[step.id for step in target_steps.values()])
step.id = step_id
if len(target_steps) > 1:
step.last = False
for target in target_steps:
r_step = step.get_result_step(target)
r_id = pipeline.add(module=step, input_ids=[step_id])
r_step.id = r_id
return StepInformation(step, pipeline)
def test_fit(self):
input_dict = {'input_data': None}
step = Step(self.module_mock, self.step_mock, file_manager=MagicMock())
step._fit(input_dict, {})
self.module_mock.fit.assert_called_once_with(**input_dict)
def test_transform(self):
input_dict = {'input_data': None}
step = Step(self.module_mock, {"x": self.step_mock}, None)
step._fit(input_dict, {})
step._transform(input_dict)
self.module_mock.transform.assert_called_once_with(**input_dict)