From ec12d261c5371584c3dfeda34a3ccb77fb8039db Mon Sep 17 00:00:00 2001
From: George Lopatenko <81328772+Lopa10ko@users.noreply.github.com>
Date: Thu, 13 Jun 2024 14:17:51 +0300
Subject: [PATCH] fix: go for forecast method instead of predict in ets (#1299)

* fix: use forecast method instead of predict

* fix: remove custom logging layer 45, defaulting to logging.INFO

* test: add unit for a smoothing+gaussian_filter+ets pipeline
---
 .../assumptions/assumptions_handler.py        |  4 +---
 .../models/ts_implementations/statsmodels.py  | 15 +++++-------
 fedot/utilities/memory.py                     |  3 +--
 test/unit/pipelines/test_pipeline.py          | 24 +++++++++++++++++++
 4 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/fedot/api/api_utils/assumptions/assumptions_handler.py b/fedot/api/api_utils/assumptions/assumptions_handler.py
index 0fe0f28f39..4a16027952 100644
--- a/fedot/api/api_utils/assumptions/assumptions_handler.py
+++ b/fedot/api/api_utils/assumptions/assumptions_handler.py
@@ -78,9 +78,7 @@ def fit_assumption_and_check_correctness(self,
             pipeline.predict(data_test)
             self.log.info('Initial pipeline was fitted successfully')
 
-            MemoryAnalytics.log(self.log,
-                                additional_info='fitting of the initial pipeline',
-                                logging_level=45)  # message logging level
+            MemoryAnalytics.log(self.log, additional_info='fitting of the initial pipeline')
 
         except Exception as ex:
             self._raise_evaluating_exception(ex)
diff --git a/fedot/core/operations/evaluation/operation_implementations/models/ts_implementations/statsmodels.py b/fedot/core/operations/evaluation/operation_implementations/models/ts_implementations/statsmodels.py
index be57649372..b2ebced32c 100644
--- a/fedot/core/operations/evaluation/operation_implementations/models/ts_implementations/statsmodels.py
+++ b/fedot/core/operations/evaluation/operation_implementations/models/ts_implementations/statsmodels.py
@@ -288,17 +288,14 @@ def fit(self, input_data):
 
     def predict(self, input_data):
         input_data = copy(input_data)
-        idx = input_data.idx
+        forecast_length = input_data.task.task_params.forecast_length
 
-        start_id = idx[0]
-        end_id = idx[-1]
-        predictions = self.model.predict(start=start_id,
-                                         end=end_id)
-        predict = predictions
-        predict = np.array(predict).reshape(1, -1)
-        new_idx = np.arange(start_id, end_id + 1)
+        start_id = input_data.idx[0]
+        end_id = start_id + forecast_length - 1
+        predictions = self.model.forecast(steps=forecast_length)
+        predict = np.array(predictions).reshape(1, -1)
 
-        input_data.idx = new_idx
+        input_data.idx = np.arange(start_id, end_id)
 
         output_data = self._convert_to_output(input_data,
                                               predict=predict,
diff --git a/fedot/utilities/memory.py b/fedot/utilities/memory.py
index 0cea1b894c..b25eb9d757 100644
--- a/fedot/utilities/memory.py
+++ b/fedot/utilities/memory.py
@@ -22,8 +22,7 @@ def finish(cls):
         """
         Finish memory monitoring session
         """
-        cls.log(additional_info='finish',
-                logging_level=45)  # message logging level
+        cls.log(additional_info='finish')
         tracemalloc.stop()
         cls.is_active = False
 
diff --git a/test/unit/pipelines/test_pipeline.py b/test/unit/pipelines/test_pipeline.py
index 21eede0588..49d56c4bb3 100644
--- a/test/unit/pipelines/test_pipeline.py
+++ b/test/unit/pipelines/test_pipeline.py
@@ -427,6 +427,30 @@ def test_ts_forecasting_pipeline_with_poly_features():
     assert prediction is not None
 
 
+def test_ts_forecasting_pipeline_with_ets():
+    """ Test for an ets-comprising pipeline, ensuring predict does not contain NaNs """
+    smoothing_node = PipelineNode('smoothing')
+    smoothing_node.parameters = {'window_size': 9}
+    gaussian_filter_node = PipelineNode('gaussian_filter', nodes_from=[smoothing_node])
+    smoothing_node.parameters = {'sigma': 4.747037682823521}
+    ets_node = PipelineNode('ets', nodes_from=[gaussian_filter_node])
+    ets_node.parameters = {
+        'error': 'mul',
+        'trend': 'mul',
+        'seasonal': 'add',
+        'damped_trend': False,
+        'seasonal_periods': 37
+    }
+    pipeline = Pipeline(ets_node)
+
+    train_data, test_data = get_ts_data(n_steps=450, forecast_length=30)
+
+    pipeline.fit(train_data)
+    prediction = pipeline.predict(test_data)
+    assert prediction is not None
+    assert not np.isnan(prediction.predict).any()
+
+
 def test_get_nodes_with_operation():
     pipeline = pipeline_first()
     actual_nodes = pipeline.get_nodes_by_name(name='rf')