Fix Enum bug (#302)

PolicyEngine · Nov 1, 2024 · cf80926 · cf80926
1 parent 6833c2b
commit cf80926
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Bug causing Enums to fail in some simulations.
diff --git a/policyengine_core/simulations/simulation.py b/policyengine_core/simulations/simulation.py
@@ -739,7 +739,15 @@ def _calculate(
                     array = holder.default_array()
 
             if variable.defined_for is not None:
-                array = np.where(mask, array, np.zeros_like(array))
+                array = np.where(mask, array, variable.default_value)
+                if variable.value_type == Enum:
+                    array = np.array(
+                        [
+                            item.index if isinstance(item, Enum) else item
+                            for item in array
+                        ]
+                    )
+                    array = EnumArray(array, variable.possible_values)
 
             array = self._cast_formula_result(array, variable)
             holder.put_in_cache(array, period, self.branch_name)