diff --git a/doctor_visits/delphi_doctor_visits/geo_maps.py b/doctor_visits/delphi_doctor_visits/geo_maps.py index ed75d699f..b7fc0fbb1 100644 --- a/doctor_visits/delphi_doctor_visits/geo_maps.py +++ b/doctor_visits/delphi_doctor_visits/geo_maps.py @@ -101,11 +101,14 @@ def county_to_megacounty(self, data, threshold_visits, threshold_len): Returns: tuple of dataframe at the daily-state resolution, and geo_id column name """ - data = self.gmpr.fips_to_megacounty(data, + all_data = self.gmpr.fips_to_megacounty(data, threshold_visits, threshold_len, fips_col="PatCountyFIPS", thr_col="Denominator", date_col="ServiceDate") - data.rename({"megafips": "PatCountyFIPS"}, axis=1, inplace=True) + all_data.rename({"megafips": "PatCountyFIPS"}, axis=1, inplace=True) + megacounties = all_data[all_data.PatCountyFIPS.str.endswith("000")] + data = pd.concat([data, megacounties]) + return data.groupby("PatCountyFIPS"), "PatCountyFIPS" diff --git a/doctor_visits/tests/test_geomap.py b/doctor_visits/tests/test_geomap.py index 3a6bdaac7..1ee054016 100644 --- a/doctor_visits/tests/test_geomap.py +++ b/doctor_visits/tests/test_geomap.py @@ -46,8 +46,13 @@ def test_county_to_megacounty(self): assert name == "PatCountyFIPS" assert set(out.groups.keys()) == { "01001", + "01003", "01005", + "01007", "01009", + "01011", + "01013", "01015", + "01017", "01000" } \ No newline at end of file