Let callers request extra variables from calculate_household_impact#287
Open
Let callers request extra variables from calculate_household_impact#287
Conversation
Adds an `extra_variables` mapping (`{entity_name: [variable_name,...]}`)
to:
- policyengine.core.Simulation
- policyengine.tax_benefit_models.us.calculate_household_impact
- policyengine.tax_benefit_models.uk.calculate_household_impact
- PolicyEngineUSLatest.run / PolicyEngineUKLatest.run
Previously both simulations computed only the hardcoded
`self.entity_variables` list (a small default curated for UI/API
views) and returned outputs missing anything else — including
adjusted_gross_income, state_agi, free_school_meals,
is_medicaid_eligible, income_tax_refundable_credits,
state_refundable_credits, state_income_tax_before_refundable_credits,
and every other variable a benchmark suite typically needs.
The fix threads a simulation-scoped override through the run path.
Extra variables are merged with the defaults (dedup; order preserved
for the defaults) so existing callers see no change, and the
returned USHouseholdOutput / UKHouseholdOutput dicts gain the
requested keys.
Unblocks the policybench migration to policyengine.py: callers can
now do
calculate_household_impact(
input,
extra_variables={"tax_unit": ["adjusted_gross_income", ...]},
)
without monkey-patching us_latest.entity_variables.
Tests: three new tests in test_household_impact.py cover the tax_unit
and household paths plus the no-op behaviour when extra_variables is
empty/omitted.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
`policyengine.tax_benefit_models.{us,uk}.calculate_household_impact` (and the underlying `Simulation.run`) only computes variables that appear in the hardcoded `entity_variables` dict on `PolicyEngineUSLatest` / `PolicyEngineUKLatest`. That list is a small curated default for UI/API views; it excludes most variables a benchmark suite needs:
`adjusted_gross_income`, `state_agi`, `free_school_meals`, `is_medicaid_eligible`, `income_tax_refundable_credits`, `state_refundable_credits`, `state_income_tax_before_refundable_credits`, …
This blocks the `policybench` migration from direct `policyengine_us` imports to the `policyengine.py` certified-bundle path.
Fix
Thread a simulation-scoped `extra_variables` mapping through `Simulation` → `PolicyEngineUSLatest.run` / `PolicyEngineUKLatest.run` → `calculate_household_impact`. The merged dict is deduped; bundle defaults keep their order. Existing callers see identical output.
```python
from policyengine.tax_benefit_models.us import (
calculate_household_impact, USHouseholdInput,
)
result = calculate_household_impact(
USHouseholdInput(people=[{"age": 35, "employment_income": 60000}], year=2026),
extra_variables={"tax_unit": ["adjusted_gross_income"]},
)
assert result.tax_unit[0]["adjusted_gross_income"] > 0
```
Test plan
🤖 Generated with Claude Code