Fix cap/floor bug in ProphetModel

CHANGELOG.md

@@ -42,7 +42,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 
 - 
 - 
-- 
+- ProphetModel doesn't work with cap and floor regressors ([#842](https://github.com/tinkoff-ai/etna/pull/842))
 - 
 - Change Docker cuda image version from 11.1 to 11.6.2 ([#838](https://github.com/tinkoff-ai/etna/pull/838))
 - 

etna/models/prophet.py

@@ -19,6 +19,8 @@
 class _ProphetAdapter(BaseAdapter):
     """Class for holding Prophet model."""
 
+    predefined_regressors_names = ("floor", "cap")
+
     def __init__(
         self,
         growth: str = "linear",
@@ -99,7 +101,8 @@ def fit(self, df: pd.DataFrame, regressors: List[str]) -> "_ProphetAdapter":
         prophet_df["ds"] = df["timestamp"]
         prophet_df[self.regressor_columns] = df[self.regressor_columns]
         for regressor in self.regressor_columns:
-            self.model.add_regressor(regressor)
+            if regressor not in self.predefined_regressors_names:
+                self.model.add_regressor(regressor)
         self.model.fit(prophet_df)
         return self
 

tests/test_models/test_prophet.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pandas as pd
 import pytest
 from prophet import Prophet
@@ -27,13 +28,15 @@ def test_run_with_reg(new_format_df, new_format_exog):
 
     regressors = new_format_exog.copy()
     regressors.columns.set_levels(["regressor_exog"], level="feature", inplace=True)
-    regressors_cap = new_format_exog.copy()
-    regressors_cap.columns.set_levels(["regressor_cap"], level="feature", inplace=True)
-    exog = pd.concat([regressors, regressors_cap], axis=1)
+    regressors_floor = new_format_exog.copy()
+    regressors_floor.columns.set_levels(["floor"], level="feature", inplace=True)
+    regressors_cap = regressors_floor.copy() + 1
+    regressors_cap.columns.set_levels(["cap"], level="feature", inplace=True)
+    exog = pd.concat([regressors, regressors_floor, regressors_cap], axis=1)
 
     ts = TSDataset(df, "1d", df_exog=exog, known_future="all")
 
-    model = ProphetModel()
+    model = ProphetModel(growth="logistic")
     model.fit(ts)
     future_ts = ts.make_future(3)
     model.forecast(future_ts)
@@ -43,6 +46,37 @@ def test_run_with_reg(new_format_df, new_format_exog):
         assert False
 
 
+def test_run_with_cap_floor():
+    cap = 101
+    floor = -1
+
+    df = pd.DataFrame(
+        {
+            "timestamp": pd.date_range(start="2020-01-01", periods=100),
+            "segment": "segment_0",
+            "target": list(range(100)),
+        }
+    )
+    df_exog = pd.DataFrame(
+        {
+            "timestamp": pd.date_range(start="2020-01-01", periods=120),
+            "segment": "segment_0",
+            "cap": cap,
+            "floor": floor,
+        }
+    )
+    ts = TSDataset(df=TSDataset.to_dataset(df), df_exog=TSDataset.to_dataset(df_exog), freq="D", known_future="all")
+
+    model = ProphetModel(growth="logistic")
+    pipeline = Pipeline(model=model, horizon=7)
+    pipeline.fit(ts)
+
+    ts_future = pipeline.forecast()
+    df_future = ts_future.to_pandas(flatten=True)
+
+    assert np.all(df_future["target"] < cap)
+
+
 def test_prediction_interval_run_insample(example_tsds):
     model = ProphetModel()
     model.fit(example_tsds)