Merge pull request #1407 from cmu-delphi/krivard/fix-hhs-prop

Refactor pop prop calculations to fix #1399

Author: krivard

hhs_hosp/delphi_hhs/run.py

@@ -109,19 +109,8 @@ def run_module(params):
                     geo_res = geo,
                     sensor = sensor,
                     smoother = smoother)
-        df = geo_mapper.add_geocode(make_signal(all_columns, sensor),
-                                    "state_id",
-                                    "state_code",
-                                    from_col="state")
-        if sensor.endswith("_prop"):
-            df=pop_proportion(df, geo_mapper)
-        df = make_geo(df, geo, geo_mapper)
-        df = smooth_values(df, smoother[0])
-        # Fix N/A MA values, see issue #1360
-        if geo == "state" and sensor.startswith(CONFIRMED_FLU):
-            ma_filter = df.val.isna() & (df.geo_id == "ma") & (df.timestamp > "08-01-2021") & \
-                (df.timestamp.dt.day_name() == "Tuesday")
-            df = df[~ma_filter]
+        df = make_signal(all_columns, sensor)
+        df = transform_signal(sensor, smoother, geo, df, geo_mapper)
         if df.empty:
             continue
         sensor_name = sensor + smoother[1]
@@ -155,17 +144,32 @@ def smooth_values(df, smoother):
     )
     return df
 
-def pop_proportion(df,geo_mapper):
-    """Get the population-proportionate variants as the dataframe val."""
-    pop_val=geo_mapper.add_population_column(df, "state_code")
-    df["val"]=round(df["val"]/pop_val["population"]*100000, 7)
-    pop_val.drop("population", axis=1, inplace=True)
+def transform_signal(sensor, smoother, geo, df, geo_mapper):
+    """Transform base df into specified geo/smoothing/prop configuration."""
+    df = geo_mapper.add_geocode(df, "state_id", "state_code", from_col="state")
+    # handling population:
+    #   add population column
+    #   sum admission counts *and* population counts during make_geo
+    #   *then* divide counts by population to get the proportion
+    if sensor.endswith("_prop"):
+        df=geo_mapper.add_population_column(df, "state_code")
+    df = make_geo(df, geo, geo_mapper)
+    if sensor.endswith("_prop"):
+        df["val"]=round(df["val"]/df["population"]*100000, 7)
+        df.drop("population", axis=1, inplace=True)
+    df = smooth_values(df, smoother[0])
+    # Fix N/A MA values, see issue #1360
+    if geo == "state" and sensor.startswith(CONFIRMED_FLU):
+        ma_filter = df.val.isna() & (df.geo_id == "ma") & (df.timestamp > "08-01-2021") & \
+            (df.timestamp.dt.day_name() == "Tuesday")
+        df = df[~ma_filter]
     return df
 
 def make_geo(state, geo, geo_mapper):
     """Transform incoming geo (state) to another geo."""
     if geo == "state":
         exported = state.rename(columns={"state": "geo_id"})
+        exported = exported.drop(columns="state_code")
     else:
         exported = geo_mapper.replace_geocode(state, "state_code", geo, new_col="geo_id")
     exported["se"] = np.nan

hhs_hosp/tests/test_run.py

@@ -5,7 +5,7 @@
 import os
 
 from delphi_hhs.run import _date_to_int, int_date_to_previous_day_datetime, generate_date_ranges, \
-    make_signal, make_geo, run_module, pop_proportion
+    make_signal, make_geo, run_module, transform_signal
 from delphi_hhs.constants import SMOOTHERS, GEOS, SIGNALS, \
     CONFIRMED, SUM_CONF_SUSP, CONFIRMED_FLU, CONFIRMED_PROP, SUM_CONF_SUSP_PROP, CONFIRMED_FLU_PROP
 from delphi_utils.geomap import GeoMapper
@@ -89,40 +89,51 @@ def test_make_signal():
     with pytest.raises(Exception):
         make_signal(data, "zig")
 
-def test_pop_proportion():
+def test_transform_signal_pop():
     geo_mapper = GeoMapper()
-    state_pop = geo_mapper.get_crosswalk("state_code", "pop")
+    state_pop = geo_mapper.get_crosswalk("state_id", "pop")
+    identity_smoother = SMOOTHERS[0]
+    hundo_k = 100000
 
     test_df = pd.DataFrame({  
-        'state': ['PA'],
-        'state_code': [42],
-        'timestamp': [datetime(year=2020, month=1, day=1)],
-        'val': [15.],})
+        'state': ['pa', 'wv'],
+        'timestamp': [datetime(year=2020, month=1, day=1)]*2,
+        'val': [15., 150.],})
 
-    pa_pop = int(state_pop.loc[state_pop.state_code == "42", "pop"])
+    pa_pop = int(state_pop.loc[state_pop.state_id == "pa", "pop"])
+    wv_pop = int(state_pop.loc[state_pop.state_id == "wv", "pop"])
     pd.testing.assert_frame_equal(
-        pop_proportion(test_df, geo_mapper),
+        transform_signal(
+            CONFIRMED_PROP,
+            identity_smoother,
+            'state',
+            test_df.copy(),
+            geo_mapper),
         pd.DataFrame({
-            'state': ['PA'],
-            'state_code': [42],
-            'timestamp': [datetime(year=2020, month=1, day=1)],
-            'val': [15/pa_pop*100000],})
+            'geo_id': ['pa', 'wv'],
+            'timestamp': [datetime(year=2020, month=1, day=1)]*2,
+            'val': [15/pa_pop*hundo_k, 150/wv_pop*hundo_k],
+            'se': [None]*2,
+            'sample_size': [None]*2,}),
+        check_dtype=False,
+        check_like=True
     )
 
-    test_df= pd.DataFrame({  
-        'state': ['WV'],
-        'state_code': [54],
-        'timestamp': [datetime(year=2020, month=1, day=1)],
-        'val': [150.],})
-
-    wv_pop = int(state_pop.loc[state_pop.state_code == "54", "pop"])
     pd.testing.assert_frame_equal(
-        pop_proportion(test_df, geo_mapper),
+        transform_signal(
+            CONFIRMED_PROP,
+            identity_smoother,
+            'nation',
+            test_df.copy(),
+            geo_mapper),
         pd.DataFrame({
-            'state': ['WV'],
-            'state_code': [54],
+            'geo_id': ['us'],
             'timestamp': [datetime(year=2020, month=1, day=1)],
-            'val': [150/wv_pop*100000],})
+            'val': [165/(pa_pop+wv_pop)*hundo_k],
+            'se': [None],
+            'sample_size': [None],}),
+        check_dtype=False,
+        check_like=True
     )
 
 def test_make_geo():