Merge pull request #311 from cmu-delphi/geo_refactor_cdccovidnet

Refactor cdc_covidnet to use geo utils

Author: krivard

cdc_covidnet/delphi_cdc_covidnet/__init__.py

@@ -10,7 +10,6 @@
 
 from . import run
 from . import api_config
-from . import geo_maps
 from . import update_sensor
 from . import covidnet
 

cdc_covidnet/delphi_cdc_covidnet/geo_maps.py

@@ -12,11 +12,10 @@
 import numpy as np
 import pandas as pd
 
-from delphi_utils import read_params
+from delphi_utils import read_params, GeoMapper
 import covidcast
 from .api_config import APIConfig
 from .covidnet import CovidNet
-from .geo_maps import GeoMaps
 from .constants import SIGNALS
 
 def write_to_csv(data: pd.DataFrame, out_name: str, output_path: str):
@@ -49,17 +48,18 @@ def write_to_csv(data: pd.DataFrame, out_name: str, output_path: str):
 
 
 def update_sensor(
-        state_files: List[str], mmwr_info: pd.DataFrame,
-        output_path: str, static_path: str,
-        start_date: datetime, end_date: datetime) -> pd.DataFrame:
+        state_files: List[str],
+        mmwr_info: pd.DataFrame,
+        output_path: str,
+        start_date: datetime,
+        end_date: datetime) -> pd.DataFrame:
     """
     Generate sensor values, and write to csv format.
 
     Args:
         state_files: List of JSON files representing COVID-NET hospitalization data for each state
         mmwr_info: Mappings from MMWR week to actual dates, as a pd.DataFrame
         output_path: Path to write the csvs to
-        static_path: Path for the static geographic fiels
         start_date: First sensor date (datetime.datetime)
         end_date: Last sensor date (datetime.datetime)
 
@@ -85,9 +85,15 @@ def update_sensor(
     ]
 
     # Set state id to two-letter abbreviation
-    geo_map = GeoMaps(static_path)
-    hosp_df = geo_map.state_name_to_abbr(hosp_df)
-
+    gmpr = GeoMapper()
+    hosp_df = gmpr.add_geocode(hosp_df,
+                               from_col=APIConfig.STATE_COL,
+                               from_code="state_name",
+                               new_code="state_id",
+                               dropna=False)
+    # To use the original column name, reassign original column and drop new one
+    hosp_df[APIConfig.STATE_COL] = hosp_df["state_id"].str.upper()
+    hosp_df.drop("state_id", axis=1, inplace=True)
     assert not hosp_df.duplicated(["date", "geo_id"]).any(), "Non-unique (date, geo_id) pairs"
     hosp_df.set_index(["date", "geo_id"], inplace=True)
 

cdc_covidnet/delphi_cdc_covidnet/update_sensor.py

@@ -75,10 +75,7 @@ def test_syn_update_sensor(self):
             end_date = datetime(year=2020, month=3, day=17)
 
             # Generate the csvs
-            hosp_df = update_sensor(
-                state_files, mmwr_info,
-                temp_dir, STATIC_DIR,
-                start_date, end_date)
+            hosp_df = update_sensor(state_files, mmwr_info, temp_dir, start_date, end_date)
 
             # Check dataframe returned
             assert hosp_df.index.nlevels == 2