Pass params to forecast cli

docs/source/commands.rst

@@ -7,17 +7,18 @@ Basic ``forecast`` usage:
 .. code-block:: console
 
         Usage: etna forecast [OPTIONS] CONFIG_PATH TARGET_PATH FREQ OUTPUT_PATH [EXOG_PATH]
-                    [RAW_OUTPUT]
+                             [FORECAST_CONFIG_PATH] [RAW_OUTPUT]
 
         Command to make forecast with etna without coding.
 
         Arguments:
-            CONFIG_PATH   path to yaml config with desired pipeline  [required]
-            TARGET_PATH   path to csv with data to forecast  [required]
-            FREQ          frequency of timestamp in files in pandas format  [required]
-            OUTPUT_PATH   where to save forecast  [required]
-            [EXOG_PATH]   path to csv with exog data
-            [RAW_OUTPUT]  by default we return only forecast without features  [default: False]
+            CONFIG_PATH             path to yaml config with desired pipeline  [required]
+            TARGET_PATH             path to csv with data to forecast  [required]
+            FREQ                    frequency of timestamp in files in pandas format  [required]
+            OUTPUT_PATH             where to save forecast  [required]
+            [EXOG_PATH]             path to csv with exog data
+            [FORECAST_CONFIG_PATH]  path to yaml config with forecast params
+            [RAW_OUTPUT]            by default we return only forecast without features [default: False]
 
 **How to create config?**
 
@@ -34,6 +35,14 @@ Example of pipeline's config:
         in_column: target
       - _target_: etna.transforms.SegmentEncoderTransform
 
+Example of forecast params config:
+
+.. code-block:: yaml
+
+    prediction_interval: true
+    quantiles: [0.025, 0.975]
+    n_folds: 3
+
 **How to prepare data?**
 
 Example of dataset with data to forecast:

docs/source/conf.py

@@ -22,8 +22,8 @@
 SOURCE_PATH = Path(os.path.dirname(__file__))  # noqa # docs source
 PROJECT_PATH = SOURCE_PATH.joinpath("../..")  # noqa # project root
 
-COMMIT_SHORT_SHA = os.environ["CI_COMMIT_SHORT_SHA"]
-WORKFLOW_NAME = os.environ["WORKFLOW_NAME"]
+COMMIT_SHORT_SHA = os.environ.get("CI_COMMIT_SHORT_SHA", None)
+WORKFLOW_NAME = os.environ.get("WORKFLOW_NAME", None)
 
 sys.path.insert(0, str(PROJECT_PATH))  # noqa
 

etna/commands/forecast_command.py

@@ -1,4 +1,6 @@
 from pathlib import Path
+from typing import Any
+from typing import Dict
 from typing import Optional
 
 import hydra_slayer
@@ -16,6 +18,7 @@ def forecast(
     freq: str = typer.Argument(..., help="frequency of timestamp in files in pandas format"),
     output_path: Path = typer.Argument(..., help="where to save forecast"),
     exog_path: Optional[Path] = typer.Argument(None, help="path to csv with exog data"),
+    forecast_config_path: Optional[Path] = typer.Argument(None, help="path to yaml config with forecast params"),
     raw_output: bool = typer.Argument(False, help="by default we return only forecast without features"),
 ):
     """Command to make forecast with etna without coding.
@@ -51,6 +54,11 @@ def forecast(
     =============  ===========  ===============  ===============
     """
     pipeline_configs = OmegaConf.to_object(OmegaConf.load(config_path))
+    if forecast_config_path:
+        forecast_params_config = OmegaConf.to_object(OmegaConf.load(forecast_config_path))
+    else:
+        forecast_params_config = {}
+    forecast_params: Dict[str, Any] = hydra_slayer.get_from_params(**forecast_params_config)
 
     df_timeseries = pd.read_csv(target_path, parse_dates=["timestamp"])
 
@@ -65,12 +73,14 @@ def forecast(
 
     pipeline: Pipeline = hydra_slayer.get_from_params(**pipeline_configs)
     pipeline.fit(tsdataset)
-    forecast = pipeline.forecast()
+    forecast = pipeline.forecast(**forecast_params)
 
+    flatten = forecast.to_pandas(flatten=True)
     if raw_output:
-        (forecast.to_pandas(True).to_csv(output_path, index=False))
+        (flatten.to_csv(output_path, index=False))
     else:
-        (forecast.to_pandas(True)[["timestamp", "segment", "target"]].to_csv(output_path, index=False))
+        quantile_columns = [column for column in flatten.columns if column.startswith("target_0.")]
+        (flatten[["timestamp", "segment", "target"] + quantile_columns].to_csv(output_path, index=False))
 
 
 if __name__ == "__main__":

tests/test_commands/conftest.py

@@ -76,3 +76,18 @@ def base_timeseries_exog_path():
     tmp.flush()
     yield Path(tmp.name)
     tmp.close()
+
+
+@pytest.fixture
+def base_forecast_omegaconf_path():
+    tmp = NamedTemporaryFile("w")
+    tmp.write(
+        """
+        prediction_interval: true
+        quantiles: [0.025, 0.975]
+        n_folds: 3
+        """
+    )
+    tmp.flush()
+    yield Path(tmp.name)
+    tmp.close()

tests/test_commands/test_forecast.py

@@ -47,3 +47,25 @@ def test_dummy_run(base_pipeline_yaml_path, base_timeseries_path):
     run(["etna", "forecast", str(base_pipeline_yaml_path), str(base_timeseries_path), "D", str(tmp_output_path)])
     df_output = pd.read_csv(tmp_output_path)
     assert len(df_output) == 2 * 4
+
+
+def test_run_with_predictive_intervals(
+    base_pipeline_yaml_path, base_timeseries_path, base_timeseries_exog_path, base_forecast_omegaconf_path
+):
+    tmp_output = NamedTemporaryFile("w")
+    tmp_output_path = Path(tmp_output.name)
+    run(
+        [
+            "etna",
+            "forecast",
+            str(base_pipeline_yaml_path),
+            str(base_timeseries_path),
+            "D",
+            str(tmp_output_path),
+            str(base_timeseries_exog_path),
+            str(base_forecast_omegaconf_path),
+        ]
+    )
+    df_output = pd.read_csv(tmp_output_path)
+    for q in [0.025, 0.975]:
+        assert f"target_{q}" in df_output.columns