Fix JHU's fips-population source file

_delphi_utils_python/delphi_utils/geomap.py

@@ -392,16 +392,17 @@ def replace_geocode(
         df = df.groupby([date_col, new_col]).sum().reset_index()
         return df
 
-    def add_population_column(self, data, geocode_type, geocode_col=None):
+    def add_population_column(self, geocode_type, data=None, geocode_col=None):
         """
-        Appends a population column to a dateframe, based on the FIPS or ZIP code.
+        Appends a population column to a dataframe, based on the FIPS or ZIP code. If no
+        dataframe is provided, the full crosswalk from geocode to population is returned.
 
         Parameters
         ---------
-        data: pd.DataFrame
-            The dataframe with a FIPS code column.
         geocode_type: {"fips", "zip"}
             The type of the geocode contained in geocode_col.
+        data: pd.DataFrame
+            The dataframe with a FIPS code column.
         geocode_col: str, default None
             The name of the column containing the geocodes. If None, uses the geocode_type
             as the name.
@@ -419,11 +420,14 @@ def add_population_column(self, data, geocode_type, geocode_col=None):
                 For other codes, aggregate those."
             )
 
+        pop_df = self._load_crosswalk(from_code=geocode_type, to_code="pop")
+
+        if data is None:
+            return pop_df.rename(columns={"pop": "population"})
+
         if not is_string_dtype(data[geocode_col]):
             data[geocode_col] = data[geocode_col].astype(str).str.zfill(5)
 
-        pop_df = self._load_crosswalk(from_code=geocode_type, to_code="pop")
-
         data_with_pop = (
             data.copy()
             .merge(pop_df, left_on=geocode_col, right_on=geocode_type, how="inner")

_delphi_utils_python/tests/test_geomap.py

@@ -277,12 +277,14 @@ def test_zip_to_state_id(self):
 
     def test_add_population_column(self):
         gmpr = GeoMapper()
-        new_data = gmpr.add_population_column(self.fips_data_3, "fips")
+        new_data = gmpr.add_population_column("fips", self.fips_data_3)
         assert new_data["population"].sum() == 274963
-        new_data = gmpr.add_population_column(self.zip_data, "zip")
+        new_data = gmpr.add_population_column("zip", self.zip_data)
         assert new_data["population"].sum() == 274902
         with pytest.raises(ValueError):
-            new_data = gmpr.add_population_column(self.zip_data, "hrr")
+            new_data = gmpr.add_population_column("hrr", self.zip_data)
+        pop_df = gmpr.add_population_column("fips")
+        assert pop_df.shape == (3274, 2)
 
     def test_add_geocode(self):
         gmpr = GeoMapper()

jhu/delphi_jhu/pull.py

@@ -69,12 +69,7 @@ def pull_jhu_data(base_url: str, metric: str, pop_df: pd.DataFrame) -> pd.DataFr
 
     gmpr = GeoMapper()
     df = gmpr.replace_geocode(df, "jhu_uid", "fips", from_col="UID", date_col="timestamp")
-
-    # Merge in population LOWERCASE, consistent across confirmed and deaths
-    # Set population as NAN for fake fips
-    pop_df.rename(columns={'FIPS':'fips'}, inplace=True)
-    pop_df['fips'] = pop_df['fips'].astype(int).\
-        astype(str).str.zfill(5)
+    # Merge in population, set population as NAN for fake fips
     df = pd.merge(df, pop_df, on="fips", how='left')
 
     # Add a dummy first row here on day before first day

jhu/delphi_jhu/run.py

@@ -17,6 +17,7 @@
     S3ArchiveDiffer,
 )
 
+from delphi_utils import GeoMapper
 from .geo import geo_map
 from .pull import pull_jhu_data
 from .smooth import (
@@ -84,10 +85,7 @@ def run_module():
     else:
         arch_diff = None
 
-    pop_df = pd.read_csv(
-        join(static_file_dir, "fips_population.csv"),
-        dtype={"fips": float, "population": float},
-    )
+    pop_df = GeoMapper().add_population_column("fips")
 
     dfs = {metric: pull_jhu_data(base_url, metric, pop_df) for metric in METRICS}
     for metric, geo_res, sensor, smoother in product(