MeanEncoderTransform generates wrong values

CHANGELOG.md

@@ -41,7 +41,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 
 - 
 - 
-- 
+- Fix working with NaN target in `MeanEncoderTransform` ([#492](https://github.com/etna-team/etna/pull/492))
 - 
 - 
 - 

etna/transforms/encoders/mean_encoder.py

@@ -214,13 +214,13 @@ def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
                     # first timestamp is NaN
                     expanding_mean = y.expanding().mean().shift()
                     # cumcount not including current timestamp
-                    cumcount = y.groupby(segment_df[self.in_column].astype(str)).agg("cumcount")
+                    cumcount = segment_df.loc[y.notna()].groupby(self.in_column, dropna=False).cumcount().reindex(y.index).replace(0, np.NaN)
                     # cumsum not including current timestamp
-                    cumsum = (
-                        y.groupby(segment_df[self.in_column].astype(str))
-                        .transform(lambda x: x.shift().cumsum())
-                        .fillna(0)
+                    cumsum = segment_df['target'].groupby(segment_df[self.in_column].astype(str), dropna=False).transform(
+                        lambda x: x.shift().fillna(0).cumsum()
                     )
+                    cumsum = cumsum.where(cumcount.notna(), np.NaN)
+
                     feature = (cumsum + expanding_mean * self.smoothing) / (cumcount + self.smoothing)
                     if self.handle_missing is MissingMode.global_mean:
                         nan_feature_index = segment_df[segment_df[self.in_column].isnull()].index
@@ -237,7 +237,7 @@ def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
                 timestamps = intersected_df.index
                 categories = pd.unique(df.loc[:, self.idx[:, self.in_column]].values.ravel())
 
-                cumstats = pd.DataFrame(data={"sum": 0, "count": 0, self.in_column: categories})
+                cumstats = pd.DataFrame(data={"sum": np.NaN, "count": np.NaN, self.in_column: categories})
                 cur_timestamp_idx = np.arange(0, len(timestamps) * n_segments, len(timestamps))
                 for _ in range(len(timestamps)):
                     timestamp_df = flatten.loc[cur_timestamp_idx]
@@ -254,8 +254,11 @@ def _transform(self, df: pd.DataFrame) -> pd.DataFrame:
                         .agg(["count", "sum"])
                         .reset_index()
                     )
+                    stats = stats.replace({"count": 0, "sum": 0}, np.NaN)
+
                     # sum current and previous statistics
                     cumstats = pd.concat([cumstats, stats]).groupby(self.in_column, as_index=False, dropna=False).sum()
+                    cumstats = cumstats.replace({"count": 0, "sum": 0}, np.NaN)
                     cur_timestamp_idx += 1
 
                 feature = (temp["cumsum"] + running_mean * self.smoothing) / (temp["cumcount"] + self.smoothing)

tests/test_transforms/test_encoders/test_mean_encoder_transform.py

@@ -27,11 +27,36 @@ def category_ts() -> TSDataset:
     return ts
 
 
+@pytest.fixture
+def mean_segment_encoder_ts() -> TSDataset:
+    df = generate_ar_df(n_segments=1, start_time="2001-01-01", periods=5)
+    df["target"] = [0, 1, np.NaN, 3, 4]
+
+    df_exog = generate_ar_df(n_segments=1, start_time="2001-01-01", periods=10)
+    df_exog.rename(columns={"target": "segment_feature"}, inplace=True)
+    df_exog["segment_feature"] = "segment_0"
+
+    ts = TSDataset(df=df, df_exog=df_exog, freq="D", known_future="all")
+
+    return ts
+
+
+@pytest.fixture
+def expected_mean_segment_encoder_ts() -> TSDataset:
+    df = generate_ar_df(n_segments=1, start_time="2001-01-01", periods=5)
+    df.rename(columns={"target": "segment_mean"}, inplace=True)
+    df["segment_mean"] = [np.NaN, 0, 0.5, 0.5, 1.33]
+
+    ts = TSDataset(df=df, freq="D")
+
+    return ts
+
+
 @pytest.fixture
 def expected_micro_category_ts() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=2)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, 1, 1.5, 1.5, 2.75, 2.25] + [np.NaN, 6.0, 6.25, 7, 7.625, 8.0]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, np.NaN, 1.5, 2.75, 2.25] + [np.NaN, np.NaN, 6.25, 7, 7.625, np.NaN]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -41,7 +66,7 @@ def expected_micro_category_ts() -> TSDataset:
 def expected_micro_global_mean_ts() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=2)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, 1, 1.5, 1.5, 2.5, 2.25] + [np.NaN, 6.0, 6.25, 7, 7.625, 8.0]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, 1.5, 1.5, 2.5, 2.25] + [np.NaN, np.NaN, 6.25, 7, 7.625, 8.0]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -61,7 +86,7 @@ def expected_micro_category_make_future_ts() -> TSDataset:
 def expected_macro_category_ts() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=2)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, 3.5, 4, 4.875, 4, 4.85] + [np.NaN, 3.5, 3.66, 4.875, 5.5, 4.275]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, np.NaN, 4.875, 4, 4.851] + [np.NaN, np.NaN, 3.66, 4.875, 5.5, 4.27]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -71,7 +96,7 @@ def expected_macro_category_ts() -> TSDataset:
 def expected_macro_global_mean_ts() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=2)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, 3.5, 4, 4.875, 5, 4.85] + [np.NaN, 3.5, 3.66, 4.875, 5.5, 5.55]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, 4, 4.875, 5, 4.85] + [np.NaN, np.NaN, 3.66, 4.875, 5.5, 5.55]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -104,7 +129,7 @@ def ts_begin_nan() -> TSDataset:
 def expected_ts_begin_nan_smooth_1() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=1)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, np.NaN, 0.5, 1.16, 1.5, 2.5]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, np.NaN, 1.75, 1.5, 2.5]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -114,7 +139,7 @@ def expected_ts_begin_nan_smooth_1() -> TSDataset:
 def expected_ts_begin_nan_smooth_2() -> TSDataset:
     df = generate_ar_df(start_time="2001-01-01", periods=6, n_segments=1)
     df.rename(columns={"target": "mean_encoded_regressor"}, inplace=True)
-    df["mean_encoded_regressor"] = [np.NaN, np.NaN, 2 / 3, 5 / 4, 5 / 3, 2.5]
+    df["mean_encoded_regressor"] = [np.NaN, np.NaN, np.NaN, 5 / 3, 5 / 3, 2.5]
 
     ts = TSDataset(df, freq="D")
     return ts
@@ -311,6 +336,22 @@ def test_ts_begin_nan_smooth_2(ts_begin_nan, expected_ts_begin_nan_smooth_2):
     )
 
 
+def test_mean_segment_encoder(mean_segment_encoder_ts, expected_mean_segment_encoder_ts):
+    mean_encoder = MeanEncoderTransform(
+        in_column="segment_feature",
+        mode="per-segment",
+        handle_missing="category",
+        smoothing=0,
+        out_column="segment_mean",
+    )
+    mean_encoder.fit_transform(mean_segment_encoder_ts)
+    assert_frame_equal(
+        mean_segment_encoder_ts.df.loc[:, pd.IndexSlice[:, "segment_mean"]],
+        expected_mean_segment_encoder_ts.df,
+        atol=0.01,
+    )
+
+
 def test_save_load(category_ts):
     mean_encoder = MeanEncoderTransform(in_column="regressor", out_column="mean_encoded_regressor")
     assert_transformation_equals_loaded_original(transform=mean_encoder, ts=category_ts)