policyengine-us 1.365.0__py3-none-any.whl → 1.365.1__py3-none-any.whl

policyengine_us/tests/utilities/test_ucgid_hierarchical.py

@@ -2,6 +2,7 @@
 Test script demonstrating UCGID hierarchical functionality.
 """
 
+import numpy as np
 from policyengine_us import Microsimulation
 from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import (
     UCGID,
@@ -18,19 +19,16 @@ class TestUCGIDHierarchical:
         # Test Congressional District (CA_23)
         ca_23 = UCGID.CA_23
         ca_23_codes = ca_23.get_hierarchical_codes()
-        print(f"CA_23 hierarchical codes: {ca_23_codes}")
         assert ca_23_codes == ["5001800US0623", "0400000US06", "0100000US"]
 
         # Test State (CA)
         ca = UCGID.CA
         ca_codes = ca.get_hierarchical_codes()
-        print(f"CA hierarchical codes: {ca_codes}")
         assert ca_codes == ["0400000US06", "0100000US"]
 
         # Test US
         us = UCGID.US
         us_codes = us.get_hierarchical_codes()
-        print(f"US hierarchical codes: {us_codes}")
         assert us_codes == ["0100000US"]
 
         # Test hierarchy matching
@@ -50,9 +48,74 @@ class TestUCGIDHierarchical:
         # Convert the string back to enum to access hierarchical methods
         ucgid_enum = UCGID[ucgid_string]
         hierarchical_codes = ucgid_enum.get_hierarchical_codes()
-        assert len(hierarchical_codes) == 1
-        assert hierarchical_codes[0] == "0100000US"
+
+        # Should have 2 hierarchical codes for state-level UCGIDs: [state, US]
+        assert len(hierarchical_codes) == 2
+        assert (
+            hierarchical_codes[1] == "0100000US"
+        )  # Second code should always be US
+        assert hierarchical_codes[0].startswith(
+            "0400000US"
+        )  # First should be state-level
 
         # Test hierarchy matching in simulation context
         is_in_us = ucgid_enum.matches_hierarchy("0100000US")
         assert is_in_us == True
+
+    def test_ucgid_str_variable(self):
+        """Test the UCGID string variable functionality."""
+        simulation = Microsimulation()
+        ucgid_str_values = simulation.calculate("ucgid_str", "2024")
+
+        # The UCGID string variable should return a string representation
+        ucgid_str = ucgid_str_values.iloc[0]
+        assert type(ucgid_str) == str
+        assert (
+            ucgid_str.startswith("0100000US")
+            | ucgid_str.startswith("0400000US")
+            | ucgid_str.startswith("5001800US")
+        )
+
+        # Create a basic simulation to test with specific input values
+        from policyengine_us import Simulation
+
+        simulation = Simulation(
+            situation={
+                "people": {"person": {}},
+                "households": {"household": {"members": ["person"]}},
+            }
+        )
+
+        # Set a specific UCGID value for testing (CA_23)
+        simulation.set_input("ucgid", 2024, UCGID.CA_23)
+
+        # Calculate the ucgid_str value
+        ucgid_str_values = simulation.calculate("ucgid_str", 2024)
+        ucgid_str = ucgid_str_values[0]
+
+        # Verify it contains all three hierarchical codes
+        assert ucgid_str == "5001800US0623,0400000US06,0100000US"
+
+    def test_ucgid_microsimulation_input_override(self):
+        """Test setting UCGID input for all households in a microsimulation."""
+        microsim = Microsimulation()
+
+        # Get initial values
+        ucgid_before = microsim.calculate("ucgid", 2024)
+        num_households = len(ucgid_before)
+
+        # Set all households to CA_23
+        input_array = np.array([UCGID.CA_23] * num_households)
+        microsim.set_input("ucgid", 2024, input_array)
+
+        # Verify the input was set
+        ucgid_after = microsim.calculate("ucgid", 2024)
+        after_values = ucgid_after.values
+        assert all(val == "CA_23" for val in after_values)
+        assert len(set(after_values)) == 1
+
+        # Test that ucgid_str now returns hierarchical codes for all households
+        ucgid_str_values = microsim.calculate("ucgid_str", 2024)
+        expected_str = "5001800US0623,0400000US06,0100000US"
+        str_values = ucgid_str_values.values
+        assert all(str_val == expected_str for str_val in str_values)

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

@@ -20,78 +20,75 @@ class ucgid(Variable):
     def formula(household, period, parameters):
         simulation: Simulation = household.simulation
 
-        # If we're running over a dataset and UCGID is provided, it should be used directly
-        # (This is handled automatically by the simulation framework)
+        # Check if this variable has been set as input
+        if simulation.get_holder("ucgid").get_known_periods():
+            return household("ucgid", period)
 
-        # For non-dataset simulations, try to derive from state_code
-        if not simulation.is_over_dataset:
-            state_code = household("state_code", period)
+        # Try to derive from state_code
+        state_code_enum = household("state_code", period)
+        state_code_strings = state_code_enum.decode_to_str()
 
-            # Map state codes to UCGID state values
-            state_mapping = {
-                "AL": UCGID.AL,
-                "AK": UCGID.AK,
-                "AZ": UCGID.AZ,
-                "AR": UCGID.AR,
-                "CA": UCGID.CA,
-                "CO": UCGID.CO,
-                "CT": UCGID.CT,
-                "DE": UCGID.DE,
-                "DC": UCGID.DC,
-                "FL": UCGID.FL,
-                "GA": UCGID.GA,
-                "HI": UCGID.HI,
-                "ID": UCGID.ID,
-                "IL": UCGID.IL,
-                "IN": UCGID.IN,
-                "IA": UCGID.IA,
-                "KS": UCGID.KS,
-                "KY": UCGID.KY,
-                "LA": UCGID.LA,
-                "ME": UCGID.ME,
-                "MD": UCGID.MD,
-                "MA": UCGID.MA,
-                "MI": UCGID.MI,
-                "MN": UCGID.MN,
-                "MS": UCGID.MS,
-                "MO": UCGID.MO,
-                "MT": UCGID.MT,
-                "NE": UCGID.NE,
-                "NV": UCGID.NV,
-                "NH": UCGID.NH,
-                "NJ": UCGID.NJ,
-                "NM": UCGID.NM,
-                "NY": UCGID.NY,
-                "NC": UCGID.NC,
-                "ND": UCGID.ND,
-                "OH": UCGID.OH,
-                "OK": UCGID.OK,
-                "OR": UCGID.OR,
-                "PA": UCGID.PA,
-                "RI": UCGID.RI,
-                "SC": UCGID.SC,
-                "SD": UCGID.SD,
-                "TN": UCGID.TN,
-                "TX": UCGID.TX,
-                "UT": UCGID.UT,
-                "VT": UCGID.VT,
-                "VA": UCGID.VA,
-                "WA": UCGID.WA,
-                "WV": UCGID.WV,
-                "WI": UCGID.WI,
-                "WY": UCGID.WY,
-            }
+        # Map state codes to UCGID state values
+        state_mapping = {
+            "AL": UCGID.AL,
+            "AK": UCGID.AK,
+            "AZ": UCGID.AZ,
+            "AR": UCGID.AR,
+            "CA": UCGID.CA,
+            "CO": UCGID.CO,
+            "CT": UCGID.CT,
+            "DE": UCGID.DE,
+            "DC": UCGID.DC,
+            "FL": UCGID.FL,
+            "GA": UCGID.GA,
+            "HI": UCGID.HI,
+            "ID": UCGID.ID,
+            "IL": UCGID.IL,
+            "IN": UCGID.IN,
+            "IA": UCGID.IA,
+            "KS": UCGID.KS,
+            "KY": UCGID.KY,
+            "LA": UCGID.LA,
+            "ME": UCGID.ME,
+            "MD": UCGID.MD,
+            "MA": UCGID.MA,
+            "MI": UCGID.MI,
+            "MN": UCGID.MN,
+            "MS": UCGID.MS,
+            "MO": UCGID.MO,
+            "MT": UCGID.MT,
+            "NE": UCGID.NE,
+            "NV": UCGID.NV,
+            "NH": UCGID.NH,
+            "NJ": UCGID.NJ,
+            "NM": UCGID.NM,
+            "NY": UCGID.NY,
+            "NC": UCGID.NC,
+            "ND": UCGID.ND,
+            "OH": UCGID.OH,
+            "OK": UCGID.OK,
+            "OR": UCGID.OR,
+            "PA": UCGID.PA,
+            "RI": UCGID.RI,
+            "SC": UCGID.SC,
+            "SD": UCGID.SD,
+            "TN": UCGID.TN,
+            "TX": UCGID.TX,
+            "UT": UCGID.UT,
+            "VT": UCGID.VT,
+            "VA": UCGID.VA,
+            "WA": UCGID.WA,
+            "WV": UCGID.WV,
+            "WI": UCGID.WI,
+            "WY": UCGID.WY,
+        }
 
-            # Convert state codes to UCGID values
-            result = np.empty(len(state_code), dtype=object)
-            for i, code in enumerate(state_code):
-                state_str = str(code) if hasattr(code, "__str__") else code
-                if state_str in state_mapping:
-                    result[i] = state_mapping[state_str]
-                else:
-                    result[i] = UCGID.US
+        # Convert state code strings to UCGID values
+        result = np.empty(len(state_code_strings), dtype=object)
+        for i, state_str in enumerate(state_code_strings):
+            if state_str in state_mapping:
+                result[i] = state_mapping[state_str]
+            else:
+                result[i] = UCGID.US
 
-            return result
-
-        # Default fallback
-        return np.full(household.count, UCGID.US, dtype=object)
+        return result

policyengine_us/variables/household/demographic/geographic/ucgid/ucgid_str.py

@@ -1,4 +1,7 @@
 from policyengine_us.model_api import *
+from policyengine_us.variables.household.demographic.geographic.ucgid.ucgid_enum import (
+    UCGID,
+)
 
 
 class ucgid_str(Variable):
@@ -9,4 +12,18 @@ class ucgid_str(Variable):
     definition_period = YEAR
 
     def formula(household, period, parameters):
-        return household("ucgid", period).decode_to_str()
+        import numpy as np
+
+        ucgid_enum_names = household("ucgid", period).decode_to_str()
+
+        # Convert each enum name to its hierarchical codes
+        result = []
+        for enum_name in ucgid_enum_names:
+            # Get the enum instance from its name
+            ucgid_enum = UCGID[enum_name]
+
+            # Get all hierarchical codes and join with commas
+            hierarchical_codes = ucgid_enum.get_hierarchical_codes()
+            result.append(",".join(hierarchical_codes))
+
+        return np.array(result)

{policyengine_us-1.365.0.dist-info → policyengine_us-1.365.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: policyengine-us
-Version: 1.365.0
+Version: 1.365.1
 Summary: Add your description here.
 Author-email: PolicyEngine <hello@policyengine.org>
 License-File: LICENSE

{policyengine_us-1.365.0.dist-info → policyengine_us-1.365.1.dist-info}/RECORD

@@ -5136,7 +5136,7 @@ policyengine_us/tests/policy/contrib/ubi_center/flat_tax.yaml,sha256=O6krfJ2rovW
 policyengine_us/tests/policy/reform/ctc_expansion.yaml,sha256=soq-LqWwbb7fNQJFphx_1gSc8uDGOhtr-5P93oY2Mpg,1596
 policyengine_us/tests/policy/reform/winship.yaml,sha256=OPfsoZ0NrlPU25hJadZXycrYLaM0wPUGXjBrn6_8gR4,1706
 policyengine_us/tests/utilities/test_load_county_fips_dataset.py,sha256=QAURDL6IMN5gafUOcTri1PUet6khLNZ_Bov6NdW_pMM,4397
-policyengine_us/tests/utilities/test_ucgid_hierarchical.py,sha256=_QBrBgX072aY-G1z7UwwNe_1pIUT9SfDMrNUhkPU-B4,2166
+policyengine_us/tests/utilities/test_ucgid_hierarchical.py,sha256=WZYqWHw6-QQAv0Wd3bGgAcgPgZ0ILhOxpC-kgo4HRkU,4532
 policyengine_us/tests/variables/gov/states/vt/tax/income/vt_eitc.yaml,sha256=gE-ZtclGqoFKs87k9VUgKEoVO2iZkFBq4k-ixGtde14,2341
 policyengine_us/tests/variables/gov/states/vt/tax/income/adjusted_gross_income/subtractions/retirement_income_exemption/vt_military_retirement_income_based_exemption.yaml,sha256=b5I5ekAoDCrojCLwEZMoBEDdUX9zDw1bM-4iQLmyTng,3267
 policyengine_us/tests/variables/gov/states/vt/tax/income/credits/vt_veteran_tax_credit.yaml,sha256=51NtThz2bNwVD7o-2VXZqplDjKUAYzIiIzyL1PfXK94,3244
@@ -7785,9 +7785,9 @@ policyengine_us/variables/household/demographic/geographic/county/county_str.py,
 policyengine_us/variables/household/demographic/geographic/county/first_county_in_state.py,sha256=yXCj-LVjxPeILIg42ijefTmi_j4NxJwX0kMtoE47xic,1578
 policyengine_us/variables/household/demographic/geographic/state/average_home_energy_use_in_state.py,sha256=aaGYAXLU7tsx-A9NW4nyLij5J5mp0WTRXsyObngekvY,319
 policyengine_us/variables/household/demographic/geographic/state/in_state.py,sha256=o3ksZCOHtlJ1JGHIXAPlqcs5-rVbTgOCcDdN1-fYqmA,551
-policyengine_us/variables/household/demographic/geographic/ucgid/ucgid.py,sha256=RX4zSK6mqUoI10Jhe8XSDhrG-H0qOv3zmKMHeSl-TQY,3270
+policyengine_us/variables/household/demographic/geographic/ucgid/ucgid.py,sha256=SwONHVRvHmUFHvPZQMYY7Y9e7gz74kisLjJICpHdyKg,2874
 policyengine_us/variables/household/demographic/geographic/ucgid/ucgid_enum.py,sha256=dkeusTBI1j4xAGmRFkfSsyT3Hq5gHZwRUkSeISeNbx0,15636
-policyengine_us/variables/household/demographic/geographic/ucgid/ucgid_str.py,sha256=ngH3tnYb6m96he1_Cj2dPtyc8nFDOsqlbjGOmh4c79A,335
+policyengine_us/variables/household/demographic/geographic/ucgid/ucgid_str.py,sha256=GnIRg2tc0376YazPX7wm7YQgMQaYPm2Au47vc1E5JcM,914
 policyengine_us/variables/household/demographic/geographic/zip_code/three_digit_zip_code.py,sha256=7q0FQRahTOINiK2vlAyHEa5xIKowM5Kgl60-kR3ZuWU,426
 policyengine_us/variables/household/demographic/geographic/zip_code/zip_code.py,sha256=PZiLb84BgKQ_Suvd6iUiEnzySu3Vtiv8FOlfB4rnagQ,1691
 policyengine_us/variables/household/demographic/household/bedrooms.py,sha256=4ltdIRLdmeuTKAskn4pdgxZuxEE-4rz0ffPnIKobc-o,164
@@ -8236,8 +8236,8 @@ policyengine_us/variables/input/farm_income.py,sha256=BEKxYmHNNnWJAAvULl5qZJigy5
 policyengine_us/variables/input/geography.py,sha256=XmBlgXhzBoLRKk6R8taVZHqUw1eU8MbNeGS9iJ7_l44,4506
 policyengine_us/variables/input/self_employment_income.py,sha256=PwsGz8R4lRikKWUYOhsC0qosNNLXq4f5SQmfw4S3mk8,511
 policyengine_us/variables/input/self_employment_income_before_lsr.py,sha256=E8fcX9Nlyqz8dziHhQv_euutdmoIwFMMWePUwbbwv_w,379
-policyengine_us-1.365.0.dist-info/METADATA,sha256=DQE2dHq4VJDvHfFPlYb1VuGGPavQDv6ZmYNVN0UP1uM,1693
-policyengine_us-1.365.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-policyengine_us-1.365.0.dist-info/entry_points.txt,sha256=MLaqNyNTbReALyKNkde85VkuFFpdPWAcy8VRG1mjczc,57
-policyengine_us-1.365.0.dist-info/licenses/LICENSE,sha256=2N5ReRelkdqkR9a-KP-y-shmcD5P62XoYiG-miLTAzo,34519
-policyengine_us-1.365.0.dist-info/RECORD,,
+policyengine_us-1.365.1.dist-info/METADATA,sha256=HOpqZAHTPNFrdBmVMwFA27fcbSixBqddzUTZwqQ9qk8,1693
+policyengine_us-1.365.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+policyengine_us-1.365.1.dist-info/entry_points.txt,sha256=MLaqNyNTbReALyKNkde85VkuFFpdPWAcy8VRG1mjczc,57
+policyengine_us-1.365.1.dist-info/licenses/LICENSE,sha256=2N5ReRelkdqkR9a-KP-y-shmcD5P62XoYiG-miLTAzo,34519
+policyengine_us-1.365.1.dist-info/RECORD,,