@@ -115,20 +115,21 @@ def test_hhs(self):
115
115
gmpr = GeoMapper ()
116
116
df = pd .DataFrame (
117
117
{
118
- "geo_id" : ["01001 " , "01009 " , "01007 " ],
118
+ "geo_id" : ["al " , "fl " , "tx " ],
119
119
"timestamp" : ["2020-02-15" , "2020-02-15" , "2020-02-15" ],
120
120
METRICS [0 ]: [10 , 15 , 2 ],
121
121
METRICS [1 ]: [100 , 20 , 45 ],
122
122
COMBINED_METRIC : [110 , 35 , 47 ],
123
123
}
124
124
)
125
125
126
- fips2hhs = gmpr .add_population_column (gmpr .get_crosswalk ("fips" , "hhs" ), "fips" )
127
- hhs_pop = fips2hhs .groupby ("hhs"
126
+ state2hhs = gmpr .add_population_column (gmpr .get_crosswalk ("state" , "state" ), "state_code" )
127
+ state2hhs = gmpr .add_geocode (state2hhs , "state_code" , "hhs" )
128
+ hhs_pop = state2hhs .groupby ("hhs"
128
129
).sum (
129
130
).reset_index (
130
131
).rename (columns = {"population" : "hhs_pop" })
131
- df_plus = df .merge (fips2hhs , left_on = "geo_id" , right_on = "fips " , how = "left"
132
+ df_plus = df .merge (state2hhs , left_on = "geo_id" , right_on = "state_id " , how = "left"
132
133
).merge (hhs_pop , on = "hhs" , how = "left"
133
134
).assign (
134
135
fractional_pop = lambda x : x .population / x .hhs_pop ,
@@ -147,7 +148,7 @@ def test_hhs(self):
147
148
new_df = geo_map (df , "hhs" ).dropna ()
148
149
149
150
assert set (new_df .keys ()) == set (df .keys ())
150
- assert set (new_df ["geo_id" ]) == set (["4" ])
151
+ assert set (new_df ["geo_id" ]) == set (["4" , "6" ])
151
152
assert new_df [METRICS [0 ]].values == pytest .approx (df_plus [METRICS [0 ]].tolist ())
152
153
assert new_df [METRICS [1 ]].values == pytest .approx (df_plus [METRICS [1 ]].tolist ())
153
154
assert new_df [COMBINED_METRIC ].values == pytest .approx (df_plus [COMBINED_METRIC ].tolist ())
@@ -156,22 +157,21 @@ def test_nation(self):
156
157
gmpr = GeoMapper ()
157
158
df = pd .DataFrame (
158
159
{
159
- "geo_id" : ["01001 " , "01009 " , "01007 " ],
160
+ "geo_id" : ["al " , "il " , "tx " ],
160
161
"timestamp" : ["2020-02-15" , "2020-02-15" , "2020-02-15" ],
161
162
METRICS [0 ]: [10 , 15 , 2 ],
162
163
METRICS [1 ]: [100 , 20 , 45 ],
163
164
COMBINED_METRIC : [110 , 35 , 47 ],
164
165
}
165
166
)
166
167
167
- fips2nation = gmpr .add_population_column (gmpr .get_crosswalk ("fips" , "hhs" ), "fips" )
168
- fips2nation .rename ({"hhs" : "nation" }, axis = 1 , inplace = True )
169
- fips2nation ["nation" ] = "nation"
170
- nation_pop = fips2nation .groupby ("nation"
168
+ state2nation = gmpr .add_population_column (gmpr .get_crosswalk ("state" , "state" ), "state_code" )
169
+ state2nation = gmpr .add_geocode (state2nation , "state_code" , "nation" )
170
+ nation_pop = state2nation .groupby ("nation"
171
171
).sum (
172
172
).reset_index (
173
173
).rename (columns = {"population" : "nation_pop" })
174
- df_plus = df .merge (fips2nation , left_on = "geo_id" , right_on = "fips " , how = "left"
174
+ df_plus = df .merge (state2nation , left_on = "geo_id" , right_on = "state_id " , how = "left"
175
175
).merge (nation_pop , on = "nation" , how = "left"
176
176
).assign (
177
177
fractional_pop = lambda x : x .population / x .nation_pop ,
@@ -190,7 +190,7 @@ def test_nation(self):
190
190
new_df = geo_map (df , "nation" ).dropna ()
191
191
192
192
assert set (new_df .keys ()) == set (df .keys ())
193
- assert set (new_df ["geo_id" ]) == set (["nation " ])
193
+ assert set (new_df ["geo_id" ]) == set (["us " ])
194
194
assert new_df [METRICS [0 ]].values == pytest .approx (df_plus [METRICS [0 ]].tolist ())
195
195
assert new_df [METRICS [1 ]].values == pytest .approx (df_plus [METRICS [1 ]].tolist ())
196
196
assert new_df [COMBINED_METRIC ].values == pytest .approx (df_plus [COMBINED_METRIC ].tolist ())
0 commit comments