Merge pull request #27 from Quantmetry/angoho_unit_tests_utils

Angoho unit tests utils

Author: JulienRoussel77

qolmat/utils/data.py

@@ -99,11 +99,9 @@ def preprocess_data(df: pd.DataFrame):
     df.set_index(["station", "datetime"], inplace=True)
     df.drop(columns=["year", "month", "day", "hour", "wd", "No"], inplace=True)
     df.sort_index(inplace=True)
-    dict_agg = {key: np.mean for key in df.columns}
-    dict_agg["RAIN"] = np.mean
     df = df.groupby(
         ["station", df.index.get_level_values("datetime").floor("d")], group_keys=False
-    ).agg(dict_agg)
+    ).mean()
     return df
 
 
@@ -142,6 +140,7 @@ def add_holes(df: pd.DataFrame, ratio_masked: float, mean_size: int):
         mask = generator.generate_mask(df)
     else:
         mask = df.groupby(groups, group_keys=False).apply(generator.generate_mask)
+
     X_with_nans = df.copy()
     X_with_nans[mask] = np.nan
     return X_with_nans
@@ -187,6 +186,7 @@ def add_station_features(df: pd.DataFrame):
     pd.DataFrame
         dataframe with missing values
     """
+    df = df.copy()
     stations = df.index.get_level_values("station")
     for station in stations.unique():
         df[f"station={station}"] = (stations == station).astype(float)
@@ -207,10 +207,11 @@ def add_datetime_features(df: pd.DataFrame):
     pd.DataFrame
         dataframe with missing values
     """
-
-    time = df.index.get_level.values("datetime")
+    df = df.copy()
+    time = df.index.get_level_values("datetime").to_series()
     days_in_year = time.dt.year.apply(
         lambda x: 366 if ((x % 4 == 0) and (x % 100 != 0)) or (x % 400 == 0) else 365
     )
-    df["time_cos"] = np.cos(2 * np.pi * time.dt.dayofyear / days_in_year)
+    time_cos = np.cos(2 * np.pi * time.dt.dayofyear / days_in_year)
+    df["time_cos"] = np.array(time_cos)
     return df

tests/utils/test_data.py

@@ -0,0 +1,121 @@
+import datetime
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from qolmat.utils import data
+
+columns = ["a", "b"]
+index = pd.MultiIndex.from_tuples(
+    [
+        ("Gucheng", datetime.datetime(2013, 3, 1)),
+        ("Gucheng", datetime.datetime(2014, 3, 1)),
+        ("Gucheng", datetime.datetime(2015, 3, 1)),
+    ],
+    names=["station", "datetime"],
+)
+df = pd.DataFrame(
+    [
+        [1, 2],
+        [3, np.nan],
+        [np.nan, 6],
+    ],
+    columns=columns,
+    index=index,
+)
+
+
+def test_preprocess_data():
+    columns_raw = [
+        "No",
+        "year",
+        "month",
+        "day",
+        "hour",
+        "a",
+        "b",
+        "wd",
+        "station",
+    ]
+    df_raw = pd.DataFrame(
+        [
+            [
+                1,
+                2013,
+                3,
+                1,
+                0,
+                1,
+                2,
+                "NW",
+                "Gucheng",
+            ],
+            [
+                2,
+                2014,
+                3,
+                1,
+                0,
+                3,
+                np.nan,
+                "NW",
+                "Gucheng",
+            ],
+            [
+                3,
+                2015,
+                3,
+                1,
+                0,
+                np.nan,
+                6,
+                "NW",
+                "Gucheng",
+            ],
+        ],
+        columns=columns_raw,
+    )
+    print(df_raw)
+    result = data.preprocess_data(df_raw)
+    print(result)
+    print(df)
+    # assert result.equals(df)
+    pd.testing.assert_frame_equal(result, df, atol=1e-3)
+
+
+def test_add_holes() -> None:
+    df_out = data.add_holes(df, 0.0, 1)
+    assert df_out.isna().sum().sum() == 2
+    df_out = data.add_holes(df, 1.0, 1)
+    assert df_out.isna().sum().sum() > 2
+
+
+def test_add_station_features() -> None:
+    columns_out = columns + ["station=Gucheng"]
+    expected = pd.DataFrame(
+        [
+            [1, 2, 1.0],
+            [3, np.nan, 1.0],
+            [np.nan, 6, 1.0],
+        ],
+        columns=columns_out,
+        index=index,
+    )
+    result = data.add_station_features(df)
+    pd.testing.assert_frame_equal(result, expected, atol=1e-3)
+
+
+def test_add_datetime_features() -> None:
+    columns_out = columns + ["time_cos"]
+    expected = pd.DataFrame(
+        [
+            [1, 2, 0.512],
+            [3, np.nan, 0.512],
+            [np.nan, 6, 0.512],
+        ],
+        columns=columns_out,
+        index=index,
+    )
+    result = data.add_datetime_features(df)
+    pd.testing.assert_frame_equal(result, expected, atol=1e-3)