@@ -110,3 +110,88 @@ def test_msa(self):
110
110
assert new_df [METRICS [0 ]].values == pytest .approx (df_plus [METRICS [0 ]].tolist ())
111
111
assert new_df [METRICS [1 ]].values == pytest .approx (df_plus [METRICS [1 ]].tolist ())
112
112
assert new_df [COMBINED_METRIC ].values == pytest .approx (df_plus [COMBINED_METRIC ].tolist ())
113
+
114
+ def test_hhs (self ):
115
+ gmpr = GeoMapper ()
116
+ df = pd .DataFrame (
117
+ {
118
+ "geo_id" : ["al" , "fl" , "tx" ],
119
+ "timestamp" : ["2020-02-15" , "2020-02-15" , "2020-02-15" ],
120
+ METRICS [0 ]: [10 , 15 , 2 ],
121
+ METRICS [1 ]: [100 , 20 , 45 ],
122
+ COMBINED_METRIC : [110 , 35 , 47 ],
123
+ }
124
+ )
125
+
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"
129
+ ).sum (
130
+ ).reset_index (
131
+ ).rename (columns = {"population" : "hhs_pop" })
132
+ df_plus = df .merge (state2hhs , left_on = "geo_id" , right_on = "state_id" , how = "left"
133
+ ).merge (hhs_pop , on = "hhs" , how = "left"
134
+ ).assign (
135
+ fractional_pop = lambda x : x .population / x .hhs_pop ,
136
+ metric_0 = lambda x : x .fractional_pop * x [METRICS [0 ]],
137
+ metric_1 = lambda x : x .fractional_pop * x [METRICS [1 ]],
138
+ combined_metric = lambda x : x .metric_0 + x .metric_1
139
+ ).groupby ("hhs"
140
+ ).sum (
141
+ ).drop (
142
+ labels = [METRICS [0 ], METRICS [1 ], COMBINED_METRIC ],
143
+ axis = "columns"
144
+ ).rename (
145
+ columns = {"metric_0" : METRICS [0 ], "metric_1" : METRICS [1 ], "combined_metric" : COMBINED_METRIC }
146
+ )
147
+
148
+ new_df = geo_map (df , "hhs" ).dropna ()
149
+
150
+ assert set (new_df .keys ()) == set (df .keys ())
151
+ assert set (new_df ["geo_id" ]) == set (["4" , "6" ])
152
+ assert new_df [METRICS [0 ]].values == pytest .approx (df_plus [METRICS [0 ]].tolist ())
153
+ assert new_df [METRICS [1 ]].values == pytest .approx (df_plus [METRICS [1 ]].tolist ())
154
+ assert new_df [COMBINED_METRIC ].values == pytest .approx (df_plus [COMBINED_METRIC ].tolist ())
155
+
156
+ def test_nation (self ):
157
+ gmpr = GeoMapper ()
158
+ df = pd .DataFrame (
159
+ {
160
+ "geo_id" : ["al" , "il" , "tx" ],
161
+ "timestamp" : ["2020-02-15" , "2020-02-15" , "2020-02-15" ],
162
+ METRICS [0 ]: [10 , 15 , 2 ],
163
+ METRICS [1 ]: [100 , 20 , 45 ],
164
+ COMBINED_METRIC : [110 , 35 , 47 ],
165
+ }
166
+ )
167
+
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
+ ).sum (
172
+ ).reset_index (
173
+ ).rename (columns = {"population" : "nation_pop" })
174
+ df_plus = df .merge (state2nation , left_on = "geo_id" , right_on = "state_id" , how = "left"
175
+ ).merge (nation_pop , on = "nation" , how = "left"
176
+ ).assign (
177
+ fractional_pop = lambda x : x .population / x .nation_pop ,
178
+ metric_0 = lambda x : x .fractional_pop * x [METRICS [0 ]],
179
+ metric_1 = lambda x : x .fractional_pop * x [METRICS [1 ]],
180
+ combined_metric = lambda x : x .metric_0 + x .metric_1
181
+ ).groupby ("nation"
182
+ ).sum (
183
+ ).drop (
184
+ labels = [METRICS [0 ], METRICS [1 ], COMBINED_METRIC ],
185
+ axis = "columns"
186
+ ).rename (
187
+ columns = {"metric_0" : METRICS [0 ], "metric_1" : METRICS [1 ], "combined_metric" : COMBINED_METRIC }
188
+ )
189
+
190
+ new_df = geo_map (df , "nation" ).dropna ()
191
+
192
+ assert set (new_df .keys ()) == set (df .keys ())
193
+ assert set (new_df ["geo_id" ]) == set (["us" ])
194
+ assert new_df [METRICS [0 ]].values == pytest .approx (df_plus [METRICS [0 ]].tolist ())
195
+ assert new_df [METRICS [1 ]].values == pytest .approx (df_plus [METRICS [1 ]].tolist ())
196
+ assert new_df [COMBINED_METRIC ].values == pytest .approx (df_plus [COMBINED_METRIC ].tolist ())
197
+
0 commit comments