Make sure the county IDs are assigned

Author: guangchenzhao

demos/models/marriage.py

@@ -196,12 +196,14 @@ def update_married_households_random(
 
     # If geoid_col is set, we copy the geoid from old households to new ones
     if module_config.geoid_col is not None:
-        new_hh_county = households.local.loc[
+        new_hh_geoid = households.local.loc[
             all_df.loc[first_index & neither_head_index, "household_id"],
             module_config.geoid_col,
         ].values
-        households.local.loc[new_hh_ids, module_config.geoid_col] = new_hh_county
-
+        households.local.loc[new_hh_ids, module_config.geoid_col] = new_hh_geoid
+    county_assignment = households.local.loc[all_df.loc[first_index & neither_head_index, "household_id"], 
+                                             "lcm_county_id"].values
+    households.local.loc[new_hh_ids, "lcm_county_id"] = county_assignment
     ## Decide who is household head in the households where the head left
     head_left_index = (all_df.relate == 0) & (all_df.household_id != all_df.new_hh_id)
     head_left_households = all_df[head_left_index]["household_id"].unique()
@@ -274,12 +276,14 @@ def update_divorce(persons, households, divorce_list, get_new_households):
 
     # If geoid_col is set, we copy the geoid from old households to new ones
     if module_config.geoid_col is not None:
-        county_assignment = households.local.loc[
+        geoid_assignment = households.local.loc[
             old_household_id, module_config.geoid_col
         ].values
         households.local.loc[new_households, module_config.geoid_col] = (
-            county_assignment
+            geoid_assignment
         )
+    county_assignment = households.local.loc[old_household_id, "lcm_county_id"].values
+    households.local.loc[new_households, "lcm_county_id"] = county_assignment
 
     ## Updates for people staying
     persons.local.loc[person_staying_index, "relate"] = 0