Delete branches earlier (#5287)

* Add branch deletion * Bug fix * Fix errors, now works * Remove print line * Fix bug in IL code that breaks subsetting * Format
PolicyEngine · Nov 1, 2024 · d9b74c5 · d9b74c5
1 parent 213fccf
commit d9b74c5
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 9 deletions.
diff --git a/changelog_entry.yaml b/changelog_entry.yaml
@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    fixed:
+    - Branch improvements.
diff --git a/policyengine_us/variables/gov/simulation/capital_gains_responses.py b/policyengine_us/variables/gov/simulation/capital_gains_responses.py
@@ -96,8 +96,6 @@ def formula(person, period, parameters):
         response_factor = np.exp(elasticity * tax_rate_change) - 1
         response = capital_gains * response_factor
 
-        print(response.sum())
-
         return response
 
 

diff --git a/policyengine_us/variables/gov/simulation/labor_supply_response/labor_supply_response.py b/policyengine_us/variables/gov/simulation/labor_supply_response/labor_supply_response.py
@@ -15,11 +15,11 @@ def formula(person, period, parameters):
         baseline_branch = simulation.get_branch("baseline").get_branch(
             "baseline_lsr_measurement"
         )
-        measurement_person = measurement_branch.populations["person"]
         baseline_person = baseline_branch.populations["person"]
         baseline_net_income = baseline_person.household(
             "household_net_income", period
         )
+        measurement_person = measurement_branch.populations["person"]
         net_income = measurement_person.household(
             "household_net_income", period
         )
@@ -47,14 +47,14 @@ class relative_wage_change(Variable):
 
     def formula(person, period, parameters):
         simulation = person.simulation
-        measurement_branch = simulation.get_branch("lsr_measurement")
         baseline_branch = simulation.get_branch("baseline").get_branch(
             "baseline_lsr_measurement"
         )
-        measurement_person = measurement_branch.populations["person"]
         baseline_person = baseline_branch.populations["person"]
         baseline_mtr = baseline_person("marginal_tax_rate", period)
         baseline_wage = 1 - baseline_mtr
+        measurement_branch = simulation.get_branch("lsr_measurement")
+        measurement_person = measurement_branch.populations["person"]
         mtr = measurement_person("marginal_tax_rate", period)
         wage_rate = 1 - mtr
         # _c suffix for "clipped"
@@ -259,6 +259,11 @@ def formula(person, period, parameters):
                 "substitution_elasticity_lsr",
             ],
         )
+        simulation = person.simulation
+        del simulation.branches["baseline"].branches[
+            "baseline_lsr_measurement"
+        ]
+        del simulation.branches["lsr_measurement"]
 
         simulation.macro_cache_read = False
         simulation.macro_cache_write = False
@@ -296,4 +301,5 @@ class self_employment_income_behavioral_response(Variable):
     def formula(person, period, parameters):
         lsr = person("labor_supply_behavioral_response", period)
         emp_response = person("employment_income_behavioral_response", period)
+
         return lsr - emp_response
diff --git a/policyengine_us/variables/gov/states/il/tax/income/exemptions/il_personal_exemption.py b/policyengine_us/variables/gov/states/il/tax/income/exemptions/il_personal_exemption.py
@@ -1,4 +1,7 @@
 from policyengine_us.model_api import *
+from policyengine_us.variables.gov.states.il.tax.income.exemptions.il_personal_exemption_eligibility_status import (
+    ILPersonalExemptionEligibilityStatus,
+)
 
 
 class il_personal_exemption(Variable):
@@ -14,18 +17,16 @@ def formula(tax_unit, period, parameters):
             "il_personal_exemption_eligibility_status", period
         )
 
-        eligibility_status = il_is_personal_exemption_eligible.possible_values
-
         personal_exemption_amount = parameters(
             period
         ).gov.states.il.tax.income.exemption.personal
 
         return personal_exemption_amount * select(
             [
                 il_is_personal_exemption_eligible
-                == eligibility_status.BOTH_ELIGIBLE,
+                == ILPersonalExemptionEligibilityStatus.BOTH_ELIGIBLE,
                 il_is_personal_exemption_eligible
-                == eligibility_status.PARTIALLY_ELIGIBLE,
+                == ILPersonalExemptionEligibilityStatus.PARTIALLY_ELIGIBLE,
             ],
             [2, 1],
             0,

diff --git a/policyengine_us/variables/household/cliff.py b/policyengine_us/variables/household/cliff.py
@@ -48,6 +48,7 @@ def formula(person, period, parameters):
             alt_person = alt_sim.person
             netinc_alt = alt_person.household("household_net_income", period)
             increase[mask] += netinc_alt[mask] - netinc_base[mask]
+            del sim.branches[f"cliff_for_adult_{adult_index}"]
         return where(increase < 0, -increase, 0)
 
 

diff --git a/policyengine_us/variables/household/marginal_tax_rate.py b/policyengine_us/variables/household/marginal_tax_rate.py
@@ -46,6 +46,7 @@ def formula(person, period, parameters):
             netinc_alt = alt_person.household("household_net_income", period)
             increase = netinc_alt - netinc_base
             mtr_values += where(mask, 1 - increase / delta, 0)
+            del sim.branches[f"mtr_for_adult_{adult_index}"]
         return mtr_values