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)