Count households once: sum the weight column unweighted

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"

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:

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