Teach OutliersTransform to ignore holidays

etna/transforms/outliers/base.py

@@ -17,17 +17,23 @@
 class OutliersTransform(ReversibleTransform, ABC):
     """Finds outliers in specific columns of DataFrame and replaces it with NaNs."""
 
-    def __init__(self, in_column: str):
+    def __init__(self, in_column: str, ignore_flag_column: Optional[str] = None):
         """
         Create instance of OutliersTransform.
 
         Parameters
         ----------
         in_column:
             name of processed column
+        ignore_flag_column:
+            name of column binary flag of holidays
         """
-        super().__init__(required_features=[in_column])
+        if ignore_flag_column:
+            super().__init__(required_features=[in_column, ignore_flag_column])
+        else:
+            super().__init__(required_features=[in_column])
         self.in_column = in_column
+        self.ignore_flag_column = ignore_flag_column
 
         self.segment_outliers: Optional[Dict[str, pd.Series]] = None
 
@@ -78,6 +84,12 @@
         :
             The fitted transform instance.
         """
+        if self.ignore_flag_column is not None and self.ignore_flag_column not in ts.regressors:
+            raise ValueError("Name ignore_flag_column not find.")
+        if self.ignore_flag_column is not None and not all(
+            ts[:, segment, self.ignore_flag_column].isin([0, 1]).all() for segment in ts.segments
+        ):
+            raise ValueError("Columns ignore_flag contain non binary value")
         self.segment_outliers = self.detect_outliers(ts)
         self._fit_segments = ts.segments
         super().fit(ts=ts)
@@ -131,8 +143,16 @@
             if segment not in segments:
                 continue
             # to locate only present indices
-            segment_outliers_timestamps = list(index_set.intersection(self.segment_outliers[segment].index.values))
+            if self.ignore_flag_column:
+                available_points = set(df[df[segment, self.ignore_flag_column] == 0].index.values)
+            else:
+                available_points = index_set
+            segment_outliers_timestamps = list(
+                available_points.intersection(self.segment_outliers[segment].index.values)
+            )
+
             df.loc[segment_outliers_timestamps, pd.IndexSlice[segment, self.in_column]] = np.NaN
+
         return df
 
     def _inverse_transform(self, df: pd.DataFrame) -> pd.DataFrame:

etna/transforms/outliers/point_outliers.py

@@ -1,6 +1,7 @@
 from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Optional
 from typing import Type
 from typing import Union
 
@@ -32,7 +33,13 @@ class MedianOutliersTransform(OutliersTransform):
     it uses information from the whole train part.
     """
 
-    def __init__(self, in_column: str, window_size: int = 10, alpha: float = 3):
+    def __init__(
+        self,
+        in_column: str,
+        window_size: int = 10,
+        alpha: float = 3,
+        ignore_flag_column: Optional[str] = None,
+    ):
         """Create instance of MedianOutliersTransform.
 
         Parameters
@@ -43,10 +50,12 @@ def __init__(self, in_column: str, window_size: int = 10, alpha: float = 3):
             number of points in the window
         alpha:
             coefficient for determining the threshold
+        ignore_flag_column:
+            name of column binary flag of holidays
         """
         self.window_size = window_size
         self.alpha = alpha
-        super().__init__(in_column=in_column)
+        super().__init__(in_column=in_column, ignore_flag_column=ignore_flag_column)
 
     def detect_outliers(self, ts: TSDataset) -> Dict[str, List[pd.Timestamp]]:
         """Call :py:func:`~etna.analysis.outliers.median_outliers.get_anomalies_median` function with self parameters.
@@ -97,6 +106,7 @@ def __init__(
         distance_coef: float = 3,
         n_neighbors: int = 3,
         distance_func: Union[Literal["absolute_difference"], Callable[[float, float], float]] = "absolute_difference",
+        ignore_flag_column: Optional[str] = None,
     ):
         """Create instance of DensityOutliersTransform.
 
@@ -118,7 +128,7 @@ def __init__(
         self.distance_coef = distance_coef
         self.n_neighbors = n_neighbors
         self.distance_func = distance_func
-        super().__init__(in_column=in_column)
+        super().__init__(in_column=in_column, ignore_flag_column=ignore_flag_column)
 
     def detect_outliers(self, ts: TSDataset) -> Dict[str, List[pd.Timestamp]]:
         """Call :py:func:`~etna.analysis.outliers.density_outliers.get_anomalies_density` function with self parameters.
@@ -169,6 +179,7 @@ def __init__(
         in_column: str,
         model: Union[Literal["prophet"], Literal["sarimax"], Type["ProphetModel"], Type["SARIMAXModel"]],
         interval_width: float = 0.95,
+        ignore_flag_column: Optional[str] = None,
         **model_kwargs,
     ):
         """Create instance of PredictionIntervalOutliersTransform.
@@ -181,7 +192,8 @@ def __init__(
             model for prediction interval estimation
         interval_width:
             width of the prediction interval
-
+        ignore_flag_column:
+            name of column binary flag of holidays
         Notes
         -----
         For not "target" column only column data will be used for learning.
@@ -190,7 +202,7 @@ def __init__(
         self.interval_width = interval_width
         self.model_kwargs = model_kwargs
         self._model_type = self._get_model_type(model)
-        super().__init__(in_column=in_column)
+        super().__init__(in_column=in_column, ignore_flag_column=ignore_flag_column)
 
     @staticmethod
     def _get_model_type(

tests/test_transforms/test_outliers/test_outliers_transform.py

@@ -9,6 +9,7 @@
 from etna.models import ProphetModel
 from etna.models import SARIMAXModel
 from etna.transforms import DensityOutliersTransform
+from etna.transforms import HolidayTransform
 from etna.transforms import MedianOutliersTransform
 from etna.transforms import PredictionIntervalOutliersTransform
 from tests.test_transforms.utils import assert_sampling_is_valid
@@ -42,6 +43,34 @@ def outliers_solid_tsds():
     return ts
 
 
+@pytest.fixture()
+def outliers_solid_tsds_with_holidays():
+    """Create TSDataset with outliers with holidays"""
+    timestamp = pd.date_range("2021-01-01", end="2021-02-20", freq="D")
+    target1 = [np.sin(i) for i in range(len(timestamp))]
+    target1[10] += 10
+
+    target2 = [np.sin(i) for i in range(len(timestamp))]
+    target2[8] += 8
+    target2[15] = 2
+    target2[26] -= 12
+
+    df1 = pd.DataFrame({"timestamp": timestamp, "target": target1, "segment": "1"})
+    df2 = pd.DataFrame({"timestamp": timestamp, "target": target2, "segment": "2"})
+    df = pd.concat([df1, df2], ignore_index=True)
+    df_exog = df.copy()
+    df_exog.columns = ["timestamp", "regressor_1", "segment"]
+    ts = TSDataset(
+        df=TSDataset.to_dataset(df).iloc[:-10],
+        df_exog=TSDataset.to_dataset(df_exog),
+        freq="D",
+        known_future="all",
+    )
+    holiday_transform = HolidayTransform(iso_code="RUS", mode="binary", out_column="is_holiday")
+    ts = holiday_transform.fit_transform(ts)
+    return ts
+
+
 @pytest.mark.parametrize("attribute_name,value_type", (("outliers_timestamps", list), ("original_values", pd.Series)))
 def test_density_outliers_deprecated_store_attributes(outliers_solid_tsds, attribute_name, value_type):
     transform = DensityOutliersTransform(in_column="target")
@@ -255,3 +284,35 @@ def test_params_to_tune(transform, outliers_solid_tsds):
     ts = outliers_solid_tsds
     assert len(transform.params_to_tune()) > 0
     assert_sampling_is_valid(transform=transform, ts=ts)
+
+
+@pytest.mark.parametrize(
+    "transform",
+    (
+        MedianOutliersTransform(in_column="target", ignore_flag_column="is_holiday"),
+        DensityOutliersTransform(in_column="target", ignore_flag_column="is_holiday"),
+        # PredictionIntervalOutliersTransform(in_column="target", model="sarimax", ignore_flag_column="is_holiday"),
+    ),
+)
+def test_correct_ignore_flag(transform, outliers_solid_tsds_with_holidays):
+    ts = outliers_solid_tsds_with_holidays
+    assert len(transform.params_to_tune()) > 0
+    transform.fit(ts)
+    ts_output = transform.transform(ts)
+    assert ts_output["2021-01-06":"2021-01-06", "1", "target"][0] != np.nan
+
+
+@pytest.mark.parametrize(
+    "transform",
+    (
+        MedianOutliersTransform(in_column="target", ignore_flag_column="is_holiday"),
+        DensityOutliersTransform(in_column="target", ignore_flag_column="is_holiday"),
+        PredictionIntervalOutliersTransform(in_column="target", model="sarimax", ignore_flag_column="is_holiday"),
+    ),
+)
+def test_incorrect_formats(transform, outliers_solid_tsds):
+    ts = outliers_solid_tsds
+    assert len(transform.params_to_tune()) > 0
+    with pytest.raises(ValueError):
+        transform.fit(ts)
+        _ = transform.transform(ts)