Fix test with pipeline tuning

aimclub · Sep 25, 2023 · 980d038 · 980d038
1 parent f089654
commit 980d038
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/test/unit/validation/test_time_series_cv.py b/test/unit/validation/test_time_series_cv.py
@@ -1,6 +1,8 @@
 import datetime
 import logging
 
+from hyperopt import hp
+
 from golem.core.log import default_log
 from golem.core.optimisers.genetic.gp_params import GPAlgorithmParameters
 from golem.core.tuning.simultaneous import SimultaneousTuner
@@ -15,6 +17,7 @@
 from fedot.core.repository.quality_metrics_repository import \
     MetricsRepository, RegressionMetricsEnum
 from fedot.core.repository.tasks import TsForecastingParams
+from fedot.core.pipelines.tuning.search_space import PipelineSearchSpace
 from fedot.core.data.cv_folds import cv_generator
 from test.unit.tasks.test_forecasting import get_simple_ts_pipeline, get_ts_data
 
@@ -69,12 +72,18 @@ def test_tuner_cv_correct():
 
     simple_pipeline = get_simple_ts_pipeline(window_size=2)
 
+    max_window = int(time_series.features.shape[0] / (folds + 1)) - (forecast_len * validation_blocks) - 1
+    ppl_ss = PipelineSearchSpace({'lagged': {'window_size': {'hyperopt-dist': hp.uniformint,
+                                                             'sampling-scope': [2, max_window],
+                                                             'type': 'discrete'}}})
+
     tuner = TunerBuilder(time_series.task)\
         .with_tuner(SimultaneousTuner)\
         .with_metric(RegressionMetricsEnum.MAE)\
         .with_cv_folds(folds) \
         .with_iterations(1) \
         .with_timeout(datetime.timedelta(minutes=1))\
+        .with_search_space(ppl_ss)\
         .build(time_series)
     _ = tuner.tune(simple_pipeline)
     is_tune_succeeded = True