Optimize select() statements by removing redundant conditions

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+      - Optimized select() statements by removing redundant conditions that match the default value, improving vectorized calculation performance.

policyengine_us/reforms/local/ny/mamdani_income_tax.py

@@ -29,21 +29,21 @@ def formula(tax_unit, period, parameters):
             filing_status = tax_unit("filing_status", period)
             filing_statuses = filing_status.possible_values
             p = parameters(period).gov.local.ny.nyc.tax.income.rates
+            # Default covers SINGLE filing status
             regular_tax = select(
                 [
-                    filing_status == filing_statuses.SINGLE,
                     filing_status == filing_statuses.JOINT,
                     filing_status == filing_statuses.HEAD_OF_HOUSEHOLD,
                     filing_status == filing_statuses.SURVIVING_SPOUSE,
                     filing_status == filing_statuses.SEPARATE,
                 ],
                 [
-                    p.single.calc(taxable_income),
                     p.joint.calc(taxable_income),
                     p.head_of_household.calc(taxable_income),
                     p.surviving_spouse.calc(taxable_income),
                     p.separate.calc(taxable_income),
                 ],
+                default=p.single.calc(taxable_income),
             )
             mamdani_tax = add(tax_unit, period, ["nyc_mamdani_income_tax"])
             return regular_tax + mamdani_tax

policyengine_us/variables/contrib/taxsim/taxsim_mstat.py

@@ -12,17 +12,15 @@ def formula(tax_unit, period, parameters):
         fstatus = filing_status.possible_values
         return select(
             [
-                filing_status == fstatus.SINGLE,
-                filing_status == fstatus.HEAD_OF_HOUSEHOLD,
                 filing_status == fstatus.JOINT,
                 filing_status == fstatus.SEPARATE,
                 filing_status == fstatus.SURVIVING_SPOUSE,
             ],
             [
-                1,
-                1,
                 2,
                 6,
                 8,
             ],
+            # Default covers SINGLE and HEAD_OF_HOUSEHOLD (both map to 1)
+            default=1,
         )

policyengine_us/variables/gov/ed/pell_grant/efc/pell_grant_efc.py

@@ -18,14 +18,15 @@ def formula(person, period, parameters):
         zero_efc_max = parameters(period).gov.ed.pell_grant.efc.automatic_zero
         automatic_zero = head_income <= zero_efc_max
         return select(
-            [formula == "A", formula == "B", formula == "C"],
+            [formula == "A", formula == "B"],
             [
                 where(
                     automatic_zero,
                     0,
                     max_(0, head_contribution + dependent_contribution),
                 ),
                 max_(0, head_contribution),
-                where(automatic_zero, 0, max_(0, head_contribution)),
             ],
+            # Default covers formula == "C".
+            default=where(automatic_zero, 0, max_(0, head_contribution)),
         )

policyengine_us/variables/gov/ed/pell_grant/sai/eligibility_type/pell_grant_household_type.py

@@ -24,13 +24,13 @@ def formula(person, period, parameters):
             [
                 dependent & ~joint,
                 dependent & joint,
-                ~dependent & ~joint,
                 ~dependent & joint,
             ],
             [
                 PellGrantHouseholdType.DEPENDENT_SINGLE,
                 PellGrantHouseholdType.DEPENDENT_NOT_SINGLE,
-                PellGrantHouseholdType.INDEPENDENT_SINGLE,
                 PellGrantHouseholdType.INDEPENDENT_NOT_SINGLE,
             ],
+            # Default covers ~dependent & ~joint (independent single).
+            default=PellGrantHouseholdType.INDEPENDENT_SINGLE,
         )

policyengine_us/variables/gov/hhs/ccdf/ccdf_market_rate.py

@@ -22,13 +22,14 @@ def formula(person, period, parameters):
         hours_per_day = person("childcare_hours_per_day", period)
         days_per_week = person("childcare_days_per_week", period)
         hours_per_week = hours_per_day * days_per_week
+        # Default covers: HOURLY duration of care
         periods_per_week = select(
             [
                 duration_of_care == durations_of_care.WEEKLY,
                 duration_of_care == durations_of_care.DAILY,
                 duration_of_care == durations_of_care.PART_DAY,
-                duration_of_care == durations_of_care.HOURLY,
             ],
-            [1, days_per_week, days_per_week, hours_per_week],
+            [1, days_per_week, days_per_week],
+            default=hours_per_week,
         )
         return rate_per_period * periods_per_week * WEEKS_IN_YEAR

policyengine_us/variables/gov/hhs/medicaid/costs/medicaid_group.py

@@ -69,16 +69,15 @@ def formula(person, period, parameters):
         il_hbi_senior = il_hbi_eligible & (age >= p_hbi.senior.min_age)
 
         # Core mapping, in precedence order:
+        # Default covers: not eligible for Medicaid
         return select(
             [
-                ~eligible,
                 disabled | il_hbi_senior,
                 non_expansion_adult,
                 expansion_adult | il_hbi_adult,
                 child | il_hbi_child,
             ],
             [
-                MedicaidGroup.NONE,
                 MedicaidGroup.AGED_DISABLED,
                 MedicaidGroup.NON_EXPANSION_ADULT,
                 MedicaidGroup.EXPANSION_ADULT,

policyengine_us/variables/gov/hhs/medicare/eligibility/part_b/income_adjusted_part_b_premium.py

@@ -38,15 +38,16 @@ def formula(person, period, parameters):
 
         p = parameters(period).gov.hhs.medicare.part_b.irmaa
 
+        # Default covers: SEPARATE filing status
         irmaa_amount = select(
-            in_status,
+            in_status[:-1],  # Exclude SEPARATE (last element)
             [
                 p.single.calc(magi),
                 p.joint.calc(magi),
                 p.head_of_household.calc(magi),
                 p.surviving_spouse.calc(magi),
-                p.separate.calc(magi),
             ],
+            default=p.separate.calc(magi),
         )
 
         # IRMAA amounts are monthly, multiply by MONTHS_IN_YEAR to get annual

policyengine_us/variables/gov/local/ny/nyc/tax/income/credits/school/nyc_school_tax_credit_rate_reduction_amount.py

@@ -25,17 +25,17 @@ def formula(tax_unit, period, parameters):
         filing_statuses = filing_status.possible_values
         return select(
             [
-                filing_status == filing_statuses.SINGLE,
                 filing_status == filing_statuses.JOINT,
                 filing_status == filing_statuses.SEPARATE,
                 filing_status == filing_statuses.HEAD_OF_HOUSEHOLD,
                 filing_status == filing_statuses.SURVIVING_SPOUSE,
             ],
             [
-                p.amount.single.calc(nyc_taxable_income),
                 p.amount.joint.calc(nyc_taxable_income),
                 p.amount.separate.calc(nyc_taxable_income),
                 p.amount.head_of_household.calc(nyc_taxable_income),
                 p.amount.surviving_spouse.calc(nyc_taxable_income),
             ],
+            # Default covers SINGLE
+            default=p.amount.single.calc(nyc_taxable_income),
         )

policyengine_us/variables/gov/local/ny/nyc/tax/income/nyc_income_tax_before_credits.py

@@ -16,17 +16,17 @@ def formula(tax_unit, period, parameters):
         rates = parameters(period).gov.local.ny.nyc.tax.income.rates
         return select(
             [
-                filing_status == filing_statuses.SINGLE,
                 filing_status == filing_statuses.JOINT,
                 filing_status == filing_statuses.HEAD_OF_HOUSEHOLD,
                 filing_status == filing_statuses.SURVIVING_SPOUSE,
                 filing_status == filing_statuses.SEPARATE,
             ],
             [
-                rates.single.calc(taxable_income),
                 rates.joint.calc(taxable_income),
                 rates.head_of_household.calc(taxable_income),
                 rates.surviving_spouse.calc(taxable_income),
                 rates.separate.calc(taxable_income),
             ],
+            # Default covers SINGLE
+            default=rates.single.calc(taxable_income),
         )

policyengine_us/variables/gov/states/ar/tax/income/deductions/itemized/post_secondary_education/ar_post_secondary_education_tuition_deduction_person.py

@@ -25,6 +25,7 @@ def formula(person, period, parameters):
                 p.weighted_average_tuition.technical_institutes,
                 p.weighted_average_tuition.four_year_college,
             ],
+            # Default covers two-year college students (and non-students)
             default=p.weighted_average_tuition.two_year_college,
         )
         return min_(uncapped, cap)

policyengine_us/variables/gov/states/ar/tax/income/low_income/ar_low_income_tax_joint.py

@@ -60,5 +60,6 @@ def formula(person, period, parameters):
                     agi_attributed_to_head, right=True
                 ),
             ],
+            # Default covers any unmatched filing status combinations
             default=0,
         )

policyengine_us/variables/gov/states/az/tax/income/az_filing_status.py

@@ -33,10 +33,11 @@ def formula(tax_unit, period, parameters):
                 us_filing_status == fsvals.SEPARATE,
             ],
             [
-                # In Arizona, surviving spouse filers are treated as heads of households.
                 ArizonaFilingStatus.JOINT,
                 ArizonaFilingStatus.SINGLE,
                 ArizonaFilingStatus.SEPARATE,
             ],
+            # Default covers HEAD_OF_HOUSEHOLD and SURVIVING_SPOUSE
+            # (in Arizona, surviving spouse filers are treated as heads of households)
             default=ArizonaFilingStatus.HEAD_OF_HOUSEHOLD,
         )

policyengine_us/variables/gov/states/ca/cdss/state_supplement/payment_standard/ca_state_supplement_aged_disabled_amount.py

@@ -21,5 +21,6 @@ def formula(spm_unit, period, parameters):
         return select(
             [aged_disabled_count > 1, aged_disabled_count == 1],
             [p.married, p.single],
+            # Default covers aged_disabled_count == 0 (no supplement)
             default=0,
         )

policyengine_us/variables/gov/states/ca/cdss/state_supplement/payment_standard/ca_state_supplement_blind_amount.py

@@ -24,6 +24,7 @@ def formula(spm_unit, period, parameters):
         blind_married_amount = select(
             [blind_count >= 2, blind_count == 1],
             [p.blind.married.two_blind, p.blind.married.one_blind],
+            # Default covers blind_count == 0 (no blind married amount)
             default=0,
         )
         total_amount_if_eligible = where(

policyengine_us/variables/gov/states/ca/cdss/tanf/child_care/reimbursement/ca_calworks_child_care_payment_factor.py

@@ -28,5 +28,6 @@ def formula(person, period, parameters):
                 p.exceptional_needs,
                 p.severely_disabled,
             ],
+            # Default covers regular child care (no special category/factor)
             default=1,
         )

policyengine_us/variables/gov/states/ca/cdss/tanf/child_care/reimbursement/ca_calworks_child_care_payment_standard.py

@@ -24,6 +24,7 @@ def formula(person, period, parameters):
         age_group = select(
             [age < p.age_threshold.lower, age > p.age_threshold.higher],
             ["younger", "older"],
+            # Default covers children in the middle age range
             default="middle",
         )
         return p.standard[age_group][provider][time][is_full_time]

policyengine_us/variables/gov/states/ca/cdss/tanf/child_care/reimbursement/ca_calworks_child_care_time_coefficient.py

@@ -27,5 +27,6 @@ def formula(person, period, parameters):
                 time_category == time_categories.WEEKLY,
             ],
             [hours_per_day * days_per_month, days_per_month, weeks_per_month],
-            default=1,  # Monthly.
+            # Default covers MONTHLY time category (coefficient of 1)
+            default=1,
         )

policyengine_us/variables/gov/states/co/tax/income/credits/ctc/co_ctc.py

@@ -26,21 +26,21 @@ def formula(tax_unit, period, parameters):
         agi = tax_unit("adjusted_gross_income", period)
         if p.ctc_matched_federal_credit:
             federal_ctc = tax_unit("co_federal_ctc", period)
+            # Default covers SINGLE filing status
             rate = select(
                 [
-                    filing_status == statuses.SINGLE,
                     filing_status == statuses.JOINT,
                     filing_status == statuses.SEPARATE,
                     filing_status == statuses.SURVIVING_SPOUSE,
                     filing_status == statuses.HEAD_OF_HOUSEHOLD,
                 ],
                 [
-                    p.rate.single.calc(agi, right=True),
                     p.rate.joint.calc(agi, right=True),
                     p.rate.separate.calc(agi, right=True),
                     p.rate.surviving_spouse.calc(agi, right=True),
                     p.rate.head_of_household.calc(agi, right=True),
                 ],
+                default=p.rate.single.calc(agi, right=True),
             )
             return rate * federal_ctc
         else:
@@ -51,20 +51,20 @@ def formula(tax_unit, period, parameters):
                 person("co_ctc_eligible_child", period) & child_age_eligible
             )
             eligible_children = tax_unit.sum(eligible_child)
+            # Default covers SINGLE filing status
             amount_per_child = select(
                 [
-                    filing_status == statuses.SINGLE,
                     filing_status == statuses.JOINT,
                     filing_status == statuses.SEPARATE,
                     filing_status == statuses.SURVIVING_SPOUSE,
                     filing_status == statuses.HEAD_OF_HOUSEHOLD,
                 ],
                 [
-                    p.amount.single.calc(agi, right=True),
                     p.amount.joint.calc(agi, right=True),
                     p.amount.separate.calc(agi, right=True),
                     p.amount.surviving_spouse.calc(agi, right=True),
                     p.amount.head_of_household.calc(agi, right=True),
                 ],
+                default=p.amount.single.calc(agi, right=True),
             )
             return amount_per_child * eligible_children

policyengine_us/variables/gov/states/ct/tax/income/credits/ct_personal_credit_rate.py

@@ -15,19 +15,19 @@ def formula(tax_unit, period, parameters):
         status = filing_status.possible_values
         p = parameters(period).gov.states.ct.tax.income.agi
 
+        # Default covers SINGLE filing status
         return select(
             [
-                filing_status == status.SINGLE,
                 filing_status == status.JOINT,
                 filing_status == status.SEPARATE,
                 filing_status == status.SURVIVING_SPOUSE,
                 filing_status == status.HEAD_OF_HOUSEHOLD,
             ],
             [
-                p.single.calc(agi, right=True),
                 p.joint.calc(agi, right=True),
                 p.separate.calc(agi, right=True),
                 p.surviving_spouse.calc(agi, right=True),
                 p.head_of_household.calc(agi, right=True),
             ],
+            default=p.single.calc(agi, right=True),
         )

policyengine_us/variables/gov/states/ct/tax/income/ct_income_tax_after_personal_credits.py

@@ -14,21 +14,21 @@ def formula(tax_unit, period, parameters):
         filing_status = tax_unit("filing_status", period)
         status = filing_status.possible_values
         p = parameters(period).gov.states.ct.tax.income.rates
+        # Default covers SINGLE filing status
         itax_before_personal_credits = select(
             [
-                filing_status == status.SINGLE,
                 filing_status == status.JOINT,
                 filing_status == status.SEPARATE,
                 filing_status == status.SURVIVING_SPOUSE,
                 filing_status == status.HEAD_OF_HOUSEHOLD,
             ],
             [
-                p.single.calc(taxable_income),
                 p.joint.calc(taxable_income),
                 p.separate.calc(taxable_income),
                 p.surviving_spouse.calc(taxable_income),
                 p.head_of_household.calc(taxable_income),
             ],
+            default=p.single.calc(taxable_income),
         )
         add_back = tax_unit("ct_income_tax_phase_out_add_back", period)
         tax_recapture = tax_unit("ct_income_tax_recapture", period)

policyengine_us/variables/gov/states/ga/tax/income/ga_income_tax_before_non_refundable_credits.py

@@ -14,19 +14,19 @@ def formula(tax_unit, period, parameters):
         filing_status = tax_unit("filing_status", period)
         status = filing_status.possible_values
         income = tax_unit("ga_taxable_income", period)
+        # Default covers SINGLE filing status
         return select(
             [
-                filing_status == status.SINGLE,
                 filing_status == status.SEPARATE,
                 filing_status == status.JOINT,
                 filing_status == status.HEAD_OF_HOUSEHOLD,
                 filing_status == status.SURVIVING_SPOUSE,
             ],
             [
-                p.single.calc(income),
                 p.separate.calc(income),
                 p.joint.calc(income),
                 p.head_of_household.calc(income),
                 p.surviving_spouse.calc(income),
             ],
+            default=p.single.calc(income),
         )

policyengine_us/variables/gov/states/hi/tax/income/credits/hi_act_115_rebate.py

@@ -15,17 +15,14 @@ def formula(tax_unit, period, parameters):
         filing_status = tax_unit("filing_status", period)
         statuses = filing_status.possible_values
         federal_agi = tax_unit("adjusted_gross_income", period)
+        # Default covers SINGLE, HEAD_OF_HOUSEHOLD, and SEPARATE (all use single rate)
         return select(
             [
                 filing_status == statuses.JOINT,
-                filing_status == statuses.HEAD_OF_HOUSEHOLD,
-                filing_status == statuses.SEPARATE,
                 filing_status == statuses.SURVIVING_SPOUSE,
             ],
             [
                 p.joint.calc(federal_agi),
-                p.single.calc(federal_agi),
-                p.single.calc(federal_agi),
                 p.joint.calc(federal_agi),
             ],
             default=p.single.calc(federal_agi),