Merge pull request #179 from ibm-granite/pandas_warnings

Updates to mitigate pandas warnings
ibm-granite · Nov 6, 2024 · 7c50f6d · 7c50f6d
2 parents d73abaf + 5dbd53d
commit 7c50f6d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 7 deletions.
diff --git a/tests/toolkit/test_time_series_preprocessor.py b/tests/toolkit/test_time_series_preprocessor.py
@@ -103,7 +103,7 @@ def test_time_series_preprocessor_scales(ts_data):
     # check scaled result
     out = tsp.preprocess(df)
     assert np.allclose(out.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.mean(x)), 0.0)
-    assert np.allclose(out.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x)), 1.0)
+    assert np.allclose(out.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x, axis=0)), 1.0)
 
     # check inverse scale result
     out_inv = tsp.inverse_scale_targets(out)
@@ -112,8 +112,8 @@ def test_time_series_preprocessor_scales(ts_data):
         == df.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.mean(x))
     )
     assert np.all(
-        out_inv.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x))
-        == df.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x))
+        out_inv.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x, axis=0))
+        == df.groupby(tsp.id_columns)[tsp.target_columns].apply(lambda x: np.std(x, axis=0))
     )
 
     # check inverse scale result, with suffix

diff --git a/tsfm_public/toolkit/time_series_preprocessor.py b/tsfm_public/toolkit/time_series_preprocessor.py
@@ -529,9 +529,8 @@ def _set_targets(self, dataset: pd.DataFrame) -> None:
     def _estimate_frequency(self, df: pd.DataFrame):
         if self.timestamp_column:
             if self.id_columns:
-                # to do: be more efficient
                 grps = df.groupby(self.id_columns)
-                _, df_subset = list(grps)[0]
+                _, df_subset = next(iter(grps))
             else:
                 df_subset = df
 
@@ -624,7 +623,7 @@ def inverse_scale_func(grp, id_columns):
         else:
             id_columns = INTERNAL_ID_COLUMN
 
-        df_inv = df.groupby(id_columns, group_keys=False).apply(
+        df_inv = df.groupby(id_columns, group_keys=False)[df.columns].apply(
             inverse_scale_func,
             id_columns=id_columns,
         )
@@ -673,7 +672,7 @@ def scale_func(grp, id_columns):
             else:
                 id_columns = INTERNAL_ID_COLUMN
 
-            df_out = df.groupby(id_columns, group_keys=False).apply(
+            df_out = df.groupby(id_columns, group_keys=False)[df.columns].apply(
                 scale_func,
                 id_columns=id_columns,
             )