Change ifelse block to a dict of functions

Author: chinandrew

doctor_visits/delphi_doctor_visits/geo_maps.py

@@ -7,6 +7,7 @@
 Created: 2020-04-18
 Last modified: 2020-04-30 by Aaron Rumack (add megacounty code)
 """
+from functools import partial
 
 import pandas as pd
 from delphi_utils.geomap import GeoMapper
@@ -20,6 +21,14 @@ class GeoMaps:
     def __init__(self):
         """Create the underlying GeoMapper."""
         self.gmpr = GeoMapper()
+        self.geo_func = {"county": partial(self.county_to_megacounty,
+                                           threshold_visits=Config.MIN_RECENT_VISITS,
+                                           threshold_len=Config.RECENT_LENGTH),
+                         "state": self.county_to_state,
+                         "msa": self.county_to_msa,
+                         "hrr": self.county_to_hrr,
+                         "hhs": self.county_to_hhs,
+                         "nation": self.county_to_nation}
 
     @staticmethod
     def convert_fips(x):

doctor_visits/delphi_doctor_visits/update_sensor.py

@@ -132,26 +132,8 @@ def update_sensor(
 
     # get right geography
     geo_map = GeoMaps()
-    if geo.lower() == "county":
-        data_groups, _ = geo_map.county_to_megacounty(
-            data, Config.MIN_RECENT_VISITS, Config.RECENT_LENGTH
-        )
-    elif geo.lower() == "state":
-        data_groups, _ = geo_map.county_to_state(data)
-    elif geo.lower() == "msa":
-        data_groups, _ = geo_map.county_to_msa(data)
-    elif geo.lower() == "hrr":
-        data_groups, _ = geo_map.county_to_hrr(data)
-    elif geo.lower() == "hhs":
-        data_groups, _ = geo_map.county_to_hhs(data)
-    elif geo.lower() == "nation":
-        data_groups, _ = geo_map.county_to_nation(data)
-    else:
-
-        logging.error(
-            f"{geo} is invalid, pick one of 'county', 'state', 'msa', 'hrr', 'hhs', 'nation'"
-        )
-        return {}
+    mapping_func = geo_map.geo_func[geo.lower()]
+    data_groups, _ = mapping_func(data)
     unique_geo_ids = list(data_groups.groups.keys())
 
     # run sensor fitting code (maybe in parallel)