From cfcdb04f3dc2992c672df2a0d6be4b9ec6eabea9 Mon Sep 17 00:00:00 2001 From: woodac Date: Wed, 4 Dec 2024 16:18:01 +0000 Subject: [PATCH] Refactoring pnp_pre_processing.py --- src/site_apportionment/pnp_pre_processing.py | 13 +++++++------ src/site_apportionment/site_apportionment_main.py | 5 +++-- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/site_apportionment/pnp_pre_processing.py b/src/site_apportionment/pnp_pre_processing.py index 833e6b742..95d770803 100644 --- a/src/site_apportionment/pnp_pre_processing.py +++ b/src/site_apportionment/pnp_pre_processing.py @@ -17,20 +17,21 @@ def pnp_pre_processing(df: pd.DataFrame) -> pd.DataFrame: # Add a column for imputation marker df = imputation_marker(df) # Add a column for a weight - df = pnp_weight(df) - + df = pnp_columns(df) + return df # Add a column for a weight -def pnp_weight(df: pd.DataFrame) -> pd.DataFrame: +def pnp_columns(df: pd.DataFrame) -> pd.DataFrame: - """Add a class column to the PNP data. + """Adding columns to run PNP data. Args: df (pd.DataFrame): The input DataFrame. Returns: - pd.DataFrame: The DataFrame with the added weight column. + pd.DataFrame: The DataFrame with the added weight and pg_numeric column. """ # Add a class column to the PNP data df["a_weight"] = 1.0 + df["pg_numeric"] = 10 - return df + return df \ No newline at end of file diff --git a/src/site_apportionment/site_apportionment_main.py b/src/site_apportionment/site_apportionment_main.py index 2300208f9..7df2d62aa 100644 --- a/src/site_apportionment/site_apportionment_main.py +++ b/src/site_apportionment/site_apportionment_main.py @@ -38,8 +38,9 @@ def run_site_apportionment( to apportion for long forms """ if config ["survey"]["survey_type"] == "PNP": - pnp_pre_processing(df) - + df = pnp_pre_processing(df) + + # Create variable for output of QA apportionment file qa_path = config["apportionment_paths"]["qa_path"]