From b411acd0a99e8655eff607fdd6f2edbd7001e8b2 Mon Sep 17 00:00:00 2001 From: PavelMakarchuk Date: Mon, 12 Oct 2026 16:56:54 -0400 Subject: [PATCH] Count households once: sum the weight column unweighted budget.households summed household_weight through the household MicroDataFrame, whose columns are already weighted by household_weight, so it returned the sum of squared weights: about 10^4 times the count on every run (state/ny 135.6bn). Count entities with a plain sum of the weight column, and test the budget on a real MicroDataFrame. Fixes #766 Co-Authored-By: Claude Opus 5.5 --- .../simulation_output_budget.py | 5 +- .../simulation_output_common.py | 16 ++++++ .../tests/test_simulation_output_budget.py | 50 +++++++++++++++++++ 3 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py index 8f7e00d34..e12028e68 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_budget.py @@ -12,6 +12,7 @@ from policyengine_simulation_executor.simulation_output_common import ( _change_output_variable, _collection_records, + _count_output_entities, _number, _sum_output_variable, ) @@ -59,8 +60,8 @@ def build_budgetary_impact(country: str, baseline, reform) -> BudgetaryImpact: state_tax_revenue_impact=state_tax_revenue_impact, benefit_spending_impact=benefit_spending_impact, budgetary_impact=tax_revenue_impact - benefit_spending_impact, - households=_sum_output_variable( - baseline, "household_weight", entity="household" + households=_count_output_entities( + baseline, entity="household", weight="household_weight" ), baseline_net_income=_sum_output_variable( baseline, "household_net_income", entity="household" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py index c031b2ee7..941b4ab72 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py @@ -8,6 +8,8 @@ from importlib import import_module from typing import Any +import numpy as np + logger = logging.getLogger(__name__) @@ -50,6 +52,20 @@ def _entity_data(simulation, entity: str): return getattr(simulation.output_dataset.data, entity) +def _count_output_entities(simulation, entity: str, weight: str) -> float: + """The weighted number of entities: the plain sum of their weights. + + Output frames are microdf ``MicroDataFrame``s weighted by the entity's + weight, so a column's ``.sum()`` is already a weighted sum. Summing the + weight column that way weights it twice and returns the sum of squared + weights, about 10^4 times the household count on the sparse US data. + Read the raw values instead. A weighted frame always holds its weight + column, so there is no aggregate fallback: that would weight it again. + """ + data = _entity_data(simulation, entity) + return float(np.asarray(data[weight], dtype=float).sum()) + + def _sum_output_variable(simulation, variable: str, entity: str) -> float: data = _entity_data(simulation, entity) if variable in data.columns: diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py new file mode 100644 index 000000000..69ff5c2fb --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_budget.py @@ -0,0 +1,50 @@ +"""Budget totals from weighted output frames, as the worker's simulations hold them.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +from microdf import MicroDataFrame + +from policyengine_simulation_executor.simulation_output_budget import ( + build_budgetary_impact, +) + + +def _simulation(columns: dict[str, list[float]]): + """A simulation whose household output is a MicroDataFrame, as policyengine builds it.""" + household = MicroDataFrame(pd.DataFrame(columns), weights="household_weight") + return SimpleNamespace( + output_dataset=SimpleNamespace(data=SimpleNamespace(household=household)) + ) + + +def test_households_counts_each_weight_once_on_weighted_frames(): + baseline = _simulation( + { + "household_weight": [1.0, 2.0], + "household_net_income": [100.0, 200.0], + "household_tax": [20.0, 40.0], + "household_benefits": [5.0, 10.0], + "household_state_income_tax": [2.0, 3.0], + } + ) + reform = _simulation( + { + "household_weight": [1.0, 2.0], + "household_net_income": [120.0, 210.0], + "household_tax": [25.0, 50.0], + "household_benefits": [4.0, 8.0], + "household_state_income_tax": [4.0, 6.0], + } + ) + + budget = build_budgetary_impact("us", baseline, reform) + + # Households weighted 1 and 2 are 3 households, not 1² + 2² = 5. + assert budget.households == 3.0 + # The other totals stay weighted sums, Σ weight × value. + assert budget.baseline_net_income == 500.0 + assert budget.tax_revenue_impact == 25.0 + assert budget.state_tax_revenue_impact == 8.0