Merge pull request #6891 from PolicyEngine/fix-state-group-enum-warnings

Fix StateGroup enum warnings for contiguous US states

Author: baogorek

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - StateGroup enum warnings for contiguous US states in state_group variable.

policyengine_us/tests/policy/baseline/household/demographic/geographic/state_group.yaml

@@ -11,3 +11,24 @@
     state_code: AK
   output:
     state_group: AK
+
+- name: Contiguous state does not cause empty array indexing error
+  period: 2024
+  input:
+    state_code: TX
+  output:
+    state_group: CONTIGUOUS_US
+
+- name: Hawaii
+  period: 2024
+  input:
+    state_code: HI
+  output:
+    state_group: HI
+
+- name: Puerto Rico
+  period: 2024
+  input:
+    state_code: PR
+  output:
+    state_group: PR

policyengine_us/variables/household/demographic/geographic/state_group.py

@@ -22,8 +22,11 @@ class state_group(Variable):
     def formula(household, period, parameters):
         NON_CONTIGUOUS_STATES = ("AK", "HI", "GU", "PR", "VI")
         state_code = household("state_code", period).decode_to_str()
-        return where(
-            np.isin(state_code, NON_CONTIGUOUS_STATES),
-            StateGroup.encode(state_code).decode(),
-            StateGroup.CONTIGUOUS_US,
-        )
+        is_non_contiguous = np.isin(state_code, NON_CONTIGUOUS_STATES)
+        encoded = np.empty(state_code.shape, dtype=object)
+        encoded[:] = StateGroup.CONTIGUOUS_US
+        if is_non_contiguous.any():
+            encoded[is_non_contiguous] = StateGroup.encode(
+                state_code[is_non_contiguous]
+            ).decode()
+        return encoded