update code and unit tests for mapping from state to hhs and nation

Author: Jingjing Tang

google_symptoms/delphi_google_symptoms/geo.py

@@ -23,24 +23,30 @@ def generate_transition_matrix(geo_res):
         The first is a data frame for HRR regions and the second are MSA
         regions.
     """
-    if geo_res != "nation":
+    if geo_res in ["hrr", "msa"]:
+        mapping_flag = "fips"
         map_df = gmpr.get_crosswalk("fips", geo_res)
+        # Add population as weights
+        map_df = gmpr.add_population_column(map_df, "fips")
     else:
-        map_df = gmpr.get_crosswalk("fips", "hhs")
-        map_df[geo_res] = "nation"
-    # Add population as weights
-    map_df = gmpr.add_population_column(map_df, "fips")
+        mapping_flag = "state_id"
+        map_df = gmpr.get_crosswalk("state", "state")
+        map_df = gmpr.add_geocode(map_df, "state_code", "hhs")
+        map_df = gmpr.add_geocode(map_df, "state_code", "nation")
+        map_df = gmpr.add_population_column(map_df, "state_code")
+
     if geo_res == "hrr":
         map_df["population"] = map_df["population"] *  map_df["weight"]
+
     aggregated_pop = map_df.groupby(geo_res).sum().reset_index()
     map_df = map_df.merge(
             aggregated_pop, on=geo_res, how="inner", suffixes=["_raw", "_groupsum"]
             )
     map_df["weight"] = map_df["population_raw"] / map_df["population_groupsum"]
 
     map_df = pd.pivot_table(
-                 map_df, values='weight', index=["fips"], columns=[geo_res]
-              ).fillna(0).reset_index().rename({"fips": "geo_id"}, axis = 1)
+                 map_df, values='weight', index=[mapping_flag], columns=[geo_res]
+              ).fillna(0).reset_index().rename({mapping_flag: "geo_id"}, axis = 1)
     return map_df
 
 def geo_map(df, geo_res):

google_symptoms/delphi_google_symptoms/run.py

@@ -87,6 +87,8 @@ def run_module(params):
     for geo_res in GEO_RESOLUTIONS:
         if geo_res == "state":
             df_pull = dfs["state"]
+        elif geo_res in ["hhs", "nation"]:
+            df_pull = geo_map(dfs["state"], geo_res)
         else:
             df_pull = geo_map(dfs["county"], geo_res)
 

google_symptoms/tests/test_geo.py

@@ -115,20 +115,21 @@ def test_hhs(self):
         gmpr = GeoMapper()
         df = pd.DataFrame(
             {
-                "geo_id": ["01001", "01009", "01007"],
+                "geo_id": ["al", "fl", "tx"],
                 "timestamp": ["2020-02-15", "2020-02-15", "2020-02-15"],
                 METRICS[0]: [10, 15, 2],
                 METRICS[1]: [100, 20, 45],
                 COMBINED_METRIC: [110, 35, 47],
             }
         )
 
-        fips2hhs = gmpr.add_population_column(gmpr.get_crosswalk("fips", "hhs"), "fips")
-        hhs_pop = fips2hhs.groupby("hhs"
+        state2hhs = gmpr.add_population_column(gmpr.get_crosswalk("state", "state"), "state_code")
+        state2hhs = gmpr.add_geocode(state2hhs, "state_code", "hhs")
+        hhs_pop = state2hhs.groupby("hhs"
             ).sum(
             ).reset_index(
             ).rename(columns={"population": "hhs_pop"})
-        df_plus = df.merge(fips2hhs, left_on="geo_id", right_on="fips", how="left"
+        df_plus = df.merge(state2hhs, left_on="geo_id", right_on="state_id", how="left"
             ).merge(hhs_pop, on="hhs", how="left"
             ).assign(
                 fractional_pop = lambda x: x.population / x.hhs_pop,
@@ -147,7 +148,7 @@ def test_hhs(self):
         new_df = geo_map(df, "hhs").dropna()
 
         assert set(new_df.keys()) == set(df.keys())
-        assert set(new_df["geo_id"]) == set(["4"])
+        assert set(new_df["geo_id"]) == set(["4", "6"])
         assert new_df[METRICS[0]].values == pytest.approx(df_plus[METRICS[0]].tolist())
         assert new_df[METRICS[1]].values == pytest.approx(df_plus[METRICS[1]].tolist())
         assert new_df[COMBINED_METRIC].values == pytest.approx(df_plus[COMBINED_METRIC].tolist())
@@ -156,22 +157,21 @@ def test_nation(self):
         gmpr = GeoMapper()
         df = pd.DataFrame(
             {
-                "geo_id": ["01001", "01009", "01007"],
+                "geo_id": ["al", "il", "tx"],
                 "timestamp": ["2020-02-15", "2020-02-15", "2020-02-15"],
                 METRICS[0]: [10, 15, 2],
                 METRICS[1]: [100, 20, 45],
                 COMBINED_METRIC: [110, 35, 47],
             }
         )
 
-        fips2nation = gmpr.add_population_column(gmpr.get_crosswalk("fips", "hhs"), "fips")
-        fips2nation.rename({"hhs": "nation"}, axis=1, inplace=True)
-        fips2nation["nation"] = "nation"
-        nation_pop = fips2nation.groupby("nation"
+        state2nation = gmpr.add_population_column(gmpr.get_crosswalk("state", "state"), "state_code")
+        state2nation = gmpr.add_geocode(state2nation, "state_code", "nation")
+        nation_pop = state2nation.groupby("nation"
             ).sum(
             ).reset_index(
             ).rename(columns={"population": "nation_pop"})
-        df_plus = df.merge(fips2nation, left_on="geo_id", right_on="fips", how="left"
+        df_plus = df.merge(state2nation, left_on="geo_id", right_on="state_id", how="left"
             ).merge(nation_pop, on="nation", how="left"
             ).assign(
                 fractional_pop = lambda x: x.population / x.nation_pop,
@@ -190,7 +190,7 @@ def test_nation(self):
         new_df = geo_map(df, "nation").dropna()
 
         assert set(new_df.keys()) == set(df.keys())
-        assert set(new_df["geo_id"]) == set(["nation"])
+        assert set(new_df["geo_id"]) == set(["us"])
         assert new_df[METRICS[0]].values == pytest.approx(df_plus[METRICS[0]].tolist())
         assert new_df[METRICS[1]].values == pytest.approx(df_plus[METRICS[1]].tolist())
         assert new_df[COMBINED_METRIC].values == pytest.approx(df_plus[COMBINED_METRIC].tolist())