Refactoring pnp_pre_processing.py

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

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"]