Geocode updates:

* update replace_geocode documentation to be clear about data columns
* add test cases for renaming columns in replace_geocode
* fix the state to state conversion dropped columns issue

Author: dshemetov

_delphi_utils_python/delphi_utils/geomap.py

@@ -77,6 +77,11 @@ class GeoMapper:
     ==========
     The main GeoMapper object loads and stores crosswalk dataframes on-demand.
 
+    When replacing geocodes with a new one an aggregation step is performed on the data columns
+    to merge entries  (i.e. in the case of a many to one mapping or a weighted mapping). This
+    requires a specification of the data columns, which are assumed to be all the columns that
+    are not the geocodes or the date column specified in date_col.
+
     Example 1: to add a new column with a new geocode, possibly with weights:
     > gmpr = GeoMapper()
     > df = gmpr.add_geocode(df, "fips", "zip", from_col="fips", new_col="geo_id",
@@ -327,8 +332,12 @@ def add_geocode(
             df = df.merge(crosswalk, left_on=from_col, right_on=from_col, how="left")
 
         # Drop extra state columns
-        if new_code in state_codes:
+        if new_code in state_codes and not from_code in state_codes:
+            state_codes.remove(new_code)
+            df.drop(columns=state_codes, inplace=True)
+        elif new_code in state_codes and from_code in state_codes:
             state_codes.remove(new_code)
+            state_codes.remove(from_code)
             df.drop(columns=state_codes, inplace=True)
 
         return df

_delphi_utils_python/tests/test_geomap.py

@@ -278,11 +278,13 @@ 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")
-        assert new_data["population"].sum() == 274963
+        assert new_data.shape == (5, 5)
         new_data = gmpr.add_population_column(self.zip_data, "zip")
-        assert new_data["population"].sum() == 274902
+        assert new_data.shape == (6, 5)
         with pytest.raises(ValueError):
             new_data = gmpr.add_population_column(self.zip_data, "hrr")
+        new_data = gmpr.add_population_column(self.fips_data_5, "fips")
+        assert new_data.shape == (4, 5)
 
     def test_add_geocode(self):
         gmpr = GeoMapper()
@@ -383,17 +385,19 @@ def test_add_geocode(self):
         assert new_data2["hhs_region_number"].unique().size == 2
 
         # state_name -> state_id
-        new_data = gmpr.add_geocode(self.zip_data, "zip", "state_name")
+        new_data = gmpr.replace_geocode(self.zip_data, "zip", "state_name")
         new_data2 = gmpr.add_geocode(new_data, "state_name", "state_id")
-        assert new_data2.shape == (12, 6)
+        assert new_data2.shape == (4, 5)
+        new_data2 = gmpr.replace_geocode(new_data, "state_name", "state_id", new_col="abbr")
+        assert "abbr" in new_data2.columns
 
         # fips -> nation
-        new_data = gmpr.replace_geocode(self.fips_data_5, "fips", "nation")
+        new_data = gmpr.replace_geocode(self.fips_data_5, "fips", "nation", new_col="NATION")
         assert new_data.equals(
             pd.DataFrame().from_dict(
                 {
                     "date": {0: pd.Timestamp("2018-01-01 00:00:00")},
-                    "nation": {0: "us"},
+                    "NATION": {0: "us"},
                     "count": {0: 10024.0},
                     "total": {0: 100006.0},
                 }
@@ -416,6 +420,23 @@ def test_add_geocode(self):
             )
         )
 
+        # hrr -> nation
+        with pytest.raises(ValueError):    
+            new_data = gmpr.replace_geocode(self.zip_data, "zip", "hrr")
+            new_data2 = gmpr.replace_geocode(new_data, "hrr", "nation")
+
+        # hrr -> nation
+        with pytest.raises(ValueError):    
+            new_data = gmpr.replace_geocode(self.zip_data, "zip", "hrr")
+            new_data2 = gmpr.replace_geocode(new_data, "hrr", "nation")
+
+        # hrr -> nation
+        with pytest.raises(ValueError):    
+            new_data = gmpr.replace_geocode(self.zip_data, "zip", "hrr")
+            new_data2 = gmpr.replace_geocode(new_data, "hrr", "nation")
+
         # fips -> hrr (dropna=True/False check)
         assert not gmpr.add_geocode(self.fips_data_3, "fips", "hrr").isna().any().any()
         assert gmpr.add_geocode(self.fips_data_3, "fips", "hrr", dropna=False).isna().any().any()
+
+TestGeoMapper().test_add_geocode()