allow option for all counties, save data files with filter suffix

Author: mauriciogtec

data-processing/Snakefile

@@ -10,7 +10,7 @@ conda: "requirements.yaml"  # Path: envs/conda.yaml
 
 
 envvars:
-    "CENSUS_API_KEY"
+    "CENSUS_API_KEY",
 
 
 # read hydra config from initialize API
@@ -21,75 +21,78 @@ with hydra.initialize(config_path="conf", version_base=None):
 
 # make sure data locations exist / user can set them to
 # symlinks separately if preferred
-os.makedirs(f"{processing_cfg.data_dir}/raw", exist_ok=True)
-os.makedirs(f"{processing_cfg.data_dir}/processed", exist_ok=True)
+data_dir = processing_cfg.data_dir
+os.makedirs(f"{data_dir}/raw", exist_ok=True)
+os.makedirs(f"{data_dir}/processed", exist_ok=True)
+
+filters = ["65000", "all"]
 
 
 rule all:
     input:
-        processing_cfg.data_dir + "/processed/endogenous_states_actions.parquet",
-        processing_cfg.data_dir + "/processed/exogenous_states.parquet",
-        processing_cfg.data_dir + "/processed/bspline_basis.parquet",
-        processing_cfg.data_dir + "/processed/confounders.parquet",
+        expand(
+            data_dir + "/processed/endogenous_states_actions_{filter}.parquet",
+            filter=filters,
+        ),
+        expand(
+            data_dir + "/processed/exogenous_states_{filter}.parquet",
+            filter=filters,
+        ),
+        expand(
+            data_dir + "/processed/bspline_basis_{filter}.parquet",
+            filter=filters,
+        ),
+        expand(
+            data_dir + "/processed/confounders_{filter}.parquet",
+            filter=filters,
+        ),
+
 
 rule merge_state_actions:
     input:
         expand(
-            processing_cfg.data_dir + "/processed/alerts/{state}.parquet",
+            data_dir + "/processed/alerts/{state}.parquet",
             state=config["states"],
         ),
-        processing_cfg.data_dir + "/processed/heatmetrics.parquet",
+        data_dir + "/processed/heatmetrics_{filter}.parquet",
     output:
-        processing_cfg.data_dir + "/processed/exogenous_states.parquet",
-        processing_cfg.data_dir + "/processed/endogenous_states_actions.parquet",
-        processing_cfg.data_dir + "/processed/budget.parquet",
+        data_dir + "/processed/exogenous_states_{filter}.parquet",
+        data_dir + "/processed/endogenous_states_actions_{filter}.parquet",
+        data_dir + "/processed/budget_{filter}.parquet",
+        data_dir + "/processed/bspline_basis_{filter}.parquet",
     log:
-        "logs/merge_state_actions.log",
+        "logs/merge_state_actions_{filter}.log",
     shell:
-        "python merge_state_actions.py &> {log}"
-
-
-# rule merge_hospitalizations:
-#     input:
-#         processing_cfg.data_dir + "/processed/exogenous_states.parquet",
-#         processing_cfg.data_dir + "/processed/endogenous_states_actions.parquet",
-#     output:
-#         processing_cfg.data_dir + "/processed/training_data.parquet",
-#     log:
-#         "logs/merge_hospitalizations.log",
-#     shell:
-#         f"""
-#         python merge_hospitalizations.py \
-#             hospitalizations.data_path={config['hosps_data_path']} 
-#             &> {{log}}
-#         """
+        "python merge_state_actions.py county_filter={wildcards.filter} &> {log}"
 
 
 rule confounders:
     output:
-        processing_cfg.data_dir + "/processed/confounders.parquet",
+        data_dir + "/processed/confounders_{filter}.parquet",
     log:
-        "logs/confounders.log",
+        "logs/confounders_{filter}.log",
     shell:
         f"""
         python confounders.py \
-            census_api_key={os.environ['CENSUS_API_KEY']} \
+            census_api_key={os.environ['CENSUS_API_KEY']} county_filter={{wildcards.filter}} \
             &> {{log}}
         """
 
 
 rule heatmetrics:
+    input:
+        data_dir + "/processed/confounders_{filter}.parquet",
     output:
-        processing_cfg.data_dir + "/processed/heatmetrics.parquet",
+        data_dir + "/processed/heatmetrics_{filter}.parquet",
     log:
-        "logs/heatmetrics.log",
+        "logs/heatmetrics_{filter}.log",
     shell:
-        "python heatmetrics.py &> {log}"
+        "python heatmetrics.py county_filter={wildcards.filter}  &> {log}"
 
 
 rule alerts:
     output:
-        processing_cfg.data_dir + "/processed/alerts/{state}.parquet",
+        data_dir + "/processed/alerts/{state}.parquet",
     log:
         "logs/alerts_{state}.log",
     shell:

data-processing/conf/config.yaml

@@ -4,6 +4,8 @@ census_api_key: ${oc.env:CENSUS_API_KEY}
 # processed_dir: ../data/processed
 data_dir: ../data
 
+county_filter: 65000  # choose 'all' for no filter
+
 heatmetrics:
   min_month: 5
   max_month: 9

data-processing/confounders.py

@@ -208,9 +208,14 @@ def main(cfg):
 
     # Keep only those places with population > 65000 and complete data cases
     print(merged_df.shape)
-    merged_df = merged_df.loc[merged_df.total_pop > 65000]
+
+    if cfg.county_filter != "all":
+        merged_df = merged_df.loc[merged_df.total_pop > float(cfg.county_filter)]
+
     merged_df = merged_df.dropna()
-    merged_df.to_parquet(f"{cfg.data_dir}/processed/confounders.parquet", index=False)
+    merged_df.to_parquet(
+        f"{cfg.data_dir}/processed/confounders_{cfg.county_filter}.parquet", index=False
+    )
 
 
 if __name__ == "__main__":

data-processing/heatmetrics.py

@@ -24,7 +24,6 @@ def transform_rds_to_parquet(rds_path, parquet_path):
 @hydra.main(config_path="conf", config_name="config", version_base=None)
 def main(cfg):
     # Download data if not already present
-    processed_path = f"{cfg.data_dir}/processed/heatmetrics.parquet"
     download_path = f"{cfg.data_dir}/raw/heatmetrics.rds"  # data is in R's native format
     url = cfg.heatmetrics.url
 
@@ -49,11 +48,12 @@ def main(cfg):
     df = df[["StCoFIPS", "Date"] + cfg.heatmetrics.cols]
     df.rename(columns={"StCoFIPS": "fips", "Date": "date"}, inplace=True)
 
-    # Keep only large fips
-    confounders = pd.read_parquet(f"{cfg.data_dir}/processed/confounders.parquet")
+    suffix = cfg.county_filter
+    confounders = pd.read_parquet(f"{cfg.data_dir}/processed/confounders_{suffix}.parquet")
     df = df[df.fips.isin(confounders.fips)]
 
     # Write to parquet
+    processed_path = f"{cfg.data_dir}/processed/heatmetrics_{suffix}.parquet"
     df.to_parquet(processed_path)
     LOGGER.info(f"Data written to {processed_path} with head\n: {df.head()}")
 

data-processing/merge_state_actions.py

@@ -18,8 +18,11 @@ def main(cfg):
     """This script merges data from the heatmetrics and heatalerts and computes
     the state space variables. The merged data is saved as a parquet file."""
 
+    # county filter suffix
+    suffix = cfg.county_filter
+
     # Load heatmetrics
-    hm = pd.read_parquet(f"{cfg.data_dir}/processed/heatmetrics.parquet")
+    hm = pd.read_parquet(f"{cfg.data_dir}/processed/heatmetrics_{suffix}.parquet")
     hm = hm.sort_values(["fips", "date"])
 
     # Load and post process heat alerts data
@@ -156,7 +159,7 @@ def main(cfg):
 
     # compute the rolling of alerts in the the entire summer
     df["rolling_alerts"] = df.groupby(["fips", "year"])["alert"].transform("cumsum")
-    df["remaiing_budget"] = df["budget"] - df["rolling_alerts"]
+    df["remaining_budget"] = df["budget"] - df["rolling_alerts"]
 
     # dos splines
     M = max(df.dos)
@@ -181,7 +184,7 @@ def main(cfg):
     # standardize and save
     bspline_basis = (bspline_basis - bspline_col_means) / bspline_col_stds
     bspline_basis.columns = [f"bspline_dos_{i}" for i in range(bspline_basis.shape[1])]
-    bspline_basis.to_parquet(f"{cfg.data_dir}/processed/bspline_basis.parquet")
+    bspline_basis.to_parquet(f"{cfg.data_dir}/processed/bspline_basis_{suffix}.parquet")
 
     # -------------------
     # save exogenous states, endogenous states, actions
@@ -200,7 +203,7 @@ def main(cfg):
         *[f"bspline_dos_{i}" for i in range(bspline_dos.shape[1])],
     ]
     exogenous_states = df[exogenous_state_vars + ["fips", "date"]]
-    exogenous_states.to_parquet(f"{cfg.data_dir}/processed/exogenous_states.parquet")
+    exogenous_states.to_parquet(f"{cfg.data_dir}/processed/exogenous_states_{suffix}.parquet")
 
     # actions and endogenous states
     action_vars = [
@@ -212,11 +215,11 @@ def main(cfg):
     ]
     action_states = df[action_vars + ["fips", "date"]]
     action_states.to_parquet(
-        f"{cfg.data_dir}/processed/endogenous_states_actions.parquet"
+        f"{cfg.data_dir}/processed/endogenous_states_actions_{suffix}.parquet"
     )
 
     # save budget
-    budget.to_parquet(f"{cfg.data_dir}/processed/budget.parquet")
+    budget.to_parquet(f"{cfg.data_dir}/processed/budget_{suffix}.parquet")
 
 
 if __name__ == "__main__":

data/.gitignore

@@ -1,5 +1,5 @@
 hospitalizations/*
 raw/nws-forecast-zones
 raw/heatmetrics.parquet
-shapefile
-!processed
+raw/shapefile
+processed/