420 flag for winsorisation

src/flag_for_winsorisation.py

@@ -0,0 +1,33 @@
+def winsorisation_flag(df, a_weight, g_weight):
+
+    """
+    Function to create a column to flag whether or not a row should have
+    winsorisation applied.
+    Function requires a_weight and g_weight columns produced
+    by the estimation method.
+
+    Parameters
+    ----------
+    df : pd.DataFrame
+        DataFrame containing a weights and g weights supplied through
+        estimation.
+
+    a_weight: float
+        Column containing the a weights.
+    g_weight: float
+        Column containing the g weights.
+
+    Returns
+    -------
+    pd.DataFrame
+        Dataframe with an additional column (nw_ag_flag) that indicates if
+        winsorisation should be applied.
+    """
+
+    df["flag_calculation"] = df[a_weight] * df[g_weight]
+
+    df["nw_ag_flag"] = df["flag_calculation"].apply(lambda x: 1 if x <= 1 else 0)
+
+    df = df.drop("flag_calculation", axis=1)
+
+    return df

tests/data/winsorisation/flag_data.csv

@@ -0,0 +1,16 @@
+a_weight,g_weight,nw_ag_flag
+1.666666667,1.023809524,0
+1.666666667,1.023809524,0
+1.666666667,1.023809524,0
+1.666666667,1.023809524,0
+1.666666667,1.023809524,0
+2.5,1.023809524,0
+2.5,1.023809524,0
+2.5,1.023809524,0
+2.5,1.023809524,0
+2.5,1.023809524,0
+0.32,0.004,1
+0.32,0.004,1
+,0.004,0
+,0.004,0
+0.5,2.0,1

tests/data/winsorisation/predicted_unit_data.csv

@@ -0,0 +1,15 @@
+strata,period,aux,sampled,a_weight,g_weight,target_variable,nw_ag_flag
+101,202401,10,0,1.666666667,1.023809524,12,0
+101,202401,23,1,1.666666667,1.023809524,20,0
+101,202401,41,1,1.666666667,1.023809524,20,0
+101,202401,53,1,1.666666667,1.023809524,40,0
+101,202401,12,0,1.666666667,1.023809524,10,0
+102,202401,50,1,2.5,1.023809524,60,0
+102,202401,40,1,2.5,1.023809524,50,0
+102,202401,45,0,2.5,1.023809524,50,0
+102,202401,70,0,2.5,1.023809524,60,0
+102,202401,86,0,2.5,1.023809524,90,0
+103,202401,20,0,0.32,0.004,90,1
+103,202401,30,0,0.32,0.004,90,1
+104,202401,20,0,,0.004,90,0
+104,202401,30,0,,0.004,90,0

tests/test_flag_for_winsorisation.py

@@ -0,0 +1,31 @@
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+from src.flag_for_winsorisation import winsorisation_flag
+
+
+@pytest.fixture(scope="class")
+def winsorisation_flag_test_data():
+    return pd.read_csv(
+        "tests/data/winsorisation/flag_data.csv",
+        low_memory=False,
+        usecols=lambda c: not c.startswith("Unnamed:"),
+    )
+
+
+class TestWinsorisationFlag:
+    def test_winsorisation_flag(self, winsorisation_flag_test_data):
+        df_expected_output = winsorisation_flag_test_data.copy()
+        df_input = df_expected_output.drop(columns=["nw_ag_flag"])
+        df_input = df_input[
+            [
+                "a_weight",
+                "g_weight",
+            ]
+        ]
+        df_output = winsorisation_flag(
+            df=df_input, a_weight="a_weight", g_weight="g_weight"
+        )
+
+        assert_frame_equal(df_output, df_expected_output)