Ensure multiple counties are present in microsims

PolicyEngine · Nov 3, 2023 · cf94f9f · cf94f9f
1 parent 33d86b2
commit cf94f9f
Showing 1 changed file with 7 additions and 2 deletions.
diff --git a/policyengine_us/variables/household/demographic/geographic/zip_code/zip_code.py b/policyengine_us/variables/household/demographic/geographic/zip_code/zip_code.py
@@ -30,10 +30,15 @@ def formula(household, period, parameters):
         else:
             household_zip_code = np.empty_like(state_code, dtype=object)
             for state in ZIP_CODE_DATASET.state.unique():
+                count_households_in_state = (state_code == state).sum()
                 household_zip_code[state_code == state] = (
                     ZIP_CODE_DATASET[ZIP_CODE_DATASET.state == state]
-                    .sample(1, weights="population")
-                    .zip_code.iloc[0]
+                    .sample(
+                        count_households_in_state,
+                        weights="population",
+                        replace=True,
+                    )
+                    .zip_code
                 )
             household_zip_code = pd.Series(household_zip_code)