333 pre calculated links (#41)

* Add functionality to use pre-calculated links * Fix if statements and add tests * Minor doc changes * Improve docstring * Pre-commit hook auto changes * Change order of if statement in test --------- Co-authored-by: Jordan Day <[email protected]>
ONSdigital · Jul 4, 2024 · 7139506 · 7139506
1 parent c8ef2a3
commit 7139506
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 16 deletions.
diff --git a/mbs_results/ratio_of_means.py b/mbs_results/ratio_of_means.py
@@ -132,6 +132,11 @@ def wrap_shift_by_strata_period(
         ),
     )
 
+    if df.columns.isin(
+        ["f_link_question", "b_link_question", "construction_link"]
+    ).all():
+        link_arguments = link_arguments[2:]
+
     for args in link_arguments:
         df = shift_by_strata_period(df, **args)
 
@@ -278,6 +283,7 @@ def ratio_of_means(
     strata: str,
     auxiliary: str,
     filters: pd.DataFrame = None,
+    imputation_links: Dict[str, str] = {},
     **kwargs
 ) -> pd.DataFrame:
     """
@@ -306,6 +312,9 @@ def ratio_of_means(
         Column name containing auxiliary information (sic).
     filters : pd.DataFrame, optional
         Dataframe with values to exclude from imputation method.
+    imputation_links : dict, optional
+        Dictionary of column name keys matching to their imputation link value
+        ("f_link_question", "b_link_question", "construction_link").
     kwargs : mapping, optional
         A dictionary of keyword arguments passed into func.
 
@@ -333,13 +342,25 @@ def ratio_of_means(
 
         df = flag_rows_to_ignore(df, filters)
 
-    # TODO: Pre calculated links
+    if all(
+        links in imputation_links.values()
+        for links in ["f_link_question", "b_link_question", "construction_link"]
+    ):
+
+        df = df.rename(columns=imputation_links).pipe(
+            wrap_shift_by_strata_period, **default_columns
+        )
+
+    else:
+
+        df = (
+            df.pipe(wrap_flag_matched_pairs, **default_columns)
+            .pipe(wrap_shift_by_strata_period, **default_columns)
+            .pipe(wrap_calculate_imputation_link, **default_columns)
+        )
 
     df = (
-        df.pipe(wrap_flag_matched_pairs, **default_columns)
-        .pipe(wrap_shift_by_strata_period, **default_columns)
-        .pipe(wrap_calculate_imputation_link, **default_columns)
-        .pipe(
+        df.pipe(
             create_impute_flags,
             **default_columns,
             predictive_auxiliary="f_predictive_auxiliary"

diff --git a/tests/test_ratio_of_means.py b/tests/test_ratio_of_means.py
@@ -26,7 +26,7 @@
     "16_BI_BI_R_NS_C_FI_FI",  # bug fixed ASAP-402
     "17_NS_R_FI_NS",
     "18_NS_BI_R_NS",
-    # "19_link_columns",  #not yet implemented ASAP-333
+    "19_link_columns",
     "20_mixed_data",
     "21_class_change_R_C_FI",
     "22_class_change_C_BI_R",
@@ -35,7 +35,7 @@
     "25_class_change_C_FI_FI",
     "26_C_FI_FI_NS_BI_BI_R_filtered",  # not yet implemented
     "27_BI_BI_R_NS_R_FI_FI_filtered",  # not yet implemented
-    # '28_link_columns_filtered', #not yet implemented + ASAP-333
+    "28_link_columns_filtered",
     "29_mixed_data_filtered",  # not yet implemented
     "30_class_change_C_C_FI_filtered",  # not yet implemented
     "31_no_response",  # bug fixed ASAP-402
@@ -73,15 +73,31 @@ def test_ratio_of_means(self, base_file_name):
             columns=["default_forward", "default_backward", "default_construction"]
         )
 
-        actual_output = ratio_of_means(
-            input_data,
-            target="question",
-            period="date",
-            reference="identifier",
-            strata="group",
-            auxiliary="other",
-            filters=filter_df,
-        )
+        if base_file_name in ["19_link_columns", "28_link_columns_filtered"]:
+            actual_output = ratio_of_means(
+                input_data,
+                target="question",
+                period="date",
+                reference="identifier",
+                strata="group",
+                auxiliary="other",
+                filters=filter_df,
+                imputation_links={
+                    "forward": "f_link_question",
+                    "backward": "b_link_question",
+                    "construction": "construction_link",
+                },
+            )
+        else:
+            actual_output = ratio_of_means(
+                input_data,
+                target="question",
+                period="date",
+                reference="identifier",
+                strata="group",
+                auxiliary="other",
+                filters=filter_df,
+            )
 
         # imputed_value is in a seperate column, remove this if otherwise
         actual_output["question"] = actual_output[["question", "imputed_value"]].agg(